diff --git a/.coveragerc b/.coveragerc index 9627cf1200..d9174acf52 100644 --- a/.coveragerc +++ b/.coveragerc @@ -1,10 +1,12 @@ [run] +concurrency = multiprocessing,thread omit = ding/utils/slurm_helper.py ding/utils/file_helper.py ding/utils/linklink_dist_helper.py ding/utils/pytorch_ddp_dist_helper.py ding/utils/k8s_helper.py + ding/utils/tests/test_k8s_launcher.py ding/utils/time_helper_cuda.py ding/utils/time_helper_base.py ding/utils/data/tests/test_dataloader.py diff --git a/.github/workflows/algo_test.yml b/.github/workflows/algo_test.yml index e61fe104aa..dbd5a95549 100644 --- a/.github/workflows/algo_test.yml +++ b/.github/workflows/algo_test.yml @@ -16,12 +16,12 @@ jobs: if: "!contains(github.event.head_commit.message, 'ci skip')" strategy: matrix: - python-version: [3.7, 3.8, 3.9] + python-version: [3.8, 3.9] steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v4 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: do_algotest @@ -31,5 +31,6 @@ jobs: run: | python -m pip install . python -m pip install ".[test,k8s]" + python -m pip install transformers ./ding/scripts/install-k8s-tools.sh make algotest diff --git a/.github/workflows/badge.yml b/.github/workflows/badge.yml index f82cf24fdf..d010d21bf8 100644 --- a/.github/workflows/badge.yml +++ b/.github/workflows/badge.yml @@ -12,13 +12,13 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: [ 3.7 ] + python-version: [ 3.8 ] env: GIST_ID: 3690cccd811e4c5f771075c2f785c7bb steps: - uses: actions/checkout@v2 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: Download cloc diff --git a/.github/workflows/deploy.yml b/.github/workflows/deploy.yml index 5246076d78..f4c6d8ce0f 100644 --- a/.github/workflows/deploy.yml +++ b/.github/workflows/deploy.yml @@ -3,11 +3,10 @@ name: deploy # deploy docker on: push: branches: [main, '*deploy*', '*docker*'] - jobs: docker_base: runs-on: ubuntu-latest - # if: "contains(github.event.head_commit.message, 'enable docker')" + if: "!contains(github.event.head_commit.message, 'ci skip')" strategy: matrix: platform: [linux/amd64] @@ -48,7 +47,7 @@ jobs: uses: docker/setup-buildx-action@v1 - name: Cache Docker layers - uses: actions/cache@v2 + uses: actions/cache@v4 with: path: /tmp/.buildx-cache key: ${{ runner.os }}-buildx-${{ github.sha }} @@ -80,11 +79,10 @@ jobs: docker_doc: runs-on: ubuntu-latest - # if: "contains(github.event.head_commit.message, 'enable docker')" + if: "contains(github.event.head_commit.message, 'doc docker')" strategy: matrix: platform: [linux/amd64] - # python-version: [3.6, 3.7, 3.8] steps: - name: Checkout uses: actions/checkout@v2 @@ -113,7 +111,7 @@ jobs: uses: docker/setup-buildx-action@v1 - name: Cache Docker layers - uses: actions/cache@v2 + uses: actions/cache@v4 with: path: /tmp/.buildx-cache key: ${{ runner.os }}-buildx-${{ github.sha }} @@ -144,7 +142,7 @@ jobs: docker_atari: runs-on: ubuntu-latest needs: docker_base - # if: "contains(github.event.head_commit.message, 'enable docker')" + if: "!contains(github.event.head_commit.message, 'ci skip')" strategy: matrix: platform: [linux/amd64] @@ -172,7 +170,7 @@ jobs: uses: docker/setup-buildx-action@v1 - name: Cache Docker layers - uses: actions/cache@v2 + uses: actions/cache@v4 with: path: /tmp/.buildx-cache key: ${{ runner.os }}-buildx-${{ github.sha }} @@ -203,7 +201,7 @@ jobs: docker_mujoco: runs-on: ubuntu-latest needs: docker_base - # if: "contains(github.event.head_commit.message, 'enable docker')" + if: "!contains(github.event.head_commit.message, 'ci skip')" strategy: matrix: platform: [linux/amd64] @@ -231,7 +229,7 @@ jobs: uses: docker/setup-buildx-action@v1 - name: Cache Docker layers - uses: actions/cache@v2 + uses: actions/cache@v4 with: path: /tmp/.buildx-cache key: ${{ runner.os }}-buildx-${{ github.sha }} @@ -261,7 +259,7 @@ jobs: docker_metaworld: runs-on: ubuntu-latest needs: docker_base - # if: "contains(github.event.head_commit.message, 'enable docker')" + if: "contains(github.event.head_commit.message, 'metaworld docker')" strategy: matrix: platform: [linux/amd64] @@ -288,7 +286,7 @@ jobs: uses: docker/setup-buildx-action@v1 - name: Cache Docker layers - uses: actions/cache@v2 + uses: actions/cache@v4 with: path: /tmp/.buildx-cache key: ${{ runner.os }}-buildx-${{ github.sha }} @@ -386,3 +384,72 @@ jobs: run: | docker buildx build -f ./docker/Dockerfile.env . -t opendilab/ding:nightly-dmc2gym --target=dmc2gym docker push opendilab/ding:nightly-dmc2gym + + docker_rpc: + runs-on: ubuntu-latest + needs: docker_base + if: "contains(github.event.head_commit.message, 'test rpc')" + strategy: + matrix: + platform: [linux/amd64] + # python-version: [3.6, 3.7, 3.8] + steps: + - name: Checkout + uses: actions/checkout@v2 + - name: Login to DockerHub + uses: docker/login-action@v1 + with: + username: ${{ secrets.DOCKERIO_USERNAME }} + password: ${{ secrets.DOCKERIO_PASSWORD }} + + - name: Build and push + id: docker_build + run: | + docker buildx build -f ./docker/Dockerfile.rpc . -t opendilab/ding:nightly-rpc-base --target=base + docker push opendilab/ding:nightly-rpc-base + + docker_evogym: + runs-on: ubuntu-latest + needs: docker_base + if: "contains(github.event.head_commit.message, 'evogym docker')" + strategy: + matrix: + platform: [linux/amd64] + # python-version: [3.6, 3.7, 3.8] + steps: + - name: Checkout + uses: actions/checkout@v2 + - name: Login to DockerHub + uses: docker/login-action@v1 + with: + username: ${{ secrets.DOCKERIO_USERNAME }} + password: ${{ secrets.DOCKERIO_PASSWORD }} + + - name: Build and push + id: docker_build + run: | + docker buildx build -f ./docker/Dockerfile.env . -t opendilab/ding:nightly-evogym --target=evogym + docker push opendilab/ding:nightly-evogym + + docker_d4rl: + runs-on: ubuntu-latest + needs: docker_mujoco + if: "contains(github.event.head_commit.message, 'd4rl docker')" + strategy: + matrix: + platform: [linux/amd64] + # python-version: [3.6, 3.7, 3.8] + steps: + - name: Checkout + uses: actions/checkout@v2 + - name: Login to DockerHub + uses: docker/login-action@v1 + with: + username: ${{ secrets.DOCKERIO_USERNAME }} + password: ${{ secrets.DOCKERIO_PASSWORD }} + + - name: Build and push + id: docker_build + run: | + docker buildx build -f ./docker/Dockerfile.env . -t opendilab/ding:nightly-d4rl --target=d4rl + docker push opendilab/ding:nightly-d4rl diff --git a/.github/workflows/doc.yml b/.github/workflows/doc.yml index 9465203300..0dc0235c70 100644 --- a/.github/workflows/doc.yml +++ b/.github/workflows/doc.yml @@ -17,7 +17,7 @@ jobs: steps: - uses: actions/checkout@v2 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: Generate diff --git a/.github/workflows/envpool_test.yml b/.github/workflows/envpool_test.yml index 66ea2f8d11..7d1e01aa8d 100644 --- a/.github/workflows/envpool_test.yml +++ b/.github/workflows/envpool_test.yml @@ -11,12 +11,12 @@ jobs: if: "!contains(github.event.head_commit.message, 'ci skip')" strategy: matrix: - python-version: [3.7, 3.8] # Envpool only supports python>=3.7 + python-version: [3.8] # Envpool only supports python>=3.7 steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v4 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: do_envpool_test @@ -24,5 +24,6 @@ jobs: python -m pip install . python -m pip install ".[test,k8s]" python -m pip install ".[envpool]" + python -m pip install transformers ./ding/scripts/install-k8s-tools.sh make envpooltest diff --git a/.github/workflows/platform_test.yml b/.github/workflows/platform_test.yml index a9ba3a2b4d..b3e92c175c 100644 --- a/.github/workflows/platform_test.yml +++ b/.github/workflows/platform_test.yml @@ -11,13 +11,13 @@ jobs: if: "!contains(github.event.head_commit.message, 'ci skip')" strategy: matrix: - os: [macos-latest, windows-latest] - python-version: [3.7, 3.8, 3.9] + os: [macos-13, windows-latest] + python-version: [3.8, 3.9] steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v4 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: do_platform_test @@ -25,5 +25,6 @@ jobs: run: | python -m pip install . python -m pip install ".[test,k8s]" + python -m pip install transformers python -m pip uninstall pytest-timeouts -y make platformtest diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 73ecaf44ff..4305e215d2 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -10,14 +10,14 @@ jobs: strategy: matrix: os: - - 'ubuntu-18.04' - python-version: [3.7] + - ubuntu-latest + python-version: [3.8] steps: - name: Checkout code - uses: actions/checkout@v2 + uses: actions/checkout@v4 - name: Set up python ${{ matrix.python-version }} - uses: actions/setup-python@v2 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: Set up python dependences diff --git a/.github/workflows/release_conda.yml b/.github/workflows/release_conda.yml index 2c4fe83191..10fff557ea 100644 --- a/.github/workflows/release_conda.yml +++ b/.github/workflows/release_conda.yml @@ -7,7 +7,7 @@ jobs: runs-on: ubuntu-latest if: github.event_name == 'push' && (startsWith(github.ref, 'refs/tags') || contains(github.event.head_commit.message, 'conda')) steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v4 - name: publish-to-conda uses: fcakyon/conda-publish-action@v1.3 with: diff --git a/.github/workflows/style.yml b/.github/workflows/style.yml index a00e690505..1ed80baa8e 100644 --- a/.github/workflows/style.yml +++ b/.github/workflows/style.yml @@ -10,16 +10,17 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: [3.7, 3.8, 3.9] + python-version: ['3.8', '3.9', '3.10'] steps: - uses: actions/checkout@v2 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: code style run: | python -m pip install "yapf==0.29.0" python -m pip install "flake8<=3.9.2" + python -m pip install "importlib-metadata<5.0.0" make format_test flake_check diff --git a/.github/workflows/unit_test.yml b/.github/workflows/unit_test.yml index c7195d820b..7bb14a2350 100644 --- a/.github/workflows/unit_test.yml +++ b/.github/workflows/unit_test.yml @@ -8,15 +8,15 @@ on: [push, pull_request] jobs: test_unittest: runs-on: ubuntu-latest - if: "!contains(github.event.head_commit.message, 'ci skip')" + if: ( !contains(github.event.head_commit.message, 'ci skip') && !contains(github.event.head_commit.message, 'ut skip')) strategy: matrix: - python-version: [3.7, 3.8, 3.9] + python-version: ['3.8', '3.9', '3.10'] steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v4 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: do_unittest @@ -25,6 +25,7 @@ jobs: python -m pip install box2d-py python -m pip install . python -m pip install ".[test,k8s]" + python -m pip install transformers ./ding/scripts/install-k8s-tools.sh make unittest - name: Upload coverage to Codecov @@ -41,17 +42,18 @@ jobs: if: "!contains(github.event.head_commit.message, 'ci skip')" strategy: matrix: - python-version: [3.7, 3.8, 3.9] + python-version: ['3.8', '3.9', '3.10'] steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v4 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: do_benchmark run: | python -m pip install . python -m pip install ".[test,k8s]" + python -m pip install transformers ./ding/scripts/install-k8s-tools.sh make benchmark diff --git a/.gitignore b/.gitignore index 3d09843579..45c01ea689 100644 --- a/.gitignore +++ b/.gitignore @@ -15,6 +15,7 @@ pkg/ src/ +impala_log/ ### CVS template /CVS/* @@ -124,6 +125,8 @@ local.properties # Graphics Interchange Format *.gif +*.mp4 +*.mpg # RAW *.raw @@ -1423,4 +1426,11 @@ events.* eval_config.py collect_demo_data_config.py !ding/**/*.py -events.* \ No newline at end of file +events.* + +evogym/* +ding/example/* +ding/framework/middleware/tests/wandb/ +ding/.style.yapf +ding/format.sh +ding/framework/middleware_v3/ diff --git a/CHANGELOG b/CHANGELOG index 36a7a4ef27..b155b976fc 100644 --- a/CHANGELOG +++ b/CHANGELOG @@ -1,3 +1,232 @@ +2024.12.23 (v0.5.3) +- env: add pistonball MARL env and its unittest/example (#833) +- env: update trading env (#831) +- env: update ppo config for better discrete action space performance (#809) +- env: remove unused config fields in MuJoCo PPO +- algo: add AWR algorithm (#828) +- algo: add encoder in MAVAC (#823) +- algo: add HPT model architecture (#841) +- algo: fix multiple model wrappers reset bug (#846) +- algo: add hybrid action space support to ActionNoiseWrapper (#829) +- algo: fix mappo adv compute bug (#812) +- feature: add resume_training option to allow the envstep and train_iter resume seamlessly (#835) +- feature: polish old/new pipeline DistributedDataParallel (DDP) implementation (#842) +- feature: adapt DingEnvWrapper to gymnasium (#817) +- fix: priority buffer delete bug (#844) +- fix: middleware collector env reset bug (#845) +- fix: fix many unittest bugs +- style: downgrade pyecharts log level to warning and polish installation doc (#838) +- style: polish necessary requirements +- style: polish api doc details +- style: polish DI-engine citation authors +- style: upgrade CI macos version from 12 to 13 + +2024.06.27(v0.5.2) +- env: add taxi env (#799) (#807) +- env: add ising model env (#782) +- env: add new Flozen Lake env (#781) +- env: optimize ppo continuous config in MuJoCo (#801) +- env: fix masac smac config multi_agent=True bug (#791) +- env: update/speed up pendulum ppo +- algo: fix gtrxl compatibility bug (#796) +- algo: fix complex obs demo for ppo pipeline (#786) +- algo: add naive PWIL demo +- algo: fix marl nstep td compatibility bug +- feature: add GPU utils (#788) +- feature: add deprecated function decorator (#778) +- style: relax flask requirement (#811) +- style: add new badge (hellogithub) in readme (#805) +- style: update discord link and badge in readme (#795) +- style: fix typo in config.py (#776) +- style: polish rl_utils api docs +- style: add constraint about numpy<2 +- style: polish macos platform test version to 12 +- style: polish ci python version + +2024.02.04(v0.5.1) +- env: add MADDPG pettingzoo example (#774) +- env: polish NGU Atari configs (#767) +- env: fix bug in cliffwalking env (#759) +- env: add PettingZoo replay video demo +- env: change default max retry in env manager from 5 to 1 +- algo: add QGPO diffusion-model related algorithm (#757) +- algo: add HAPPO multi-agent algorithm (#717) +- algo: add DreamerV3 + MiniGrid adaption (#725) +- algo: fix hppo entropy_weight to avoid nan error in log_prob (#761) +- algo: fix structured action bug (#760) +- algo: polish Decision Transformer entry (#754) +- algo: fix EDAC policy/model bug +- fix: env typos +- fix: pynng requirements bug +- fix: communication module unittest bug +- style: polish policy API doc (#762) (#764) (#768) +- style: add agent API doc (#758) +- style: polish torch_utils/utils API doc (#745) (#747) (#752) (#755) (#763) + +2023.11.06(v0.5.0) +- env: add tabmwp env (#667) +- env: polish anytrading env issues (#731) +- algo: add PromptPG algorithm (#667) +- algo: add Plan Diffuser algorithm (#700) +- algo: add new pipeline implementation of IMPALA algorithm (#713) +- algo: add dropout layers to DQN-style algorithms (#712) +- feature: add new pipeline agent for sac/ddpg/a2c/ppo and Hugging Face support (#637) (#730) (#737) +- feature: add more unittest cases for model (#728) +- feature: add collector logging in new pipeline (#735) +- fix: logger middleware problems (#715) +- fix: ppo parallel bug (#709) +- fix: typo in optimizer_helper.py (#726) +- fix: mlp dropout if condition bug +- fix: drex collecting data unittest bugs +- style: polish env manager/wrapper comments and API doc (#742) +- style: polish model comments and API doc (#722) (#729) (#734) (#736) (#741) +- style: polish policy comments and API doc (#732) +- style: polish rl_utils comments and API doc (#724) +- style: polish torch_utils comments and API doc (#738) +- style: update README.md and Colab demo (#733) +- style: update metaworld docker image + +2023.08.23(v0.4.9) +- env: add cliffwalking env (#677) +- env: add lunarlander ppo config and example +- algo: add BCQ offline RL algorithm (#640) +- algo: add Dreamerv3 model-based RL algorithm (#652) +- algo: add tensor stream merge network tools (#673) +- algo: add scatter connection model (#680) +- algo: refactor Decision Transformer in new pipeline and support img input and discrete output (#693) +- algo: add three variants of Bilinear classes and a FiLM class (#703) +- feature: polish offpolicy RL multi-gpu DDP training (#679) +- feature: add middleware for Ape-X distributed pipeline (#696) +- feature: add example for evaluating trained DQN (#706) +- fix: to_ndarray fails to assign dtype for scalars (#708) +- fix: evaluator return episode_info compatibility bug +- fix: cql example entry wrong config bug +- fix: enable_save_figure env interface +- fix: redundant env info bug in evaluator +- fix: to_item unittest bug +- style: polish and simplify requirements (#672) +- style: add Hugging Face Model Zoo badge (#674) +- style: add openxlab Model Zoo badge (#675) +- style: fix py37 macos ci bug and update default pytorch from 1.7.1 to 1.12.1 (#678) +- style: fix mujoco-py compatibility issue for cython<3 (#711) +- style: fix type spell error (#704) +- style: fix pypi release actions ubuntu 18.04 bug +- style: update contact information (e.g. wechat) +- style: polish algorithm doc tables + +2023.05.25(v0.4.8) +- env: fix gym hybrid reward dtype bug (#664) +- env: fix atari env id noframeskip bug (#655) +- env: fix typo in gym any_trading env (#654) +- env: update td3bc d4rl config (#659) +- env: polish bipedalwalker config +- algo: add EDAC offline RL algorithm (#639) +- algo: add LN and GN norm_type support in ResBlock (#660) +- algo: add normal value norm baseline for PPOF (#658) +- algo: polish last layer init/norm in MLP (#650) +- algo: polish TD3 monitor variable +- feature: add MAPPO/MASAC task example (#661) +- feature: add PPO example for complex env observation (#644) +- feature: add barrier middleware (#570) +- fix: abnormal collector log and add record_random_collect option (#662) +- fix: to_item compatibility bug (#646) +- fix: trainer dtype transform compatibility bug +- fix: pettingzoo 1.23.0 compatibility bug +- fix: ensemble head unittest bug +- style: fix incompatible gym version bug in Dockerfile.env (#653) +- style: add more algorithm docs + +2023.04.11(v0.4.7) +- env: add dmc2gym env support and baseline (#451) +- env: update pettingzoo to the latest version (#597) +- env: polish icm/rnd+onppo config bugs and add app_door_to_key env (#564) +- env: add lunarlander continuous TD3/SAC config +- env: polish lunarlander discrete C51 config +- algo: add Procedure Cloning (PC) imitation learning algorithm (#514) +- algo: add Munchausen Reinforcement Learning (MDQN) algorithm (#590) +- algo: add reward/value norm methods: popart & value rescale & symlog (#605) +- algo: polish reward model config and training pipeline (#624) +- algo: add PPOF reward space demo support (#608) +- algo: add PPOF Atari demo support (#589) +- algo: polish dqn default config and env examples (#611) +- algo: polish comment and clean code about SAC +- feature: add language model (e.g. GPT) training utils (#625) +- feature: remove policy cfg sub fields requirements (#620) +- feature: add full wandb support (#579) +- fix: confusing shallow copy operation about next_obs (#641) +- fix: unsqueeze action_args in PDQN when shape is 1 (#599) +- fix: evaluator return_info tensor type bug (#592) +- fix: deque buffer wrapper PER bug (#586) +- fix: reward model save method compatibility bug +- fix: logger assertion and unittest bug +- fix: bfs test py3.9 compatibility bug +- fix: zergling collector unittest bug +- style: add DI-engine torch-rpc p2p communication docker (#628) +- style: add D4RL docker (#591) +- style: correct typo in task (#617) +- style: correct typo in time_helper (#602) +- style: polish readme and add treetensor example +- style: update contributing doc + +2023.02.16(v0.4.6) +- env: add metadrive env and related ppo config (#574) +- env: add acrobot env and related dqn config (#577) +- env: add carracing in box2d (#575) +- env: add new gym hybrid viz (#563) +- env: update cartpole IL config (#578) +- algo: add BDQ algorithm (#558) +- algo: add procedure cloning model (#573) +- feature: add simplified PPOF (PPO × Family) interface (#567) (#568) (#581) (#582) +- fix: to_device and prev_state bug when using ttorch (#571) +- fix: py38 and numpy unittest bugs (#565) +- fix: typo in contrastive_loss.py (#572) +- fix: dizoo envs pkg installation bugs +- fix: multi_trainer middleware unittest bug +- style: add evogym docker (#580) +- style: fix metaworld docker bug +- style: fix setuptools high version incompatibility bug +- style: extend treetensor lowest version + +2022.12.13(v0.4.5) +- env: add beergame supply chain optimization env (#512) +- env: add env gym_pybullet_drones (#526) +- env: rename eval reward to episode return (#536) +- algo: add policy gradient algo implementation (#544) +- algo: add MADDPG algo implementation (#550) +- algo: add IMPALA continuous algo implementation (#551) +- algo: add MADQN algo implementation (#540) +- feature: add new task IMPALA-type distributed training scheme (#321) +- feature: add load and save method for replaybuffer (#542) +- feature: add more DingEnvWrapper example (#525) +- feature: add evaluator more info viz support (#538) +- feature: add trackback log for subprocess env manager (#534) +- fix: halfcheetah td3 config file (#537) +- fix: mujoco action_clip args compatibility bug (#535) +- fix: atari a2c config entry bug +- fix: drex unittest compatibility bug +- style: add Roadmap issue of DI-engine (#548) +- style: update related project link and new env doc + +2022.10.31(v0.4.4) +- env: add modified gym-hybrid including moving, sliding and hardmove (#505) (#519) +- env: add evogym support (#495) (#527) +- env: add save_replay_gif option (#506) +- env: adapt minigrid_env and related config to latest MiniGrid v2.0.0 (#500) +- algo: add pcgrad optimizer (#489) +- algo: add some features in MLP and ResBlock (#511) +- algo: delete mcts related modules (#518) +- feature: add wandb middleware and demo (#488) (#523) (#528) +- feature: add new properties in Context (#499) +- feature: add single env policy wrapper for policy deployment +- feature: add custom model demo and doc +- fix: build logger args and unittests (#522) +- fix: total_loss calculation in PDQN (#504) +- fix: save gif function bug +- fix: level sample unittest bug +- style: update contact email address (#503) +- style: polish env log and resblock name +- style: add details button in readme + 2022.09.23(v0.4.3) - env: add rule-based gomoku expert (#465) - algo: fix a2c policy batch size bug (#481) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 3600608c22..ecfc1e0670 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -1,5 +1,7 @@ -[git guide](https://opendilab.github.io/DI-engine/guide/git_guide.html) +[Git Guide](https://di-engine-docs.readthedocs.io/en/latest/24_cooperation/git_guide.html) -[doc contribution guide](https://opendilab.github.io/DI-engine/guide/doc_contribution.html) +[GitHub Cooperation Guide](https://di-engine-docs.readthedocs.io/en/latest/24_cooperation/issue_pr.html) -[code review guide](https://opendilab.github.io/DI-engine/guide/code_comment.html) + - [Code Style](https://di-engine-docs.readthedocs.io/en/latest/21_code_style/index.html) + - [Unit Test](https://di-engine-docs.readthedocs.io/en/latest/22_test/index.html) + - [Code Review](https://di-engine-docs.readthedocs.io/en/latest/24_cooperation/issue_pr.html#pr-s-code-review) diff --git a/README.md b/README.md index eee4a3f3b3..738e7b3e59 100644 --- a/README.md +++ b/README.md @@ -15,14 +15,13 @@ ![Comments](https://img.shields.io/endpoint?url=https://gist.githubusercontent.com/HansBug/3690cccd811e4c5f771075c2f785c7bb/raw/comments.json) ![Style](https://github.com/opendilab/DI-engine/actions/workflows/style.yml/badge.svg) -![Docs](https://github.com/opendilab/DI-engine/actions/workflows/doc.yml/badge.svg) +[![Read en Docs](https://github.com/opendilab/DI-engine/actions/workflows/doc.yml/badge.svg)](https://di-engine-docs.readthedocs.io/en/latest) +[![Read zh_CN Docs](https://img.shields.io/readthedocs/di-engine-docs?label=%E4%B8%AD%E6%96%87%E6%96%87%E6%A1%A3)](https://di-engine-docs.readthedocs.io/zh_CN/latest) ![Unittest](https://github.com/opendilab/DI-engine/actions/workflows/unit_test.yml/badge.svg) ![Algotest](https://github.com/opendilab/DI-engine/actions/workflows/algo_test.yml/badge.svg) ![deploy](https://github.com/opendilab/DI-engine/actions/workflows/deploy.yml/badge.svg) [![codecov](https://codecov.io/gh/opendilab/DI-engine/branch/main/graph/badge.svg?token=B0Q15JI301)](https://codecov.io/gh/opendilab/DI-engine) - - ![GitHub Org's stars](https://img.shields.io/github/stars/opendilab) [![GitHub stars](https://img.shields.io/github/stars/opendilab/DI-engine)](https://github.com/opendilab/DI-engine/stargazers) [![GitHub forks](https://img.shields.io/github/forks/opendilab/DI-engine)](https://github.com/opendilab/DI-engine/network) @@ -31,50 +30,99 @@ [![GitHub pulls](https://img.shields.io/github/issues-pr/opendilab/DI-engine)](https://github.com/opendilab/DI-engine/pulls) [![Contributors](https://img.shields.io/github/contributors/opendilab/DI-engine)](https://github.com/opendilab/DI-engine/graphs/contributors) [![GitHub license](https://img.shields.io/github/license/opendilab/DI-engine)](https://github.com/opendilab/DI-engine/blob/master/LICENSE) +[![Hugging Face](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Models-yellow)](https://huggingface.co/OpenDILabCommunity) +[![Open in OpenXLab](https://cdn-static.openxlab.org.cn/header/openxlab_models.svg)](https://openxlab.org.cn/models?search=opendilab) +[![discord badge](https://dcbadge.vercel.app/api/server/dkZS2JF56X?style=flat)](https://discord.gg/dkZS2JF56X) +[![slack badge](https://img.shields.io/badge/Slack-join-blueviolet?logo=slack&)](https://join.slack.com/t/opendilab/shared_invite/zt-v9tmv4fp-nUBAQEH1_Kuyu_q4plBssQ) -Updated on 2022.09.23 DI-engine-v0.4.3 +
+ + Featured｜HelloGitHub + +
+
+Updated on 2024.12.23 DI-engine-v0.5.3 ## Introduction to DI-engine -[DI-engine doc](https://di-engine-docs.readthedocs.io/en/latest/) | [中文文档](https://di-engine-docs.readthedocs.io/zh_CN/latest/) -**DI-engine** is a generalized decision intelligence engine. It supports **various [deep reinforcement learning](https://di-engine-docs.readthedocs.io/en/latest/10_concepts/index.html) algorithms** ([link](https://di-engine-docs.readthedocs.io/en/latest/12_policies/index.html)): +[Documentation](https://di-engine-docs.readthedocs.io/en/latest/) | [中文文档](https://di-engine-docs.readthedocs.io/zh_CN/latest/) | [Tutorials](https://di-engine-docs.readthedocs.io/en/latest/01_quickstart/index.html) | [Feature](#feature) | [Task & Middleware](https://di-engine-docs.readthedocs.io/en/latest/03_system/index.html) | [TreeTensor](#general-data-container-treetensor) | [Roadmap](https://github.com/opendilab/DI-engine/issues/548) + +**DI-engine** is a generalized decision intelligence engine for PyTorch and JAX. -- Most basic DRL algorithms, such as DQN, PPO, SAC, R2D2, IMPALA -- Multi-agent RL algorithms like QMIX, MAPPO -- Imitation learning algorithms (BC/IRL/GAIL) , such as GAIL, SQIL, Guided Cost Learning, Implicit Behavioral Cloning -- Exploration algorithms like HER, RND, ICM, NGU -- Offline RL algorithms: CQL, TD3BC, Decision Transformer -- Model-based RL algorithms: SVG, MVE, STEVE / MBPO, DDPPO +It provides **python-first** and **asynchronous-native** task and middleware abstractions, and modularly integrates several of the most important decision-making concepts: Env, Policy and Model. Based on the above mechanisms, DI-engine supports **various [deep reinforcement learning](https://di-engine-docs.readthedocs.io/en/latest/10_concepts/index.html) algorithms** with superior performance, high efficiency, well-organized [documentation](https://di-engine-docs.readthedocs.io/en/latest/) and [unittest](https://github.com/opendilab/DI-engine/actions): -**DI-engine** aims to **standardize different Decision Intelligence enviroments and applications**. Various training pipelines and customized decision AI applications are also supported. +- Most basic DRL algorithms: such as DQN, Rainbow, PPO, TD3, SAC, R2D2, IMPALA +- Multi-agent RL algorithms: such as QMIX, WQMIX, MAPPO, HAPPO, ACE +- Imitation learning algorithms (BC/IRL/GAIL): such as GAIL, SQIL, Guided Cost Learning, Implicit BC +- Offline RL algorithms: BCQ, CQL, TD3BC, Decision Transformer, EDAC, Diffuser, Decision Diffuser, SO2 +- Model-based RL algorithms: SVG, STEVE, MBPO, DDPPO, DreamerV3 +- Exploration algorithms: HER, RND, ICM, NGU +- LLM + RL Algorithms: PPO-max, DPO, PromptPG, PromptAWR +- Other algorithms: such as PER, PLR, PCGrad +- MCTS + RL algorithms: AlphaZero, MuZero, please refer to [LightZero](https://github.com/opendilab/LightZero) +- Generative Model + RL algorithms: Diffusion-QL, QGPO, SRPO, please refer to [GenerativeRL](https://github.com/opendilab/GenerativeRL) + + +**DI-engine** aims to **standardize different Decision Intelligence environments and applications**, supporting both academic research and prototype applications. Various training pipelines and customized decision AI applications are also supported: + +
+(Click to Collapse) - Traditional academic environments - - [DI-zoo](https://github.com/opendilab/DI-engine#environment-versatility) + - [DI-zoo](https://github.com/opendilab/DI-engine#environment-versatility): various decision intelligence demonstrations and benchmark environments with DI-engine. +- Tutorial courses + - [PPOxFamily](https://github.com/opendilab/PPOxFamily): PPO x Family DRL Tutorial Course - Real world decision AI applications - [DI-star](https://github.com/opendilab/DI-star): Decision AI in StarCraftII + - [PsyDI](https://github.com/opendilab/PsyDI): Towards a Multi-Modal and Interactive Chatbot for Psychological Assessments - [DI-drive](https://github.com/opendilab/DI-drive): Auto-driving platform - - [GoBigger](https://github.com/opendilab/GoBigger): Multi-Agent Decision Intelligence Environment - [DI-sheep](https://github.com/opendilab/DI-sheep): Decision AI in 3 Tiles Game - [DI-smartcross](https://github.com/opendilab/DI-smartcross): Decision AI in Traffic Light Control - [DI-bioseq](https://github.com/opendilab/DI-bioseq): Decision AI in Biological Sequence Prediction and Searching + - [DI-1024](https://github.com/opendilab/DI-1024): Deep Reinforcement Learning + 1024 Game - Research paper - - [InterFuser](https://github.com/opendilab/InterFuser): (CoRL 2022) Safety-Enhanced Autonomous Driving Using Interpretable Sensor Fusion Transformer -- General nested data lib - - [treevalue](https://github.com/opendilab/treevalue): Tree-nested data structure - - [DI-treetensor](https://github.com/opendilab/DI-treetensor): Tree-nested PyTorch tensor Lib + - [InterFuser](https://github.com/opendilab/InterFuser): [CoRL 2022] Safety-Enhanced Autonomous Driving Using Interpretable Sensor Fusion Transformer + - [ACE](https://github.com/opendilab/ACE): [AAAI 2023] ACE: Cooperative Multi-agent Q-learning with Bidirectional Action-Dependency + - [GoBigger](https://github.com/opendilab/GoBigger): [ICLR 2023] Multi-Agent Decision Intelligence Environment + - [DOS](https://github.com/opendilab/DOS): [CVPR 2023] ReasonNet: End-to-End Driving with Temporal and Global Reasoning + - [LightZero](https://github.com/opendilab/LightZero): [NeurIPS 2023 Spotlight] A lightweight and efficient MCTS/AlphaZero/MuZero algorithm toolkit + - [SO2](https://github.com/opendilab/SO2): [AAAI 2024] A Perspective of Q-value Estimation on Offline-to-Online Reinforcement Learning + - [LMDrive](https://github.com/opendilab/LMDrive): [CVPR 2024] LMDrive: Closed-Loop End-to-End Driving with Large Language Models + - [SmartRefine](https://github.com/opendilab/SmartRefine): [CVPR 2024] SmartRefine: A Scenario-Adaptive Refinement Framework for Efficient Motion Prediction + - [ReZero](https://github.com/opendilab/LightZero): Boosting MCTS-based Algorithms by Backward-view and Entire-buffer Reanalyze + - [UniZero](https://github.com/opendilab/LightZero): Generalized and Efficient Planning with Scalable Latent World Models - Docs and Tutorials - - [DI-engine-docs](https://github.com/opendilab/DI-engine-docs) + - [DI-engine-docs](https://github.com/opendilab/DI-engine-docs): Tutorials, best practice and the API reference. - [awesome-model-based-RL](https://github.com/opendilab/awesome-model-based-RL): A curated list of awesome Model-Based RL resources - [awesome-exploration-RL](https://github.com/opendilab/awesome-exploration-rl): A curated list of awesome exploration RL resources - [awesome-decision-transformer](https://github.com/opendilab/awesome-decision-transformer): A curated list of Decision Transformer resources + - [awesome-RLHF](https://github.com/opendilab/awesome-RLHF): A curated list of reinforcement learning with human feedback resources + - [awesome-multi-modal-reinforcement-learning](https://github.com/opendilab/awesome-multi-modal-reinforcement-learning): A curated list of Multi-Modal Reinforcement Learning resources + - [awesome-diffusion-model-in-rl](https://github.com/opendilab/awesome-diffusion-model-in-rl): A curated list of Diffusion Model in RL resources + - [awesome-ui-agents](https://github.com/opendilab/awesome-ui-agents): A curated list of of awesome UI agents resources, encompassing Web, App, OS, and beyond + - [awesome-AI-based-protein-design](https://github.com/opendilab/awesome-AI-based-protein-design): a collection of research papers for AI-based protein design + - [awesome-end-to-end-autonomous-driving](https://github.com/opendilab/awesome-end-to-end-autonomous-driving): A curated list of awesome End-to-End Autonomous Driving resources + - [awesome-driving-behavior-prediction](https://github.com/opendilab/awesome-driving-behavior-prediction): A collection of research papers for Driving Behavior Prediction + +
-**DI-engine** also has some **system optimization and design** for efficient and robust large-scale RL training: +On the low-level end, DI-engine comes with a set of highly re-usable modules, including [RL optimization functions](https://github.com/opendilab/DI-engine/tree/main/ding/rl_utils), [PyTorch utilities](https://github.com/opendilab/DI-engine/tree/main/ding/torch_utils) and [auxiliary tools](https://github.com/opendilab/DI-engine/tree/main/ding/utils). +BTW, **DI-engine** also has some special **system optimization and design** for efficient and robust large-scale RL training: + +
+(Click for Details) + +- [treevalue](https://github.com/opendilab/treevalue): Tree-nested data structure +- [DI-treetensor](https://github.com/opendilab/DI-treetensor): Tree-nested PyTorch tensor Lib +- [DI-toolkit](https://github.com/opendilab/DI-toolkit): A simple toolkit package for decision intelligence - [DI-orchestrator](https://github.com/opendilab/DI-orchestrator): RL Kubernetes Custom Resource and Operator Lib - [DI-hpc](https://github.com/opendilab/DI-hpc): RL HPC OP Lib - [DI-store](https://github.com/opendilab/DI-store): RL Object Store +
+ Have fun with exploration and exploitation. ## Outline @@ -84,38 +132,44 @@ Have fun with exploration and exploitation. - [Installation](#installation) - [Quick Start](#quick-start) - [Feature](#feature) - - [↳ Algorithm Versatility](#algorithm-versatility) - - [↳ Environment Versatility](#environment-versatility) + - [Algorithm Versatility](#algorithm-versatility) + - [Environment Versatility](#environment-versatility) + - [General Data Container: TreeTensor](#general-data-container-treetensor) - [Feedback and Contribution](#feedback-and-contribution) - [Supporters](#supporters) - - [↳ Stargazers](#-stargazers) - - [↳ Forkers](#-forkers) + - [↳ Stargazers](#-stargazers) + - [↳ Forkers](#-forkers) - [Citation](#citation) - [License](#license) ## Installation You can simply install DI-engine from PyPI with the following command: -```bash -pip install DI-engine -``` -If you use Anaconda or Miniconda, you can install DI-engine from conda-forge through the following command: ```bash -conda install -c opendilab di-engine +pip install DI-engine ``` For more information about installation, you can refer to [installation](https://di-engine-docs.readthedocs.io/en/latest/01_quickstart/installation.html). And our dockerhub repo can be found [here](https://hub.docker.com/repository/docker/opendilab/ding)，we prepare `base image` and `env image` with common RL environments. +
+(Click for Details) + - base: opendilab/ding:nightly +- rpc: opendilab/ding:nightly-rpc - atari: opendilab/ding:nightly-atari - mujoco: opendilab/ding:nightly-mujoco - dmc: opendilab/ding:nightly-dmc2gym - metaworld: opendilab/ding:nightly-metaworld - smac: opendilab/ding:nightly-smac - grf: opendilab/ding:nightly-grf +- cityflow: opendilab/ding:nightly-cityflow +- evogym: opendilab/ding:nightly-evogym +- d4rl: opendilab/ding:nightly-d4rl + +
The detailed documentation are hosted on [doc](https://di-engine-docs.readthedocs.io/en/latest/) | [中文文档](https://di-engine-docs.readthedocs.io/zh_CN/latest/). @@ -123,23 +177,30 @@ The detailed documentation are hosted on [doc](https://di-engine-docs.readthedoc [3 Minutes Kickoff](https://di-engine-docs.readthedocs.io/en/latest/01_quickstart/first_rl_program.html) -[3 Minutes Kickoff (colab)](https://colab.research.google.com/drive/1K3DGi3dOT9fhFqa6bBtinwCDdWkOM3zE?usp=sharing) +[3 Minutes Kickoff (colab)](https://colab.research.google.com/drive/1_7L-QFDfeCvMvLJzRyBRUW5_Q6ESXcZ4) + +[DI-engine Huggingface Kickoff (colab)](https://colab.research.google.com/drive/1UH1GQOjcHrmNSaW77hnLGxFJrLSLwCOk) [How to migrate a new **RL Env**](https://di-engine-docs.readthedocs.io/en/latest/11_dizoo/index.html) | [如何迁移一个新的**强化学习环境**](https://di-engine-docs.readthedocs.io/zh_CN/latest/11_dizoo/index_zh.html) -**Bonus: Train RL agent in one line code:** +[How to customize the neural network model](https://di-engine-docs.readthedocs.io/en/latest/04_best_practice/custom_model.html) | [如何定制策略使用的**神经网络模型**](https://di-engine-docs.readthedocs.io/zh_CN/latest/04_best_practice/custom_model_zh.html) -```bash -ding -m serial -e cartpole -p dqn -s 0 -``` +[测试/部署 **强化学习策略** 的样例](https://github.com/opendilab/DI-engine/blob/main/dizoo/classic_control/cartpole/entry/cartpole_c51_deploy.py) + +[新老 pipeline 的异同对比](https://di-engine-docs.readthedocs.io/zh_CN/latest/04_best_practice/diff_in_new_pipeline_zh.html) ## Feature + ### Algorithm Versatility -![discrete](https://img.shields.io/badge/-discrete-brightgreen)  discrete means discrete action space, which is only label in normal DRL algorithms (1-18) -![continuous](https://img.shields.io/badge/-continous-green)  means continuous action space, which is only label in normal DRL algorithms (1-18) +
+(Click to Collapse) + +![discrete](https://img.shields.io/badge/-discrete-brightgreen)  discrete means discrete action space, which is only label in normal DRL algorithms (1-23) -![hybrid](https://img.shields.io/badge/-hybrid-darkgreen)  means hybrid (discrete + continuous) action space (1-18) +![continuous](https://img.shields.io/badge/-continous-green)  means continuous action space, which is only label in normal DRL algorithms (1-23) + +![hybrid](https://img.shields.io/badge/-hybrid-darkgreen)  means hybrid (discrete + continuous) action space (1-23) ![dist](https://img.shields.io/badge/-distributed-blue)  [Distributed Reinforcement Learning](https://di-engine-docs.readthedocs.io/en/latest/02_algo/distributed_rl.html)｜[分布式强化学习](https://di-engine-docs.readthedocs.io/zh_CN/latest/02_algo/distributed_rl_zh.html) @@ -151,99 +212,124 @@ ding -m serial -e cartpole -p dqn -s 0 ![offline](https://img.shields.io/badge/-offlineRL-darkblue)  [Offiline Reinforcement Learning](https://di-engine-docs.readthedocs.io/en/latest/02_algo/offline_rl.html)｜[离线强化学习](https://di-engine-docs.readthedocs.io/zh_CN/latest/02_algo/offline_rl_zh.html) - ![mbrl](https://img.shields.io/badge/-ModelBasedRL-lightblue)  [Model-Based Reinforcement Learning](https://di-engine-docs.readthedocs.io/en/latest/02_algo/model_based_rl.html)｜[基于模型的强化学习](https://di-engine-docs.readthedocs.io/zh_CN/latest/02_algo/model_based_rl_zh.html) -![other](https://img.shields.io/badge/-other-lightgrey)  means other sub-direction algorithm, usually as plugin-in in the whole pipeline +![other](https://img.shields.io/badge/-other-lightgrey)  means other sub-direction algorithms, usually as plugin-in in the whole pipeline P.S: The `.py` file in `Runnable Demo` can be found in `dizoo` +| No. | Algorithm | Label | Doc and Implementation | Runnable Demo | +| :-: | :---------------------------------------------------------------------------------------------------------------------------------: | :---------------------------------------------------------------------------------------------------------------------------------------------------------------------------: | :----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------: | :--------------------------------------------------------------------------------------: | +| 1 | [DQN](https://storage.googleapis.com/deepmind-media/dqn/DQNNaturePaper.pdf) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | [DQN doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/dqn.html)
[DQN中文文档](https://di-engine-docs.readthedocs.io/zh_CN/latest/12_policies/dqn_zh.html)
[policy/dqn](https://github.com/opendilab/DI-engine/blob/main/ding/policy/dqn.py) | python3 -u cartpole_dqn_main.py / ding -m serial -c cartpole_dqn_config.py -s 0 | +| 2 | [C51](https://arxiv.org/pdf/1707.06887.pdf) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | [C51 doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/c51.html)
[policy/c51](https://github.com/opendilab/DI-engine/blob/main/ding/policy/c51.py) | ding -m serial -c cartpole_c51_config.py -s 0 | +| 3 | [QRDQN](https://arxiv.org/pdf/1710.10044.pdf) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | [QRDQN doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/qrdqn.html)
[policy/qrdqn](https://github.com/opendilab/DI-engine/blob/main/ding/policy/qrdqn.py) | ding -m serial -c cartpole_qrdqn_config.py -s 0 | +| 4 | [IQN](https://arxiv.org/pdf/1806.06923.pdf) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | [IQN doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/iqn.html)
[policy/iqn](https://github.com/opendilab/DI-engine/blob/main/ding/policy/iqn.py) | ding -m serial -c cartpole_iqn_config.py -s 0 | +| 5 | [FQF](https://arxiv.org/pdf/1911.02140.pdf) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | [FQF doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/fqf.html)
[policy/fqf](https://github.com/opendilab/DI-engine/blob/main/ding/policy/fqf.py) | ding -m serial -c cartpole_fqf_config.py -s 0 | +| 6 | [Rainbow](https://arxiv.org/pdf/1710.02298.pdf) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | [Rainbow doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/rainbow.html)
[policy/rainbow](https://github.com/opendilab/DI-engine/blob/main/ding/policy/rainbow.py) | ding -m serial -c cartpole_rainbow_config.py -s 0 | +| 7 | [SQL](https://arxiv.org/pdf/1702.08165.pdf) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen)![continuous](https://img.shields.io/badge/-continous-green) | [SQL doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/sql.html)
[policy/sql](https://github.com/opendilab/DI-engine/blob/main/ding/policy/sql.py) | ding -m serial -c cartpole_sql_config.py -s 0 | +| 8 | [R2D2](https://openreview.net/forum?id=r1lyTjAqYX) | ![dist](https://img.shields.io/badge/-distributed-blue)![discrete](https://img.shields.io/badge/-discrete-brightgreen) | [R2D2 doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/r2d2.html)
[policy/r2d2](https://github.com/opendilab/DI-engine/blob/main/ding/policy/r2d2.py) | ding -m serial -c cartpole_r2d2_config.py -s 0 | +| 9 | [PG](https://proceedings.neurips.cc/paper/1999/file/464d828b85b0bed98e80ade0a5c43b0f-Paper.pdf) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | [PG doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/a2c.html)
[policy/pg](https://github.com/opendilab/DI-engine/blob/main/ding/policy/pg.py) | ding -m serial -c cartpole_pg_config.py -s 0 | +| 10 | [PromptPG](https://arxiv.org/abs/2209.14610) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | [policy/prompt_pg](https://github.com/opendilab/DI-engine/blob/main/ding/policy/prompt_pg.py) | ding -m serial_onpolicy -c tabmwp_pg_config.py -s 0 | +| 11 | [A2C](https://arxiv.org/pdf/1602.01783.pdf) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | [A2C doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/a2c.html)
[policy/a2c](https://github.com/opendilab/DI-engine/blob/main/ding/policy/a2c.py) | ding -m serial -c cartpole_a2c_config.py -s 0 | +| 12 | [PPO](https://arxiv.org/abs/1707.06347)/[MAPPO](https://arxiv.org/pdf/2103.01955.pdf) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen)![continuous](https://img.shields.io/badge/-continous-green)![MARL](https://img.shields.io/badge/-MARL-yellow) | [PPO doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/ppo.html)
[policy/ppo](https://github.com/opendilab/DI-engine/blob/main/ding/policy/ppo.py) | python3 -u cartpole_ppo_main.py / ding -m serial_onpolicy -c cartpole_ppo_config.py -s 0 | +| 13 | [PPG](https://arxiv.org/pdf/2009.04416.pdf) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | [PPG doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/ppg.html)
[policy/ppg](https://github.com/opendilab/DI-engine/blob/main/ding/policy/ppg.py) | python3 -u cartpole_ppg_main.py | +| 14 | [ACER](https://arxiv.org/pdf/1611.01224.pdf) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen)![continuous](https://img.shields.io/badge/-continous-green) | [ACER doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/acer.html)
[policy/acer](https://github.com/opendilab/DI-engine/blob/main/ding/policy/acer.py) | ding -m serial -c cartpole_acer_config.py -s 0 | +| 15 | [IMPALA](https://arxiv.org/abs/1802.01561) | ![dist](https://img.shields.io/badge/-distributed-blue)![discrete](https://img.shields.io/badge/-discrete-brightgreen) | [IMPALA doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/impala.html)
[policy/impala](https://github.com/opendilab/DI-engine/blob/main/ding/policy/impala.py) | ding -m serial -c cartpole_impala_config.py -s 0 | +| 16 | [DDPG](https://arxiv.org/pdf/1509.02971.pdf)/[PADDPG](https://arxiv.org/pdf/1511.04143.pdf) | ![continuous](https://img.shields.io/badge/-continous-green)![hybrid](https://img.shields.io/badge/-hybrid-darkgreen) | [DDPG doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/ddpg.html)
[policy/ddpg](https://github.com/opendilab/DI-engine/blob/main/ding/policy/ddpg.py) | ding -m serial -c pendulum_ddpg_config.py -s 0 | +| 17 | [TD3](https://arxiv.org/pdf/1802.09477.pdf) | ![continuous](https://img.shields.io/badge/-continous-green)![hybrid](https://img.shields.io/badge/-hybrid-darkgreen) | [TD3 doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/td3.html)
[policy/td3](https://github.com/opendilab/DI-engine/blob/main/ding/policy/td3.py) | python3 -u pendulum_td3_main.py / ding -m serial -c pendulum_td3_config.py -s 0 | +| 18 | [D4PG](https://arxiv.org/pdf/1804.08617.pdf) | ![continuous](https://img.shields.io/badge/-continous-green) | [D4PG doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/d4pg.html)
[policy/d4pg](https://github.com/opendilab/DI-engine/blob/main/ding/policy/d4pg.py) | python3 -u pendulum_d4pg_config.py | +| 19 | [SAC](https://arxiv.org/abs/1801.01290)/[MASAC] | ![discrete](https://img.shields.io/badge/-discrete-brightgreen)![continuous](https://img.shields.io/badge/-continous-green)![MARL](https://img.shields.io/badge/-MARL-yellow) | [SAC doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/sac.html)
[policy/sac](https://github.com/opendilab/DI-engine/blob/main/ding/policy/sac.py) | ding -m serial -c pendulum_sac_config.py -s 0 | +| 20 | [PDQN](https://arxiv.org/pdf/1810.06394.pdf) | ![hybrid](https://img.shields.io/badge/-hybrid-darkgreen) | [policy/pdqn](https://github.com/opendilab/DI-engine/blob/main/ding/policy/pdqn.py) | ding -m serial -c gym_hybrid_pdqn_config.py -s 0 | +| 21 | [MPDQN](https://arxiv.org/pdf/1905.04388.pdf) | ![hybrid](https://img.shields.io/badge/-hybrid-darkgreen) | [policy/pdqn](https://github.com/opendilab/DI-engine/blob/main/ding/policy/pdqn.py) | ding -m serial -c gym_hybrid_mpdqn_config.py -s 0 | +| 22 | [HPPO](https://arxiv.org/pdf/1903.01344.pdf) | ![hybrid](https://img.shields.io/badge/-hybrid-darkgreen) | [policy/ppo](https://github.com/opendilab/DI-engine/blob/main/ding/policy/ppo.py) | ding -m serial_onpolicy -c gym_hybrid_hppo_config.py -s 0 | +| 23 | [BDQ](https://arxiv.org/pdf/1711.08946.pdf) | ![hybrid](https://img.shields.io/badge/-hybrid-darkgreen) | [policy/bdq](https://github.com/opendilab/DI-engine/blob/main/ding/policy/dqn.py) | python3 -u hopper_bdq_config.py | +| 24 | [MDQN](https://arxiv.org/abs/2007.14430) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | [policy/mdqn](https://github.com/opendilab/DI-engine/blob/main/ding/policy/mdqn.py) | python3 -u asterix_mdqn_config.py | +| 25 | [QMIX](https://arxiv.org/pdf/1803.11485.pdf) | ![MARL](https://img.shields.io/badge/-MARL-yellow) | [QMIX doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/qmix.html)
[policy/qmix](https://github.com/opendilab/DI-engine/blob/main/ding/policy/qmix.py) | ding -m serial -c smac_3s5z_qmix_config.py -s 0 | +| 26 | [COMA](https://arxiv.org/pdf/1705.08926.pdf) | ![MARL](https://img.shields.io/badge/-MARL-yellow) | [COMA doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/coma.html)
[policy/coma](https://github.com/opendilab/DI-engine/blob/main/ding/policy/coma.py) | ding -m serial -c smac_3s5z_coma_config.py -s 0 | +| 27 | [QTran](https://arxiv.org/abs/1905.05408) | ![MARL](https://img.shields.io/badge/-MARL-yellow) | [policy/qtran](https://github.com/opendilab/DI-engine/blob/main/ding/policy/qtran.py) | ding -m serial -c smac_3s5z_qtran_config.py -s 0 | +| 28 | [WQMIX](https://arxiv.org/abs/2006.10800) | ![MARL](https://img.shields.io/badge/-MARL-yellow) | [WQMIX doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/wqmix.html)
[policy/wqmix](https://github.com/opendilab/DI-engine/blob/main/ding/policy/wqmix.py) | ding -m serial -c smac_3s5z_wqmix_config.py -s 0 | +| 29 | [CollaQ](https://arxiv.org/pdf/2010.08531.pdf) | ![MARL](https://img.shields.io/badge/-MARL-yellow) | [CollaQ doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/collaq.html)
[policy/collaq](https://github.com/opendilab/DI-engine/blob/main/ding/policy/collaq.py) | ding -m serial -c smac_3s5z_collaq_config.py -s 0 | +| 30 | [MADDPG](https://arxiv.org/pdf/1706.02275.pdf) | ![MARL](https://img.shields.io/badge/-MARL-yellow) | [MADDPG doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/ddpg.html)
[policy/ddpg](https://github.com/opendilab/DI-engine/blob/main/ding/policy/ddpg.py) | ding -m serial -c ptz_simple_spread_maddpg_config.py -s 0 | +| 31 | [GAIL](https://arxiv.org/pdf/1606.03476.pdf) | ![IL](https://img.shields.io/badge/-IL-purple) | [GAIL doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/gail.html)
[reward_model/gail](https://github.com/opendilab/DI-engine/blob/main/ding/reward_model/gail_irl_model.py) | ding -m serial_gail -c cartpole_dqn_gail_config.py -s 0 | +| 32 | [SQIL](https://arxiv.org/pdf/1905.11108.pdf) | ![IL](https://img.shields.io/badge/-IL-purple) | [SQIL doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/sqil.html)
[entry/sqil](https://github.com/opendilab/DI-engine/blob/main/ding/entry/serial_entry_sqil.py) | ding -m serial_sqil -c cartpole_sqil_config.py -s 0 | +| 33 | [DQFD](https://arxiv.org/pdf/1704.03732.pdf) | ![IL](https://img.shields.io/badge/-IL-purple) | [DQFD doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/dqfd.html)
[policy/dqfd](https://github.com/opendilab/DI-engine/blob/main/ding/policy/dqfd.py) | ding -m serial_dqfd -c cartpole_dqfd_config.py -s 0 | +| 34 | [R2D3](https://arxiv.org/pdf/1909.01387.pdf) | ![IL](https://img.shields.io/badge/-IL-purple) | [R2D3 doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/r2d3.html)
[R2D3中文文档](https://di-engine-docs.readthedocs.io/zh_CN/latest/12_policies/r2d3_zh.html)
[policy/r2d3](https://di-engine-docs.readthedocs.io/zh_CN/latest/12_policies/r2d3_zh.html) | python3 -u pong_r2d3_r2d2expert_config.py | +| 35 | [Guided Cost Learning](https://arxiv.org/pdf/1603.00448.pdf) | ![IL](https://img.shields.io/badge/-IL-purple) | [Guided Cost Learning中文文档](https://di-engine-docs.readthedocs.io/zh_CN/latest/12_policies/guided_cost_zh.html)
[reward_model/guided_cost](https://github.com/opendilab/DI-engine/blob/main/ding/reward_model/guided_cost_reward_model.py) | python3 lunarlander_gcl_config.py | +| 36 | [TREX](https://arxiv.org/abs/1904.06387) | ![IL](https://img.shields.io/badge/-IL-purple) | [TREX doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/trex.html)
[reward_model/trex](https://github.com/opendilab/DI-engine/blob/main/ding/reward_model/trex_reward_model.py) | python3 mujoco_trex_main.py | +| 37 | [Implicit Behavorial Cloning](https://implicitbc.github.io/) (DFO+MCMC) | ![IL](https://img.shields.io/badge/-IL-purple) | [policy/ibc](https://github.com/opendilab/DI-engine/blob/main/ding/policy/ibc.py)
[model/template/ebm](https://github.com/opendilab/DI-engine/blob/main/ding/model/template/ebm.py) | python3 d4rl_ibc_main.py -s 0 -c pen_human_ibc_mcmc_config.py | +| 38 | [BCO](https://arxiv.org/pdf/1805.01954.pdf) | ![IL](https://img.shields.io/badge/-IL-purple) | [entry/bco](https://github.com/opendilab/DI-engine/blob/main/ding/entry/serial_entry_bco.py) | python3 -u cartpole_bco_config.py | +| 39 | [HER](https://arxiv.org/pdf/1707.01495.pdf) | ![exp](https://img.shields.io/badge/-exploration-orange) | [HER doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/her.html)
[reward_model/her](https://github.com/opendilab/DI-engine/blob/main/ding/reward_model/her_reward_model.py) | python3 -u bitflip_her_dqn.py | +| 40 | [RND](https://arxiv.org/abs/1810.12894) | ![exp](https://img.shields.io/badge/-exploration-orange) | [RND doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/rnd.html)
[reward_model/rnd](https://github.com/opendilab/DI-engine/blob/main/ding/reward_model/rnd_reward_model.py) | python3 -u cartpole_rnd_onppo_config.py | +| 41 | [ICM](https://arxiv.org/pdf/1705.05363.pdf) | ![exp](https://img.shields.io/badge/-exploration-orange) | [ICM doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/icm.html)
[ICM中文文档](https://di-engine-docs.readthedocs.io/zh_CN/latest/12_policies/icm_zh.html)
[reward_model/icm](https://github.com/opendilab/DI-engine/blob/main/ding/reward_model/icm_reward_model.py) | python3 -u cartpole_ppo_icm_config.py | +| 42 | [CQL](https://arxiv.org/pdf/2006.04779.pdf) | ![offline](https://img.shields.io/badge/-offlineRL-darkblue) | [CQL doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/cql.html)
[policy/cql](https://github.com/opendilab/DI-engine/blob/main/ding/policy/cql.py) | python3 -u d4rl_cql_main.py | +| 43 | [TD3BC](https://arxiv.org/pdf/2106.06860.pdf) | ![offline](https://img.shields.io/badge/-offlineRL-darkblue) | [TD3BC doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/td3_bc.html)
[policy/td3_bc](https://github.com/opendilab/DI-engine/blob/main/ding/policy/td3_bc.py) | python3 -u d4rl_td3_bc_main.py | +| 44 | [Decision Transformer](https://arxiv.org/pdf/2106.01345.pdf) | ![offline](https://img.shields.io/badge/-offlineRL-darkblue) | [policy/dt](https://github.com/opendilab/DI-engine/blob/main/ding/policy/dt.py) | python3 -u d4rl_dt_mujoco.py | +| 45 | [EDAC](https://arxiv.org/pdf/2110.01548.pdf) | ![offline](https://img.shields.io/badge/-offlineRL-darkblue) | [EDAC doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/edac.html)
[policy/edac](https://github.com/opendilab/DI-engine/blob/main/ding/policy/edac.py) | python3 -u d4rl_edac_main.py | +| 46 | [QGPO](https://arxiv.org/pdf/2304.12824.pdf) | ![offline](https://img.shields.io/badge/-offlineRL-darkblue) | [QGPO doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/qgpo.html)
[policy/qgpo](https://github.com/opendilab/DI-engine/blob/main/ding/policy/qgpo.py) | python3 -u ding/example/qgpo.py | +| 47 | MBSAC([SAC](https://arxiv.org/abs/1801.01290)+[MVE](https://arxiv.org/abs/1803.00101)+[SVG](https://arxiv.org/abs/1510.09142)) | ![continuous](https://img.shields.io/badge/-continous-green)![mbrl](https://img.shields.io/badge/-ModelBasedRL-lightblue) | [policy/mbpolicy/mbsac](https://github.com/opendilab/DI-engine/blob/main/ding/policy/mbpolicy/mbsac.py) | python3 -u pendulum_mbsac_mbpo_config.py \ python3 -u pendulum_mbsac_ddppo_config.py | +| 48 | STEVESAC([SAC](https://arxiv.org/abs/1801.01290)+[STEVE](https://arxiv.org/abs/1807.01675)+[SVG](https://arxiv.org/abs/1510.09142)) | ![continuous](https://img.shields.io/badge/-continous-green)![mbrl](https://img.shields.io/badge/-ModelBasedRL-lightblue) | [policy/mbpolicy/mbsac](https://github.com/opendilab/DI-engine/blob/main/ding/policy/mbpolicy/mbsac.py) | python3 -u pendulum_stevesac_mbpo_config.py | +| 49 | [MBPO](https://arxiv.org/pdf/1906.08253.pdf) | ![mbrl](https://img.shields.io/badge/-ModelBasedRL-lightblue) | [MBPO doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/mbpo.html)
[world_model/mbpo](https://github.com/opendilab/DI-engine/blob/main/ding/world_model/mbpo.py) | python3 -u pendulum_sac_mbpo_config.py | +| 50 | [DDPPO](https://openreview.net/forum?id=rzvOQrnclO0) | ![mbrl](https://img.shields.io/badge/-ModelBasedRL-lightblue) | [world_model/ddppo](https://github.com/opendilab/DI-engine/blob/main/ding/world_model/ddppo.py) | python3 -u pendulum_mbsac_ddppo_config.py | +| 51 | [DreamerV3](https://arxiv.org/pdf/2301.04104.pdf) | ![mbrl](https://img.shields.io/badge/-ModelBasedRL-lightblue) | [world_model/dreamerv3](https://github.com/opendilab/DI-engine/blob/main/ding/world_model/dreamerv3.py) | python3 -u cartpole_balance_dreamer_config.py | +| 52 | [PER](https://arxiv.org/pdf/1511.05952.pdf) | ![other](https://img.shields.io/badge/-other-lightgrey) | [worker/replay_buffer](https://github.com/opendilab/DI-engine/blob/main/ding/worker/replay_buffer/advanced_buffer.py) | `rainbow demo` | +| 53 | [GAE](https://arxiv.org/pdf/1506.02438.pdf) | ![other](https://img.shields.io/badge/-other-lightgrey) | [rl_utils/gae](https://github.com/opendilab/DI-engine/blob/main/ding/rl_utils/gae.py) | `ppo demo` | +| 54 | [ST-DIM](https://arxiv.org/pdf/1906.08226.pdf) | ![other](https://img.shields.io/badge/-other-lightgrey) | [torch_utils/loss/contrastive_loss](https://github.com/opendilab/DI-engine/blob/main/ding/torch_utils/loss/contrastive_loss.py) | ding -m serial -c cartpole_dqn_stdim_config.py -s 0 | +| 55 | [PLR](https://arxiv.org/pdf/2010.03934.pdf) | ![other](https://img.shields.io/badge/-other-lightgrey) | [PLR doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/plr.html)
[data/level_replay/level_sampler](https://github.com/opendilab/DI-engine/blob/main/ding/data/level_replay/level_sampler.py) | python3 -u bigfish_plr_config.py -s 0 | +| 56 | [PCGrad](https://arxiv.org/pdf/2001.06782.pdf) | ![other](https://img.shields.io/badge/-other-lightgrey) | [torch_utils/optimizer_helper/PCGrad](https://github.com/opendilab/DI-engine/blob/main/ding/data/torch_utils/optimizer_helper.py) | python3 -u multi_mnist_pcgrad_main.py -s 0 | +| 57 | [AWR](https://arxiv.org/pdf/1910.00177) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | [policy/ibc](https://github.com/opendilab/DI-engine/blob/main/ding/policy/prompt_awr.py) | python3 -u tabmwp_awr_config.py | + +
+### Environment Versatility - -| No. | Algorithm | Label | Doc and Implementation | Runnable Demo | -| :--: | :----------------------------------------------------------: | :----------------------------------------------------------: | :----------------------------------------------------------: | :----------------------------------------------------------: | -| 1 | [DQN](https://storage.googleapis.com/deepmind-media/dqn/DQNNaturePaper.pdf) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | [DQN doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/dqn.html)
[DQN中文文档](https://di-engine-docs.readthedocs.io/zh_CN/latest/12_policies/dqn_zh.html)
[policy/dqn](https://github.com/opendilab/DI-engine/blob/main/ding/policy/dqn.py) | python3 -u cartpole_dqn_main.py / ding -m serial -c cartpole_dqn_config.py -s 0 | -| 2 | [C51](https://arxiv.org/pdf/1707.06887.pdf) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | [C51 doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/c51.html)
[policy/c51](https://github.com/opendilab/DI-engine/blob/main/ding/policy/c51.py) | ding -m serial -c cartpole_c51_config.py -s 0 | -| 3 | [QRDQN](https://arxiv.org/pdf/1710.10044.pdf) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | [QRDQN doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/qrdqn.html)
[policy/qrdqn](https://github.com/opendilab/DI-engine/blob/main/ding/policy/qrdqn.py) | ding -m serial -c cartpole_qrdqn_config.py -s 0 | -| 4 | [IQN](https://arxiv.org/pdf/1806.06923.pdf) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | [IQN doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/iqn.html)
[policy/iqn](https://github.com/opendilab/DI-engine/blob/main/ding/policy/iqn.py) | ding -m serial -c cartpole_iqn_config.py -s 0 | -| 5 | [FQF](https://arxiv.org/pdf/1911.02140.pdf) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | [FQF doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/fqf.html)
[policy/fqf](https://github.com/opendilab/DI-engine/blob/main/ding/policy/fqf.py) | ding -m serial -c cartpole_fqf_config.py -s 0 | -| 6 | [Rainbow](https://arxiv.org/pdf/1710.02298.pdf) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | [Rainbow doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/rainbow.html)
[policy/rainbow](https://github.com/opendilab/DI-engine/blob/main/ding/policy/rainbow.py) | ding -m serial -c cartpole_rainbow_config.py -s 0 | -| 7 | [SQL](https://arxiv.org/pdf/1702.08165.pdf) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen)![continuous](https://img.shields.io/badge/-continous-green) | [SQL doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/sql.html)
[policy/sql](https://github.com/opendilab/DI-engine/blob/main/ding/policy/sql.py) | ding -m serial -c cartpole_sql_config.py -s 0 | -| 8 | [R2D2](https://openreview.net/forum?id=r1lyTjAqYX) | ![dist](https://img.shields.io/badge/-distributed-blue)![discrete](https://img.shields.io/badge/-discrete-brightgreen) | [R2D2 doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/r2d2.html)
[policy/r2d2](https://github.com/opendilab/DI-engine/blob/main/ding/policy/r2d2.py) | ding -m serial -c cartpole_r2d2_config.py -s 0 | -| 9 | [A2C](https://arxiv.org/pdf/1602.01783.pdf) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | [A2C doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/a2c.html)
[policy/a2c](https://github.com/opendilab/DI-engine/blob/main/ding/policy/a2c.py) | ding -m serial -c cartpole_a2c_config.py -s 0 | -| 10 | [PPO](https://arxiv.org/abs/1707.06347)/[MAPPO](https://arxiv.org/pdf/2103.01955.pdf) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen)![continuous](https://img.shields.io/badge/-continous-green)![MARL](https://img.shields.io/badge/-MARL-yellow) | [PPO doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/ppo.html)
[policy/ppo](https://github.com/opendilab/DI-engine/blob/main/ding/policy/ppo.py) | python3 -u cartpole_ppo_main.py / ding -m serial_onpolicy -c cartpole_ppo_config.py -s 0 | -| 11 | [PPG](https://arxiv.org/pdf/2009.04416.pdf) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | [PPG doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/ppg.html)
[policy/ppg](https://github.com/opendilab/DI-engine/blob/main/ding/policy/ppg.py) | python3 -u cartpole_ppg_main.py | -| 12 | [ACER](https://arxiv.org/pdf/1611.01224.pdf) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen)![continuous](https://img.shields.io/badge/-continous-green) | [ACER doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/acer.html)
[policy/acer](https://github.com/opendilab/DI-engine/blob/main/ding/policy/acer.py) | ding -m serial -c cartpole_acer_config.py -s 0 | -| 13 | [IMPALA](https://arxiv.org/abs/1802.01561) | ![dist](https://img.shields.io/badge/-distributed-blue)![discrete](https://img.shields.io/badge/-discrete-brightgreen) | [IMPALA doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/impala.html)
[policy/impala](https://github.com/opendilab/DI-engine/blob/main/ding/policy/impala.py) | ding -m serial -c cartpole_impala_config.py -s 0 | -| 14 | [DDPG](https://arxiv.org/pdf/1509.02971.pdf)/[PADDPG](https://arxiv.org/pdf/1511.04143.pdf) | ![continuous](https://img.shields.io/badge/-continous-green)![hybrid](https://img.shields.io/badge/-hybrid-darkgreen) | [DDPG doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/ddpg.html)
[policy/ddpg](https://github.com/opendilab/DI-engine/blob/main/ding/policy/ddpg.py) | ding -m serial -c pendulum_ddpg_config.py -s 0 | -| 15 | [TD3](https://arxiv.org/pdf/1802.09477.pdf) | ![continuous](https://img.shields.io/badge/-continous-green)![hybrid](https://img.shields.io/badge/-hybrid-darkgreen) | [TD3 doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/td3.html)
[policy/td3](https://github.com/opendilab/DI-engine/blob/main/ding/policy/td3.py) | python3 -u pendulum_td3_main.py / ding -m serial -c pendulum_td3_config.py -s 0 | -| 16 | [D4PG](https://arxiv.org/pdf/1804.08617.pdf) | ![continuous](https://img.shields.io/badge/-continous-green) | [D4PG doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/d4pg.html)
[policy/d4pg](https://github.com/opendilab/DI-engine/blob/main/ding/policy/d4pg.py) | python3 -u pendulum_d4pg_config.py | -| 17 | [SAC](https://arxiv.org/abs/1801.01290)/[MASAC] | ![discrete](https://img.shields.io/badge/-discrete-brightgreen)![continuous](https://img.shields.io/badge/-continous-green)![MARL](https://img.shields.io/badge/-MARL-yellow) | [SAC doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/sac.html)
[policy/sac](https://github.com/opendilab/DI-engine/blob/main/ding/policy/sac.py) | ding -m serial -c pendulum_sac_config.py -s 0 | -| 18 | [PDQN](https://arxiv.org/pdf/1810.06394.pdf) | ![hybrid](https://img.shields.io/badge/-hybrid-darkgreen) | [policy/pdqn](https://github.com/opendilab/DI-engine/blob/main/ding/policy/pdqn.py) | ding -m serial -c gym_hybrid_pdqn_config.py -s 0 | -| 19 | [MPDQN](https://arxiv.org/pdf/1905.04388.pdf) | ![hybrid](https://img.shields.io/badge/-hybrid-darkgreen) | [policy/pdqn](https://github.com/opendilab/DI-engine/blob/main/ding/policy/pdqn.py) | ding -m serial -c gym_hybrid_mpdqn_config.py -s 0 | -| 20 | [HPPO](https://arxiv.org/pdf/1903.01344.pdf) | ![hybrid](https://img.shields.io/badge/-hybrid-darkgreen) | [policy/ppo](https://github.com/opendilab/DI-engine/blob/main/ding/policy/ppo.py) | ding -m serial_onpolicy -c gym_hybrid_hppo_config.py -s 0 | -| 21 | [QMIX](https://arxiv.org/pdf/1803.11485.pdf) | ![MARL](https://img.shields.io/badge/-MARL-yellow) | [QMIX doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/qmix.html)
[policy/qmix](https://github.com/opendilab/DI-engine/blob/main/ding/policy/qmix.py) | ding -m serial -c smac_3s5z_qmix_config.py -s 0 | -| 22 | [COMA](https://arxiv.org/pdf/1705.08926.pdf) | ![MARL](https://img.shields.io/badge/-MARL-yellow) | [COMA doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/coma.html)
[policy/coma](https://github.com/opendilab/DI-engine/blob/main/ding/policy/coma.py) | ding -m serial -c smac_3s5z_coma_config.py -s 0 | -| 23 | [QTran](https://arxiv.org/abs/1905.05408) | ![MARL](https://img.shields.io/badge/-MARL-yellow) | [policy/qtran](https://github.com/opendilab/DI-engine/blob/main/ding/policy/qtran.py) | ding -m serial -c smac_3s5z_qtran_config.py -s 0 | -| 24 | [WQMIX](https://arxiv.org/abs/2006.10800) | ![MARL](https://img.shields.io/badge/-MARL-yellow) | [WQMIX doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/wqmix.html)
[policy/wqmix](https://github.com/opendilab/DI-engine/blob/main/ding/policy/wqmix.py) | ding -m serial -c smac_3s5z_wqmix_config.py -s 0 | -| 25 | [CollaQ](https://arxiv.org/pdf/2010.08531.pdf) | ![MARL](https://img.shields.io/badge/-MARL-yellow) | [CollaQ doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/collaq.html)
[policy/collaq](https://github.com/opendilab/DI-engine/blob/main/ding/policy/collaq.py) | ding -m serial -c smac_3s5z_collaq_config.py -s 0 | -| 26 | [GAIL](https://arxiv.org/pdf/1606.03476.pdf) | ![IL](https://img.shields.io/badge/-IL-purple) | [GAIL doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/gail.html)
[reward_model/gail](https://github.com/opendilab/DI-engine/blob/main/ding/reward_model/gail_irl_model.py) | ding -m serial_gail -c cartpole_dqn_gail_config.py -s 0 | -| 27 | [SQIL](https://arxiv.org/pdf/1905.11108.pdf) | ![IL](https://img.shields.io/badge/-IL-purple) | [SQIL doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/sqil.html)
[entry/sqil](https://github.com/opendilab/DI-engine/blob/main/ding/entry/serial_entry_sqil.py) | ding -m serial_sqil -c cartpole_sqil_config.py -s 0 | -| 28 | [DQFD](https://arxiv.org/pdf/1704.03732.pdf) | ![IL](https://img.shields.io/badge/-IL-purple) | [DQFD doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/dqfd.html)
[policy/dqfd](https://github.com/opendilab/DI-engine/blob/main/ding/policy/dqfd.py) | ding -m serial_dqfd -c cartpole_dqfd_config.py -s 0 | -| 29 | [R2D3](https://arxiv.org/pdf/1909.01387.pdf) | ![IL](https://img.shields.io/badge/-IL-purple) | [R2D3 doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/r2d3.html)
[R2D3中文文档](https://di-engine-docs.readthedocs.io/zh_CN/latest/12_policies/r2d3_zh.html)
[policy/r2d3](https://di-engine-docs.readthedocs.io/zh_CN/latest/12_policies/r2d3_zh.html) | python3 -u pong_r2d3_r2d2expert_config.py | -| 30 | [Guided Cost Learning](https://arxiv.org/pdf/1603.00448.pdf) | ![IL](https://img.shields.io/badge/-IL-purple) | [Guided Cost Learning中文文档](https://di-engine-docs.readthedocs.io/zh_CN/latest/12_policies/guided_cost_zh.html)
[reward_model/guided_cost](https://github.com/opendilab/DI-engine/blob/main/ding/reward_model/guided_cost_reward_model.py) | python3 lunarlander_gcl_config.py | -| 31 | [TREX](https://arxiv.org/abs/1904.06387) | ![IL](https://img.shields.io/badge/-IL-purple) | [TREX doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/trex.html)
[reward_model/trex](https://github.com/opendilab/DI-engine/blob/main/ding/reward_model/trex_reward_model.py) | python3 mujoco_trex_main.py | -| 32 | [Implicit Behavorial Cloning](https://implicitbc.github.io/) (DFO+MCMC) | ![IL](https://img.shields.io/badge/-IL-purple) | [policy/ibc](https://github.com/opendilab/DI-engine/blob/main/ding/policy/ibc.py) & [model/template/ebm](https://github.com/opendilab/DI-engine/blob/main/ding/model/template/ebm.py) | python3 d4rl_ibc_main.py -s 0 -c pen_human_ibc_mcmc_config.py | -| 33 | [BCO](https://arxiv.org/pdf/1805.01954.pdf) | ![IL](https://img.shields.io/badge/-IL-purple) | [entry/bco](https://github.com/opendilab/DI-engine/blob/main/ding/entry/serial_entry_bco.py) | python3 -u cartpole_bco_config.py | -| 34 | [HER](https://arxiv.org/pdf/1707.01495.pdf) | ![exp](https://img.shields.io/badge/-exploration-orange) | [HER doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/her.html)
[reward_model/her](https://github.com/opendilab/DI-engine/blob/main/ding/reward_model/her_reward_model.py) | python3 -u bitflip_her_dqn.py | -| 35 | [RND](https://arxiv.org/abs/1810.12894) | ![exp](https://img.shields.io/badge/-exploration-orange) | [RND doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/rnd.html)
[reward_model/rnd](https://github.com/opendilab/DI-engine/blob/main/ding/reward_model/rnd_reward_model.py) | python3 -u cartpole_rnd_onppo_config.py | -| 36 | [ICM](https://arxiv.org/pdf/1705.05363.pdf) | ![exp](https://img.shields.io/badge/-exploration-orange) | [ICM doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/icm.html)
[ICM中文文档](https://di-engine-docs.readthedocs.io/zh_CN/latest/12_policies/icm_zh.html)
[reward_model/icm](https://github.com/opendilab/DI-engine/blob/main/ding/reward_model/icm_reward_model.py) | python3 -u cartpole_ppo_icm_config.py | -| 37 | [CQL](https://arxiv.org/pdf/2006.04779.pdf) | ![offline](https://img.shields.io/badge/-offlineRL-darkblue) | [CQL doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/cql.html)
[policy/cql](https://github.com/opendilab/DI-engine/blob/main/ding/policy/cql.py) | python3 -u d4rl_cql_main.py | -| 38 | [TD3BC](https://arxiv.org/pdf/2106.06860.pdf) | ![offline](https://img.shields.io/badge/-offlineRL-darkblue) | [TD3BC doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/td3_bc.html)
[policy/td3_bc](https://github.com/opendilab/DI-engine/blob/main/ding/policy/td3_bc.py) | python3 -u mujoco_td3_bc_main.py | -| 39 | MBSAC([SAC](https://arxiv.org/abs/1801.01290)+[MVE](https://arxiv.org/abs/1803.00101)+[SVG](https://arxiv.org/abs/1510.09142)) | ![continuous](https://img.shields.io/badge/-continous-green)![mbrl](https://img.shields.io/badge/-ModelBasedRL-lightblue) | [policy/mbpolicy/mbsac](https://github.com/opendilab/DI-engine/blob/main/ding/policy/mbpolicy/mbsac.py) | python3 -u pendulum_mbsac_mbpo_config.py \ python3 -u pendulum_mbsac_ddppo_config.py | -| 40 | STEVESAC([SAC](https://arxiv.org/abs/1801.01290)+[STEVE](https://arxiv.org/abs/1807.01675)+[SVG](https://arxiv.org/abs/1510.09142)) | ![continuous](https://img.shields.io/badge/-continous-green)![mbrl](https://img.shields.io/badge/-ModelBasedRL-lightblue) | [policy/mbpolicy/mbsac](https://github.com/opendilab/DI-engine/blob/main/ding/policy/mbpolicy/mbsac.py) | python3 -u pendulum_stevesac_mbpo_config.py | -| 41 | [MBPO](https://arxiv.org/pdf/1906.08253.pdf) | ![mbrl](https://img.shields.io/badge/-ModelBasedRL-lightblue) | [MBPO doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/mbpo.html)
[world_model/mbpo](https://github.com/opendilab/DI-engine/blob/main/ding/world_model/mbpo.py) | python3 -u pendulum_sac_mbpo_config.py | -| 42 | [DDPPO](https://openreview.net/forum?id=rzvOQrnclO0) | ![mbrl](https://img.shields.io/badge/-ModelBasedRL-lightblue) | [world_model/ddppo](https://github.com/opendilab/DI-engine/blob/main/ding/world_model/ddppo.py) | python3 -u pendulum_mbsac_ddppo_config.py | -| 43 | [PER](https://arxiv.org/pdf/1511.05952.pdf) | ![other](https://img.shields.io/badge/-other-lightgrey) | [worker/replay_buffer](https://github.com/opendilab/DI-engine/blob/main/ding/worker/replay_buffer/advanced_buffer.py) | `rainbow demo` | -| 44 | [GAE](https://arxiv.org/pdf/1506.02438.pdf) | ![other](https://img.shields.io/badge/-other-lightgrey) | [rl_utils/gae](https://github.com/opendilab/DI-engine/blob/main/ding/rl_utils/gae.py) | `ppo demo` | -| 45 | [ST-DIM](https://arxiv.org/pdf/1906.08226.pdf) | ![other](https://img.shields.io/badge/-other-lightgrey) | [torch_utils/loss/contrastive_loss](https://github.com/opendilab/DI-engine/blob/main/ding/torch_utils/loss/contrastive_loss.py) | ding -m serial -c cartpole_dqn_stdim_config.py -s 0 | -| 46 | [PLR](https://arxiv.org/pdf/2010.03934.pdf) | ![other](https://img.shields.io/badge/-other-lightgrey) | [PLR doc](https://di-engine-docs.readthedocs.io/en/latest/12_policies/plr.html)
[data/level_replay/level_sampler](https://github.com/opendilab/DI-engine/blob/main/ding/data/level_replay/level_sampler.py) | python3 -u bigfish_plr_config.py -s 0 | -| 47 | [PCGrad](https://arxiv.org/pdf/2001.06782.pdf) | ![other](https://img.shields.io/badge/-other-lightgrey) | [torch_utils/optimizer_helper/PCGrad](https://github.com/opendilab/DI-engine/blob/main/ding/data/torch_utils/optimizer_helper.py) | python3 -u multi_mnist_pcgrad_main.py -s 0 | +
+(Click to Collapse) + +| No | Environment | Label | Visualization | Code and Doc Links | +| :-: | :--------------------------------------------------------------------------------------------: | :---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------: | :----------------------------------------------------------------------------------------------------: | :-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------: | +| 1 | [Atari](https://ale.farama.org) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | ![original](./dizoo/atari/atari.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/atari/envs)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/atari.html)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/atari_zh.html) | +| 2 | [box2d/bipedalwalker](https://github.com/openai/gym/tree/master/gym/envs/box2d) | ![continuous](https://img.shields.io/badge/-continous-green) | ![original](./dizoo/box2d/bipedalwalker/original.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/box2d/bipedalwalker/envs)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/bipedalwalker.html)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/bipedalwalker_zh.html) | +| 3 | [box2d/lunarlander](https://github.com/openai/gym/tree/master/gym/envs/box2d) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | ![original](./dizoo/box2d/lunarlander/lunarlander.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/box2d/lunarlander/envs)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/lunarlander.html)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/lunarlander_zh.html) | +| 4 | [classic_control/cartpole](https://github.com/openai/gym/tree/master/gym/envs/classic_control) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | ![original](./dizoo/classic_control/cartpole/cartpole.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/classic_control/cartpole/envs)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/cartpole.html)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/cartpole_zh.html) | +| 5 | [classic_control/pendulum](https://github.com/openai/gym/tree/master/gym/envs/classic_control) | ![continuous](https://img.shields.io/badge/-continous-green) | ![original](./dizoo/classic_control/pendulum/pendulum.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/classic_control/pendulum/envs)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/pendulum.html)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/pendulum_zh.html) | +| 6 | [competitive_rl](https://github.com/cuhkrlcourse/competitive-rl) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) ![selfplay](https://img.shields.io/badge/-selfplay-blue) | ![original](./dizoo/competitive_rl/competitive_rl.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo.classic_control)
[环境指南](https://di-engine-docs.readthedocs.io/en/latest/13_envs/competitive_rl_zh.html) | +| 7 | [gfootball](https://github.com/google-research/football) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen)![sparse](https://img.shields.io/badge/-sparse%20reward-orange)![selfplay](https://img.shields.io/badge/-selfplay-blue) | ![original](./dizoo/gfootball/gfootball.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo.gfootball/envs)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/gfootball.html)
[环境指南](https://di-engine-docs.readthedocs.io/en/latest/13_envs/gfootball_zh.html) | +| 8 | [minigrid](https://github.com/maximecb/gym-minigrid) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen)![sparse](https://img.shields.io/badge/-sparse%20reward-orange) | ![original](./dizoo/minigrid/minigrid.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/minigrid/envs)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/minigrid.html)
[环境指南](https://di-engine-docs.readthedocs.io/en/latest/13_envs/minigrid_zh.html) | +| 9 | [MuJoCo](https://github.com/openai/gym/tree/master/gym/envs/mujoco) | ![continuous](https://img.shields.io/badge/-continous-green) | ![original](./dizoo/mujoco/mujoco.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/majoco/envs)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/mujoco.html)
[环境指南](https://di-engine-docs.readthedocs.io/en/latest/13_envs/mujoco_zh.html) | +| 10 | [PettingZoo](https://github.com/Farama-Foundation/PettingZoo) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) ![continuous](https://img.shields.io/badge/-continous-green) ![marl](https://img.shields.io/badge/-MARL-yellow) | ![original](./dizoo/petting_zoo/petting_zoo_mpe_simple_spread.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/petting_zoo/envs)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/pettingzoo.html)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/pettingzoo_zh.html) | +| 11 | [overcooked](https://github.com/HumanCompatibleAI/overcooked-demo) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) ![marl](https://img.shields.io/badge/-MARL-yellow) | ![original](./dizoo/overcooked/overcooked.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/overcooded/envs)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/overcooked.html) | +| 12 | [procgen](https://github.com/openai/procgen) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | ![original](./dizoo/procgen/coinrun.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/procgen)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/procgen.html)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/procgen_zh.html) | +| 13 | [pybullet](https://github.com/benelot/pybullet-gym) | ![continuous](https://img.shields.io/badge/-continous-green) | ![original](./dizoo/pybullet/pybullet.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/pybullet/envs)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/pybullet_zh.html) | +| 14 | [smac](https://github.com/oxwhirl/smac) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) ![marl](https://img.shields.io/badge/-MARL-yellow)![selfplay](https://img.shields.io/badge/-selfplay-blue)![sparse](https://img.shields.io/badge/-sparse%20reward-orange) | ![original](./dizoo/smac/smac.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/smac/envs)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/smac.html)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/smac_zh.html) | +| 15 | [d4rl](https://github.com/rail-berkeley/d4rl) | ![offline](https://img.shields.io/badge/-offlineRL-darkblue) | ![ori](dizoo/d4rl/d4rl.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/d4rl)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/d4rl_zh.html) | +| 16 | league_demo | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) ![selfplay](https://img.shields.io/badge/-selfplay-blue) | ![original](./dizoo/league_demo/league_demo.png) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/league_demo/envs) | +| 17 | pomdp atari | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/pomdp/envs) | +| 18 | [bsuite](https://github.com/deepmind/bsuite) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | ![original](./dizoo/bsuite/bsuite.png) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/bsuite/envs)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs//bsuite.html)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/bsuite_zh.html) | +| 19 | [ImageNet](https://www.image-net.org/) | ![IL](https://img.shields.io/badge/-IL/SL-purple) | ![original](./dizoo/image_classification/imagenet.png) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/image_classification)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/image_cls_zh.html) | +| 20 | [slime_volleyball](https://github.com/hardmaru/slimevolleygym) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen)![selfplay](https://img.shields.io/badge/-selfplay-blue) | ![ori](dizoo/slime_volley/slime_volley.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/slime_volley)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/slime_volleyball.html)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/slime_volleyball_zh.html) | +| 21 | [gym_hybrid](https://github.com/thomashirtz/gym-hybrid) | ![hybrid](https://img.shields.io/badge/-hybrid-darkgreen) | ![ori](dizoo/gym_hybrid/moving_v0.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/gym_hybrid)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/gym_hybrid.html)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/gym_hybrid_zh.html) | +| 22 | [GoBigger](https://github.com/opendilab/GoBigger) | ![hybrid](https://img.shields.io/badge/-hybrid-darkgreen)![marl](https://img.shields.io/badge/-MARL-yellow)![selfplay](https://img.shields.io/badge/-selfplay-blue) | ![ori](./dizoo/gobigger_overview.gif) | [dizoo link](https://github.com/opendilab/GoBigger-Challenge-2021/tree/main/di_baseline)
[env tutorial](https://gobigger.readthedocs.io/en/latest/index.html)
[环境指南](https://gobigger.readthedocs.io/zh_CN/latest/) | +| 23 | [gym_soccer](https://github.com/openai/gym-soccer) | ![hybrid](https://img.shields.io/badge/-hybrid-darkgreen) | ![ori](dizoo/gym_soccer/half_offensive.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/gym_soccer)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/gym_soccer_zh.html) | +| 24 | [multiagent_mujoco](https://github.com/schroederdewitt/multiagent_mujoco) | ![continuous](https://img.shields.io/badge/-continous-green) ![marl](https://img.shields.io/badge/-MARL-yellow) | ![original](./dizoo/mujoco/mujoco.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/multiagent_mujoco/envs)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/mujoco_zh.html) | +| 25 | bitflip | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) ![sparse](https://img.shields.io/badge/-sparse%20reward-orange) | ![original](./dizoo/bitflip/bitflip.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/bitflip/envs)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/bitflip_zh.html) | +| 26 | [sokoban](https://github.com/mpSchrader/gym-sokoban) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | ![Game 2](https://github.com/mpSchrader/gym-sokoban/raw/default/docs/Animations/solved_4.gif?raw=true) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/sokoban/envs)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/sokoban.html)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/sokoban_zh.html) | +| 27 | [gym_anytrading](https://github.com/AminHP/gym-anytrading) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | ![original](./dizoo/gym_anytrading/envs/position.png) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/gym_anytrading)
[env tutorial](https://github.com/opendilab/DI-engine/blob/main/dizoo/gym_anytrading/envs/README.md) | +| 28 | [mario](https://github.com/Kautenja/gym-super-mario-bros) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | ![original](./dizoo/mario/mario.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/mario)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/gym_super_mario_bros.html)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/gym_super_mario_bros_zh.html) | +| 29 | [dmc2gym](https://github.com/denisyarats/dmc2gym) | ![continuous](https://img.shields.io/badge/-continous-green) | ![original](./dizoo/dmc2gym/dmc2gym_cheetah.png) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/dmc2gym)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/dmc2gym.html)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/dmc2gym_zh.html) | +| 30 | [evogym](https://github.com/EvolutionGym/evogym) | ![continuous](https://img.shields.io/badge/-continous-green) | ![original](./dizoo/evogym/evogym.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/evogym/envs)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/evogym.html)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/Evogym_zh.html) | +| 31 | [gym-pybullet-drones](https://github.com/utiasDSL/gym-pybullet-drones) | ![continuous](https://img.shields.io/badge/-continous-green) | ![original](./dizoo/gym_pybullet_drones/gym_pybullet_drones.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/gym_pybullet_drones/envs)
环境指南 | +| 32 | [beergame](https://github.com/OptMLGroup/DeepBeerInventory-RL) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | ![original](./dizoo/beergame/beergame.png) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/beergame/envs)
环境指南 | +| 33 | [classic_control/acrobot](https://github.com/openai/gym/tree/master/gym/envs/classic_control) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | ![original](./dizoo/classic_control/acrobot/acrobot.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/classic_control/acrobot/envs)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/acrobot_zh.html) | +| 34 | [box2d/car_racing](https://github.com/openai/gym/blob/master/gym/envs/box2d/car_racing.py) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen)
![continuous](https://img.shields.io/badge/-continous-green) | ![original](./dizoo/box2d/carracing/car_racing.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/box2d/carracing/envs)
环境指南 | +| 35 | [metadrive](https://github.com/metadriverse/metadrive) | ![continuous](https://img.shields.io/badge/-continous-green) | ![original](./dizoo/metadrive/metadrive_env.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/metadrive/env)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/metadrive_zh.html) | +| 36 | [cliffwalking](https://github.com/openai/gym/blob/master/gym/envs/toy_text/cliffwalking.py) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | ![original](./dizoo/cliffwalking/cliff_walking.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/cliffwalking/envs)
env tutorial
环境指南 | +| 37 | [tabmwp](https://promptpg.github.io/explore.html) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | ![original](./dizoo/tabmwp/tabmwp.jpeg) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/tabmwp)
env tutorial
环境指南 | +| 38 | [frozen_lake](https://gymnasium.farama.org/environments/toy_text/frozen_lake) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | ![original](./dizoo/frozen_lake/FrozenLake.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/frozen_lake)
env tutorial
环境指南 | +| 39 | [ising_model](https://github.com/mlii/mfrl/tree/master/examples/ising_model) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) ![marl](https://img.shields.io/badge/-MARL-yellow) | ![original](./dizoo/ising_env/ising_env.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/ising_env)
env tutorial
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/ising_model_zh.html) | +| 40 | [taxi](https://www.gymlibrary.dev/environments/toy_text/taxi/) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | ![original](./dizoo/taxi/Taxi-v3_episode_0.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/taxi/envs)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/taxi.html)
[环境指南](https://di-engine-docs.readthedocs.io/zh-cn/latest/13_envs/taxi_zh.html) | -### Environment Versatility -| No | Environment | Label | Visualization | Code and Doc Links | -| :--: | :--------------------------------------: | :---------------------------------: | :--------------------------------:|:---------------------------------------------------------: | -| 1 | [atari](https://github.com/openai/gym/tree/master/gym/envs/atari) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | ![original](./dizoo/atari/atari.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/atari/envs)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/atari.html)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/atari_zh.html) | -| 2 | [box2d/bipedalwalker](https://github.com/openai/gym/tree/master/gym/envs/box2d) | ![continuous](https://img.shields.io/badge/-continous-green) | ![original](./dizoo/box2d/bipedalwalker/original.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/box2d/bipedalwalker/envs)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/bipedalwalker.html)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/bipedalwalker_zh.html) | -| 3 | [box2d/lunarlander](https://github.com/openai/gym/tree/master/gym/envs/box2d) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | ![original](./dizoo/box2d/lunarlander/lunarlander.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/box2d/lunarlander/envs)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/lunarlander.html)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/lunarlander_zh.html) | -| 4 | [classic_control/cartpole](https://github.com/openai/gym/tree/master/gym/envs/classic_control) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | ![original](./dizoo/classic_control/cartpole/cartpole.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/classic_control/cartpole/envs)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/cartpole.html)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/cartpole_zh.html) | -| 5 | [classic_control/pendulum](https://github.com/openai/gym/tree/master/gym/envs/classic_control) | ![continuous](https://img.shields.io/badge/-continous-green) | ![original](./dizoo/classic_control/pendulum/pendulum.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/classic_control/pendulum/envs)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/pendulum.html)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/pendulum_zh.html) | -| 6 | [competitive_rl](https://github.com/cuhkrlcourse/competitive-rl) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) ![selfplay](https://img.shields.io/badge/-selfplay-blue) | ![original](./dizoo/competitive_rl/competitive_rl.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo.classic_control)
[环境指南](https://di-engine-docs.readthedocs.io/en/latest/13_envs/competitive_rl_zh.html) | -| 7 | [gfootball](https://github.com/google-research/football) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen)![sparse](https://img.shields.io/badge/-sparse%20reward-orange)![selfplay](https://img.shields.io/badge/-selfplay-blue) | ![original](./dizoo/gfootball/gfootball.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo.gfootball/envs)
[环境指南](https://di-engine-docs.readthedocs.io/en/latest/13_envs/gfootball_zh.html) | -| 8 | [minigrid](https://github.com/maximecb/gym-minigrid) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen)![sparse](https://img.shields.io/badge/-sparse%20reward-orange) | ![original](./dizoo/minigrid/minigrid.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/minigrid/envs)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/minigrid.html)
[环境指南](https://di-engine-docs.readthedocs.io/en/latest/13_envs/minigrid_zh.html) | -| 9 | [mujoco](https://github.com/openai/gym/tree/master/gym/envs/mujoco) | ![continuous](https://img.shields.io/badge/-continous-green) | ![original](./dizoo/mujoco/mujoco.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/majoco/envs)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/mujoco.html)
[环境指南](https://di-engine-docs.readthedocs.io/en/latest/13_envs/mujoco_zh.html) | -| 10 | [PettingZoo](https://github.com/Farama-Foundation/PettingZoo) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) ![continuous](https://img.shields.io/badge/-continous-green) ![marl](https://img.shields.io/badge/-MARL-yellow) | ![original](./dizoo/petting_zoo/petting_zoo_mpe_simple_spread.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/petting_zoo/envs)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/pettingzoo_zh.html) | -| 11 | [overcooked](https://github.com/HumanCompatibleAI/overcooked-demo) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) ![marl](https://img.shields.io/badge/-MARL-yellow) | ![original](./dizoo/overcooked/overcooked.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/overcooded/envs)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/overcooked.html) | -| 12 | [procgen](https://github.com/openai/procgen) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | ![original](./dizoo/procgen/coinrun.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/procgen)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/procgen.html)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/procgen_zh.html) | -| 13 | [pybullet](https://github.com/benelot/pybullet-gym) | ![continuous](https://img.shields.io/badge/-continous-green) | ![original](./dizoo/pybullet/pybullet.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/pybullet/envs)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/pybullet_zh.html) | -| 14 | [smac](https://github.com/oxwhirl/smac) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) ![marl](https://img.shields.io/badge/-MARL-yellow)![selfplay](https://img.shields.io/badge/-selfplay-blue)![sparse](https://img.shields.io/badge/-sparse%20reward-orange) | ![original](./dizoo/smac/smac.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/smac/envs)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/smac.html)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/smac_zh.html) | -| 15 | [d4rl](https://github.com/rail-berkeley/d4rl) | ![offline](https://img.shields.io/badge/-offlineRL-darkblue) | ![ori](dizoo/d4rl/d4rl.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/d4rl)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/d4rl_zh.html) | -| 16 | league_demo | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) ![selfplay](https://img.shields.io/badge/-selfplay-blue) | ![original](./dizoo/league_demo/league_demo.png) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/league_demo/envs) | -| 17 | pomdp atari | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/pomdp/envs) | -| 18 | [bsuite](https://github.com/deepmind/bsuite) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | ![original](./dizoo/bsuite/bsuite.png) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/bsuite/envs)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs//bsuite.html) | -| 19 | [ImageNet](https://www.image-net.org/) | ![IL](https://img.shields.io/badge/-IL/SL-purple) | ![original](./dizoo/image_classification/imagenet.png) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/image_classification)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/image_cls_zh.html) | -| 20 | [slime_volleyball](https://github.com/hardmaru/slimevolleygym) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen)![selfplay](https://img.shields.io/badge/-selfplay-blue) | ![ori](dizoo/slime_volley/slime_volley.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/slime_volley)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/slime_volleyball.html)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/slime_volleyball_zh.html) | -| 21 | [gym_hybrid](https://github.com/thomashirtz/gym-hybrid) | ![hybrid](https://img.shields.io/badge/-hybrid-darkgreen) | ![ori](dizoo/gym_hybrid/moving_v0.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/gym_hybrid)
[env tutorial](https://di-engine-docs.readthedocs.io/en/latest/13_envs/gym_hybrid.html)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/gym_hybrid_zh.html) | -| 22 | [GoBigger](https://github.com/opendilab/GoBigger) | ![hybrid](https://img.shields.io/badge/-hybrid-darkgreen)![marl](https://img.shields.io/badge/-MARL-yellow)![selfplay](https://img.shields.io/badge/-selfplay-blue) | ![ori](./dizoo/gobigger_overview.gif) | [dizoo link](https://github.com/opendilab/GoBigger-Challenge-2021/tree/main/di_baseline)
[env tutorial](https://gobigger.readthedocs.io/en/latest/index.html)
[环境指南](https://gobigger.readthedocs.io/zh_CN/latest/) | -| 23 | [gym_soccer](https://github.com/openai/gym-soccer) | ![hybrid](https://img.shields.io/badge/-hybrid-darkgreen) | ![ori](dizoo/gym_soccer/half_offensive.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/gym_soccer)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/gym_soccer_zh.html) | -| 24 |[multiagent_mujoco](https://github.com/schroederdewitt/multiagent_mujoco) | ![continuous](https://img.shields.io/badge/-continous-green) ![marl](https://img.shields.io/badge/-MARL-yellow) | ![original](./dizoo/mujoco/mujoco.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/multiagent_mujoco/envs)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/mujoco_zh.html) | -| 25 |bitflip | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) ![sparse](https://img.shields.io/badge/-sparse%20reward-orange) | ![original](./dizoo/bitflip/bitflip.gif) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/bitflip/envs)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/bitflip_zh.html) | -| 26 |[sokoban](https://github.com/mpSchrader/gym-sokoban) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | ![Game 2](https://github.com/mpSchrader/gym-sokoban/raw/default/docs/Animations/solved_4.gif?raw=true) | [dizoo link](https://github.com/opendilab/DI-engine/tree/main/dizoo/sokoban/envs)
[环境指南](https://di-engine-docs.readthedocs.io/zh_CN/latest/13_envs/sokoban_zh.html) | -| 27 |[gym_anytrading](https://github.com/AminHP/gym-anytrading) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | ![original](./dizoo/gym_anytrading/envs/position.png) | dizoo link
环境指南 | -| 28 |[mario](https://github.com/Kautenja/gym-super-mario-bros) | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) | ![original](./dizoo/mario/mario.gif) | dizoo link
环境指南 | ![discrete](https://img.shields.io/badge/-discrete-brightgreen) means discrete action space @@ -261,19 +347,127 @@ P.S: The `.py` file in `Runnable Demo` can be found in `dizoo` ![selfplay](https://img.shields.io/badge/-selfplay-blue) means environment that allows agent VS agent battle -P.S. some enviroments in Atari, such as **MontezumaRevenge**, are also sparse reward type - +P.S. some enviroments in Atari, such as **MontezumaRevenge**, are also the sparse reward type. + +
+ +### General Data Container: TreeTensor + +DI-engine utilizes [TreeTensor](https://github.com/opendilab/DI-treetensor) as the basic data container in various components, which is ease of use and consistent across different code modules such as environment definition, data processing and DRL optimization. Here are some concrete code examples: + +- TreeTensor can easily extend all the operations of `torch.Tensor` to nested data: + +
+ (Click for Details) + + ```python + import treetensor.torch as ttorch + + + # create random tensor + data = ttorch.randn({'a': (3, 2), 'b': {'c': (3, )}}) + # clone+detach tensor + data_clone = data.clone().detach() + # access tree structure like attribute + a = data.a + c = data.b.c + # stack/cat/split + stacked_data = ttorch.stack([data, data_clone], 0) + cat_data = ttorch.cat([data, data_clone], 0) + data, data_clone = ttorch.split(stacked_data, 1) + # reshape + data = data.unsqueeze(-1) + data = data.squeeze(-1) + flatten_data = data.view(-1) + # indexing + data_0 = data[0] + data_1to2 = data[1:2] + # execute math calculations + data = data.sin() + data.b.c.cos_().clamp_(-1, 1) + data += data ** 2 + # backward + data.requires_grad_(True) + loss = data.arctan().mean() + loss.backward() + # print shape + print(data.shape) + # result + # + # ├── 'a' --> torch.Size([1, 3, 2]) + # └── 'b' --> + # └── 'c' --> torch.Size([1, 3]) + ``` + +
+- TreeTensor can make it simple yet effective to implement classic deep reinforcement learning pipeline + +
+ (Click for Details) + + ```diff + import torch + import treetensor.torch as ttorch + + B = 4 + + + def get_item(): + return { + 'obs': { + 'scalar': torch.randn(12), + 'image': torch.randn(3, 32, 32), + }, + 'action': torch.randint(0, 10, size=(1,)), + 'reward': torch.rand(1), + 'done': False, + } + + + data = [get_item() for _ in range(B)] + + + # execute `stack` op + - def stack(data, dim): + - elem = data[0] + - if isinstance(elem, torch.Tensor): + - return torch.stack(data, dim) + - elif isinstance(elem, dict): + - return {k: stack([item[k] for item in data], dim) for k in elem.keys()} + - elif isinstance(elem, bool): + - return torch.BoolTensor(data) + - else: + - raise TypeError("not support elem type: {}".format(type(elem))) + - stacked_data = stack(data, dim=0) + + data = [ttorch.tensor(d) for d in data] + + stacked_data = ttorch.stack(data, dim=0) + + # validate + - assert stacked_data['obs']['image'].shape == (B, 3, 32, 32) + - assert stacked_data['action'].shape == (B, 1) + - assert stacked_data['reward'].shape == (B, 1) + - assert stacked_data['done'].shape == (B,) + - assert stacked_data['done'].dtype == torch.bool + + assert stacked_data.obs.image.shape == (B, 3, 32, 32) + + assert stacked_data.action.shape == (B, 1) + + assert stacked_data.reward.shape == (B, 1) + + assert stacked_data.done.shape == (B,) + + assert stacked_data.done.dtype == torch.bool + ``` + +
## Feedback and Contribution - [File an issue](https://github.com/opendilab/DI-engine/issues/new/choose) on Github - Open or participate in our [forum](https://github.com/opendilab/DI-engine/discussions) +- Discuss on DI-engine [discord server](https://discord.gg/dkZS2JF56X) - Discuss on DI-engine [slack communication channel](https://join.slack.com/t/opendilab/shared_invite/zt-v9tmv4fp-nUBAQEH1_Kuyu_q4plBssQ) -- Discuss on DI-engine's QQ group (700157520) or add us on WeChat - - ![WeChat](https://github.com/opendilab/DI-engine/blob/main/assets/wechat.png) -- Contact our email (opendilab.contact@gmail.com) -- Contributes to our future plan [Roadmap](https://github.com/opendilab/DI-engine/projects) +- Discuss on DI-engine's WeChat group (i.e. add us on WeChat: ding314assist) + + +- Contact our email (opendilab@pjlab.org.cn) +- Contributes to our future plan [Roadmap](https://github.com/opendilab/DI-engine/issues/548) We appreciate all the feedbacks and contributions to improve DI-engine, both algorithms and system designs. And `CONTRIBUTING.md` offers some necessary information. @@ -287,17 +481,18 @@ We appreciate all the feedbacks and contributions to improve DI-engine, both alg [![Forkers repo roster for @opendilab/DI-engine](https://reporoster.com/forks/opendilab/DI-engine)](https://github.com/opendilab/DI-engine/network/members) - ## Citation + ```latex @misc{ding, - title={{DI-engine: OpenDILab} Decision Intelligence Engine}, - author={DI-engine Contributors}, - publisher = {GitHub}, - howpublished = {\url{https://github.com/opendilab/DI-engine}}, + title={DI-engine: A Universal AI System/Engine for Decision Intelligence}, + author={Niu, Yazhe and Xu, Jingxin and Pu, Yuan and Nie, Yunpeng and Zhang, Jinouwen and Hu, Shuai and Zhao, Liangxuan and Zhang, Ming and Liu, Yu}, + publisher={GitHub}, + howpublished={\url{https://github.com/opendilab/DI-engine}}, year={2021}, } ``` ## License + DI-engine released under the Apache 2.0 license. diff --git a/SECURITY.md b/SECURITY.md new file mode 100644 index 0000000000..034e848032 --- /dev/null +++ b/SECURITY.md @@ -0,0 +1,21 @@ +# Security Policy + +## Supported Versions + +Use this section to tell people about which versions of your project are +currently being supported with security updates. + +| Version | Supported | +| ------- | ------------------ | +| 5.1.x | :white_check_mark: | +| 5.0.x | :x: | +| 4.0.x | :white_check_mark: | +| < 4.0 | :x: | + +## Reporting a Vulnerability + +Use this section to tell people how to report a vulnerability. + +Tell them where to go, how often they can expect to get an update on a +reported vulnerability, what to expect if the vulnerability is accepted or +declined, etc. diff --git a/assets/wechat.jpeg b/assets/wechat.jpeg new file mode 100644 index 0000000000..d4285c1f26 Binary files /dev/null and b/assets/wechat.jpeg differ diff --git a/assets/wechat.png b/assets/wechat.png deleted file mode 100644 index 9861f6ebb1..0000000000 Binary files a/assets/wechat.png and /dev/null differ diff --git a/conda/meta.yaml b/conda/meta.yaml index ff4e91d0d2..c6dfcb7c90 100644 --- a/conda/meta.yaml +++ b/conda/meta.yaml @@ -1,7 +1,7 @@ {% set data = load_setup_py_data() %} package: name: di-engine - version: v0.4.3 + version: v0.5.3 source: path: .. diff --git a/ding/__init__.py b/ding/__init__.py index 0844b0e105..e3e1e4554b 100644 --- a/ding/__init__.py +++ b/ding/__init__.py @@ -1,7 +1,7 @@ import os __TITLE__ = 'DI-engine' -__VERSION__ = 'v0.4.3' +__VERSION__ = 'v0.5.3' __DESCRIPTION__ = 'Decision AI Engine' __AUTHOR__ = "OpenDILab Contributors" __AUTHOR_EMAIL__ = "opendilab@pjlab.org.cn" diff --git a/ding/bonus/__init__.py b/ding/bonus/__init__.py new file mode 100644 index 0000000000..3329830b39 --- /dev/null +++ b/ding/bonus/__init__.py @@ -0,0 +1,132 @@ +import ding.config +from .a2c import A2CAgent +from .c51 import C51Agent +from .ddpg import DDPGAgent +from .dqn import DQNAgent +from .pg import PGAgent +from .ppof import PPOF +from .ppo_offpolicy import PPOOffPolicyAgent +from .sac import SACAgent +from .sql import SQLAgent +from .td3 import TD3Agent + +supported_algo = dict( + A2C=A2CAgent, + C51=C51Agent, + DDPG=DDPGAgent, + DQN=DQNAgent, + PG=PGAgent, + PPOF=PPOF, + PPOOffPolicy=PPOOffPolicyAgent, + SAC=SACAgent, + SQL=SQLAgent, + TD3=TD3Agent, +) + +supported_algo_list = list(supported_algo.keys()) + + +def env_supported(algo: str = None) -> list: + """ + return list of the envs that supported by di-engine. + """ + + if algo is not None: + if algo.upper() == "A2C": + return list(ding.config.example.A2C.supported_env.keys()) + elif algo.upper() == "C51": + return list(ding.config.example.C51.supported_env.keys()) + elif algo.upper() == "DDPG": + return list(ding.config.example.DDPG.supported_env.keys()) + elif algo.upper() == "DQN": + return list(ding.config.example.DQN.supported_env.keys()) + elif algo.upper() == "PG": + return list(ding.config.example.PG.supported_env.keys()) + elif algo.upper() == "PPOF": + return list(ding.config.example.PPOF.supported_env.keys()) + elif algo.upper() == "PPOOFFPOLICY": + return list(ding.config.example.PPOOffPolicy.supported_env.keys()) + elif algo.upper() == "SAC": + return list(ding.config.example.SAC.supported_env.keys()) + elif algo.upper() == "SQL": + return list(ding.config.example.SQL.supported_env.keys()) + elif algo.upper() == "TD3": + return list(ding.config.example.TD3.supported_env.keys()) + else: + raise ValueError("The algo {} is not supported by di-engine.".format(algo)) + else: + supported_env = set() + supported_env.update(ding.config.example.A2C.supported_env.keys()) + supported_env.update(ding.config.example.C51.supported_env.keys()) + supported_env.update(ding.config.example.DDPG.supported_env.keys()) + supported_env.update(ding.config.example.DQN.supported_env.keys()) + supported_env.update(ding.config.example.PG.supported_env.keys()) + supported_env.update(ding.config.example.PPOF.supported_env.keys()) + supported_env.update(ding.config.example.PPOOffPolicy.supported_env.keys()) + supported_env.update(ding.config.example.SAC.supported_env.keys()) + supported_env.update(ding.config.example.SQL.supported_env.keys()) + supported_env.update(ding.config.example.TD3.supported_env.keys()) + # return the list of the envs + return list(supported_env) + + +supported_env = env_supported() + + +def algo_supported(env_id: str = None) -> list: + """ + return list of the algos that supported by di-engine. + """ + if env_id is not None: + algo = [] + if env_id.upper() in [item.upper() for item in ding.config.example.A2C.supported_env.keys()]: + algo.append("A2C") + if env_id.upper() in [item.upper() for item in ding.config.example.C51.supported_env.keys()]: + algo.append("C51") + if env_id.upper() in [item.upper() for item in ding.config.example.DDPG.supported_env.keys()]: + algo.append("DDPG") + if env_id.upper() in [item.upper() for item in ding.config.example.DQN.supported_env.keys()]: + algo.append("DQN") + if env_id.upper() in [item.upper() for item in ding.config.example.PG.supported_env.keys()]: + algo.append("PG") + if env_id.upper() in [item.upper() for item in ding.config.example.PPOF.supported_env.keys()]: + algo.append("PPOF") + if env_id.upper() in [item.upper() for item in ding.config.example.PPOOffPolicy.supported_env.keys()]: + algo.append("PPOOffPolicy") + if env_id.upper() in [item.upper() for item in ding.config.example.SAC.supported_env.keys()]: + algo.append("SAC") + if env_id.upper() in [item.upper() for item in ding.config.example.SQL.supported_env.keys()]: + algo.append("SQL") + if env_id.upper() in [item.upper() for item in ding.config.example.TD3.supported_env.keys()]: + algo.append("TD3") + + if len(algo) == 0: + raise ValueError("The env {} is not supported by di-engine.".format(env_id)) + return algo + else: + return supported_algo_list + + +def is_supported(env_id: str = None, algo: str = None) -> bool: + """ + Check if the env-algo pair is supported by di-engine. + """ + if env_id is not None and env_id.upper() in [item.upper() for item in supported_env.keys()]: + if algo is not None and algo.upper() in supported_algo_list: + if env_id.upper() in env_supported(algo): + return True + else: + return False + elif algo is None: + return True + else: + return False + elif env_id is None: + if algo is not None and algo.upper() in supported_algo_list: + return True + elif algo is None: + raise ValueError("Please specify the env or algo.") + else: + return False + else: + return False diff --git a/ding/bonus/a2c.py b/ding/bonus/a2c.py new file mode 100644 index 0000000000..4533664f76 --- /dev/null +++ b/ding/bonus/a2c.py @@ -0,0 +1,460 @@ +from typing import Optional, Union, List +from ditk import logging +from easydict import EasyDict +import os +import numpy as np +import torch +import treetensor.torch as ttorch +from ding.framework import task, OnlineRLContext +from ding.framework.middleware import CkptSaver, trainer, \ + wandb_online_logger, offline_data_saver, termination_checker, interaction_evaluator, StepCollector, \ + gae_estimator, final_ctx_saver +from ding.envs import BaseEnv +from ding.envs import setup_ding_env_manager +from ding.policy import A2CPolicy +from ding.utils import set_pkg_seed +from ding.utils import get_env_fps, render +from ding.config import save_config_py, compile_config +from ding.model import VAC +from ding.model import model_wrap +from ding.bonus.common import TrainingReturn, EvalReturn +from ding.config.example.A2C import supported_env_cfg +from ding.config.example.A2C import supported_env + + +class A2CAgent: + """ + Overview: + Class of agent for training, evaluation and deployment of Reinforcement learning algorithm \ + Advantage Actor Critic(A2C). + For more information about the system design of RL agent, please refer to \ + . + Interface: + ``__init__``, ``train``, ``deploy``, ``collect_data``, ``batch_evaluate``, ``best`` + """ + supported_env_list = list(supported_env_cfg.keys()) + """ + Overview: + List of supported envs. + Examples: + >>> from ding.bonus.a2c import A2CAgent + >>> print(A2CAgent.supported_env_list) + """ + + def __init__( + self, + env_id: str = None, + env: BaseEnv = None, + seed: int = 0, + exp_name: str = None, + model: Optional[torch.nn.Module] = None, + cfg: Optional[Union[EasyDict, dict]] = None, + policy_state_dict: str = None, + ) -> None: + """ + Overview: + Initialize agent for A2C algorithm. + Arguments: + - env_id (:obj:`str`): The environment id, which is a registered environment name in gym or gymnasium. \ + If ``env_id`` is not specified, ``env_id`` in ``cfg.env`` must be specified. \ + If ``env_id`` is specified, ``env_id`` in ``cfg.env`` will be ignored. \ + ``env_id`` should be one of the supported envs, which can be found in ``supported_env_list``. + - env (:obj:`BaseEnv`): The environment instance for training and evaluation. \ + If ``env`` is not specified, `env_id`` or ``cfg.env.env_id`` must be specified. \ + ``env_id`` or ``cfg.env.env_id`` will be used to create environment instance. \ + If ``env`` is specified, ``env_id`` and ``cfg.env.env_id`` will be ignored. + - seed (:obj:`int`): The random seed, which is set before running the program. \ + Default to 0. + - exp_name (:obj:`str`): The name of this experiment, which will be used to create the folder to save \ + log data. Default to None. If not specified, the folder name will be ``env_id``-``algorithm``. + - model (:obj:`torch.nn.Module`): The model of A2C algorithm, which should be an instance of class \ + :class:`ding.model.VAC`. \ + If not specified, a default model will be generated according to the configuration. + - cfg (:obj:`Union[EasyDict, dict]`): The configuration of A2C algorithm, which is a dict. \ + Default to None. If not specified, the default configuration will be used. \ + The default configuration can be found in ``ding/config/example/A2C/gym_lunarlander_v2.py``. + - policy_state_dict (:obj:`str`): The path of policy state dict saved by PyTorch a in local file. \ + If specified, the policy will be loaded from this file. Default to None. + + .. note:: + An RL Agent Instance can be initialized in two basic ways. \ + For example, we have an environment with id ``LunarLanderContinuous-v2`` registered in gym, \ + and we want to train an agent with A2C algorithm with default configuration. \ + Then we can initialize the agent in the following ways: + >>> agent = A2CAgent(env_id='LunarLanderContinuous-v2') + or, if we want can specify the env_id in the configuration: + >>> cfg = {'env': {'env_id': 'LunarLanderContinuous-v2'}, 'policy': ...... } + >>> agent = A2CAgent(cfg=cfg) + There are also other arguments to specify the agent when initializing. + For example, if we want to specify the environment instance: + >>> env = CustomizedEnv('LunarLanderContinuous-v2') + >>> agent = A2CAgent(cfg=cfg, env=env) + or, if we want to specify the model: + >>> model = VAC(**cfg.policy.model) + >>> agent = A2CAgent(cfg=cfg, model=model) + or, if we want to reload the policy from a saved policy state dict: + >>> agent = A2CAgent(cfg=cfg, policy_state_dict='LunarLanderContinuous-v2.pth.tar') + Make sure that the configuration is consistent with the saved policy state dict. + """ + + assert env_id is not None or cfg is not None, "Please specify env_id or cfg." + + if cfg is not None and not isinstance(cfg, EasyDict): + cfg = EasyDict(cfg) + + if env_id is not None: + assert env_id in A2CAgent.supported_env_list, "Please use supported envs: {}".format( + A2CAgent.supported_env_list + ) + if cfg is None: + cfg = supported_env_cfg[env_id] + else: + assert cfg.env.env_id == env_id, "env_id in cfg should be the same as env_id in args." + else: + assert hasattr(cfg.env, "env_id"), "Please specify env_id in cfg." + assert cfg.env.env_id in A2CAgent.supported_env_list, "Please use supported envs: {}".format( + A2CAgent.supported_env_list + ) + default_policy_config = EasyDict({"policy": A2CPolicy.default_config()}) + default_policy_config.update(cfg) + cfg = default_policy_config + + if exp_name is not None: + cfg.exp_name = exp_name + self.cfg = compile_config(cfg, policy=A2CPolicy) + self.exp_name = self.cfg.exp_name + if env is None: + self.env = supported_env[cfg.env.env_id](cfg=cfg.env) + else: + assert isinstance(env, BaseEnv), "Please use BaseEnv as env data type." + self.env = env + + logging.getLogger().setLevel(logging.INFO) + self.seed = seed + set_pkg_seed(self.seed, use_cuda=self.cfg.policy.cuda) + if not os.path.exists(self.exp_name): + os.makedirs(self.exp_name) + save_config_py(self.cfg, os.path.join(self.exp_name, 'policy_config.py')) + if model is None: + model = VAC(**self.cfg.policy.model) + self.policy = A2CPolicy(self.cfg.policy, model=model) + if policy_state_dict is not None: + self.policy.learn_mode.load_state_dict(policy_state_dict) + self.checkpoint_save_dir = os.path.join(self.exp_name, "ckpt") + + def train( + self, + step: int = int(1e7), + collector_env_num: int = 4, + evaluator_env_num: int = 4, + n_iter_log_show: int = 500, + n_iter_save_ckpt: int = 1000, + context: Optional[str] = None, + debug: bool = False, + wandb_sweep: bool = False, + ) -> TrainingReturn: + """ + Overview: + Train the agent with A2C algorithm for ``step`` iterations with ``collector_env_num`` collector \ + environments and ``evaluator_env_num`` evaluator environments. Information during training will be \ + recorded and saved by wandb. + Arguments: + - step (:obj:`int`): The total training environment steps of all collector environments. Default to 1e7. + - collector_env_num (:obj:`int`): The collector environment number. Default to None. \ + If not specified, it will be set according to the configuration. + - evaluator_env_num (:obj:`int`): The evaluator environment number. Default to None. \ + If not specified, it will be set according to the configuration. + - n_iter_save_ckpt (:obj:`int`): The frequency of saving checkpoint every training iteration. \ + Default to 1000. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + - wandb_sweep (:obj:`bool`): Whether to use wandb sweep, \ + which is a hyper-parameter optimization process for seeking the best configurations. \ + Default to False. If True, the wandb sweep id will be used as the experiment name. + Returns: + - (:obj:`TrainingReturn`): The training result, of which the attributions are: + - wandb_url (:obj:`str`): The weight & biases (wandb) project url of the trainning experiment. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + logging.debug(self.policy._model) + # define env and policy + collector_env = self._setup_env_manager(collector_env_num, context, debug, 'collector') + evaluator_env = self._setup_env_manager(evaluator_env_num, context, debug, 'evaluator') + + with task.start(ctx=OnlineRLContext()): + task.use( + interaction_evaluator( + self.cfg, + self.policy.eval_mode, + evaluator_env, + render=self.cfg.policy.eval.render if hasattr(self.cfg.policy.eval, "render") else False + ) + ) + task.use(CkptSaver(policy=self.policy, save_dir=self.checkpoint_save_dir, train_freq=n_iter_save_ckpt)) + task.use( + StepCollector( + self.cfg, + self.policy.collect_mode, + collector_env, + random_collect_size=self.cfg.policy.random_collect_size + if hasattr(self.cfg.policy, 'random_collect_size') else 0, + ) + ) + task.use(gae_estimator(self.cfg, self.policy.collect_mode)) + task.use(trainer(self.cfg, self.policy.learn_mode)) + task.use( + wandb_online_logger( + metric_list=self.policy._monitor_vars_learn(), + model=self.policy._model, + anonymous=True, + project_name=self.exp_name, + wandb_sweep=wandb_sweep, + ) + ) + task.use(termination_checker(max_env_step=step)) + task.use(final_ctx_saver(name=self.exp_name)) + task.run() + + return TrainingReturn(wandb_url=task.ctx.wandb_url) + + def deploy( + self, + enable_save_replay: bool = False, + concatenate_all_replay: bool = False, + replay_save_path: str = None, + seed: Optional[Union[int, List]] = None, + debug: bool = False + ) -> EvalReturn: + """ + Overview: + Deploy the agent with A2C algorithm by interacting with the environment, during which the replay video \ + can be saved if ``enable_save_replay`` is True. The evaluation result will be returned. + Arguments: + - enable_save_replay (:obj:`bool`): Whether to save the replay video. Default to False. + - concatenate_all_replay (:obj:`bool`): Whether to concatenate all replay videos into one video. \ + Default to False. If ``enable_save_replay`` is False, this argument will be ignored. \ + If ``enable_save_replay`` is True and ``concatenate_all_replay`` is False, \ + the replay video of each episode will be saved separately. + - replay_save_path (:obj:`str`): The path to save the replay video. Default to None. \ + If not specified, the video will be saved in ``exp_name/videos``. + - seed (:obj:`Union[int, List]`): The random seed, which is set before running the program. \ + Default to None. If not specified, ``self.seed`` will be used. \ + If ``seed`` is an integer, the agent will be deployed once. \ + If ``seed`` is a list of integers, the agent will be deployed once for each seed in the list. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + Returns: + - (:obj:`EvalReturn`): The evaluation result, of which the attributions are: + - eval_value (:obj:`np.float32`): The mean of evaluation return. + - eval_value_std (:obj:`np.float32`): The standard deviation of evaluation return. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + # define env and policy + env = self.env.clone(caller='evaluator') + + if seed is not None and isinstance(seed, int): + seeds = [seed] + elif seed is not None and isinstance(seed, list): + seeds = seed + else: + seeds = [self.seed] + + returns = [] + images = [] + if enable_save_replay: + replay_save_path = os.path.join(self.exp_name, 'videos') if replay_save_path is None else replay_save_path + env.enable_save_replay(replay_path=replay_save_path) + else: + logging.warning('No video would be generated during the deploy.') + if concatenate_all_replay: + logging.warning('concatenate_all_replay is set to False because enable_save_replay is False.') + concatenate_all_replay = False + + def single_env_forward_wrapper(forward_fn, cuda=True): + + if self.cfg.policy.action_space == 'continuous': + forward_fn = model_wrap(forward_fn, wrapper_name='deterministic_sample').forward + elif self.cfg.policy.action_space == 'discrete': + forward_fn = model_wrap(forward_fn, wrapper_name='argmax_sample').forward + else: + raise NotImplementedError + + def _forward(obs): + # unsqueeze means add batch dim, i.e. (O, ) -> (1, O) + obs = ttorch.as_tensor(obs).unsqueeze(0) + if cuda and torch.cuda.is_available(): + obs = obs.cuda() + action = forward_fn(obs, mode='compute_actor')["action"] + # squeeze means delete batch dim, i.e. (1, A) -> (A, ) + action = action.squeeze(0).detach().cpu().numpy() + return action + + return _forward + + forward_fn = single_env_forward_wrapper(self.policy._model, self.cfg.policy.cuda) + + # reset first to make sure the env is in the initial state + # env will be reset again in the main loop + env.reset() + + for seed in seeds: + env.seed(seed, dynamic_seed=False) + return_ = 0. + step = 0 + obs = env.reset() + images.append(render(env)[None]) if concatenate_all_replay else None + while True: + action = forward_fn(obs) + obs, rew, done, info = env.step(action) + images.append(render(env)[None]) if concatenate_all_replay else None + return_ += rew + step += 1 + if done: + break + logging.info(f'DQN deploy is finished, final episode return with {step} steps is: {return_}') + returns.append(return_) + + env.close() + + if concatenate_all_replay: + images = np.concatenate(images, axis=0) + import imageio + imageio.mimwrite(os.path.join(replay_save_path, 'deploy.mp4'), images, fps=get_env_fps(env)) + + return EvalReturn(eval_value=np.mean(returns), eval_value_std=np.std(returns)) + + def collect_data( + self, + env_num: int = 8, + save_data_path: Optional[str] = None, + n_sample: Optional[int] = None, + n_episode: Optional[int] = None, + context: Optional[str] = None, + debug: bool = False + ) -> None: + """ + Overview: + Collect data with A2C algorithm for ``n_episode`` episodes with ``env_num`` collector environments. \ + The collected data will be saved in ``save_data_path`` if specified, otherwise it will be saved in \ + ``exp_name/demo_data``. + Arguments: + - env_num (:obj:`int`): The number of collector environments. Default to 8. + - save_data_path (:obj:`str`): The path to save the collected data. Default to None. \ + If not specified, the data will be saved in ``exp_name/demo_data``. + - n_sample (:obj:`int`): The number of samples to collect. Default to None. \ + If not specified, ``n_episode`` must be specified. + - n_episode (:obj:`int`): The number of episodes to collect. Default to None. \ + If not specified, ``n_sample`` must be specified. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + if n_episode is not None: + raise NotImplementedError + # define env and policy + env_num = env_num if env_num else self.cfg.env.collector_env_num + env = setup_ding_env_manager(self.env, env_num, context, debug, 'collector') + + if save_data_path is None: + save_data_path = os.path.join(self.exp_name, 'demo_data') + + # main execution task + with task.start(ctx=OnlineRLContext()): + task.use( + StepCollector( + self.cfg, self.policy.collect_mode, env, random_collect_size=self.cfg.policy.random_collect_size + ) + ) + task.use(offline_data_saver(save_data_path, data_type='hdf5')) + task.run(max_step=1) + logging.info( + f'A2C collecting is finished, more than {n_sample} samples are collected and saved in `{save_data_path}`' + ) + + def batch_evaluate( + self, + env_num: int = 4, + n_evaluator_episode: int = 4, + context: Optional[str] = None, + debug: bool = False + ) -> EvalReturn: + """ + Overview: + Evaluate the agent with A2C algorithm for ``n_evaluator_episode`` episodes with ``env_num`` evaluator \ + environments. The evaluation result will be returned. + The difference between methods ``batch_evaluate`` and ``deploy`` is that ``batch_evaluate`` will create \ + multiple evaluator environments to evaluate the agent to get an average performance, while ``deploy`` \ + will only create one evaluator environment to evaluate the agent and save the replay video. + Arguments: + - env_num (:obj:`int`): The number of evaluator environments. Default to 4. + - n_evaluator_episode (:obj:`int`): The number of episodes to evaluate. Default to 4. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + Returns: + - (:obj:`EvalReturn`): The evaluation result, of which the attributions are: + - eval_value (:obj:`np.float32`): The mean of evaluation return. + - eval_value_std (:obj:`np.float32`): The standard deviation of evaluation return. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + # define env and policy + env_num = env_num if env_num else self.cfg.env.evaluator_env_num + env = setup_ding_env_manager(self.env, env_num, context, debug, 'evaluator') + + # reset first to make sure the env is in the initial state + # env will be reset again in the main loop + env.launch() + env.reset() + + evaluate_cfg = self.cfg + evaluate_cfg.env.n_evaluator_episode = n_evaluator_episode + + # main execution task + with task.start(ctx=OnlineRLContext()): + task.use(interaction_evaluator(self.cfg, self.policy.eval_mode, env)) + task.run(max_step=1) + + return EvalReturn(eval_value=task.ctx.eval_value, eval_value_std=task.ctx.eval_value_std) + + @property + def best(self) -> 'A2CAgent': + """ + Overview: + Load the best model from the checkpoint directory, \ + which is by default in folder ``exp_name/ckpt/eval.pth.tar``. \ + The return value is the agent with the best model. + Returns: + - (:obj:`A2CAgent`): The agent with the best model. + Examples: + >>> agent = A2CAgent(env_id='LunarLanderContinuous-v2') + >>> agent.train() + >>> agent = agent.best + + .. note:: + The best model is the model with the highest evaluation return. If this method is called, the current \ + model will be replaced by the best model. + """ + + best_model_file_path = os.path.join(self.checkpoint_save_dir, "eval.pth.tar") + # Load best model if it exists + if os.path.exists(best_model_file_path): + policy_state_dict = torch.load(best_model_file_path, map_location=torch.device("cpu")) + self.policy.learn_mode.load_state_dict(policy_state_dict) + return self diff --git a/ding/bonus/c51.py b/ding/bonus/c51.py new file mode 100644 index 0000000000..0fe2cc64e7 --- /dev/null +++ b/ding/bonus/c51.py @@ -0,0 +1,458 @@ +from typing import Optional, Union, List +from ditk import logging +from easydict import EasyDict +import os +import numpy as np +import torch +import treetensor.torch as ttorch +from ding.framework import task, OnlineRLContext +from ding.framework.middleware import CkptSaver, \ + wandb_online_logger, offline_data_saver, termination_checker, interaction_evaluator, StepCollector, data_pusher, \ + OffPolicyLearner, final_ctx_saver, eps_greedy_handler, nstep_reward_enhancer +from ding.envs import BaseEnv +from ding.envs import setup_ding_env_manager +from ding.policy import C51Policy +from ding.utils import set_pkg_seed +from ding.utils import get_env_fps, render +from ding.config import save_config_py, compile_config +from ding.model import C51DQN +from ding.model import model_wrap +from ding.data import DequeBuffer +from ding.bonus.common import TrainingReturn, EvalReturn +from ding.config.example.C51 import supported_env_cfg +from ding.config.example.C51 import supported_env + + +class C51Agent: + """ + Overview: + Class of agent for training, evaluation and deployment of Reinforcement learning algorithm C51. + For more information about the system design of RL agent, please refer to \ + . + Interface: + ``__init__``, ``train``, ``deploy``, ``collect_data``, ``batch_evaluate``, ``best`` + """ + supported_env_list = list(supported_env_cfg.keys()) + """ + Overview: + List of supported envs. + Examples: + >>> from ding.bonus.c51 import C51Agent + >>> print(C51Agent.supported_env_list) + """ + + def __init__( + self, + env_id: str = None, + env: BaseEnv = None, + seed: int = 0, + exp_name: str = None, + model: Optional[torch.nn.Module] = None, + cfg: Optional[Union[EasyDict, dict]] = None, + policy_state_dict: str = None, + ) -> None: + """ + Overview: + Initialize agent for C51 algorithm. + Arguments: + - env_id (:obj:`str`): The environment id, which is a registered environment name in gym or gymnasium. \ + If ``env_id`` is not specified, ``env_id`` in ``cfg.env`` must be specified. \ + If ``env_id`` is specified, ``env_id`` in ``cfg.env`` will be ignored. \ + ``env_id`` should be one of the supported envs, which can be found in ``supported_env_list``. + - env (:obj:`BaseEnv`): The environment instance for training and evaluation. \ + If ``env`` is not specified, `env_id`` or ``cfg.env.env_id`` must be specified. \ + ``env_id`` or ``cfg.env.env_id`` will be used to create environment instance. \ + If ``env`` is specified, ``env_id`` and ``cfg.env.env_id`` will be ignored. + - seed (:obj:`int`): The random seed, which is set before running the program. \ + Default to 0. + - exp_name (:obj:`str`): The name of this experiment, which will be used to create the folder to save \ + log data. Default to None. If not specified, the folder name will be ``env_id``-``algorithm``. + - model (:obj:`torch.nn.Module`): The model of C51 algorithm, which should be an instance of class \ + :class:`ding.model.C51DQN`. If not specified, a default model will be generated according to the config. + - cfg (:obj:`Union[EasyDict, dict]`): The configuration of C51 algorithm, which is a dict. \ + Default to None. If not specified, the default configuration will be used. \ + The default configuration can be found in ``ding/config/example/C51/gym_lunarlander_v2.py``. + - policy_state_dict (:obj:`str`): The path of policy state dict saved by PyTorch a in local file. \ + If specified, the policy will be loaded from this file. Default to None. + + .. note:: + An RL Agent Instance can be initialized in two basic ways. \ + For example, we have an environment with id ``LunarLander-v2`` registered in gym, \ + and we want to train an agent with C51 algorithm with default configuration. \ + Then we can initialize the agent in the following ways: + >>> agent = C51Agent(env_id='LunarLander-v2') + or, if we want can specify the env_id in the configuration: + >>> cfg = {'env': {'env_id': 'LunarLander-v2'}, 'policy': ...... } + >>> agent = C51Agent(cfg=cfg) + There are also other arguments to specify the agent when initializing. + For example, if we want to specify the environment instance: + >>> env = CustomizedEnv('LunarLander-v2') + >>> agent = C51Agent(cfg=cfg, env=env) + or, if we want to specify the model: + >>> model = C51DQN(**cfg.policy.model) + >>> agent = C51Agent(cfg=cfg, model=model) + or, if we want to reload the policy from a saved policy state dict: + >>> agent = C51Agent(cfg=cfg, policy_state_dict='LunarLander-v2.pth.tar') + Make sure that the configuration is consistent with the saved policy state dict. + """ + + assert env_id is not None or cfg is not None, "Please specify env_id or cfg." + + if cfg is not None and not isinstance(cfg, EasyDict): + cfg = EasyDict(cfg) + + if env_id is not None: + assert env_id in C51Agent.supported_env_list, "Please use supported envs: {}".format( + C51Agent.supported_env_list + ) + if cfg is None: + cfg = supported_env_cfg[env_id] + else: + assert cfg.env.env_id == env_id, "env_id in cfg should be the same as env_id in args." + else: + assert hasattr(cfg.env, "env_id"), "Please specify env_id in cfg." + assert cfg.env.env_id in C51Agent.supported_env_list, "Please use supported envs: {}".format( + C51Agent.supported_env_list + ) + default_policy_config = EasyDict({"policy": C51Policy.default_config()}) + default_policy_config.update(cfg) + cfg = default_policy_config + + if exp_name is not None: + cfg.exp_name = exp_name + self.cfg = compile_config(cfg, policy=C51Policy) + self.exp_name = self.cfg.exp_name + if env is None: + self.env = supported_env[cfg.env.env_id](cfg=cfg.env) + else: + assert isinstance(env, BaseEnv), "Please use BaseEnv as env data type." + self.env = env + + logging.getLogger().setLevel(logging.INFO) + self.seed = seed + set_pkg_seed(self.seed, use_cuda=self.cfg.policy.cuda) + if not os.path.exists(self.exp_name): + os.makedirs(self.exp_name) + save_config_py(self.cfg, os.path.join(self.exp_name, 'policy_config.py')) + if model is None: + model = C51DQN(**self.cfg.policy.model) + self.buffer_ = DequeBuffer(size=self.cfg.policy.other.replay_buffer.replay_buffer_size) + self.policy = C51Policy(self.cfg.policy, model=model) + if policy_state_dict is not None: + self.policy.learn_mode.load_state_dict(policy_state_dict) + self.checkpoint_save_dir = os.path.join(self.exp_name, "ckpt") + + def train( + self, + step: int = int(1e7), + collector_env_num: int = None, + evaluator_env_num: int = None, + n_iter_save_ckpt: int = 1000, + context: Optional[str] = None, + debug: bool = False, + wandb_sweep: bool = False, + ) -> TrainingReturn: + """ + Overview: + Train the agent with C51 algorithm for ``step`` iterations with ``collector_env_num`` collector \ + environments and ``evaluator_env_num`` evaluator environments. Information during training will be \ + recorded and saved by wandb. + Arguments: + - step (:obj:`int`): The total training environment steps of all collector environments. Default to 1e7. + - collector_env_num (:obj:`int`): The collector environment number. Default to None. \ + If not specified, it will be set according to the configuration. + - evaluator_env_num (:obj:`int`): The evaluator environment number. Default to None. \ + If not specified, it will be set according to the configuration. + - n_iter_save_ckpt (:obj:`int`): The frequency of saving checkpoint every training iteration. \ + Default to 1000. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + - wandb_sweep (:obj:`bool`): Whether to use wandb sweep, \ + which is a hyper-parameter optimization process for seeking the best configurations. \ + Default to False. If True, the wandb sweep id will be used as the experiment name. + Returns: + - (:obj:`TrainingReturn`): The training result, of which the attributions are: + - wandb_url (:obj:`str`): The weight & biases (wandb) project url of the trainning experiment. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + logging.debug(self.policy._model) + # define env and policy + collector_env_num = collector_env_num if collector_env_num else self.cfg.env.collector_env_num + evaluator_env_num = evaluator_env_num if evaluator_env_num else self.cfg.env.evaluator_env_num + collector_env = setup_ding_env_manager(self.env, collector_env_num, context, debug, 'collector') + evaluator_env = setup_ding_env_manager(self.env, evaluator_env_num, context, debug, 'evaluator') + + with task.start(ctx=OnlineRLContext()): + task.use( + interaction_evaluator( + self.cfg, + self.policy.eval_mode, + evaluator_env, + render=self.cfg.policy.eval.render if hasattr(self.cfg.policy.eval, "render") else False + ) + ) + task.use(CkptSaver(policy=self.policy, save_dir=self.checkpoint_save_dir, train_freq=n_iter_save_ckpt)) + task.use(eps_greedy_handler(self.cfg)) + task.use( + StepCollector( + self.cfg, + self.policy.collect_mode, + collector_env, + random_collect_size=self.cfg.policy.random_collect_size + if hasattr(self.cfg.policy, 'random_collect_size') else 0, + ) + ) + task.use(nstep_reward_enhancer(self.cfg)) + task.use(data_pusher(self.cfg, self.buffer_)) + task.use(OffPolicyLearner(self.cfg, self.policy.learn_mode, self.buffer_)) + task.use( + wandb_online_logger( + metric_list=self.policy._monitor_vars_learn(), + model=self.policy._model, + anonymous=True, + project_name=self.exp_name, + wandb_sweep=wandb_sweep, + ) + ) + task.use(termination_checker(max_env_step=step)) + task.use(final_ctx_saver(name=self.exp_name)) + task.run() + + return TrainingReturn(wandb_url=task.ctx.wandb_url) + + def deploy( + self, + enable_save_replay: bool = False, + concatenate_all_replay: bool = False, + replay_save_path: str = None, + seed: Optional[Union[int, List]] = None, + debug: bool = False + ) -> EvalReturn: + """ + Overview: + Deploy the agent with C51 algorithm by interacting with the environment, during which the replay video \ + can be saved if ``enable_save_replay`` is True. The evaluation result will be returned. + Arguments: + - enable_save_replay (:obj:`bool`): Whether to save the replay video. Default to False. + - concatenate_all_replay (:obj:`bool`): Whether to concatenate all replay videos into one video. \ + Default to False. If ``enable_save_replay`` is False, this argument will be ignored. \ + If ``enable_save_replay`` is True and ``concatenate_all_replay`` is False, \ + the replay video of each episode will be saved separately. + - replay_save_path (:obj:`str`): The path to save the replay video. Default to None. \ + If not specified, the video will be saved in ``exp_name/videos``. + - seed (:obj:`Union[int, List]`): The random seed, which is set before running the program. \ + Default to None. If not specified, ``self.seed`` will be used. \ + If ``seed`` is an integer, the agent will be deployed once. \ + If ``seed`` is a list of integers, the agent will be deployed once for each seed in the list. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + Returns: + - (:obj:`EvalReturn`): The evaluation result, of which the attributions are: + - eval_value (:obj:`np.float32`): The mean of evaluation return. + - eval_value_std (:obj:`np.float32`): The standard deviation of evaluation return. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + # define env and policy + env = self.env.clone(caller='evaluator') + + if seed is not None and isinstance(seed, int): + seeds = [seed] + elif seed is not None and isinstance(seed, list): + seeds = seed + else: + seeds = [self.seed] + + returns = [] + images = [] + if enable_save_replay: + replay_save_path = os.path.join(self.exp_name, 'videos') if replay_save_path is None else replay_save_path + env.enable_save_replay(replay_path=replay_save_path) + else: + logging.warning('No video would be generated during the deploy.') + if concatenate_all_replay: + logging.warning('concatenate_all_replay is set to False because enable_save_replay is False.') + concatenate_all_replay = False + + def single_env_forward_wrapper(forward_fn, cuda=True): + + forward_fn = model_wrap(forward_fn, wrapper_name='argmax_sample').forward + + def _forward(obs): + # unsqueeze means add batch dim, i.e. (O, ) -> (1, O) + obs = ttorch.as_tensor(obs).unsqueeze(0) + if cuda and torch.cuda.is_available(): + obs = obs.cuda() + action = forward_fn(obs)["action"] + # squeeze means delete batch dim, i.e. (1, A) -> (A, ) + action = action.squeeze(0).detach().cpu().numpy() + return action + + return _forward + + forward_fn = single_env_forward_wrapper(self.policy._model, self.cfg.policy.cuda) + + # reset first to make sure the env is in the initial state + # env will be reset again in the main loop + env.reset() + + for seed in seeds: + env.seed(seed, dynamic_seed=False) + return_ = 0. + step = 0 + obs = env.reset() + images.append(render(env)[None]) if concatenate_all_replay else None + while True: + action = forward_fn(obs) + obs, rew, done, info = env.step(action) + images.append(render(env)[None]) if concatenate_all_replay else None + return_ += rew + step += 1 + if done: + break + logging.info(f'C51 deploy is finished, final episode return with {step} steps is: {return_}') + returns.append(return_) + + env.close() + + if concatenate_all_replay: + images = np.concatenate(images, axis=0) + import imageio + imageio.mimwrite(os.path.join(replay_save_path, 'deploy.mp4'), images, fps=get_env_fps(env)) + + return EvalReturn(eval_value=np.mean(returns), eval_value_std=np.std(returns)) + + def collect_data( + self, + env_num: int = 8, + save_data_path: Optional[str] = None, + n_sample: Optional[int] = None, + n_episode: Optional[int] = None, + context: Optional[str] = None, + debug: bool = False + ) -> None: + """ + Overview: + Collect data with C51 algorithm for ``n_episode`` episodes with ``env_num`` collector environments. \ + The collected data will be saved in ``save_data_path`` if specified, otherwise it will be saved in \ + ``exp_name/demo_data``. + Arguments: + - env_num (:obj:`int`): The number of collector environments. Default to 8. + - save_data_path (:obj:`str`): The path to save the collected data. Default to None. \ + If not specified, the data will be saved in ``exp_name/demo_data``. + - n_sample (:obj:`int`): The number of samples to collect. Default to None. \ + If not specified, ``n_episode`` must be specified. + - n_episode (:obj:`int`): The number of episodes to collect. Default to None. \ + If not specified, ``n_sample`` must be specified. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + if n_episode is not None: + raise NotImplementedError + # define env and policy + env_num = env_num if env_num else self.cfg.env.collector_env_num + env = setup_ding_env_manager(self.env, env_num, context, debug, 'collector') + + if save_data_path is None: + save_data_path = os.path.join(self.exp_name, 'demo_data') + + # main execution task + with task.start(ctx=OnlineRLContext()): + task.use( + StepCollector( + self.cfg, self.policy.collect_mode, env, random_collect_size=self.cfg.policy.random_collect_size + ) + ) + task.use(offline_data_saver(save_data_path, data_type='hdf5')) + task.run(max_step=1) + logging.info( + f'C51 collecting is finished, more than {n_sample} samples are collected and saved in `{save_data_path}`' + ) + + def batch_evaluate( + self, + env_num: int = 4, + n_evaluator_episode: int = 4, + context: Optional[str] = None, + debug: bool = False + ) -> EvalReturn: + """ + Overview: + Evaluate the agent with C51 algorithm for ``n_evaluator_episode`` episodes with ``env_num`` evaluator \ + environments. The evaluation result will be returned. + The difference between methods ``batch_evaluate`` and ``deploy`` is that ``batch_evaluate`` will create \ + multiple evaluator environments to evaluate the agent to get an average performance, while ``deploy`` \ + will only create one evaluator environment to evaluate the agent and save the replay video. + Arguments: + - env_num (:obj:`int`): The number of evaluator environments. Default to 4. + - n_evaluator_episode (:obj:`int`): The number of episodes to evaluate. Default to 4. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + Returns: + - (:obj:`EvalReturn`): The evaluation result, of which the attributions are: + - eval_value (:obj:`np.float32`): The mean of evaluation return. + - eval_value_std (:obj:`np.float32`): The standard deviation of evaluation return. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + # define env and policy + env_num = env_num if env_num else self.cfg.env.evaluator_env_num + env = setup_ding_env_manager(self.env, env_num, context, debug, 'evaluator') + + # reset first to make sure the env is in the initial state + # env will be reset again in the main loop + env.launch() + env.reset() + + evaluate_cfg = self.cfg + evaluate_cfg.env.n_evaluator_episode = n_evaluator_episode + + # main execution task + with task.start(ctx=OnlineRLContext()): + task.use(interaction_evaluator(self.cfg, self.policy.eval_mode, env)) + task.run(max_step=1) + + return EvalReturn(eval_value=task.ctx.eval_value, eval_value_std=task.ctx.eval_value_std) + + @property + def best(self) -> 'C51Agent': + """ + Overview: + Load the best model from the checkpoint directory, \ + which is by default in folder ``exp_name/ckpt/eval.pth.tar``. \ + The return value is the agent with the best model. + Returns: + - (:obj:`C51Agent`): The agent with the best model. + Examples: + >>> agent = C51Agent(env_id='LunarLander-v2') + >>> agent.train() + >>> agent = agent.best + + .. note:: + The best model is the model with the highest evaluation return. If this method is called, the current \ + model will be replaced by the best model. + """ + + best_model_file_path = os.path.join(self.checkpoint_save_dir, "eval.pth.tar") + # Load best model if it exists + if os.path.exists(best_model_file_path): + policy_state_dict = torch.load(best_model_file_path, map_location=torch.device("cpu")) + self.policy.learn_mode.load_state_dict(policy_state_dict) + return self diff --git a/ding/bonus/common.py b/ding/bonus/common.py new file mode 100644 index 0000000000..1d4ddfc711 --- /dev/null +++ b/ding/bonus/common.py @@ -0,0 +1,22 @@ +from dataclasses import dataclass +import numpy as np + + +@dataclass +class TrainingReturn: + ''' + Attributions + wandb_url: The weight & biases (wandb) project url of the trainning experiment. + ''' + wandb_url: str + + +@dataclass +class EvalReturn: + ''' + Attributions + eval_value: The mean of evaluation return. + eval_value_std: The standard deviation of evaluation return. + ''' + eval_value: np.float32 + eval_value_std: np.float32 diff --git a/ding/bonus/config.py b/ding/bonus/config.py new file mode 100644 index 0000000000..285eff6586 --- /dev/null +++ b/ding/bonus/config.py @@ -0,0 +1,326 @@ +from easydict import EasyDict +import os +import gym +from ding.envs import BaseEnv, DingEnvWrapper +from ding.envs.env_wrappers import MaxAndSkipWrapper, WarpFrameWrapper, ScaledFloatFrameWrapper, FrameStackWrapper, \ + EvalEpisodeReturnWrapper, TransposeWrapper, TimeLimitWrapper, FlatObsWrapper, GymToGymnasiumWrapper +from ding.policy import PPOFPolicy + + +def get_instance_config(env_id: str, algorithm: str) -> EasyDict: + if algorithm == 'PPOF': + cfg = PPOFPolicy.default_config() + if env_id == 'LunarLander-v2': + cfg.n_sample = 512 + cfg.value_norm = 'popart' + cfg.entropy_weight = 1e-3 + elif env_id == 'LunarLanderContinuous-v2': + cfg.action_space = 'continuous' + cfg.n_sample = 400 + elif env_id == 'BipedalWalker-v3': + cfg.learning_rate = 1e-3 + cfg.action_space = 'continuous' + cfg.n_sample = 1024 + elif env_id == 'Pendulum-v1': + cfg.action_space = 'continuous' + cfg.n_sample = 400 + elif env_id == 'acrobot': + cfg.learning_rate = 1e-4 + cfg.n_sample = 400 + elif env_id == 'rocket_landing': + cfg.n_sample = 2048 + cfg.adv_norm = False + cfg.model = dict( + encoder_hidden_size_list=[64, 64, 128], + actor_head_hidden_size=128, + critic_head_hidden_size=128, + ) + elif env_id == 'drone_fly': + cfg.action_space = 'continuous' + cfg.adv_norm = False + cfg.epoch_per_collect = 5 + cfg.learning_rate = 5e-5 + cfg.n_sample = 640 + elif env_id == 'hybrid_moving': + cfg.action_space = 'hybrid' + cfg.n_sample = 3200 + cfg.entropy_weight = 0.03 + cfg.batch_size = 320 + cfg.adv_norm = False + cfg.model = dict( + encoder_hidden_size_list=[256, 128, 64, 64], + sigma_type='fixed', + fixed_sigma_value=0.3, + bound_type='tanh', + ) + elif env_id == 'evogym_carrier': + cfg.action_space = 'continuous' + cfg.n_sample = 2048 + cfg.batch_size = 256 + cfg.epoch_per_collect = 10 + cfg.learning_rate = 3e-3 + elif env_id == 'mario': + cfg.n_sample = 256 + cfg.batch_size = 64 + cfg.epoch_per_collect = 2 + cfg.learning_rate = 1e-3 + cfg.model = dict( + encoder_hidden_size_list=[64, 64, 128], + critic_head_hidden_size=128, + actor_head_hidden_size=128, + ) + elif env_id == 'di_sheep': + cfg.n_sample = 3200 + cfg.batch_size = 320 + cfg.epoch_per_collect = 10 + cfg.learning_rate = 3e-4 + cfg.adv_norm = False + cfg.entropy_weight = 0.001 + elif env_id == 'procgen_bigfish': + cfg.n_sample = 16384 + cfg.batch_size = 16384 + cfg.epoch_per_collect = 10 + cfg.learning_rate = 5e-4 + cfg.model = dict( + encoder_hidden_size_list=[64, 128, 256], + critic_head_hidden_size=256, + actor_head_hidden_size=256, + ) + elif env_id in ['KangarooNoFrameskip-v4', 'BowlingNoFrameskip-v4']: + cfg.n_sample = 1024 + cfg.batch_size = 128 + cfg.epoch_per_collect = 10 + cfg.learning_rate = 0.0001 + cfg.model = dict( + encoder_hidden_size_list=[32, 64, 64, 128], + actor_head_hidden_size=128, + critic_head_hidden_size=128, + critic_head_layer_num=2, + ) + elif env_id == 'PongNoFrameskip-v4': + cfg.n_sample = 3200 + cfg.batch_size = 320 + cfg.epoch_per_collect = 10 + cfg.learning_rate = 3e-4 + cfg.model = dict( + encoder_hidden_size_list=[64, 64, 128], + actor_head_hidden_size=128, + critic_head_hidden_size=128, + ) + elif env_id == 'SpaceInvadersNoFrameskip-v4': + cfg.n_sample = 320 + cfg.batch_size = 320 + cfg.epoch_per_collect = 1 + cfg.learning_rate = 1e-3 + cfg.entropy_weight = 0.01 + cfg.lr_scheduler = (2000, 0.1) + cfg.model = dict( + encoder_hidden_size_list=[64, 64, 128], + actor_head_hidden_size=128, + critic_head_hidden_size=128, + ) + elif env_id == 'QbertNoFrameskip-v4': + cfg.n_sample = 3200 + cfg.batch_size = 320 + cfg.epoch_per_collect = 10 + cfg.learning_rate = 5e-4 + cfg.lr_scheduler = (1000, 0.1) + cfg.model = dict( + encoder_hidden_size_list=[64, 64, 128], + actor_head_hidden_size=128, + critic_head_hidden_size=128, + ) + elif env_id == 'minigrid_fourroom': + cfg.n_sample = 3200 + cfg.batch_size = 320 + cfg.learning_rate = 3e-4 + cfg.epoch_per_collect = 10 + cfg.entropy_weight = 0.001 + elif env_id == 'metadrive': + cfg.learning_rate = 3e-4 + cfg.action_space = 'continuous' + cfg.entropy_weight = 0.001 + cfg.n_sample = 3000 + cfg.epoch_per_collect = 10 + cfg.learning_rate = 0.0001 + cfg.model = dict( + encoder_hidden_size_list=[32, 64, 64, 128], + actor_head_hidden_size=128, + critic_head_hidden_size=128, + critic_head_layer_num=2, + ) + elif env_id == 'Hopper-v3': + cfg.action_space = "continuous" + cfg.n_sample = 3200 + cfg.batch_size = 320 + cfg.epoch_per_collect = 10 + cfg.learning_rate = 3e-4 + elif env_id == 'HalfCheetah-v3': + cfg.action_space = "continuous" + cfg.n_sample = 3200 + cfg.batch_size = 320 + cfg.epoch_per_collect = 10 + cfg.learning_rate = 3e-4 + elif env_id == 'Walker2d-v3': + cfg.action_space = "continuous" + cfg.n_sample = 3200 + cfg.batch_size = 320 + cfg.epoch_per_collect = 10 + cfg.learning_rate = 3e-4 + else: + raise KeyError("not supported env type: {}".format(env_id)) + else: + raise KeyError("not supported algorithm type: {}".format(algorithm)) + + return cfg + + +def get_instance_env(env_id: str) -> BaseEnv: + if env_id == 'LunarLander-v2': + return DingEnvWrapper(gym.make('LunarLander-v2')) + elif env_id == 'LunarLanderContinuous-v2': + return DingEnvWrapper(gym.make('LunarLanderContinuous-v2', continuous=True)) + elif env_id == 'BipedalWalker-v3': + return DingEnvWrapper(gym.make('BipedalWalker-v3'), cfg={'act_scale': True, 'rew_clip': True}) + elif env_id == 'Pendulum-v1': + return DingEnvWrapper(gym.make('Pendulum-v1'), cfg={'act_scale': True}) + elif env_id == 'acrobot': + return DingEnvWrapper(gym.make('Acrobot-v1')) + elif env_id == 'rocket_landing': + from dizoo.rocket.envs import RocketEnv + cfg = EasyDict({ + 'task': 'landing', + 'max_steps': 800, + }) + return RocketEnv(cfg) + elif env_id == 'drone_fly': + from dizoo.gym_pybullet_drones.envs import GymPybulletDronesEnv + cfg = EasyDict({ + 'env_id': 'flythrugate-aviary-v0', + 'action_type': 'VEL', + }) + return GymPybulletDronesEnv(cfg) + elif env_id == 'hybrid_moving': + import gym_hybrid + return DingEnvWrapper(gym.make('Moving-v0')) + elif env_id == 'evogym_carrier': + import evogym.envs + from evogym import sample_robot, WorldObject + path = os.path.join(os.path.dirname(__file__), '../../dizoo/evogym/envs/world_data/carry_bot.json') + robot_object = WorldObject.from_json(path) + body = robot_object.get_structure() + return DingEnvWrapper( + gym.make('Carrier-v0', body=body), + cfg={ + 'env_wrapper': [ + lambda env: TimeLimitWrapper(env, max_limit=300), + lambda env: EvalEpisodeReturnWrapper(env), + ] + } + ) + elif env_id == 'mario': + import gym_super_mario_bros + from nes_py.wrappers import JoypadSpace + return DingEnvWrapper( + JoypadSpace(gym_super_mario_bros.make("SuperMarioBros-1-1-v1"), [["right"], ["right", "A"]]), + cfg={ + 'env_wrapper': [ + lambda env: MaxAndSkipWrapper(env, skip=4), + lambda env: WarpFrameWrapper(env, size=84), + lambda env: ScaledFloatFrameWrapper(env), + lambda env: FrameStackWrapper(env, n_frames=4), + lambda env: TimeLimitWrapper(env, max_limit=200), + lambda env: EvalEpisodeReturnWrapper(env), + ] + } + ) + elif env_id == 'di_sheep': + from sheep_env import SheepEnv + return DingEnvWrapper(SheepEnv(level=9)) + elif env_id == 'procgen_bigfish': + return DingEnvWrapper( + gym.make('procgen:procgen-bigfish-v0', start_level=0, num_levels=1), + cfg={ + 'env_wrapper': [ + lambda env: TransposeWrapper(env), + lambda env: ScaledFloatFrameWrapper(env), + lambda env: EvalEpisodeReturnWrapper(env), + ] + }, + seed_api=False, + ) + elif env_id == 'Hopper-v3': + cfg = EasyDict( + env_id='Hopper-v3', + env_wrapper='mujoco_default', + act_scale=True, + rew_clip=True, + ) + return DingEnvWrapper(gym.make('Hopper-v3'), cfg=cfg) + elif env_id == 'HalfCheetah-v3': + cfg = EasyDict( + env_id='HalfCheetah-v3', + env_wrapper='mujoco_default', + act_scale=True, + rew_clip=True, + ) + return DingEnvWrapper(gym.make('HalfCheetah-v3'), cfg=cfg) + elif env_id == 'Walker2d-v3': + cfg = EasyDict( + env_id='Walker2d-v3', + env_wrapper='mujoco_default', + act_scale=True, + rew_clip=True, + ) + return DingEnvWrapper(gym.make('Walker2d-v3'), cfg=cfg) + + elif env_id in [ + 'BowlingNoFrameskip-v4', + 'BreakoutNoFrameskip-v4', + 'GopherNoFrameskip-v4' + 'KangarooNoFrameskip-v4', + 'PongNoFrameskip-v4', + 'QbertNoFrameskip-v4', + 'SpaceInvadersNoFrameskip-v4', + ]: + + cfg = EasyDict({ + 'env_id': env_id, + 'env_wrapper': 'atari_default', + }) + ding_env_atari = DingEnvWrapper(gym.make(env_id), cfg=cfg) + return ding_env_atari + elif env_id == 'minigrid_fourroom': + import gymnasium + return DingEnvWrapper( + gymnasium.make('MiniGrid-FourRooms-v0'), + cfg={ + 'env_wrapper': [ + lambda env: GymToGymnasiumWrapper(env), + lambda env: FlatObsWrapper(env), + lambda env: TimeLimitWrapper(env, max_limit=300), + lambda env: EvalEpisodeReturnWrapper(env), + ] + } + ) + elif env_id == 'metadrive': + from dizoo.metadrive.env.drive_env import MetaDrivePPOOriginEnv + from dizoo.metadrive.env.drive_wrapper import DriveEnvWrapper + cfg = dict( + map='XSOS', + horizon=4000, + out_of_road_penalty=40.0, + crash_vehicle_penalty=40.0, + out_of_route_done=True, + ) + cfg = EasyDict(cfg) + return DriveEnvWrapper(MetaDrivePPOOriginEnv(cfg)) + else: + raise KeyError("not supported env type: {}".format(env_id)) + + +def get_hybrid_shape(action_space) -> EasyDict: + return EasyDict({ + 'action_type_shape': action_space[0].n, + 'action_args_shape': action_space[1].shape, + }) diff --git a/ding/bonus/ddpg.py b/ding/bonus/ddpg.py new file mode 100644 index 0000000000..88110c7598 --- /dev/null +++ b/ding/bonus/ddpg.py @@ -0,0 +1,456 @@ +from typing import Optional, Union, List +from ditk import logging +from easydict import EasyDict +import os +import numpy as np +import torch +import treetensor.torch as ttorch +from ding.framework import task, OnlineRLContext +from ding.framework.middleware import CkptSaver, \ + wandb_online_logger, offline_data_saver, termination_checker, interaction_evaluator, StepCollector, data_pusher, \ + OffPolicyLearner, final_ctx_saver +from ding.envs import BaseEnv +from ding.envs import setup_ding_env_manager +from ding.policy import DDPGPolicy +from ding.utils import set_pkg_seed +from ding.utils import get_env_fps, render +from ding.config import save_config_py, compile_config +from ding.model import ContinuousQAC +from ding.data import DequeBuffer +from ding.bonus.common import TrainingReturn, EvalReturn +from ding.config.example.DDPG import supported_env_cfg +from ding.config.example.DDPG import supported_env + + +class DDPGAgent: + """ + Overview: + Class of agent for training, evaluation and deployment of Reinforcement learning algorithm \ + Deep Deterministic Policy Gradient(DDPG). + For more information about the system design of RL agent, please refer to \ + . + Interface: + ``__init__``, ``train``, ``deploy``, ``collect_data``, ``batch_evaluate``, ``best`` + """ + supported_env_list = list(supported_env_cfg.keys()) + """ + Overview: + List of supported envs. + Examples: + >>> from ding.bonus.ddpg import DDPGAgent + >>> print(DDPGAgent.supported_env_list) + """ + + def __init__( + self, + env_id: str = None, + env: BaseEnv = None, + seed: int = 0, + exp_name: str = None, + model: Optional[torch.nn.Module] = None, + cfg: Optional[Union[EasyDict, dict]] = None, + policy_state_dict: str = None, + ) -> None: + """ + Overview: + Initialize agent for DDPG algorithm. + Arguments: + - env_id (:obj:`str`): The environment id, which is a registered environment name in gym or gymnasium. \ + If ``env_id`` is not specified, ``env_id`` in ``cfg.env`` must be specified. \ + If ``env_id`` is specified, ``env_id`` in ``cfg.env`` will be ignored. \ + ``env_id`` should be one of the supported envs, which can be found in ``supported_env_list``. + - env (:obj:`BaseEnv`): The environment instance for training and evaluation. \ + If ``env`` is not specified, `env_id`` or ``cfg.env.env_id`` must be specified. \ + ``env_id`` or ``cfg.env.env_id`` will be used to create environment instance. \ + If ``env`` is specified, ``env_id`` and ``cfg.env.env_id`` will be ignored. + - seed (:obj:`int`): The random seed, which is set before running the program. \ + Default to 0. + - exp_name (:obj:`str`): The name of this experiment, which will be used to create the folder to save \ + log data. Default to None. If not specified, the folder name will be ``env_id``-``algorithm``. + - model (:obj:`torch.nn.Module`): The model of DDPG algorithm, which should be an instance of class \ + :class:`ding.model.ContinuousQAC`. \ + If not specified, a default model will be generated according to the configuration. + - cfg (:obj:`Union[EasyDict, dict]`): The configuration of DDPG algorithm, which is a dict. \ + Default to None. If not specified, the default configuration will be used. \ + The default configuration can be found in ``ding/config/example/DDPG/gym_lunarlander_v2.py``. + - policy_state_dict (:obj:`str`): The path of policy state dict saved by PyTorch a in local file. \ + If specified, the policy will be loaded from this file. Default to None. + + .. note:: + An RL Agent Instance can be initialized in two basic ways. \ + For example, we have an environment with id ``LunarLanderContinuous-v2`` registered in gym, \ + and we want to train an agent with DDPG algorithm with default configuration. \ + Then we can initialize the agent in the following ways: + >>> agent = DDPGAgent(env_id='LunarLanderContinuous-v2') + or, if we want can specify the env_id in the configuration: + >>> cfg = {'env': {'env_id': 'LunarLanderContinuous-v2'}, 'policy': ...... } + >>> agent = DDPGAgent(cfg=cfg) + There are also other arguments to specify the agent when initializing. + For example, if we want to specify the environment instance: + >>> env = CustomizedEnv('LunarLanderContinuous-v2') + >>> agent = DDPGAgent(cfg=cfg, env=env) + or, if we want to specify the model: + >>> model = ContinuousQAC(**cfg.policy.model) + >>> agent = DDPGAgent(cfg=cfg, model=model) + or, if we want to reload the policy from a saved policy state dict: + >>> agent = DDPGAgent(cfg=cfg, policy_state_dict='LunarLanderContinuous-v2.pth.tar') + Make sure that the configuration is consistent with the saved policy state dict. + """ + + assert env_id is not None or cfg is not None, "Please specify env_id or cfg." + + if cfg is not None and not isinstance(cfg, EasyDict): + cfg = EasyDict(cfg) + + if env_id is not None: + assert env_id in DDPGAgent.supported_env_list, "Please use supported envs: {}".format( + DDPGAgent.supported_env_list + ) + if cfg is None: + cfg = supported_env_cfg[env_id] + else: + assert cfg.env.env_id == env_id, "env_id in cfg should be the same as env_id in args." + else: + assert hasattr(cfg.env, "env_id"), "Please specify env_id in cfg." + assert cfg.env.env_id in DDPGAgent.supported_env_list, "Please use supported envs: {}".format( + DDPGAgent.supported_env_list + ) + default_policy_config = EasyDict({"policy": DDPGPolicy.default_config()}) + default_policy_config.update(cfg) + cfg = default_policy_config + + if exp_name is not None: + cfg.exp_name = exp_name + self.cfg = compile_config(cfg, policy=DDPGPolicy) + self.exp_name = self.cfg.exp_name + if env is None: + self.env = supported_env[cfg.env.env_id](cfg=cfg.env) + else: + assert isinstance(env, BaseEnv), "Please use BaseEnv as env data type." + self.env = env + + logging.getLogger().setLevel(logging.INFO) + self.seed = seed + set_pkg_seed(self.seed, use_cuda=self.cfg.policy.cuda) + if not os.path.exists(self.exp_name): + os.makedirs(self.exp_name) + save_config_py(self.cfg, os.path.join(self.exp_name, 'policy_config.py')) + if model is None: + model = ContinuousQAC(**self.cfg.policy.model) + self.buffer_ = DequeBuffer(size=self.cfg.policy.other.replay_buffer.replay_buffer_size) + self.policy = DDPGPolicy(self.cfg.policy, model=model) + if policy_state_dict is not None: + self.policy.learn_mode.load_state_dict(policy_state_dict) + self.checkpoint_save_dir = os.path.join(self.exp_name, "ckpt") + + def train( + self, + step: int = int(1e7), + collector_env_num: int = None, + evaluator_env_num: int = None, + n_iter_log_show: int = 500, + n_iter_save_ckpt: int = 1000, + context: Optional[str] = None, + debug: bool = False, + wandb_sweep: bool = False, + ) -> TrainingReturn: + """ + Overview: + Train the agent with DDPG algorithm for ``step`` iterations with ``collector_env_num`` collector \ + environments and ``evaluator_env_num`` evaluator environments. Information during training will be \ + recorded and saved by wandb. + Arguments: + - step (:obj:`int`): The total training environment steps of all collector environments. Default to 1e7. + - collector_env_num (:obj:`int`): The collector environment number. Default to None. \ + If not specified, it will be set according to the configuration. + - evaluator_env_num (:obj:`int`): The evaluator environment number. Default to None. \ + If not specified, it will be set according to the configuration. + - n_iter_save_ckpt (:obj:`int`): The frequency of saving checkpoint every training iteration. \ + Default to 1000. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + - wandb_sweep (:obj:`bool`): Whether to use wandb sweep, \ + which is a hyper-parameter optimization process for seeking the best configurations. \ + Default to False. If True, the wandb sweep id will be used as the experiment name. + Returns: + - (:obj:`TrainingReturn`): The training result, of which the attributions are: + - wandb_url (:obj:`str`): The weight & biases (wandb) project url of the trainning experiment. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + logging.debug(self.policy._model) + # define env and policy + collector_env_num = collector_env_num if collector_env_num else self.cfg.env.collector_env_num + evaluator_env_num = evaluator_env_num if evaluator_env_num else self.cfg.env.evaluator_env_num + collector_env = setup_ding_env_manager(self.env, collector_env_num, context, debug, 'collector') + evaluator_env = setup_ding_env_manager(self.env, evaluator_env_num, context, debug, 'evaluator') + + with task.start(ctx=OnlineRLContext()): + task.use( + interaction_evaluator( + self.cfg, + self.policy.eval_mode, + evaluator_env, + render=self.cfg.policy.eval.render if hasattr(self.cfg.policy.eval, "render") else False + ) + ) + task.use(CkptSaver(policy=self.policy, save_dir=self.checkpoint_save_dir, train_freq=n_iter_save_ckpt)) + task.use( + StepCollector( + self.cfg, + self.policy.collect_mode, + collector_env, + random_collect_size=self.cfg.policy.random_collect_size + if hasattr(self.cfg.policy, 'random_collect_size') else 0, + ) + ) + task.use(data_pusher(self.cfg, self.buffer_)) + task.use(OffPolicyLearner(self.cfg, self.policy.learn_mode, self.buffer_)) + task.use( + wandb_online_logger( + metric_list=self.policy._monitor_vars_learn(), + model=self.policy._model, + anonymous=True, + project_name=self.exp_name, + wandb_sweep=wandb_sweep, + ) + ) + task.use(termination_checker(max_env_step=step)) + task.use(final_ctx_saver(name=self.exp_name)) + task.run() + + return TrainingReturn(wandb_url=task.ctx.wandb_url) + + def deploy( + self, + enable_save_replay: bool = False, + concatenate_all_replay: bool = False, + replay_save_path: str = None, + seed: Optional[Union[int, List]] = None, + debug: bool = False + ) -> EvalReturn: + """ + Overview: + Deploy the agent with DDPG algorithm by interacting with the environment, during which the replay video \ + can be saved if ``enable_save_replay`` is True. The evaluation result will be returned. + Arguments: + - enable_save_replay (:obj:`bool`): Whether to save the replay video. Default to False. + - concatenate_all_replay (:obj:`bool`): Whether to concatenate all replay videos into one video. \ + Default to False. If ``enable_save_replay`` is False, this argument will be ignored. \ + If ``enable_save_replay`` is True and ``concatenate_all_replay`` is False, \ + the replay video of each episode will be saved separately. + - replay_save_path (:obj:`str`): The path to save the replay video. Default to None. \ + If not specified, the video will be saved in ``exp_name/videos``. + - seed (:obj:`Union[int, List]`): The random seed, which is set before running the program. \ + Default to None. If not specified, ``self.seed`` will be used. \ + If ``seed`` is an integer, the agent will be deployed once. \ + If ``seed`` is a list of integers, the agent will be deployed once for each seed in the list. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + Returns: + - (:obj:`EvalReturn`): The evaluation result, of which the attributions are: + - eval_value (:obj:`np.float32`): The mean of evaluation return. + - eval_value_std (:obj:`np.float32`): The standard deviation of evaluation return. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + # define env and policy + env = self.env.clone(caller='evaluator') + + if seed is not None and isinstance(seed, int): + seeds = [seed] + elif seed is not None and isinstance(seed, list): + seeds = seed + else: + seeds = [self.seed] + + returns = [] + images = [] + if enable_save_replay: + replay_save_path = os.path.join(self.exp_name, 'videos') if replay_save_path is None else replay_save_path + env.enable_save_replay(replay_path=replay_save_path) + else: + logging.warning('No video would be generated during the deploy.') + if concatenate_all_replay: + logging.warning('concatenate_all_replay is set to False because enable_save_replay is False.') + concatenate_all_replay = False + + def single_env_forward_wrapper(forward_fn, cuda=True): + + def _forward(obs): + # unsqueeze means add batch dim, i.e. (O, ) -> (1, O) + obs = ttorch.as_tensor(obs).unsqueeze(0) + if cuda and torch.cuda.is_available(): + obs = obs.cuda() + action = forward_fn(obs, mode='compute_actor')["action"] + # squeeze means delete batch dim, i.e. (1, A) -> (A, ) + action = action.squeeze(0).detach().cpu().numpy() + return action + + return _forward + + forward_fn = single_env_forward_wrapper(self.policy._model, self.cfg.policy.cuda) + + # reset first to make sure the env is in the initial state + # env will be reset again in the main loop + env.reset() + + for seed in seeds: + env.seed(seed, dynamic_seed=False) + return_ = 0. + step = 0 + obs = env.reset() + images.append(render(env)[None]) if concatenate_all_replay else None + while True: + action = forward_fn(obs) + obs, rew, done, info = env.step(action) + images.append(render(env)[None]) if concatenate_all_replay else None + return_ += rew + step += 1 + if done: + break + logging.info(f'DDPG deploy is finished, final episode return with {step} steps is: {return_}') + returns.append(return_) + + env.close() + + if concatenate_all_replay: + images = np.concatenate(images, axis=0) + import imageio + imageio.mimwrite(os.path.join(replay_save_path, 'deploy.mp4'), images, fps=get_env_fps(env)) + + return EvalReturn(eval_value=np.mean(returns), eval_value_std=np.std(returns)) + + def collect_data( + self, + env_num: int = 8, + save_data_path: Optional[str] = None, + n_sample: Optional[int] = None, + n_episode: Optional[int] = None, + context: Optional[str] = None, + debug: bool = False + ) -> None: + """ + Overview: + Collect data with DDPG algorithm for ``n_episode`` episodes with ``env_num`` collector environments. \ + The collected data will be saved in ``save_data_path`` if specified, otherwise it will be saved in \ + ``exp_name/demo_data``. + Arguments: + - env_num (:obj:`int`): The number of collector environments. Default to 8. + - save_data_path (:obj:`str`): The path to save the collected data. Default to None. \ + If not specified, the data will be saved in ``exp_name/demo_data``. + - n_sample (:obj:`int`): The number of samples to collect. Default to None. \ + If not specified, ``n_episode`` must be specified. + - n_episode (:obj:`int`): The number of episodes to collect. Default to None. \ + If not specified, ``n_sample`` must be specified. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + if n_episode is not None: + raise NotImplementedError + # define env and policy + env_num = env_num if env_num else self.cfg.env.collector_env_num + env = setup_ding_env_manager(self.env, env_num, context, debug, 'collector') + + if save_data_path is None: + save_data_path = os.path.join(self.exp_name, 'demo_data') + + # main execution task + with task.start(ctx=OnlineRLContext()): + task.use( + StepCollector( + self.cfg, self.policy.collect_mode, env, random_collect_size=self.cfg.policy.random_collect_size + ) + ) + task.use(offline_data_saver(save_data_path, data_type='hdf5')) + task.run(max_step=1) + logging.info( + f'DDPG collecting is finished, more than {n_sample} samples are collected and saved in `{save_data_path}`' + ) + + def batch_evaluate( + self, + env_num: int = 4, + n_evaluator_episode: int = 4, + context: Optional[str] = None, + debug: bool = False + ) -> EvalReturn: + """ + Overview: + Evaluate the agent with DDPG algorithm for ``n_evaluator_episode`` episodes with ``env_num`` evaluator \ + environments. The evaluation result will be returned. + The difference between methods ``batch_evaluate`` and ``deploy`` is that ``batch_evaluate`` will create \ + multiple evaluator environments to evaluate the agent to get an average performance, while ``deploy`` \ + will only create one evaluator environment to evaluate the agent and save the replay video. + Arguments: + - env_num (:obj:`int`): The number of evaluator environments. Default to 4. + - n_evaluator_episode (:obj:`int`): The number of episodes to evaluate. Default to 4. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + Returns: + - (:obj:`EvalReturn`): The evaluation result, of which the attributions are: + - eval_value (:obj:`np.float32`): The mean of evaluation return. + - eval_value_std (:obj:`np.float32`): The standard deviation of evaluation return. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + # define env and policy + env_num = env_num if env_num else self.cfg.env.evaluator_env_num + env = setup_ding_env_manager(self.env, env_num, context, debug, 'evaluator') + + # reset first to make sure the env is in the initial state + # env will be reset again in the main loop + env.launch() + env.reset() + + evaluate_cfg = self.cfg + evaluate_cfg.env.n_evaluator_episode = n_evaluator_episode + + # main execution task + with task.start(ctx=OnlineRLContext()): + task.use(interaction_evaluator(self.cfg, self.policy.eval_mode, env)) + task.run(max_step=1) + + return EvalReturn(eval_value=task.ctx.eval_value, eval_value_std=task.ctx.eval_value_std) + + @property + def best(self) -> 'DDPGAgent': + """ + Overview: + Load the best model from the checkpoint directory, \ + which is by default in folder ``exp_name/ckpt/eval.pth.tar``. \ + The return value is the agent with the best model. + Returns: + - (:obj:`DDPGAgent`): The agent with the best model. + Examples: + >>> agent = DDPGAgent(env_id='LunarLanderContinuous-v2') + >>> agent.train() + >>> agent = agent.best + + .. note:: + The best model is the model with the highest evaluation return. If this method is called, the current \ + model will be replaced by the best model. + """ + + best_model_file_path = os.path.join(self.checkpoint_save_dir, "eval.pth.tar") + # Load best model if it exists + if os.path.exists(best_model_file_path): + policy_state_dict = torch.load(best_model_file_path, map_location=torch.device("cpu")) + self.policy.learn_mode.load_state_dict(policy_state_dict) + return self diff --git a/ding/bonus/dqn.py b/ding/bonus/dqn.py new file mode 100644 index 0000000000..3207c30490 --- /dev/null +++ b/ding/bonus/dqn.py @@ -0,0 +1,460 @@ +from typing import Optional, Union, List +from ditk import logging +from easydict import EasyDict +import os +import numpy as np +import torch +import treetensor.torch as ttorch +from ding.framework import task, OnlineRLContext +from ding.framework.middleware import CkptSaver, \ + wandb_online_logger, offline_data_saver, termination_checker, interaction_evaluator, StepCollector, data_pusher, \ + OffPolicyLearner, final_ctx_saver, nstep_reward_enhancer, eps_greedy_handler +from ding.envs import BaseEnv +from ding.envs import setup_ding_env_manager +from ding.policy import DQNPolicy +from ding.utils import set_pkg_seed +from ding.utils import get_env_fps, render +from ding.config import save_config_py, compile_config +from ding.model import DQN +from ding.model import model_wrap +from ding.data import DequeBuffer +from ding.bonus.common import TrainingReturn, EvalReturn +from ding.config.example.DQN import supported_env_cfg +from ding.config.example.DQN import supported_env + + +class DQNAgent: + """ + Overview: + Class of agent for training, evaluation and deployment of Reinforcement learning algorithm Deep Q-Learning(DQN). + For more information about the system design of RL agent, please refer to \ + . + Interface: + ``__init__``, ``train``, ``deploy``, ``collect_data``, ``batch_evaluate``, ``best`` + """ + supported_env_list = list(supported_env_cfg.keys()) + """ + Overview: + List of supported envs. + Examples: + >>> from ding.bonus.dqn import DQNAgent + >>> print(DQNAgent.supported_env_list) + """ + + def __init__( + self, + env_id: str = None, + env: BaseEnv = None, + seed: int = 0, + exp_name: str = None, + model: Optional[torch.nn.Module] = None, + cfg: Optional[Union[EasyDict, dict]] = None, + policy_state_dict: str = None, + ) -> None: + """ + Overview: + Initialize agent for DQN algorithm. + Arguments: + - env_id (:obj:`str`): The environment id, which is a registered environment name in gym or gymnasium. \ + If ``env_id`` is not specified, ``env_id`` in ``cfg.env`` must be specified. \ + If ``env_id`` is specified, ``env_id`` in ``cfg.env`` will be ignored. \ + ``env_id`` should be one of the supported envs, which can be found in ``supported_env_list``. + - env (:obj:`BaseEnv`): The environment instance for training and evaluation. \ + If ``env`` is not specified, `env_id`` or ``cfg.env.env_id`` must be specified. \ + ``env_id`` or ``cfg.env.env_id`` will be used to create environment instance. \ + If ``env`` is specified, ``env_id`` and ``cfg.env.env_id`` will be ignored. + - seed (:obj:`int`): The random seed, which is set before running the program. \ + Default to 0. + - exp_name (:obj:`str`): The name of this experiment, which will be used to create the folder to save \ + log data. Default to None. If not specified, the folder name will be ``env_id``-``algorithm``. + - model (:obj:`torch.nn.Module`): The model of DQN algorithm, which should be an instance of class \ + :class:`ding.model.DQN`. \ + If not specified, a default model will be generated according to the configuration. + - cfg (:obj:`Union[EasyDict, dict]`): The configuration of DQN algorithm, which is a dict. \ + Default to None. If not specified, the default configuration will be used. \ + The default configuration can be found in ``ding/config/example/DQN/gym_lunarlander_v2.py``. + - policy_state_dict (:obj:`str`): The path of policy state dict saved by PyTorch a in local file. \ + If specified, the policy will be loaded from this file. Default to None. + + .. note:: + An RL Agent Instance can be initialized in two basic ways. \ + For example, we have an environment with id ``LunarLander-v2`` registered in gym, \ + and we want to train an agent with DQN algorithm with default configuration. \ + Then we can initialize the agent in the following ways: + >>> agent = DQNAgent(env_id='LunarLander-v2') + or, if we want can specify the env_id in the configuration: + >>> cfg = {'env': {'env_id': 'LunarLander-v2'}, 'policy': ...... } + >>> agent = DQNAgent(cfg=cfg) + There are also other arguments to specify the agent when initializing. + For example, if we want to specify the environment instance: + >>> env = CustomizedEnv('LunarLander-v2') + >>> agent = DQNAgent(cfg=cfg, env=env) + or, if we want to specify the model: + >>> model = DQN(**cfg.policy.model) + >>> agent = DQNAgent(cfg=cfg, model=model) + or, if we want to reload the policy from a saved policy state dict: + >>> agent = DQNAgent(cfg=cfg, policy_state_dict='LunarLander-v2.pth.tar') + Make sure that the configuration is consistent with the saved policy state dict. + """ + + assert env_id is not None or cfg is not None, "Please specify env_id or cfg." + + if cfg is not None and not isinstance(cfg, EasyDict): + cfg = EasyDict(cfg) + + if env_id is not None: + assert env_id in DQNAgent.supported_env_list, "Please use supported envs: {}".format( + DQNAgent.supported_env_list + ) + if cfg is None: + cfg = supported_env_cfg[env_id] + else: + assert cfg.env.env_id == env_id, "env_id in cfg should be the same as env_id in args." + else: + assert hasattr(cfg.env, "env_id"), "Please specify env_id in cfg." + assert cfg.env.env_id in DQNAgent.supported_env_list, "Please use supported envs: {}".format( + DQNAgent.supported_env_list + ) + default_policy_config = EasyDict({"policy": DQNPolicy.default_config()}) + default_policy_config.update(cfg) + cfg = default_policy_config + + if exp_name is not None: + cfg.exp_name = exp_name + self.cfg = compile_config(cfg, policy=DQNPolicy) + self.exp_name = self.cfg.exp_name + if env is None: + self.env = supported_env[cfg.env.env_id](cfg=cfg.env) + else: + assert isinstance(env, BaseEnv), "Please use BaseEnv as env data type." + self.env = env + + logging.getLogger().setLevel(logging.INFO) + self.seed = seed + set_pkg_seed(self.seed, use_cuda=self.cfg.policy.cuda) + if not os.path.exists(self.exp_name): + os.makedirs(self.exp_name) + save_config_py(self.cfg, os.path.join(self.exp_name, 'policy_config.py')) + if model is None: + model = DQN(**self.cfg.policy.model) + self.buffer_ = DequeBuffer(size=self.cfg.policy.other.replay_buffer.replay_buffer_size) + self.policy = DQNPolicy(self.cfg.policy, model=model) + if policy_state_dict is not None: + self.policy.learn_mode.load_state_dict(policy_state_dict) + self.checkpoint_save_dir = os.path.join(self.exp_name, "ckpt") + + def train( + self, + step: int = int(1e7), + collector_env_num: int = None, + evaluator_env_num: int = None, + n_iter_save_ckpt: int = 1000, + context: Optional[str] = None, + debug: bool = False, + wandb_sweep: bool = False, + ) -> TrainingReturn: + """ + Overview: + Train the agent with DQN algorithm for ``step`` iterations with ``collector_env_num`` collector \ + environments and ``evaluator_env_num`` evaluator environments. Information during training will be \ + recorded and saved by wandb. + Arguments: + - step (:obj:`int`): The total training environment steps of all collector environments. Default to 1e7. + - collector_env_num (:obj:`int`): The collector environment number. Default to None. \ + If not specified, it will be set according to the configuration. + - evaluator_env_num (:obj:`int`): The evaluator environment number. Default to None. \ + If not specified, it will be set according to the configuration. + - n_iter_save_ckpt (:obj:`int`): The frequency of saving checkpoint every training iteration. \ + Default to 1000. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + - wandb_sweep (:obj:`bool`): Whether to use wandb sweep, \ + which is a hyper-parameter optimization process for seeking the best configurations. \ + Default to False. If True, the wandb sweep id will be used as the experiment name. + Returns: + - (:obj:`TrainingReturn`): The training result, of which the attributions are: + - wandb_url (:obj:`str`): The weight & biases (wandb) project url of the trainning experiment. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + logging.debug(self.policy._model) + # define env and policy + collector_env_num = collector_env_num if collector_env_num else self.cfg.env.collector_env_num + evaluator_env_num = evaluator_env_num if evaluator_env_num else self.cfg.env.evaluator_env_num + collector_env = setup_ding_env_manager(self.env, collector_env_num, context, debug, 'collector') + evaluator_env = setup_ding_env_manager(self.env, evaluator_env_num, context, debug, 'evaluator') + + with task.start(ctx=OnlineRLContext()): + task.use( + interaction_evaluator( + self.cfg, + self.policy.eval_mode, + evaluator_env, + render=self.cfg.policy.eval.render if hasattr(self.cfg.policy.eval, "render") else False + ) + ) + task.use(CkptSaver(policy=self.policy, save_dir=self.checkpoint_save_dir, train_freq=n_iter_save_ckpt)) + task.use(eps_greedy_handler(self.cfg)) + task.use( + StepCollector( + self.cfg, + self.policy.collect_mode, + collector_env, + random_collect_size=self.cfg.policy.random_collect_size + if hasattr(self.cfg.policy, 'random_collect_size') else 0, + ) + ) + if "nstep" in self.cfg.policy and self.cfg.policy.nstep > 1: + task.use(nstep_reward_enhancer(self.cfg)) + task.use(data_pusher(self.cfg, self.buffer_)) + task.use(OffPolicyLearner(self.cfg, self.policy.learn_mode, self.buffer_)) + task.use( + wandb_online_logger( + metric_list=self.policy._monitor_vars_learn(), + model=self.policy._model, + anonymous=True, + project_name=self.exp_name, + wandb_sweep=wandb_sweep, + ) + ) + task.use(termination_checker(max_env_step=step)) + task.use(final_ctx_saver(name=self.exp_name)) + task.run() + + return TrainingReturn(wandb_url=task.ctx.wandb_url) + + def deploy( + self, + enable_save_replay: bool = False, + concatenate_all_replay: bool = False, + replay_save_path: str = None, + seed: Optional[Union[int, List]] = None, + debug: bool = False + ) -> EvalReturn: + """ + Overview: + Deploy the agent with DQN algorithm by interacting with the environment, during which the replay video \ + can be saved if ``enable_save_replay`` is True. The evaluation result will be returned. + Arguments: + - enable_save_replay (:obj:`bool`): Whether to save the replay video. Default to False. + - concatenate_all_replay (:obj:`bool`): Whether to concatenate all replay videos into one video. \ + Default to False. If ``enable_save_replay`` is False, this argument will be ignored. \ + If ``enable_save_replay`` is True and ``concatenate_all_replay`` is False, \ + the replay video of each episode will be saved separately. + - replay_save_path (:obj:`str`): The path to save the replay video. Default to None. \ + If not specified, the video will be saved in ``exp_name/videos``. + - seed (:obj:`Union[int, List]`): The random seed, which is set before running the program. \ + Default to None. If not specified, ``self.seed`` will be used. \ + If ``seed`` is an integer, the agent will be deployed once. \ + If ``seed`` is a list of integers, the agent will be deployed once for each seed in the list. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + Returns: + - (:obj:`EvalReturn`): The evaluation result, of which the attributions are: + - eval_value (:obj:`np.float32`): The mean of evaluation return. + - eval_value_std (:obj:`np.float32`): The standard deviation of evaluation return. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + # define env and policy + env = self.env.clone(caller='evaluator') + + if seed is not None and isinstance(seed, int): + seeds = [seed] + elif seed is not None and isinstance(seed, list): + seeds = seed + else: + seeds = [self.seed] + + returns = [] + images = [] + if enable_save_replay: + replay_save_path = os.path.join(self.exp_name, 'videos') if replay_save_path is None else replay_save_path + env.enable_save_replay(replay_path=replay_save_path) + else: + logging.warning('No video would be generated during the deploy.') + if concatenate_all_replay: + logging.warning('concatenate_all_replay is set to False because enable_save_replay is False.') + concatenate_all_replay = False + + def single_env_forward_wrapper(forward_fn, cuda=True): + + forward_fn = model_wrap(forward_fn, wrapper_name='argmax_sample').forward + + def _forward(obs): + # unsqueeze means add batch dim, i.e. (O, ) -> (1, O) + obs = ttorch.as_tensor(obs).unsqueeze(0) + if cuda and torch.cuda.is_available(): + obs = obs.cuda() + action = forward_fn(obs)["action"] + # squeeze means delete batch dim, i.e. (1, A) -> (A, ) + action = action.squeeze(0).detach().cpu().numpy() + return action + + return _forward + + forward_fn = single_env_forward_wrapper(self.policy._model, self.cfg.policy.cuda) + + # reset first to make sure the env is in the initial state + # env will be reset again in the main loop + env.reset() + + for seed in seeds: + env.seed(seed, dynamic_seed=False) + return_ = 0. + step = 0 + obs = env.reset() + images.append(render(env)[None]) if concatenate_all_replay else None + while True: + action = forward_fn(obs) + obs, rew, done, info = env.step(action) + images.append(render(env)[None]) if concatenate_all_replay else None + return_ += rew + step += 1 + if done: + break + logging.info(f'DQN deploy is finished, final episode return with {step} steps is: {return_}') + returns.append(return_) + + env.close() + + if concatenate_all_replay: + images = np.concatenate(images, axis=0) + import imageio + imageio.mimwrite(os.path.join(replay_save_path, 'deploy.mp4'), images, fps=get_env_fps(env)) + + return EvalReturn(eval_value=np.mean(returns), eval_value_std=np.std(returns)) + + def collect_data( + self, + env_num: int = 8, + save_data_path: Optional[str] = None, + n_sample: Optional[int] = None, + n_episode: Optional[int] = None, + context: Optional[str] = None, + debug: bool = False + ) -> None: + """ + Overview: + Collect data with DQN algorithm for ``n_episode`` episodes with ``env_num`` collector environments. \ + The collected data will be saved in ``save_data_path`` if specified, otherwise it will be saved in \ + ``exp_name/demo_data``. + Arguments: + - env_num (:obj:`int`): The number of collector environments. Default to 8. + - save_data_path (:obj:`str`): The path to save the collected data. Default to None. \ + If not specified, the data will be saved in ``exp_name/demo_data``. + - n_sample (:obj:`int`): The number of samples to collect. Default to None. \ + If not specified, ``n_episode`` must be specified. + - n_episode (:obj:`int`): The number of episodes to collect. Default to None. \ + If not specified, ``n_sample`` must be specified. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + if n_episode is not None: + raise NotImplementedError + # define env and policy + env_num = env_num if env_num else self.cfg.env.collector_env_num + env = setup_ding_env_manager(self.env, env_num, context, debug, 'collector') + + if save_data_path is None: + save_data_path = os.path.join(self.exp_name, 'demo_data') + + # main execution task + with task.start(ctx=OnlineRLContext()): + task.use( + StepCollector( + self.cfg, self.policy.collect_mode, env, random_collect_size=self.cfg.policy.random_collect_size + ) + ) + task.use(offline_data_saver(save_data_path, data_type='hdf5')) + task.run(max_step=1) + logging.info( + f'DQN collecting is finished, more than {n_sample} samples are collected and saved in `{save_data_path}`' + ) + + def batch_evaluate( + self, + env_num: int = 4, + n_evaluator_episode: int = 4, + context: Optional[str] = None, + debug: bool = False + ) -> EvalReturn: + """ + Overview: + Evaluate the agent with DQN algorithm for ``n_evaluator_episode`` episodes with ``env_num`` evaluator \ + environments. The evaluation result will be returned. + The difference between methods ``batch_evaluate`` and ``deploy`` is that ``batch_evaluate`` will create \ + multiple evaluator environments to evaluate the agent to get an average performance, while ``deploy`` \ + will only create one evaluator environment to evaluate the agent and save the replay video. + Arguments: + - env_num (:obj:`int`): The number of evaluator environments. Default to 4. + - n_evaluator_episode (:obj:`int`): The number of episodes to evaluate. Default to 4. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + Returns: + - (:obj:`EvalReturn`): The evaluation result, of which the attributions are: + - eval_value (:obj:`np.float32`): The mean of evaluation return. + - eval_value_std (:obj:`np.float32`): The standard deviation of evaluation return. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + # define env and policy + env_num = env_num if env_num else self.cfg.env.evaluator_env_num + env = setup_ding_env_manager(self.env, env_num, context, debug, 'evaluator') + + # reset first to make sure the env is in the initial state + # env will be reset again in the main loop + env.launch() + env.reset() + + evaluate_cfg = self.cfg + evaluate_cfg.env.n_evaluator_episode = n_evaluator_episode + + # main execution task + with task.start(ctx=OnlineRLContext()): + task.use(interaction_evaluator(self.cfg, self.policy.eval_mode, env)) + task.run(max_step=1) + + return EvalReturn(eval_value=task.ctx.eval_value, eval_value_std=task.ctx.eval_value_std) + + @property + def best(self) -> 'DQNAgent': + """ + Overview: + Load the best model from the checkpoint directory, \ + which is by default in folder ``exp_name/ckpt/eval.pth.tar``. \ + The return value is the agent with the best model. + Returns: + - (:obj:`DQNAgent`): The agent with the best model. + Examples: + >>> agent = DQNAgent(env_id='LunarLander-v2') + >>> agent.train() + >>> agent = agent.best + + .. note:: + The best model is the model with the highest evaluation return. If this method is called, the current \ + model will be replaced by the best model. + """ + + best_model_file_path = os.path.join(self.checkpoint_save_dir, "eval.pth.tar") + # Load best model if it exists + if os.path.exists(best_model_file_path): + policy_state_dict = torch.load(best_model_file_path, map_location=torch.device("cpu")) + self.policy.learn_mode.load_state_dict(policy_state_dict) + return self diff --git a/ding/bonus/model.py b/ding/bonus/model.py new file mode 100644 index 0000000000..d33fa4c779 --- /dev/null +++ b/ding/bonus/model.py @@ -0,0 +1,245 @@ +from typing import Union, Optional +from easydict import EasyDict +import torch +import torch.nn as nn +import treetensor.torch as ttorch +from copy import deepcopy +from ding.utils import SequenceType, squeeze +from ding.model.common import ReparameterizationHead, RegressionHead, MultiHead, \ + FCEncoder, ConvEncoder, IMPALAConvEncoder, PopArtVHead +from ding.torch_utils import MLP, fc_block + + +class DiscretePolicyHead(nn.Module): + + def __init__( + self, + hidden_size: int, + output_size: int, + layer_num: int = 1, + activation: Optional[nn.Module] = nn.ReLU(), + norm_type: Optional[str] = None, + ) -> None: + super(DiscretePolicyHead, self).__init__() + self.main = nn.Sequential( + MLP( + hidden_size, + hidden_size, + hidden_size, + layer_num, + layer_fn=nn.Linear, + activation=activation, + norm_type=norm_type + ), fc_block(hidden_size, output_size) + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.main(x) + + +class PPOFModel(nn.Module): + mode = ['compute_actor', 'compute_critic', 'compute_actor_critic'] + + def __init__( + self, + obs_shape: Union[int, SequenceType], + action_shape: Union[int, SequenceType, EasyDict], + action_space: str = 'discrete', + share_encoder: bool = True, + encoder_hidden_size_list: SequenceType = [128, 128, 64], + actor_head_hidden_size: int = 64, + actor_head_layer_num: int = 1, + critic_head_hidden_size: int = 64, + critic_head_layer_num: int = 1, + activation: Optional[nn.Module] = nn.ReLU(), + norm_type: Optional[str] = None, + sigma_type: Optional[str] = 'independent', + fixed_sigma_value: Optional[int] = 0.3, + bound_type: Optional[str] = None, + encoder: Optional[torch.nn.Module] = None, + popart_head=False, + ) -> None: + super(PPOFModel, self).__init__() + obs_shape = squeeze(obs_shape) + action_shape = squeeze(action_shape) + self.obs_shape, self.action_shape = obs_shape, action_shape + self.share_encoder = share_encoder + + # Encoder Type + def new_encoder(outsize): + if isinstance(obs_shape, int) or len(obs_shape) == 1: + return FCEncoder( + obs_shape=obs_shape, + hidden_size_list=encoder_hidden_size_list, + activation=activation, + norm_type=norm_type + ) + elif len(obs_shape) == 3: + return ConvEncoder( + obs_shape=obs_shape, + hidden_size_list=encoder_hidden_size_list, + activation=activation, + norm_type=norm_type + ) + else: + raise RuntimeError( + "not support obs_shape for pre-defined encoder: {}, please customize your own encoder". + format(obs_shape) + ) + + if self.share_encoder: + assert actor_head_hidden_size == critic_head_hidden_size, \ + "actor and critic network head should have same size." + if encoder: + if isinstance(encoder, torch.nn.Module): + self.encoder = encoder + else: + raise ValueError("illegal encoder instance.") + else: + self.encoder = new_encoder(actor_head_hidden_size) + else: + if encoder: + if isinstance(encoder, torch.nn.Module): + self.actor_encoder = encoder + self.critic_encoder = deepcopy(encoder) + else: + raise ValueError("illegal encoder instance.") + else: + self.actor_encoder = new_encoder(actor_head_hidden_size) + self.critic_encoder = new_encoder(critic_head_hidden_size) + + # Head Type + if not popart_head: + self.critic_head = RegressionHead( + critic_head_hidden_size, 1, critic_head_layer_num, activation=activation, norm_type=norm_type + ) + else: + self.critic_head = PopArtVHead( + critic_head_hidden_size, 1, critic_head_layer_num, activation=activation, norm_type=norm_type + ) + + self.action_space = action_space + assert self.action_space in ['discrete', 'continuous', 'hybrid'], self.action_space + if self.action_space == 'continuous': + self.multi_head = False + self.actor_head = ReparameterizationHead( + actor_head_hidden_size, + action_shape, + actor_head_layer_num, + sigma_type=sigma_type, + activation=activation, + norm_type=norm_type, + bound_type=bound_type + ) + elif self.action_space == 'discrete': + actor_head_cls = DiscretePolicyHead + multi_head = not isinstance(action_shape, int) + self.multi_head = multi_head + if multi_head: + self.actor_head = MultiHead( + actor_head_cls, + actor_head_hidden_size, + action_shape, + layer_num=actor_head_layer_num, + activation=activation, + norm_type=norm_type + ) + else: + self.actor_head = actor_head_cls( + actor_head_hidden_size, + action_shape, + actor_head_layer_num, + activation=activation, + norm_type=norm_type + ) + elif self.action_space == 'hybrid': # HPPO + # hybrid action space: action_type(discrete) + action_args(continuous), + # such as {'action_type_shape': torch.LongTensor([0]), 'action_args_shape': torch.FloatTensor([0.1, -0.27])} + action_shape.action_args_shape = squeeze(action_shape.action_args_shape) + action_shape.action_type_shape = squeeze(action_shape.action_type_shape) + actor_action_args = ReparameterizationHead( + actor_head_hidden_size, + action_shape.action_args_shape, + actor_head_layer_num, + sigma_type=sigma_type, + fixed_sigma_value=fixed_sigma_value, + activation=activation, + norm_type=norm_type, + bound_type=bound_type, + ) + actor_action_type = DiscretePolicyHead( + actor_head_hidden_size, + action_shape.action_type_shape, + actor_head_layer_num, + activation=activation, + norm_type=norm_type, + ) + self.actor_head = nn.ModuleList([actor_action_type, actor_action_args]) + + # must use list, not nn.ModuleList + if self.share_encoder: + self.actor = [self.encoder, self.actor_head] + self.critic = [self.encoder, self.critic_head] + else: + self.actor = [self.actor_encoder, self.actor_head] + self.critic = [self.critic_encoder, self.critic_head] + # Convenient for calling some apis (e.g. self.critic.parameters()), + # but may cause misunderstanding when `print(self)` + self.actor = nn.ModuleList(self.actor) + self.critic = nn.ModuleList(self.critic) + + def forward(self, inputs: ttorch.Tensor, mode: str) -> ttorch.Tensor: + assert mode in self.mode, "not support forward mode: {}/{}".format(mode, self.mode) + return getattr(self, mode)(inputs) + + def compute_actor(self, x: ttorch.Tensor) -> ttorch.Tensor: + if self.share_encoder: + x = self.encoder(x) + else: + x = self.actor_encoder(x) + + if self.action_space == 'discrete': + return self.actor_head(x) + elif self.action_space == 'continuous': + x = self.actor_head(x) # mu, sigma + return ttorch.as_tensor(x) + elif self.action_space == 'hybrid': + action_type = self.actor_head[0](x) + action_args = self.actor_head[1](x) + return ttorch.as_tensor({'action_type': action_type, 'action_args': action_args}) + + def compute_critic(self, x: ttorch.Tensor) -> ttorch.Tensor: + if self.share_encoder: + x = self.encoder(x) + else: + x = self.critic_encoder(x) + x = self.critic_head(x) + return x + + def compute_actor_critic(self, x: ttorch.Tensor) -> ttorch.Tensor: + if self.share_encoder: + actor_embedding = critic_embedding = self.encoder(x) + else: + actor_embedding = self.actor_encoder(x) + critic_embedding = self.critic_encoder(x) + + value = self.critic_head(critic_embedding) + + if self.action_space == 'discrete': + logit = self.actor_head(actor_embedding) + return ttorch.as_tensor({'logit': logit, 'value': value['pred']}) + elif self.action_space == 'continuous': + x = self.actor_head(actor_embedding) + return ttorch.as_tensor({'logit': x, 'value': value['pred']}) + elif self.action_space == 'hybrid': + action_type = self.actor_head[0](actor_embedding) + action_args = self.actor_head[1](actor_embedding) + return ttorch.as_tensor( + { + 'logit': { + 'action_type': action_type, + 'action_args': action_args + }, + 'value': value['pred'] + } + ) diff --git a/ding/bonus/pg.py b/ding/bonus/pg.py new file mode 100644 index 0000000000..1e7cb8fb99 --- /dev/null +++ b/ding/bonus/pg.py @@ -0,0 +1,453 @@ +from typing import Optional, Union, List +from ditk import logging +from easydict import EasyDict +import os +import numpy as np +import torch +import treetensor.torch as ttorch +from ding.framework import task, OnlineRLContext +from ding.framework.middleware import CkptSaver, trainer, \ + wandb_online_logger, offline_data_saver, termination_checker, interaction_evaluator, StepCollector, \ + montecarlo_return_estimator, final_ctx_saver, EpisodeCollector +from ding.envs import BaseEnv +from ding.envs import setup_ding_env_manager +from ding.policy import PGPolicy +from ding.utils import set_pkg_seed +from ding.utils import get_env_fps, render +from ding.config import save_config_py, compile_config +from ding.model import PG +from ding.bonus.common import TrainingReturn, EvalReturn +from ding.config.example.PG import supported_env_cfg +from ding.config.example.PG import supported_env + + +class PGAgent: + """ + Overview: + Class of agent for training, evaluation and deployment of Reinforcement learning algorithm Policy Gradient(PG). + For more information about the system design of RL agent, please refer to \ + . + Interface: + ``__init__``, ``train``, ``deploy``, ``collect_data``, ``batch_evaluate``, ``best`` + """ + supported_env_list = list(supported_env_cfg.keys()) + """ + Overview: + List of supported envs. + Examples: + >>> from ding.bonus.pg import PGAgent + >>> print(PGAgent.supported_env_list) + """ + + def __init__( + self, + env_id: str = None, + env: BaseEnv = None, + seed: int = 0, + exp_name: str = None, + model: Optional[torch.nn.Module] = None, + cfg: Optional[Union[EasyDict, dict]] = None, + policy_state_dict: str = None, + ) -> None: + """ + Overview: + Initialize agent for PG algorithm. + Arguments: + - env_id (:obj:`str`): The environment id, which is a registered environment name in gym or gymnasium. \ + If ``env_id`` is not specified, ``env_id`` in ``cfg.env`` must be specified. \ + If ``env_id`` is specified, ``env_id`` in ``cfg.env`` will be ignored. \ + ``env_id`` should be one of the supported envs, which can be found in ``supported_env_list``. + - env (:obj:`BaseEnv`): The environment instance for training and evaluation. \ + If ``env`` is not specified, `env_id`` or ``cfg.env.env_id`` must be specified. \ + ``env_id`` or ``cfg.env.env_id`` will be used to create environment instance. \ + If ``env`` is specified, ``env_id`` and ``cfg.env.env_id`` will be ignored. + - seed (:obj:`int`): The random seed, which is set before running the program. \ + Default to 0. + - exp_name (:obj:`str`): The name of this experiment, which will be used to create the folder to save \ + log data. Default to None. If not specified, the folder name will be ``env_id``-``algorithm``. + - model (:obj:`torch.nn.Module`): The model of PG algorithm, which should be an instance of class \ + :class:`ding.model.PG`. \ + If not specified, a default model will be generated according to the configuration. + - cfg (:obj:`Union[EasyDict, dict]`): The configuration of PG algorithm, which is a dict. \ + Default to None. If not specified, the default configuration will be used. \ + The default configuration can be found in ``ding/config/example/PG/gym_lunarlander_v2.py``. + - policy_state_dict (:obj:`str`): The path of policy state dict saved by PyTorch a in local file. \ + If specified, the policy will be loaded from this file. Default to None. + + .. note:: + An RL Agent Instance can be initialized in two basic ways. \ + For example, we have an environment with id ``LunarLanderContinuous-v2`` registered in gym, \ + and we want to train an agent with PG algorithm with default configuration. \ + Then we can initialize the agent in the following ways: + >>> agent = PGAgent(env_id='LunarLanderContinuous-v2') + or, if we want can specify the env_id in the configuration: + >>> cfg = {'env': {'env_id': 'LunarLanderContinuous-v2'}, 'policy': ...... } + >>> agent = PGAgent(cfg=cfg) + There are also other arguments to specify the agent when initializing. + For example, if we want to specify the environment instance: + >>> env = CustomizedEnv('LunarLanderContinuous-v2') + >>> agent = PGAgent(cfg=cfg, env=env) + or, if we want to specify the model: + >>> model = PG(**cfg.policy.model) + >>> agent = PGAgent(cfg=cfg, model=model) + or, if we want to reload the policy from a saved policy state dict: + >>> agent = PGAgent(cfg=cfg, policy_state_dict='LunarLanderContinuous-v2.pth.tar') + Make sure that the configuration is consistent with the saved policy state dict. + """ + + assert env_id is not None or cfg is not None, "Please specify env_id or cfg." + + if cfg is not None and not isinstance(cfg, EasyDict): + cfg = EasyDict(cfg) + + if env_id is not None: + assert env_id in PGAgent.supported_env_list, "Please use supported envs: {}".format( + PGAgent.supported_env_list + ) + if cfg is None: + cfg = supported_env_cfg[env_id] + else: + assert cfg.env.env_id == env_id, "env_id in cfg should be the same as env_id in args." + else: + assert hasattr(cfg.env, "env_id"), "Please specify env_id in cfg." + assert cfg.env.env_id in PGAgent.supported_env_list, "Please use supported envs: {}".format( + PGAgent.supported_env_list + ) + default_policy_config = EasyDict({"policy": PGPolicy.default_config()}) + default_policy_config.update(cfg) + cfg = default_policy_config + + if exp_name is not None: + cfg.exp_name = exp_name + self.cfg = compile_config(cfg, policy=PGPolicy) + self.exp_name = self.cfg.exp_name + if env is None: + self.env = supported_env[cfg.env.env_id](cfg=cfg.env) + else: + assert isinstance(env, BaseEnv), "Please use BaseEnv as env data type." + self.env = env + + logging.getLogger().setLevel(logging.INFO) + self.seed = seed + set_pkg_seed(self.seed, use_cuda=self.cfg.policy.cuda) + if not os.path.exists(self.exp_name): + os.makedirs(self.exp_name) + save_config_py(self.cfg, os.path.join(self.exp_name, 'policy_config.py')) + if model is None: + model = PG(**self.cfg.policy.model) + self.policy = PGPolicy(self.cfg.policy, model=model) + if policy_state_dict is not None: + self.policy.learn_mode.load_state_dict(policy_state_dict) + self.checkpoint_save_dir = os.path.join(self.exp_name, "ckpt") + + def train( + self, + step: int = int(1e7), + collector_env_num: int = None, + evaluator_env_num: int = None, + n_iter_save_ckpt: int = 1000, + context: Optional[str] = None, + debug: bool = False, + wandb_sweep: bool = False, + ) -> TrainingReturn: + """ + Overview: + Train the agent with PG algorithm for ``step`` iterations with ``collector_env_num`` collector \ + environments and ``evaluator_env_num`` evaluator environments. Information during training will be \ + recorded and saved by wandb. + Arguments: + - step (:obj:`int`): The total training environment steps of all collector environments. Default to 1e7. + - collector_env_num (:obj:`int`): The collector environment number. Default to None. \ + If not specified, it will be set according to the configuration. + - evaluator_env_num (:obj:`int`): The evaluator environment number. Default to None. \ + If not specified, it will be set according to the configuration. + - n_iter_save_ckpt (:obj:`int`): The frequency of saving checkpoint every training iteration. \ + Default to 1000. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + - wandb_sweep (:obj:`bool`): Whether to use wandb sweep, \ + which is a hyper-parameter optimization process for seeking the best configurations. \ + Default to False. If True, the wandb sweep id will be used as the experiment name. + Returns: + - (:obj:`TrainingReturn`): The training result, of which the attributions are: + - wandb_url (:obj:`str`): The weight & biases (wandb) project url of the trainning experiment. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + logging.debug(self.policy._model) + # define env and policy + collector_env_num = collector_env_num if collector_env_num else self.cfg.env.collector_env_num + evaluator_env_num = evaluator_env_num if evaluator_env_num else self.cfg.env.evaluator_env_num + collector_env = setup_ding_env_manager(self.env, collector_env_num, context, debug, 'collector') + evaluator_env = setup_ding_env_manager(self.env, evaluator_env_num, context, debug, 'evaluator') + + with task.start(ctx=OnlineRLContext()): + task.use( + interaction_evaluator( + self.cfg, + self.policy.eval_mode, + evaluator_env, + render=self.cfg.policy.eval.render if hasattr(self.cfg.policy.eval, "render") else False + ) + ) + task.use(CkptSaver(policy=self.policy, save_dir=self.checkpoint_save_dir, train_freq=n_iter_save_ckpt)) + task.use(EpisodeCollector(self.cfg, self.policy.collect_mode, collector_env)) + task.use(montecarlo_return_estimator(self.policy)) + task.use(trainer(self.cfg, self.policy.learn_mode)) + task.use( + wandb_online_logger( + metric_list=self.policy._monitor_vars_learn(), + model=self.policy._model, + anonymous=True, + project_name=self.exp_name, + wandb_sweep=wandb_sweep, + ) + ) + task.use(termination_checker(max_env_step=step)) + task.use(final_ctx_saver(name=self.exp_name)) + task.run() + + return TrainingReturn(wandb_url=task.ctx.wandb_url) + + def deploy( + self, + enable_save_replay: bool = False, + concatenate_all_replay: bool = False, + replay_save_path: str = None, + seed: Optional[Union[int, List]] = None, + debug: bool = False + ) -> EvalReturn: + """ + Overview: + Deploy the agent with PG algorithm by interacting with the environment, during which the replay video \ + can be saved if ``enable_save_replay`` is True. The evaluation result will be returned. + Arguments: + - enable_save_replay (:obj:`bool`): Whether to save the replay video. Default to False. + - concatenate_all_replay (:obj:`bool`): Whether to concatenate all replay videos into one video. \ + Default to False. If ``enable_save_replay`` is False, this argument will be ignored. \ + If ``enable_save_replay`` is True and ``concatenate_all_replay`` is False, \ + the replay video of each episode will be saved separately. + - replay_save_path (:obj:`str`): The path to save the replay video. Default to None. \ + If not specified, the video will be saved in ``exp_name/videos``. + - seed (:obj:`Union[int, List]`): The random seed, which is set before running the program. \ + Default to None. If not specified, ``self.seed`` will be used. \ + If ``seed`` is an integer, the agent will be deployed once. \ + If ``seed`` is a list of integers, the agent will be deployed once for each seed in the list. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + Returns: + - (:obj:`EvalReturn`): The evaluation result, of which the attributions are: + - eval_value (:obj:`np.float32`): The mean of evaluation return. + - eval_value_std (:obj:`np.float32`): The standard deviation of evaluation return. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + # define env and policy + env = self.env.clone(caller='evaluator') + + if seed is not None and isinstance(seed, int): + seeds = [seed] + elif seed is not None and isinstance(seed, list): + seeds = seed + else: + seeds = [self.seed] + + returns = [] + images = [] + if enable_save_replay: + replay_save_path = os.path.join(self.exp_name, 'videos') if replay_save_path is None else replay_save_path + env.enable_save_replay(replay_path=replay_save_path) + else: + logging.warning('No video would be generated during the deploy.') + if concatenate_all_replay: + logging.warning('concatenate_all_replay is set to False because enable_save_replay is False.') + concatenate_all_replay = False + + def single_env_forward_wrapper(forward_fn, cuda=True): + + def _forward(obs): + # unsqueeze means add batch dim, i.e. (O, ) -> (1, O) + obs = ttorch.as_tensor(obs).unsqueeze(0) + if cuda and torch.cuda.is_available(): + obs = obs.cuda() + output = forward_fn(obs) + if self.policy._cfg.deterministic_eval: + if self.policy._cfg.action_space == 'discrete': + output['action'] = output['logit'].argmax(dim=-1) + elif self.policy._cfg.action_space == 'continuous': + output['action'] = output['logit']['mu'] + else: + raise KeyError("invalid action_space: {}".format(self.policy._cfg.action_space)) + else: + output['action'] = output['dist'].sample() + # squeeze means delete batch dim, i.e. (1, A) -> (A, ) + action = output['action'].squeeze(0).detach().cpu().numpy() + return action + + return _forward + + forward_fn = single_env_forward_wrapper(self.policy._model, self.cfg.policy.cuda) + + # reset first to make sure the env is in the initial state + # env will be reset again in the main loop + env.reset() + + for seed in seeds: + env.seed(seed, dynamic_seed=False) + return_ = 0. + step = 0 + obs = env.reset() + images.append(render(env)[None]) if concatenate_all_replay else None + while True: + action = forward_fn(obs) + obs, rew, done, info = env.step(action) + images.append(render(env)[None]) if concatenate_all_replay else None + return_ += rew + step += 1 + if done: + break + logging.info(f'DQN deploy is finished, final episode return with {step} steps is: {return_}') + returns.append(return_) + + env.close() + + if concatenate_all_replay: + images = np.concatenate(images, axis=0) + import imageio + imageio.mimwrite(os.path.join(replay_save_path, 'deploy.mp4'), images, fps=get_env_fps(env)) + + return EvalReturn(eval_value=np.mean(returns), eval_value_std=np.std(returns)) + + def collect_data( + self, + env_num: int = 8, + save_data_path: Optional[str] = None, + n_sample: Optional[int] = None, + n_episode: Optional[int] = None, + context: Optional[str] = None, + debug: bool = False + ) -> None: + """ + Overview: + Collect data with PG algorithm for ``n_episode`` episodes with ``env_num`` collector environments. \ + The collected data will be saved in ``save_data_path`` if specified, otherwise it will be saved in \ + ``exp_name/demo_data``. + Arguments: + - env_num (:obj:`int`): The number of collector environments. Default to 8. + - save_data_path (:obj:`str`): The path to save the collected data. Default to None. \ + If not specified, the data will be saved in ``exp_name/demo_data``. + - n_sample (:obj:`int`): The number of samples to collect. Default to None. \ + If not specified, ``n_episode`` must be specified. + - n_episode (:obj:`int`): The number of episodes to collect. Default to None. \ + If not specified, ``n_sample`` must be specified. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + if n_episode is not None: + raise NotImplementedError + # define env and policy + env_num = env_num if env_num else self.cfg.env.collector_env_num + env = setup_ding_env_manager(self.env, env_num, context, debug, 'collector') + + if save_data_path is None: + save_data_path = os.path.join(self.exp_name, 'demo_data') + + # main execution task + with task.start(ctx=OnlineRLContext()): + task.use( + StepCollector( + self.cfg, self.policy.collect_mode, env, random_collect_size=self.cfg.policy.random_collect_size + ) + ) + task.use(offline_data_saver(save_data_path, data_type='hdf5')) + task.run(max_step=1) + logging.info( + f'PG collecting is finished, more than {n_sample} samples are collected and saved in `{save_data_path}`' + ) + + def batch_evaluate( + self, + env_num: int = 4, + n_evaluator_episode: int = 4, + context: Optional[str] = None, + debug: bool = False + ) -> EvalReturn: + """ + Overview: + Evaluate the agent with PG algorithm for ``n_evaluator_episode`` episodes with ``env_num`` evaluator \ + environments. The evaluation result will be returned. + The difference between methods ``batch_evaluate`` and ``deploy`` is that ``batch_evaluate`` will create \ + multiple evaluator environments to evaluate the agent to get an average performance, while ``deploy`` \ + will only create one evaluator environment to evaluate the agent and save the replay video. + Arguments: + - env_num (:obj:`int`): The number of evaluator environments. Default to 4. + - n_evaluator_episode (:obj:`int`): The number of episodes to evaluate. Default to 4. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + Returns: + - (:obj:`EvalReturn`): The evaluation result, of which the attributions are: + - eval_value (:obj:`np.float32`): The mean of evaluation return. + - eval_value_std (:obj:`np.float32`): The standard deviation of evaluation return. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + # define env and policy + env_num = env_num if env_num else self.cfg.env.evaluator_env_num + env = setup_ding_env_manager(self.env, env_num, context, debug, 'evaluator') + + # reset first to make sure the env is in the initial state + # env will be reset again in the main loop + env.launch() + env.reset() + + evaluate_cfg = self.cfg + evaluate_cfg.env.n_evaluator_episode = n_evaluator_episode + + # main execution task + with task.start(ctx=OnlineRLContext()): + task.use(interaction_evaluator(self.cfg, self.policy.eval_mode, env)) + task.run(max_step=1) + + return EvalReturn(eval_value=task.ctx.eval_value, eval_value_std=task.ctx.eval_value_std) + + @property + def best(self) -> 'PGAgent': + """ + Overview: + Load the best model from the checkpoint directory, \ + which is by default in folder ``exp_name/ckpt/eval.pth.tar``. \ + The return value is the agent with the best model. + Returns: + - (:obj:`PGAgent`): The agent with the best model. + Examples: + >>> agent = PGAgent(env_id='LunarLanderContinuous-v2') + >>> agent.train() + >>> agent = agent.best + + .. note:: + The best model is the model with the highest evaluation return. If this method is called, the current \ + model will be replaced by the best model. + """ + + best_model_file_path = os.path.join(self.checkpoint_save_dir, "eval.pth.tar") + # Load best model if it exists + if os.path.exists(best_model_file_path): + policy_state_dict = torch.load(best_model_file_path, map_location=torch.device("cpu")) + self.policy.learn_mode.load_state_dict(policy_state_dict) + return self diff --git a/ding/bonus/ppo_offpolicy.py b/ding/bonus/ppo_offpolicy.py new file mode 100644 index 0000000000..5aef20593f --- /dev/null +++ b/ding/bonus/ppo_offpolicy.py @@ -0,0 +1,471 @@ +from typing import Optional, Union, List +from ditk import logging +from easydict import EasyDict +import os +import numpy as np +import torch +import treetensor.torch as ttorch +from ding.framework import task, OnlineRLContext +from ding.framework.middleware import CkptSaver, final_ctx_saver, OffPolicyLearner, StepCollector, \ + wandb_online_logger, offline_data_saver, termination_checker, interaction_evaluator, gae_estimator +from ding.envs import BaseEnv +from ding.envs import setup_ding_env_manager +from ding.policy import PPOOffPolicy +from ding.utils import set_pkg_seed +from ding.utils import get_env_fps, render +from ding.config import save_config_py, compile_config +from ding.model import VAC +from ding.model import model_wrap +from ding.data import DequeBuffer +from ding.bonus.common import TrainingReturn, EvalReturn +from ding.config.example.PPOOffPolicy import supported_env_cfg +from ding.config.example.PPOOffPolicy import supported_env + + +class PPOOffPolicyAgent: + """ + Overview: + Class of agent for training, evaluation and deployment of Reinforcement learning algorithm \ + Proximal Policy Optimization(PPO) in an off-policy style. + For more information about the system design of RL agent, please refer to \ + . + Interface: + ``__init__``, ``train``, ``deploy``, ``collect_data``, ``batch_evaluate``, ``best`` + """ + supported_env_list = list(supported_env_cfg.keys()) + """ + Overview: + List of supported envs. + Examples: + >>> from ding.bonus.ppo_offpolicy import PPOOffPolicyAgent + >>> print(PPOOffPolicyAgent.supported_env_list) + """ + + def __init__( + self, + env_id: str = None, + env: BaseEnv = None, + seed: int = 0, + exp_name: str = None, + model: Optional[torch.nn.Module] = None, + cfg: Optional[Union[EasyDict, dict]] = None, + policy_state_dict: str = None, + ) -> None: + """ + Overview: + Initialize agent for PPO (offpolicy) algorithm. + Arguments: + - env_id (:obj:`str`): The environment id, which is a registered environment name in gym or gymnasium. \ + If ``env_id`` is not specified, ``env_id`` in ``cfg.env`` must be specified. \ + If ``env_id`` is specified, ``env_id`` in ``cfg.env`` will be ignored. \ + ``env_id`` should be one of the supported envs, which can be found in ``supported_env_list``. + - env (:obj:`BaseEnv`): The environment instance for training and evaluation. \ + If ``env`` is not specified, `env_id`` or ``cfg.env.env_id`` must be specified. \ + ``env_id`` or ``cfg.env.env_id`` will be used to create environment instance. \ + If ``env`` is specified, ``env_id`` and ``cfg.env.env_id`` will be ignored. + - seed (:obj:`int`): The random seed, which is set before running the program. \ + Default to 0. + - exp_name (:obj:`str`): The name of this experiment, which will be used to create the folder to save \ + log data. Default to None. If not specified, the folder name will be ``env_id``-``algorithm``. + - model (:obj:`torch.nn.Module`): The model of PPO (offpolicy) algorithm, \ + which should be an instance of class :class:`ding.model.VAC`. \ + If not specified, a default model will be generated according to the configuration. + - cfg (:obj:`Union[EasyDict, dict]`): The configuration of PPO (offpolicy) algorithm, which is a dict. \ + Default to None. If not specified, the default configuration will be used. \ + The default configuration can be found in ``ding/config/example/PPO (offpolicy)/gym_lunarlander_v2.py``. + - policy_state_dict (:obj:`str`): The path of policy state dict saved by PyTorch a in local file. \ + If specified, the policy will be loaded from this file. Default to None. + + .. note:: + An RL Agent Instance can be initialized in two basic ways. \ + For example, we have an environment with id ``LunarLander-v2`` registered in gym, \ + and we want to train an agent with PPO (offpolicy) algorithm with default configuration. \ + Then we can initialize the agent in the following ways: + >>> agent = PPOOffPolicyAgent(env_id='LunarLander-v2') + or, if we want can specify the env_id in the configuration: + >>> cfg = {'env': {'env_id': 'LunarLander-v2'}, 'policy': ...... } + >>> agent = PPOOffPolicyAgent(cfg=cfg) + There are also other arguments to specify the agent when initializing. + For example, if we want to specify the environment instance: + >>> env = CustomizedEnv('LunarLander-v2') + >>> agent = PPOOffPolicyAgent(cfg=cfg, env=env) + or, if we want to specify the model: + >>> model = VAC(**cfg.policy.model) + >>> agent = PPOOffPolicyAgent(cfg=cfg, model=model) + or, if we want to reload the policy from a saved policy state dict: + >>> agent = PPOOffPolicyAgent(cfg=cfg, policy_state_dict='LunarLander-v2.pth.tar') + Make sure that the configuration is consistent with the saved policy state dict. + """ + + assert env_id is not None or cfg is not None, "Please specify env_id or cfg." + + if cfg is not None and not isinstance(cfg, EasyDict): + cfg = EasyDict(cfg) + + if env_id is not None: + assert env_id in PPOOffPolicyAgent.supported_env_list, "Please use supported envs: {}".format( + PPOOffPolicyAgent.supported_env_list + ) + if cfg is None: + cfg = supported_env_cfg[env_id] + else: + assert cfg.env.env_id == env_id, "env_id in cfg should be the same as env_id in args." + else: + assert hasattr(cfg.env, "env_id"), "Please specify env_id in cfg." + assert cfg.env.env_id in PPOOffPolicyAgent.supported_env_list, "Please use supported envs: {}".format( + PPOOffPolicyAgent.supported_env_list + ) + default_policy_config = EasyDict({"policy": PPOOffPolicy.default_config()}) + default_policy_config.update(cfg) + cfg = default_policy_config + + if exp_name is not None: + cfg.exp_name = exp_name + self.cfg = compile_config(cfg, policy=PPOOffPolicy) + self.exp_name = self.cfg.exp_name + if env is None: + self.env = supported_env[cfg.env.env_id](cfg=cfg.env) + else: + assert isinstance(env, BaseEnv), "Please use BaseEnv as env data type." + self.env = env + + logging.getLogger().setLevel(logging.INFO) + self.seed = seed + set_pkg_seed(self.seed, use_cuda=self.cfg.policy.cuda) + if not os.path.exists(self.exp_name): + os.makedirs(self.exp_name) + save_config_py(self.cfg, os.path.join(self.exp_name, 'policy_config.py')) + if model is None: + model = VAC(**self.cfg.policy.model) + self.buffer_ = DequeBuffer(size=self.cfg.policy.other.replay_buffer.replay_buffer_size) + self.policy = PPOOffPolicy(self.cfg.policy, model=model) + if policy_state_dict is not None: + self.policy.learn_mode.load_state_dict(policy_state_dict) + self.checkpoint_save_dir = os.path.join(self.exp_name, "ckpt") + + def train( + self, + step: int = int(1e7), + collector_env_num: int = None, + evaluator_env_num: int = None, + n_iter_save_ckpt: int = 1000, + context: Optional[str] = None, + debug: bool = False, + wandb_sweep: bool = False, + ) -> TrainingReturn: + """ + Overview: + Train the agent with PPO (offpolicy) algorithm for ``step`` iterations with ``collector_env_num`` \ + collector environments and ``evaluator_env_num`` evaluator environments. \ + Information during training will be recorded and saved by wandb. + Arguments: + - step (:obj:`int`): The total training environment steps of all collector environments. Default to 1e7. + - collector_env_num (:obj:`int`): The collector environment number. Default to None. \ + If not specified, it will be set according to the configuration. + - evaluator_env_num (:obj:`int`): The evaluator environment number. Default to None. \ + If not specified, it will be set according to the configuration. + - n_iter_save_ckpt (:obj:`int`): The frequency of saving checkpoint every training iteration. \ + Default to 1000. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + - wandb_sweep (:obj:`bool`): Whether to use wandb sweep, \ + which is a hyper-parameter optimization process for seeking the best configurations. \ + Default to False. If True, the wandb sweep id will be used as the experiment name. + Returns: + - (:obj:`TrainingReturn`): The training result, of which the attributions are: + - wandb_url (:obj:`str`): The weight & biases (wandb) project url of the trainning experiment. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + logging.debug(self.policy._model) + # define env and policy + collector_env_num = collector_env_num if collector_env_num else self.cfg.env.collector_env_num + evaluator_env_num = evaluator_env_num if evaluator_env_num else self.cfg.env.evaluator_env_num + collector_env = setup_ding_env_manager(self.env, collector_env_num, context, debug, 'collector') + evaluator_env = setup_ding_env_manager(self.env, evaluator_env_num, context, debug, 'evaluator') + + with task.start(ctx=OnlineRLContext()): + task.use( + interaction_evaluator( + self.cfg, + self.policy.eval_mode, + evaluator_env, + render=self.cfg.policy.eval.render if hasattr(self.cfg.policy.eval, "render") else False + ) + ) + task.use(CkptSaver(policy=self.policy, save_dir=self.checkpoint_save_dir, train_freq=n_iter_save_ckpt)) + task.use( + StepCollector( + self.cfg, + self.policy.collect_mode, + collector_env, + random_collect_size=self.cfg.policy.random_collect_size + if hasattr(self.cfg.policy, 'random_collect_size') else 0, + ) + ) + task.use(gae_estimator(self.cfg, self.policy.collect_mode, self.buffer_)) + task.use(OffPolicyLearner(self.cfg, self.policy.learn_mode, self.buffer_)) + task.use( + wandb_online_logger( + cfg=self.cfg.wandb_logger, + exp_config=self.cfg, + metric_list=self.policy._monitor_vars_learn(), + model=self.policy._model, + anonymous=True, + project_name=self.exp_name, + wandb_sweep=wandb_sweep, + ) + ) + task.use(termination_checker(max_env_step=step)) + task.use(final_ctx_saver(name=self.exp_name)) + task.run() + + return TrainingReturn(wandb_url=task.ctx.wandb_url) + + def deploy( + self, + enable_save_replay: bool = False, + concatenate_all_replay: bool = False, + replay_save_path: str = None, + seed: Optional[Union[int, List]] = None, + debug: bool = False + ) -> EvalReturn: + """ + Overview: + Deploy the agent with PPO (offpolicy) algorithm by interacting with the environment, \ + during which the replay video can be saved if ``enable_save_replay`` is True. \ + The evaluation result will be returned. + Arguments: + - enable_save_replay (:obj:`bool`): Whether to save the replay video. Default to False. + - concatenate_all_replay (:obj:`bool`): Whether to concatenate all replay videos into one video. \ + Default to False. If ``enable_save_replay`` is False, this argument will be ignored. \ + If ``enable_save_replay`` is True and ``concatenate_all_replay`` is False, \ + the replay video of each episode will be saved separately. + - replay_save_path (:obj:`str`): The path to save the replay video. Default to None. \ + If not specified, the video will be saved in ``exp_name/videos``. + - seed (:obj:`Union[int, List]`): The random seed, which is set before running the program. \ + Default to None. If not specified, ``self.seed`` will be used. \ + If ``seed`` is an integer, the agent will be deployed once. \ + If ``seed`` is a list of integers, the agent will be deployed once for each seed in the list. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + Returns: + - (:obj:`EvalReturn`): The evaluation result, of which the attributions are: + - eval_value (:obj:`np.float32`): The mean of evaluation return. + - eval_value_std (:obj:`np.float32`): The standard deviation of evaluation return. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + # define env and policy + env = self.env.clone(caller='evaluator') + + if seed is not None and isinstance(seed, int): + seeds = [seed] + elif seed is not None and isinstance(seed, list): + seeds = seed + else: + seeds = [self.seed] + + returns = [] + images = [] + if enable_save_replay: + replay_save_path = os.path.join(self.exp_name, 'videos') if replay_save_path is None else replay_save_path + env.enable_save_replay(replay_path=replay_save_path) + else: + logging.warning('No video would be generated during the deploy.') + if concatenate_all_replay: + logging.warning('concatenate_all_replay is set to False because enable_save_replay is False.') + concatenate_all_replay = False + + def single_env_forward_wrapper(forward_fn, cuda=True): + + if self.cfg.policy.action_space == 'discrete': + forward_fn = model_wrap(forward_fn, wrapper_name='argmax_sample').forward + elif self.cfg.policy.action_space == 'continuous': + forward_fn = model_wrap(forward_fn, wrapper_name='deterministic_sample').forward + elif self.cfg.policy.action_space == 'hybrid': + forward_fn = model_wrap(forward_fn, wrapper_name='hybrid_deterministic_argmax_sample').forward + elif self.cfg.policy.action_space == 'general': + forward_fn = model_wrap(forward_fn, wrapper_name='base').forward + else: + raise NotImplementedError + + def _forward(obs): + # unsqueeze means add batch dim, i.e. (O, ) -> (1, O) + obs = ttorch.as_tensor(obs).unsqueeze(0) + if cuda and torch.cuda.is_available(): + obs = obs.cuda() + action = forward_fn(obs, mode='compute_actor')["action"] + # squeeze means delete batch dim, i.e. (1, A) -> (A, ) + action = action.squeeze(0).detach().cpu().numpy() + return action + + return _forward + + forward_fn = single_env_forward_wrapper(self.policy._model, self.cfg.policy.cuda) + + # reset first to make sure the env is in the initial state + # env will be reset again in the main loop + env.reset() + + for seed in seeds: + env.seed(seed, dynamic_seed=False) + return_ = 0. + step = 0 + obs = env.reset() + images.append(render(env)[None]) if concatenate_all_replay else None + while True: + action = forward_fn(obs) + obs, rew, done, info = env.step(action) + images.append(render(env)[None]) if concatenate_all_replay else None + return_ += rew + step += 1 + if done: + break + logging.info(f'PPO (offpolicy) deploy is finished, final episode return with {step} steps is: {return_}') + returns.append(return_) + + env.close() + + if concatenate_all_replay: + images = np.concatenate(images, axis=0) + import imageio + imageio.mimwrite(os.path.join(replay_save_path, 'deploy.mp4'), images, fps=get_env_fps(env)) + + return EvalReturn(eval_value=np.mean(returns), eval_value_std=np.std(returns)) + + def collect_data( + self, + env_num: int = 8, + save_data_path: Optional[str] = None, + n_sample: Optional[int] = None, + n_episode: Optional[int] = None, + context: Optional[str] = None, + debug: bool = False + ) -> None: + """ + Overview: + Collect data with PPO (offpolicy) algorithm for ``n_episode`` episodes \ + with ``env_num`` collector environments. \ + The collected data will be saved in ``save_data_path`` if specified, otherwise it will be saved in \ + ``exp_name/demo_data``. + Arguments: + - env_num (:obj:`int`): The number of collector environments. Default to 8. + - save_data_path (:obj:`str`): The path to save the collected data. Default to None. \ + If not specified, the data will be saved in ``exp_name/demo_data``. + - n_sample (:obj:`int`): The number of samples to collect. Default to None. \ + If not specified, ``n_episode`` must be specified. + - n_episode (:obj:`int`): The number of episodes to collect. Default to None. \ + If not specified, ``n_sample`` must be specified. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + if n_episode is not None: + raise NotImplementedError + # define env and policy + env_num = env_num if env_num else self.cfg.env.collector_env_num + env = setup_ding_env_manager(self.env, env_num, context, debug, 'collector') + + if save_data_path is None: + save_data_path = os.path.join(self.exp_name, 'demo_data') + + # main execution task + with task.start(ctx=OnlineRLContext()): + task.use( + StepCollector( + self.cfg, self.policy.collect_mode, env, random_collect_size=self.cfg.policy.random_collect_size + ) + ) + task.use(offline_data_saver(save_data_path, data_type='hdf5')) + task.run(max_step=1) + logging.info( + f'PPOOffPolicy collecting is finished, more than {n_sample} \ + samples are collected and saved in `{save_data_path}`' + ) + + def batch_evaluate( + self, + env_num: int = 4, + n_evaluator_episode: int = 4, + context: Optional[str] = None, + debug: bool = False + ) -> EvalReturn: + """ + Overview: + Evaluate the agent with PPO (offpolicy) algorithm for ``n_evaluator_episode`` episodes \ + with ``env_num`` evaluator environments. The evaluation result will be returned. + The difference between methods ``batch_evaluate`` and ``deploy`` is that ``batch_evaluate`` will create \ + multiple evaluator environments to evaluate the agent to get an average performance, while ``deploy`` \ + will only create one evaluator environment to evaluate the agent and save the replay video. + Arguments: + - env_num (:obj:`int`): The number of evaluator environments. Default to 4. + - n_evaluator_episode (:obj:`int`): The number of episodes to evaluate. Default to 4. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + Returns: + - (:obj:`EvalReturn`): The evaluation result, of which the attributions are: + - eval_value (:obj:`np.float32`): The mean of evaluation return. + - eval_value_std (:obj:`np.float32`): The standard deviation of evaluation return. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + # define env and policy + env_num = env_num if env_num else self.cfg.env.evaluator_env_num + env = setup_ding_env_manager(self.env, env_num, context, debug, 'evaluator') + + # reset first to make sure the env is in the initial state + # env will be reset again in the main loop + env.launch() + env.reset() + + evaluate_cfg = self.cfg + evaluate_cfg.env.n_evaluator_episode = n_evaluator_episode + + # main execution task + with task.start(ctx=OnlineRLContext()): + task.use(interaction_evaluator(self.cfg, self.policy.eval_mode, env)) + task.run(max_step=1) + + return EvalReturn(eval_value=task.ctx.eval_value, eval_value_std=task.ctx.eval_value_std) + + @property + def best(self) -> 'PPOOffPolicyAgent': + """ + Overview: + Load the best model from the checkpoint directory, \ + which is by default in folder ``exp_name/ckpt/eval.pth.tar``. \ + The return value is the agent with the best model. + Returns: + - (:obj:`PPOOffPolicyAgent`): The agent with the best model. + Examples: + >>> agent = PPOOffPolicyAgent(env_id='LunarLander-v2') + >>> agent.train() + >>> agent.best + + .. note:: + The best model is the model with the highest evaluation return. If this method is called, the current \ + model will be replaced by the best model. + """ + + best_model_file_path = os.path.join(self.checkpoint_save_dir, "eval.pth.tar") + # Load best model if it exists + if os.path.exists(best_model_file_path): + policy_state_dict = torch.load(best_model_file_path, map_location=torch.device("cpu")) + self.policy.learn_mode.load_state_dict(policy_state_dict) + return self diff --git a/ding/bonus/ppof.py b/ding/bonus/ppof.py new file mode 100644 index 0000000000..88d0b43e1e --- /dev/null +++ b/ding/bonus/ppof.py @@ -0,0 +1,509 @@ +from typing import Optional, Union, List +from ditk import logging +from easydict import EasyDict +from functools import partial +import os +import gym +import gymnasium +import numpy as np +import torch +from ding.framework import task, OnlineRLContext +from ding.framework.middleware import interaction_evaluator_ttorch, PPOFStepCollector, multistep_trainer, CkptSaver, \ + wandb_online_logger, offline_data_saver, termination_checker, ppof_adv_estimator +from ding.envs import BaseEnv, BaseEnvManagerV2, SubprocessEnvManagerV2 +from ding.policy import PPOFPolicy, single_env_forward_wrapper_ttorch +from ding.utils import set_pkg_seed +from ding.utils import get_env_fps, render +from ding.config import save_config_py +from .model import PPOFModel +from .config import get_instance_config, get_instance_env, get_hybrid_shape +from ding.bonus.common import TrainingReturn, EvalReturn + + +class PPOF: + """ + Overview: + Class of agent for training, evaluation and deployment of Reinforcement learning algorithm \ + Proximal Policy Optimization(PPO). + For more information about the system design of RL agent, please refer to \ + . + Interface: + ``__init__``, ``train``, ``deploy``, ``collect_data``, ``batch_evaluate``, ``best`` + """ + + supported_env_list = [ + # common + 'LunarLander-v2', + 'LunarLanderContinuous-v2', + 'BipedalWalker-v3', + 'Pendulum-v1', + 'acrobot', + # ch2: action + 'rocket_landing', + 'drone_fly', + 'hybrid_moving', + # ch3: obs + 'evogym_carrier', + 'mario', + 'di_sheep', + 'procgen_bigfish', + # ch4: reward + 'minigrid_fourroom', + 'metadrive', + # atari + 'BowlingNoFrameskip-v4', + 'BreakoutNoFrameskip-v4', + 'GopherNoFrameskip-v4' + 'KangarooNoFrameskip-v4', + 'PongNoFrameskip-v4', + 'QbertNoFrameskip-v4', + 'SpaceInvadersNoFrameskip-v4', + # mujoco + 'Hopper-v3', + 'HalfCheetah-v3', + 'Walker2d-v3', + ] + """ + Overview: + List of supported envs. + Examples: + >>> from ding.bonus.ppof import PPOF + >>> print(PPOF.supported_env_list) + """ + + def __init__( + self, + env_id: str = None, + env: BaseEnv = None, + seed: int = 0, + exp_name: str = None, + model: Optional[torch.nn.Module] = None, + cfg: Optional[Union[EasyDict, dict]] = None, + policy_state_dict: str = None + ) -> None: + """ + Overview: + Initialize agent for PPO algorithm. + Arguments: + - env_id (:obj:`str`): The environment id, which is a registered environment name in gym or gymnasium. \ + If ``env_id`` is not specified, ``env_id`` in ``cfg`` must be specified. \ + If ``env_id`` is specified, ``env_id`` in ``cfg`` will be ignored. \ + ``env_id`` should be one of the supported envs, which can be found in ``PPOF.supported_env_list``. + - env (:obj:`BaseEnv`): The environment instance for training and evaluation. \ + If ``env`` is not specified, ``env_id`` or ``cfg.env_id`` must be specified. \ + ``env_id`` or ``cfg.env_id`` will be used to create environment instance. \ + If ``env`` is specified, ``env_id`` and ``cfg.env_id`` will be ignored. + - seed (:obj:`int`): The random seed, which is set before running the program. \ + Default to 0. + - exp_name (:obj:`str`): The name of this experiment, which will be used to create the folder to save \ + log data. Default to None. If not specified, the folder name will be ``env_id``-``algorithm``. + - model (:obj:`torch.nn.Module`): The model of PPO algorithm, which should be an instance of class \ + ``ding.model.PPOFModel``. \ + If not specified, a default model will be generated according to the configuration. + - cfg (:obj:`Union[EasyDict, dict]`): The configuration of PPO algorithm, which is a dict. \ + Default to None. If not specified, the default configuration will be used. + - policy_state_dict (:obj:`str`): The path of policy state dict saved by PyTorch a in local file. \ + If specified, the policy will be loaded from this file. Default to None. + + .. note:: + An RL Agent Instance can be initialized in two basic ways. \ + For example, we have an environment with id ``LunarLander-v2`` registered in gym, \ + and we want to train an agent with PPO algorithm with default configuration. \ + Then we can initialize the agent in the following ways: + >>> agent = PPOF(env_id='LunarLander-v2') + or, if we want can specify the env_id in the configuration: + >>> cfg = {'env': {'env_id': 'LunarLander-v2'}, 'policy': ...... } + >>> agent = PPOF(cfg=cfg) + There are also other arguments to specify the agent when initializing. + For example, if we want to specify the environment instance: + >>> env = CustomizedEnv('LunarLander-v2') + >>> agent = PPOF(cfg=cfg, env=env) + or, if we want to specify the model: + >>> model = VAC(**cfg.policy.model) + >>> agent = PPOF(cfg=cfg, model=model) + or, if we want to reload the policy from a saved policy state dict: + >>> agent = PPOF(cfg=cfg, policy_state_dict='LunarLander-v2.pth.tar') + Make sure that the configuration is consistent with the saved policy state dict. + """ + + assert env_id is not None or cfg is not None, "Please specify env_id or cfg." + + if cfg is not None and not isinstance(cfg, EasyDict): + cfg = EasyDict(cfg) + + if env_id is not None: + assert env_id in PPOF.supported_env_list, "Please use supported envs: {}".format(PPOF.supported_env_list) + if cfg is None: + cfg = get_instance_config(env_id, algorithm="PPOF") + + if not hasattr(cfg, "env_id"): + cfg.env_id = env_id + assert cfg.env_id == env_id, "env_id in cfg should be the same as env_id in args." + else: + assert hasattr(cfg, "env_id"), "Please specify env_id in cfg." + assert cfg.env_id in PPOF.supported_env_list, "Please use supported envs: {}".format( + PPOF.supported_env_list + ) + + if exp_name is not None: + cfg.exp_name = exp_name + elif not hasattr(cfg, "exp_name"): + cfg.exp_name = "{}-{}".format(cfg.env_id, "PPO") + self.cfg = cfg + self.exp_name = self.cfg.exp_name + + if env is None: + self.env = get_instance_env(self.cfg.env_id) + else: + self.env = env + + logging.getLogger().setLevel(logging.INFO) + self.seed = seed + set_pkg_seed(self.seed, use_cuda=self.cfg.cuda) + + if not os.path.exists(self.exp_name): + os.makedirs(self.exp_name) + save_config_py(self.cfg, os.path.join(self.exp_name, 'policy_config.py')) + + action_space = self.env.action_space + if isinstance(action_space, (gym.spaces.Discrete, gymnasium.spaces.Discrete)): + action_shape = int(action_space.n) + elif isinstance(action_space, (gym.spaces.Tuple, gymnasium.spaces.Tuple)): + action_shape = get_hybrid_shape(action_space) + else: + action_shape = action_space.shape + + # Three types of value normalization is supported currently + assert self.cfg.value_norm in ['popart', 'value_rescale', 'symlog', 'baseline'] + if model is None: + if self.cfg.value_norm != 'popart': + model = PPOFModel( + self.env.observation_space.shape, + action_shape, + action_space=self.cfg.action_space, + **self.cfg.model + ) + else: + model = PPOFModel( + self.env.observation_space.shape, + action_shape, + action_space=self.cfg.action_space, + popart_head=True, + **self.cfg.model + ) + self.policy = PPOFPolicy(self.cfg, model=model) + if policy_state_dict is not None: + self.policy.load_state_dict(policy_state_dict) + self.checkpoint_save_dir = os.path.join(self.exp_name, "ckpt") + + def train( + self, + step: int = int(1e7), + collector_env_num: int = 4, + evaluator_env_num: int = 4, + n_iter_log_show: int = 500, + n_iter_save_ckpt: int = 1000, + context: Optional[str] = None, + reward_model: Optional[str] = None, + debug: bool = False, + wandb_sweep: bool = False, + ) -> TrainingReturn: + """ + Overview: + Train the agent with PPO algorithm for ``step`` iterations with ``collector_env_num`` collector \ + environments and ``evaluator_env_num`` evaluator environments. Information during training will be \ + recorded and saved by wandb. + Arguments: + - step (:obj:`int`): The total training environment steps of all collector environments. Default to 1e7. + - collector_env_num (:obj:`int`): The number of collector environments. Default to 4. + - evaluator_env_num (:obj:`int`): The number of evaluator environments. Default to 4. + - n_iter_log_show (:obj:`int`): The frequency of logging every training iteration. Default to 500. + - n_iter_save_ckpt (:obj:`int`): The frequency of saving checkpoint every training iteration. \ + Default to 1000. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - reward_model (:obj:`str`): The reward model name. Default to None. This argument is not supported yet. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + - wandb_sweep (:obj:`bool`): Whether to use wandb sweep, \ + which is a hyper-parameter optimization process for seeking the best configurations. \ + Default to False. If True, the wandb sweep id will be used as the experiment name. + Returns: + - (:obj:`TrainingReturn`): The training result, of which the attributions are: + - wandb_url (:obj:`str`): The weight & biases (wandb) project url of the trainning experiment. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + logging.debug(self.policy._model) + # define env and policy + collector_env = self._setup_env_manager(collector_env_num, context, debug, 'collector') + evaluator_env = self._setup_env_manager(evaluator_env_num, context, debug, 'evaluator') + + if reward_model is not None: + # self.reward_model = create_reward_model(reward_model, self.cfg.reward_model) + pass + + with task.start(ctx=OnlineRLContext()): + task.use(interaction_evaluator_ttorch(self.seed, self.policy, evaluator_env)) + task.use(CkptSaver(self.policy, save_dir=self.checkpoint_save_dir, train_freq=n_iter_save_ckpt)) + task.use(PPOFStepCollector(self.seed, self.policy, collector_env, self.cfg.n_sample)) + task.use(ppof_adv_estimator(self.policy)) + task.use(multistep_trainer(self.policy, log_freq=n_iter_log_show)) + task.use( + wandb_online_logger( + metric_list=self.policy.monitor_vars(), + model=self.policy._model, + anonymous=True, + project_name=self.exp_name, + wandb_sweep=wandb_sweep, + ) + ) + task.use(termination_checker(max_env_step=step)) + task.run() + + return TrainingReturn(wandb_url=task.ctx.wandb_url) + + def deploy( + self, + enable_save_replay: bool = False, + concatenate_all_replay: bool = False, + replay_save_path: str = None, + seed: Optional[Union[int, List]] = None, + debug: bool = False + ) -> EvalReturn: + """ + Overview: + Deploy the agent with PPO algorithm by interacting with the environment, during which the replay video \ + can be saved if ``enable_save_replay`` is True. The evaluation result will be returned. + Arguments: + - enable_save_replay (:obj:`bool`): Whether to save the replay video. Default to False. + - concatenate_all_replay (:obj:`bool`): Whether to concatenate all replay videos into one video. \ + Default to False. If ``enable_save_replay`` is False, this argument will be ignored. \ + If ``enable_save_replay`` is True and ``concatenate_all_replay`` is False, \ + the replay video of each episode will be saved separately. + - replay_save_path (:obj:`str`): The path to save the replay video. Default to None. \ + If not specified, the video will be saved in ``exp_name/videos``. + - seed (:obj:`Union[int, List]`): The random seed, which is set before running the program. \ + Default to None. If not specified, ``self.seed`` will be used. \ + If ``seed`` is an integer, the agent will be deployed once. \ + If ``seed`` is a list of integers, the agent will be deployed once for each seed in the list. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + Returns: + - (:obj:`EvalReturn`): The evaluation result, of which the attributions are: + - eval_value (:obj:`np.float32`): The mean of evaluation return. + - eval_value_std (:obj:`np.float32`): The standard deviation of evaluation return. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + # define env and policy + env = self.env.clone(caller='evaluator') + + if seed is not None and isinstance(seed, int): + seeds = [seed] + elif seed is not None and isinstance(seed, list): + seeds = seed + else: + seeds = [self.seed] + + returns = [] + images = [] + if enable_save_replay: + replay_save_path = os.path.join(self.exp_name, 'videos') if replay_save_path is None else replay_save_path + env.enable_save_replay(replay_path=replay_save_path) + else: + logging.warning('No video would be generated during the deploy.') + if concatenate_all_replay: + logging.warning('concatenate_all_replay is set to False because enable_save_replay is False.') + concatenate_all_replay = False + + forward_fn = single_env_forward_wrapper_ttorch(self.policy.eval, self.cfg.cuda) + + # reset first to make sure the env is in the initial state + # env will be reset again in the main loop + env.reset() + + for seed in seeds: + env.seed(seed, dynamic_seed=False) + return_ = 0. + step = 0 + obs = env.reset() + images.append(render(env)[None]) if concatenate_all_replay else None + while True: + action = forward_fn(obs) + obs, rew, done, info = env.step(action) + images.append(render(env)[None]) if concatenate_all_replay else None + return_ += rew + step += 1 + if done: + break + logging.info(f'DQN deploy is finished, final episode return with {step} steps is: {return_}') + returns.append(return_) + + env.close() + + if concatenate_all_replay: + images = np.concatenate(images, axis=0) + import imageio + imageio.mimwrite(os.path.join(replay_save_path, 'deploy.mp4'), images, fps=get_env_fps(env)) + + return EvalReturn(eval_value=np.mean(returns), eval_value_std=np.std(returns)) + + def collect_data( + self, + env_num: int = 8, + save_data_path: Optional[str] = None, + n_sample: Optional[int] = None, + n_episode: Optional[int] = None, + context: Optional[str] = None, + debug: bool = False + ) -> None: + """ + Overview: + Collect data with PPO algorithm for ``n_episode`` episodes with ``env_num`` collector environments. \ + The collected data will be saved in ``save_data_path`` if specified, otherwise it will be saved in \ + ``exp_name/demo_data``. + Arguments: + - env_num (:obj:`int`): The number of collector environments. Default to 8. + - save_data_path (:obj:`str`): The path to save the collected data. Default to None. \ + If not specified, the data will be saved in ``exp_name/demo_data``. + - n_sample (:obj:`int`): The number of samples to collect. Default to None. \ + If not specified, ``n_episode`` must be specified. + - n_episode (:obj:`int`): The number of episodes to collect. Default to None. \ + If not specified, ``n_sample`` must be specified. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + if n_episode is not None: + raise NotImplementedError + # define env and policy + env = self._setup_env_manager(env_num, context, debug, 'collector') + if save_data_path is None: + save_data_path = os.path.join(self.exp_name, 'demo_data') + + # main execution task + with task.start(ctx=OnlineRLContext()): + task.use(PPOFStepCollector(self.seed, self.policy, env, n_sample)) + task.use(offline_data_saver(save_data_path, data_type='hdf5')) + task.run(max_step=1) + logging.info( + f'PPOF collecting is finished, more than {n_sample} samples are collected and saved in `{save_data_path}`' + ) + + def batch_evaluate( + self, + env_num: int = 4, + n_evaluator_episode: int = 4, + context: Optional[str] = None, + debug: bool = False, + ) -> EvalReturn: + """ + Overview: + Evaluate the agent with PPO algorithm for ``n_evaluator_episode`` episodes with ``env_num`` evaluator \ + environments. The evaluation result will be returned. + The difference between methods ``batch_evaluate`` and ``deploy`` is that ``batch_evaluate`` will create \ + multiple evaluator environments to evaluate the agent to get an average performance, while ``deploy`` \ + will only create one evaluator environment to evaluate the agent and save the replay video. + Arguments: + - env_num (:obj:`int`): The number of evaluator environments. Default to 4. + - n_evaluator_episode (:obj:`int`): The number of episodes to evaluate. Default to 4. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + Returns: + - (:obj:`EvalReturn`): The evaluation result, of which the attributions are: + - eval_value (:obj:`np.float32`): The mean of evaluation return. + - eval_value_std (:obj:`np.float32`): The standard deviation of evaluation return. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + # define env and policy + env = self._setup_env_manager(env_num, context, debug, 'evaluator') + + # reset first to make sure the env is in the initial state + # env will be reset again in the main loop + env.launch() + env.reset() + + # main execution task + with task.start(ctx=OnlineRLContext()): + task.use(interaction_evaluator_ttorch( + self.seed, + self.policy, + env, + n_evaluator_episode, + )) + task.run(max_step=1) + + return EvalReturn(eval_value=task.ctx.eval_value, eval_value_std=task.ctx.eval_value_std) + + def _setup_env_manager( + self, + env_num: int, + context: Optional[str] = None, + debug: bool = False, + caller: str = 'collector' + ) -> BaseEnvManagerV2: + """ + Overview: + Setup the environment manager. The environment manager is used to manage multiple environments. + Arguments: + - env_num (:obj:`int`): The number of environments. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + - caller (:obj:`str`): The caller of the environment manager. Default to 'collector'. + Returns: + - (:obj:`BaseEnvManagerV2`): The environment manager. + """ + assert caller in ['evaluator', 'collector'] + if debug: + env_cls = BaseEnvManagerV2 + manager_cfg = env_cls.default_config() + else: + env_cls = SubprocessEnvManagerV2 + manager_cfg = env_cls.default_config() + if context is not None: + manager_cfg.context = context + return env_cls([partial(self.env.clone, caller) for _ in range(env_num)], manager_cfg) + + @property + def best(self) -> 'PPOF': + """ + Overview: + Load the best model from the checkpoint directory, \ + which is by default in folder ``exp_name/ckpt/eval.pth.tar``. \ + The return value is the agent with the best model. + Returns: + - (:obj:`PPOF`): The agent with the best model. + Examples: + >>> agent = PPOF(env_id='LunarLander-v2') + >>> agent.train() + >>> agent = agent.best() + + .. note:: + The best model is the model with the highest evaluation return. If this method is called, the current \ + model will be replaced by the best model. + """ + + best_model_file_path = os.path.join(self.checkpoint_save_dir, "eval.pth.tar") + # Load best model if it exists + if os.path.exists(best_model_file_path): + policy_state_dict = torch.load(best_model_file_path, map_location=torch.device("cpu")) + self.policy.learn_mode.load_state_dict(policy_state_dict) + return self diff --git a/ding/bonus/sac.py b/ding/bonus/sac.py new file mode 100644 index 0000000000..f635d5db37 --- /dev/null +++ b/ding/bonus/sac.py @@ -0,0 +1,457 @@ +from typing import Optional, Union, List +from ditk import logging +from easydict import EasyDict +import os +import numpy as np +import torch +import treetensor.torch as ttorch +from ding.framework import task, OnlineRLContext +from ding.framework.middleware import CkptSaver, \ + wandb_online_logger, offline_data_saver, termination_checker, interaction_evaluator, StepCollector, data_pusher, \ + OffPolicyLearner, final_ctx_saver +from ding.envs import BaseEnv +from ding.envs import setup_ding_env_manager +from ding.policy import SACPolicy +from ding.utils import set_pkg_seed +from ding.utils import get_env_fps, render +from ding.config import save_config_py, compile_config +from ding.model import ContinuousQAC +from ding.model import model_wrap +from ding.data import DequeBuffer +from ding.bonus.common import TrainingReturn, EvalReturn +from ding.config.example.SAC import supported_env_cfg +from ding.config.example.SAC import supported_env + + +class SACAgent: + """ + Overview: + Class of agent for training, evaluation and deployment of Reinforcement learning algorithm \ + Soft Actor-Critic(SAC). + For more information about the system design of RL agent, please refer to \ + . + Interface: + ``__init__``, ``train``, ``deploy``, ``collect_data``, ``batch_evaluate``, ``best`` + """ + supported_env_list = list(supported_env_cfg.keys()) + """ + Overview: + List of supported envs. + Examples: + >>> from ding.bonus.sac import SACAgent + >>> print(SACAgent.supported_env_list) + """ + + def __init__( + self, + env_id: str = None, + env: BaseEnv = None, + seed: int = 0, + exp_name: str = None, + model: Optional[torch.nn.Module] = None, + cfg: Optional[Union[EasyDict, dict]] = None, + policy_state_dict: str = None, + ) -> None: + """ + Overview: + Initialize agent for SAC algorithm. + Arguments: + - env_id (:obj:`str`): The environment id, which is a registered environment name in gym or gymnasium. \ + If ``env_id`` is not specified, ``env_id`` in ``cfg.env`` must be specified. \ + If ``env_id`` is specified, ``env_id`` in ``cfg.env`` will be ignored. \ + ``env_id`` should be one of the supported envs, which can be found in ``supported_env_list``. + - env (:obj:`BaseEnv`): The environment instance for training and evaluation. \ + If ``env`` is not specified, `env_id`` or ``cfg.env.env_id`` must be specified. \ + ``env_id`` or ``cfg.env.env_id`` will be used to create environment instance. \ + If ``env`` is specified, ``env_id`` and ``cfg.env.env_id`` will be ignored. + - seed (:obj:`int`): The random seed, which is set before running the program. \ + Default to 0. + - exp_name (:obj:`str`): The name of this experiment, which will be used to create the folder to save \ + log data. Default to None. If not specified, the folder name will be ``env_id``-``algorithm``. + - model (:obj:`torch.nn.Module`): The model of SAC algorithm, which should be an instance of class \ + :class:`ding.model.ContinuousQAC`. \ + If not specified, a default model will be generated according to the configuration. + - cfg (:obj:`Union[EasyDict, dict]`): The configuration of SAC algorithm, which is a dict. \ + Default to None. If not specified, the default configuration will be used. \ + The default configuration can be found in ``ding/config/example/SAC/gym_lunarlander_v2.py``. + - policy_state_dict (:obj:`str`): The path of policy state dict saved by PyTorch a in local file. \ + If specified, the policy will be loaded from this file. Default to None. + + .. note:: + An RL Agent Instance can be initialized in two basic ways. \ + For example, we have an environment with id ``LunarLanderContinuous-v2`` registered in gym, \ + and we want to train an agent with SAC algorithm with default configuration. \ + Then we can initialize the agent in the following ways: + >>> agent = SACAgent(env_id='LunarLanderContinuous-v2') + or, if we want can specify the env_id in the configuration: + >>> cfg = {'env': {'env_id': 'LunarLanderContinuous-v2'}, 'policy': ...... } + >>> agent = SACAgent(cfg=cfg) + There are also other arguments to specify the agent when initializing. + For example, if we want to specify the environment instance: + >>> env = CustomizedEnv('LunarLanderContinuous-v2') + >>> agent = SACAgent(cfg=cfg, env=env) + or, if we want to specify the model: + >>> model = ContinuousQAC(**cfg.policy.model) + >>> agent = SACAgent(cfg=cfg, model=model) + or, if we want to reload the policy from a saved policy state dict: + >>> agent = SACAgent(cfg=cfg, policy_state_dict='LunarLanderContinuous-v2.pth.tar') + Make sure that the configuration is consistent with the saved policy state dict. + """ + + assert env_id is not None or cfg is not None, "Please specify env_id or cfg." + + if cfg is not None and not isinstance(cfg, EasyDict): + cfg = EasyDict(cfg) + + if env_id is not None: + assert env_id in SACAgent.supported_env_list, "Please use supported envs: {}".format( + SACAgent.supported_env_list + ) + if cfg is None: + cfg = supported_env_cfg[env_id] + else: + assert cfg.env.env_id == env_id, "env_id in cfg should be the same as env_id in args." + else: + assert hasattr(cfg.env, "env_id"), "Please specify env_id in cfg." + assert cfg.env.env_id in SACAgent.supported_env_list, "Please use supported envs: {}".format( + SACAgent.supported_env_list + ) + default_policy_config = EasyDict({"policy": SACPolicy.default_config()}) + default_policy_config.update(cfg) + cfg = default_policy_config + + if exp_name is not None: + cfg.exp_name = exp_name + self.cfg = compile_config(cfg, policy=SACPolicy) + self.exp_name = self.cfg.exp_name + if env is None: + self.env = supported_env[cfg.env.env_id](cfg=cfg.env) + else: + assert isinstance(env, BaseEnv), "Please use BaseEnv as env data type." + self.env = env + + logging.getLogger().setLevel(logging.INFO) + self.seed = seed + set_pkg_seed(self.seed, use_cuda=self.cfg.policy.cuda) + if not os.path.exists(self.exp_name): + os.makedirs(self.exp_name) + save_config_py(self.cfg, os.path.join(self.exp_name, 'policy_config.py')) + if model is None: + model = ContinuousQAC(**self.cfg.policy.model) + self.buffer_ = DequeBuffer(size=self.cfg.policy.other.replay_buffer.replay_buffer_size) + self.policy = SACPolicy(self.cfg.policy, model=model) + if policy_state_dict is not None: + self.policy.learn_mode.load_state_dict(policy_state_dict) + self.checkpoint_save_dir = os.path.join(self.exp_name, "ckpt") + + def train( + self, + step: int = int(1e7), + collector_env_num: int = None, + evaluator_env_num: int = None, + n_iter_save_ckpt: int = 1000, + context: Optional[str] = None, + debug: bool = False, + wandb_sweep: bool = False, + ) -> TrainingReturn: + """ + Overview: + Train the agent with SAC algorithm for ``step`` iterations with ``collector_env_num`` collector \ + environments and ``evaluator_env_num`` evaluator environments. Information during training will be \ + recorded and saved by wandb. + Arguments: + - step (:obj:`int`): The total training environment steps of all collector environments. Default to 1e7. + - collector_env_num (:obj:`int`): The collector environment number. Default to None. \ + If not specified, it will be set according to the configuration. + - evaluator_env_num (:obj:`int`): The evaluator environment number. Default to None. \ + If not specified, it will be set according to the configuration. + - n_iter_save_ckpt (:obj:`int`): The frequency of saving checkpoint every training iteration. \ + Default to 1000. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + - wandb_sweep (:obj:`bool`): Whether to use wandb sweep, \ + which is a hyper-parameter optimization process for seeking the best configurations. \ + Default to False. If True, the wandb sweep id will be used as the experiment name. + Returns: + - (:obj:`TrainingReturn`): The training result, of which the attributions are: + - wandb_url (:obj:`str`): The weight & biases (wandb) project url of the trainning experiment. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + logging.debug(self.policy._model) + # define env and policy + collector_env_num = collector_env_num if collector_env_num else self.cfg.env.collector_env_num + evaluator_env_num = evaluator_env_num if evaluator_env_num else self.cfg.env.evaluator_env_num + collector_env = setup_ding_env_manager(self.env, collector_env_num, context, debug, 'collector') + evaluator_env = setup_ding_env_manager(self.env, evaluator_env_num, context, debug, 'evaluator') + + with task.start(ctx=OnlineRLContext()): + task.use( + interaction_evaluator( + self.cfg, + self.policy.eval_mode, + evaluator_env, + render=self.cfg.policy.eval.render if hasattr(self.cfg.policy.eval, "render") else False + ) + ) + task.use(CkptSaver(policy=self.policy, save_dir=self.checkpoint_save_dir, train_freq=n_iter_save_ckpt)) + task.use( + StepCollector( + self.cfg, + self.policy.collect_mode, + collector_env, + random_collect_size=self.cfg.policy.random_collect_size + if hasattr(self.cfg.policy, 'random_collect_size') else 0, + ) + ) + task.use(data_pusher(self.cfg, self.buffer_)) + task.use(OffPolicyLearner(self.cfg, self.policy.learn_mode, self.buffer_)) + task.use( + wandb_online_logger( + metric_list=self.policy._monitor_vars_learn(), + model=self.policy._model, + anonymous=True, + project_name=self.exp_name, + wandb_sweep=wandb_sweep, + ) + ) + task.use(termination_checker(max_env_step=step)) + task.use(final_ctx_saver(name=self.exp_name)) + task.run() + + return TrainingReturn(wandb_url=task.ctx.wandb_url) + + def deploy( + self, + enable_save_replay: bool = False, + concatenate_all_replay: bool = False, + replay_save_path: str = None, + seed: Optional[Union[int, List]] = None, + debug: bool = False + ) -> EvalReturn: + """ + Overview: + Deploy the agent with SAC algorithm by interacting with the environment, during which the replay video \ + can be saved if ``enable_save_replay`` is True. The evaluation result will be returned. + Arguments: + - enable_save_replay (:obj:`bool`): Whether to save the replay video. Default to False. + - concatenate_all_replay (:obj:`bool`): Whether to concatenate all replay videos into one video. \ + Default to False. If ``enable_save_replay`` is False, this argument will be ignored. \ + If ``enable_save_replay`` is True and ``concatenate_all_replay`` is False, \ + the replay video of each episode will be saved separately. + - replay_save_path (:obj:`str`): The path to save the replay video. Default to None. \ + If not specified, the video will be saved in ``exp_name/videos``. + - seed (:obj:`Union[int, List]`): The random seed, which is set before running the program. \ + Default to None. If not specified, ``self.seed`` will be used. \ + If ``seed`` is an integer, the agent will be deployed once. \ + If ``seed`` is a list of integers, the agent will be deployed once for each seed in the list. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + Returns: + - (:obj:`EvalReturn`): The evaluation result, of which the attributions are: + - eval_value (:obj:`np.float32`): The mean of evaluation return. + - eval_value_std (:obj:`np.float32`): The standard deviation of evaluation return. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + # define env and policy + env = self.env.clone(caller='evaluator') + + if seed is not None and isinstance(seed, int): + seeds = [seed] + elif seed is not None and isinstance(seed, list): + seeds = seed + else: + seeds = [self.seed] + + returns = [] + images = [] + if enable_save_replay: + replay_save_path = os.path.join(self.exp_name, 'videos') if replay_save_path is None else replay_save_path + env.enable_save_replay(replay_path=replay_save_path) + else: + logging.warning('No video would be generated during the deploy.') + if concatenate_all_replay: + logging.warning('concatenate_all_replay is set to False because enable_save_replay is False.') + concatenate_all_replay = False + + def single_env_forward_wrapper(forward_fn, cuda=True): + + forward_fn = model_wrap(forward_fn, wrapper_name='base').forward + + def _forward(obs): + # unsqueeze means add batch dim, i.e. (O, ) -> (1, O) + obs = ttorch.as_tensor(obs).unsqueeze(0) + if cuda and torch.cuda.is_available(): + obs = obs.cuda() + (mu, sigma) = forward_fn(obs, mode='compute_actor')['logit'] + action = torch.tanh(mu).detach().cpu().numpy()[0] # deterministic_eval + return action + + return _forward + + forward_fn = single_env_forward_wrapper(self.policy._model, self.cfg.policy.cuda) + + # reset first to make sure the env is in the initial state + # env will be reset again in the main loop + env.reset() + + for seed in seeds: + env.seed(seed, dynamic_seed=False) + return_ = 0. + step = 0 + obs = env.reset() + images.append(render(env)[None]) if concatenate_all_replay else None + while True: + action = forward_fn(obs) + obs, rew, done, info = env.step(action) + images.append(render(env)[None]) if concatenate_all_replay else None + return_ += rew + step += 1 + if done: + break + logging.info(f'DQN deploy is finished, final episode return with {step} steps is: {return_}') + returns.append(return_) + + env.close() + + if concatenate_all_replay: + images = np.concatenate(images, axis=0) + import imageio + imageio.mimwrite(os.path.join(replay_save_path, 'deploy.mp4'), images, fps=get_env_fps(env)) + + return EvalReturn(eval_value=np.mean(returns), eval_value_std=np.std(returns)) + + def collect_data( + self, + env_num: int = 8, + save_data_path: Optional[str] = None, + n_sample: Optional[int] = None, + n_episode: Optional[int] = None, + context: Optional[str] = None, + debug: bool = False + ) -> None: + """ + Overview: + Collect data with SAC algorithm for ``n_episode`` episodes with ``env_num`` collector environments. \ + The collected data will be saved in ``save_data_path`` if specified, otherwise it will be saved in \ + ``exp_name/demo_data``. + Arguments: + - env_num (:obj:`int`): The number of collector environments. Default to 8. + - save_data_path (:obj:`str`): The path to save the collected data. Default to None. \ + If not specified, the data will be saved in ``exp_name/demo_data``. + - n_sample (:obj:`int`): The number of samples to collect. Default to None. \ + If not specified, ``n_episode`` must be specified. + - n_episode (:obj:`int`): The number of episodes to collect. Default to None. \ + If not specified, ``n_sample`` must be specified. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + if n_episode is not None: + raise NotImplementedError + # define env and policy + env_num = env_num if env_num else self.cfg.env.collector_env_num + env = setup_ding_env_manager(self.env, env_num, context, debug, 'collector') + + if save_data_path is None: + save_data_path = os.path.join(self.exp_name, 'demo_data') + + # main execution task + with task.start(ctx=OnlineRLContext()): + task.use( + StepCollector( + self.cfg, self.policy.collect_mode, env, random_collect_size=self.cfg.policy.random_collect_size + ) + ) + task.use(offline_data_saver(save_data_path, data_type='hdf5')) + task.run(max_step=1) + logging.info( + f'SAC collecting is finished, more than {n_sample} samples are collected and saved in `{save_data_path}`' + ) + + def batch_evaluate( + self, + env_num: int = 4, + n_evaluator_episode: int = 4, + context: Optional[str] = None, + debug: bool = False + ) -> EvalReturn: + """ + Overview: + Evaluate the agent with SAC algorithm for ``n_evaluator_episode`` episodes with ``env_num`` evaluator \ + environments. The evaluation result will be returned. + The difference between methods ``batch_evaluate`` and ``deploy`` is that ``batch_evaluate`` will create \ + multiple evaluator environments to evaluate the agent to get an average performance, while ``deploy`` \ + will only create one evaluator environment to evaluate the agent and save the replay video. + Arguments: + - env_num (:obj:`int`): The number of evaluator environments. Default to 4. + - n_evaluator_episode (:obj:`int`): The number of episodes to evaluate. Default to 4. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + Returns: + - (:obj:`EvalReturn`): The evaluation result, of which the attributions are: + - eval_value (:obj:`np.float32`): The mean of evaluation return. + - eval_value_std (:obj:`np.float32`): The standard deviation of evaluation return. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + # define env and policy + env_num = env_num if env_num else self.cfg.env.evaluator_env_num + env = setup_ding_env_manager(self.env, env_num, context, debug, 'evaluator') + + # reset first to make sure the env is in the initial state + # env will be reset again in the main loop + env.launch() + env.reset() + + evaluate_cfg = self.cfg + evaluate_cfg.env.n_evaluator_episode = n_evaluator_episode + + # main execution task + with task.start(ctx=OnlineRLContext()): + task.use(interaction_evaluator(self.cfg, self.policy.eval_mode, env)) + task.run(max_step=1) + + return EvalReturn(eval_value=task.ctx.eval_value, eval_value_std=task.ctx.eval_value_std) + + @property + def best(self) -> 'SACAgent': + """ + Overview: + Load the best model from the checkpoint directory, \ + which is by default in folder ``exp_name/ckpt/eval.pth.tar``. \ + The return value is the agent with the best model. + Returns: + - (:obj:`SACAgent`): The agent with the best model. + Examples: + >>> agent = SACAgent(env_id='LunarLanderContinuous-v2') + >>> agent.train() + >>> agent = agent.best + + .. note:: + The best model is the model with the highest evaluation return. If this method is called, the current \ + model will be replaced by the best model. + """ + + best_model_file_path = os.path.join(self.checkpoint_save_dir, "eval.pth.tar") + # Load best model if it exists + if os.path.exists(best_model_file_path): + policy_state_dict = torch.load(best_model_file_path, map_location=torch.device("cpu")) + self.policy.learn_mode.load_state_dict(policy_state_dict) + return self diff --git a/ding/bonus/sql.py b/ding/bonus/sql.py new file mode 100644 index 0000000000..63d26acce2 --- /dev/null +++ b/ding/bonus/sql.py @@ -0,0 +1,461 @@ +from typing import Optional, Union, List +from ditk import logging +from easydict import EasyDict +import os +import numpy as np +import torch +import treetensor.torch as ttorch +from ding.framework import task, OnlineRLContext +from ding.framework.middleware import CkptSaver, \ + wandb_online_logger, offline_data_saver, termination_checker, interaction_evaluator, StepCollector, data_pusher, \ + OffPolicyLearner, final_ctx_saver, nstep_reward_enhancer, eps_greedy_handler +from ding.envs import BaseEnv +from ding.envs import setup_ding_env_manager +from ding.policy import SQLPolicy +from ding.utils import set_pkg_seed +from ding.utils import get_env_fps, render +from ding.config import save_config_py, compile_config +from ding.model import DQN +from ding.model import model_wrap +from ding.data import DequeBuffer +from ding.bonus.common import TrainingReturn, EvalReturn +from ding.config.example.SQL import supported_env_cfg +from ding.config.example.SQL import supported_env + + +class SQLAgent: + """ + Overview: + Class of agent for training, evaluation and deployment of Reinforcement learning algorithm \ + Soft Q-Learning(SQL). + For more information about the system design of RL agent, please refer to \ + . + Interface: + ``__init__``, ``train``, ``deploy``, ``collect_data``, ``batch_evaluate``, ``best`` + """ + supported_env_list = list(supported_env_cfg.keys()) + """ + Overview: + List of supported envs. + Examples: + >>> from ding.bonus.sql import SQLAgent + >>> print(SQLAgent.supported_env_list) + """ + + def __init__( + self, + env_id: str = None, + env: BaseEnv = None, + seed: int = 0, + exp_name: str = None, + model: Optional[torch.nn.Module] = None, + cfg: Optional[Union[EasyDict, dict]] = None, + policy_state_dict: str = None, + ) -> None: + """ + Overview: + Initialize agent for SQL algorithm. + Arguments: + - env_id (:obj:`str`): The environment id, which is a registered environment name in gym or gymnasium. \ + If ``env_id`` is not specified, ``env_id`` in ``cfg.env`` must be specified. \ + If ``env_id`` is specified, ``env_id`` in ``cfg.env`` will be ignored. \ + ``env_id`` should be one of the supported envs, which can be found in ``supported_env_list``. + - env (:obj:`BaseEnv`): The environment instance for training and evaluation. \ + If ``env`` is not specified, `env_id`` or ``cfg.env.env_id`` must be specified. \ + ``env_id`` or ``cfg.env.env_id`` will be used to create environment instance. \ + If ``env`` is specified, ``env_id`` and ``cfg.env.env_id`` will be ignored. + - seed (:obj:`int`): The random seed, which is set before running the program. \ + Default to 0. + - exp_name (:obj:`str`): The name of this experiment, which will be used to create the folder to save \ + log data. Default to None. If not specified, the folder name will be ``env_id``-``algorithm``. + - model (:obj:`torch.nn.Module`): The model of SQL algorithm, which should be an instance of class \ + :class:`ding.model.DQN`. \ + If not specified, a default model will be generated according to the configuration. + - cfg (:obj:Union[EasyDict, dict]): The configuration of SQL algorithm, which is a dict. \ + Default to None. If not specified, the default configuration will be used. \ + The default configuration can be found in ``ding/config/example/SQL/gym_lunarlander_v2.py``. + - policy_state_dict (:obj:`str`): The path of policy state dict saved by PyTorch a in local file. \ + If specified, the policy will be loaded from this file. Default to None. + + .. note:: + An RL Agent Instance can be initialized in two basic ways. \ + For example, we have an environment with id ``LunarLander-v2`` registered in gym, \ + and we want to train an agent with SQL algorithm with default configuration. \ + Then we can initialize the agent in the following ways: + >>> agent = SQLAgent(env_id='LunarLander-v2') + or, if we want can specify the env_id in the configuration: + >>> cfg = {'env': {'env_id': 'LunarLander-v2'}, 'policy': ...... } + >>> agent = SQLAgent(cfg=cfg) + There are also other arguments to specify the agent when initializing. + For example, if we want to specify the environment instance: + >>> env = CustomizedEnv('LunarLander-v2') + >>> agent = SQLAgent(cfg=cfg, env=env) + or, if we want to specify the model: + >>> model = DQN(**cfg.policy.model) + >>> agent = SQLAgent(cfg=cfg, model=model) + or, if we want to reload the policy from a saved policy state dict: + >>> agent = SQLAgent(cfg=cfg, policy_state_dict='LunarLander-v2.pth.tar') + Make sure that the configuration is consistent with the saved policy state dict. + """ + + assert env_id is not None or cfg is not None, "Please specify env_id or cfg." + + if cfg is not None and not isinstance(cfg, EasyDict): + cfg = EasyDict(cfg) + + if env_id is not None: + assert env_id in SQLAgent.supported_env_list, "Please use supported envs: {}".format( + SQLAgent.supported_env_list + ) + if cfg is None: + cfg = supported_env_cfg[env_id] + else: + assert cfg.env.env_id == env_id, "env_id in cfg should be the same as env_id in args." + else: + assert hasattr(cfg.env, "env_id"), "Please specify env_id in cfg." + assert cfg.env.env_id in SQLAgent.supported_env_list, "Please use supported envs: {}".format( + SQLAgent.supported_env_list + ) + default_policy_config = EasyDict({"policy": SQLPolicy.default_config()}) + default_policy_config.update(cfg) + cfg = default_policy_config + + if exp_name is not None: + cfg.exp_name = exp_name + self.cfg = compile_config(cfg, policy=SQLPolicy) + self.exp_name = self.cfg.exp_name + if env is None: + self.env = supported_env[cfg.env.env_id](cfg=cfg.env) + else: + assert isinstance(env, BaseEnv), "Please use BaseEnv as env data type." + self.env = env + + logging.getLogger().setLevel(logging.INFO) + self.seed = seed + set_pkg_seed(self.seed, use_cuda=self.cfg.policy.cuda) + if not os.path.exists(self.exp_name): + os.makedirs(self.exp_name) + save_config_py(self.cfg, os.path.join(self.exp_name, 'policy_config.py')) + if model is None: + model = DQN(**self.cfg.policy.model) + self.buffer_ = DequeBuffer(size=self.cfg.policy.other.replay_buffer.replay_buffer_size) + self.policy = SQLPolicy(self.cfg.policy, model=model) + if policy_state_dict is not None: + self.policy.learn_mode.load_state_dict(policy_state_dict) + self.checkpoint_save_dir = os.path.join(self.exp_name, "ckpt") + + def train( + self, + step: int = int(1e7), + collector_env_num: int = None, + evaluator_env_num: int = None, + n_iter_save_ckpt: int = 1000, + context: Optional[str] = None, + debug: bool = False, + wandb_sweep: bool = False, + ) -> TrainingReturn: + """ + Overview: + Train the agent with SQL algorithm for ``step`` iterations with ``collector_env_num`` collector \ + environments and ``evaluator_env_num`` evaluator environments. Information during training will be \ + recorded and saved by wandb. + Arguments: + - step (:obj:`int`): The total training environment steps of all collector environments. Default to 1e7. + - collector_env_num (:obj:`int`): The collector environment number. Default to None. \ + If not specified, it will be set according to the configuration. + - evaluator_env_num (:obj:`int`): The evaluator environment number. Default to None. \ + If not specified, it will be set according to the configuration. + - n_iter_save_ckpt (:obj:`int`): The frequency of saving checkpoint every training iteration. \ + Default to 1000. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + - wandb_sweep (:obj:`bool`): Whether to use wandb sweep, \ + which is a hyper-parameter optimization process for seeking the best configurations. \ + Default to False. If True, the wandb sweep id will be used as the experiment name. + Returns: + - (:obj:`TrainingReturn`): The training result, of which the attributions are: + - wandb_url (:obj:`str`): The weight & biases (wandb) project url of the trainning experiment. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + logging.debug(self.policy._model) + # define env and policy + collector_env_num = collector_env_num if collector_env_num else self.cfg.env.collector_env_num + evaluator_env_num = evaluator_env_num if evaluator_env_num else self.cfg.env.evaluator_env_num + collector_env = setup_ding_env_manager(self.env, collector_env_num, context, debug, 'collector') + evaluator_env = setup_ding_env_manager(self.env, evaluator_env_num, context, debug, 'evaluator') + + with task.start(ctx=OnlineRLContext()): + task.use( + interaction_evaluator( + self.cfg, + self.policy.eval_mode, + evaluator_env, + render=self.cfg.policy.eval.render if hasattr(self.cfg.policy.eval, "render") else False + ) + ) + task.use(CkptSaver(policy=self.policy, save_dir=self.checkpoint_save_dir, train_freq=n_iter_save_ckpt)) + task.use(eps_greedy_handler(self.cfg)) + task.use( + StepCollector( + self.cfg, + self.policy.collect_mode, + collector_env, + random_collect_size=self.cfg.policy.random_collect_size + if hasattr(self.cfg.policy, 'random_collect_size') else 0, + ) + ) + if "nstep" in self.cfg.policy and self.cfg.policy.nstep > 1: + task.use(nstep_reward_enhancer(self.cfg)) + task.use(data_pusher(self.cfg, self.buffer_)) + task.use(OffPolicyLearner(self.cfg, self.policy.learn_mode, self.buffer_)) + task.use( + wandb_online_logger( + metric_list=self.policy._monitor_vars_learn(), + model=self.policy._model, + anonymous=True, + project_name=self.exp_name, + wandb_sweep=wandb_sweep, + ) + ) + task.use(termination_checker(max_env_step=step)) + task.use(final_ctx_saver(name=self.exp_name)) + task.run() + + return TrainingReturn(wandb_url=task.ctx.wandb_url) + + def deploy( + self, + enable_save_replay: bool = False, + concatenate_all_replay: bool = False, + replay_save_path: str = None, + seed: Optional[Union[int, List]] = None, + debug: bool = False + ) -> EvalReturn: + """ + Overview: + Deploy the agent with SQL algorithm by interacting with the environment, during which the replay video \ + can be saved if ``enable_save_replay`` is True. The evaluation result will be returned. + Arguments: + - enable_save_replay (:obj:`bool`): Whether to save the replay video. Default to False. + - concatenate_all_replay (:obj:`bool`): Whether to concatenate all replay videos into one video. \ + Default to False. If ``enable_save_replay`` is False, this argument will be ignored. \ + If ``enable_save_replay`` is True and ``concatenate_all_replay`` is False, \ + the replay video of each episode will be saved separately. + - replay_save_path (:obj:`str`): The path to save the replay video. Default to None. \ + If not specified, the video will be saved in ``exp_name/videos``. + - seed (:obj:`Union[int, List]`): The random seed, which is set before running the program. \ + Default to None. If not specified, ``self.seed`` will be used. \ + If ``seed`` is an integer, the agent will be deployed once. \ + If ``seed`` is a list of integers, the agent will be deployed once for each seed in the list. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + Returns: + - (:obj:`EvalReturn`): The evaluation result, of which the attributions are: + - eval_value (:obj:`np.float32`): The mean of evaluation return. + - eval_value_std (:obj:`np.float32`): The standard deviation of evaluation return. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + # define env and policy + env = self.env.clone(caller='evaluator') + + if seed is not None and isinstance(seed, int): + seeds = [seed] + elif seed is not None and isinstance(seed, list): + seeds = seed + else: + seeds = [self.seed] + + returns = [] + images = [] + if enable_save_replay: + replay_save_path = os.path.join(self.exp_name, 'videos') if replay_save_path is None else replay_save_path + env.enable_save_replay(replay_path=replay_save_path) + else: + logging.warning('No video would be generated during the deploy.') + if concatenate_all_replay: + logging.warning('concatenate_all_replay is set to False because enable_save_replay is False.') + concatenate_all_replay = False + + def single_env_forward_wrapper(forward_fn, cuda=True): + + forward_fn = model_wrap(forward_fn, wrapper_name='argmax_sample').forward + + def _forward(obs): + # unsqueeze means add batch dim, i.e. (O, ) -> (1, O) + obs = ttorch.as_tensor(obs).unsqueeze(0) + if cuda and torch.cuda.is_available(): + obs = obs.cuda() + action = forward_fn(obs)["action"] + # squeeze means delete batch dim, i.e. (1, A) -> (A, ) + action = action.squeeze(0).detach().cpu().numpy() + return action + + return _forward + + forward_fn = single_env_forward_wrapper(self.policy._model, self.cfg.policy.cuda) + + # reset first to make sure the env is in the initial state + # env will be reset again in the main loop + env.reset() + + for seed in seeds: + env.seed(seed, dynamic_seed=False) + return_ = 0. + step = 0 + obs = env.reset() + images.append(render(env)[None]) if concatenate_all_replay else None + while True: + action = forward_fn(obs) + obs, rew, done, info = env.step(action) + images.append(render(env)[None]) if concatenate_all_replay else None + return_ += rew + step += 1 + if done: + break + logging.info(f'SQL deploy is finished, final episode return with {step} steps is: {return_}') + returns.append(return_) + + env.close() + + if concatenate_all_replay: + images = np.concatenate(images, axis=0) + import imageio + imageio.mimwrite(os.path.join(replay_save_path, 'deploy.mp4'), images, fps=get_env_fps(env)) + + return EvalReturn(eval_value=np.mean(returns), eval_value_std=np.std(returns)) + + def collect_data( + self, + env_num: int = 8, + save_data_path: Optional[str] = None, + n_sample: Optional[int] = None, + n_episode: Optional[int] = None, + context: Optional[str] = None, + debug: bool = False + ) -> None: + """ + Overview: + Collect data with SQL algorithm for ``n_episode`` episodes with ``env_num`` collector environments. \ + The collected data will be saved in ``save_data_path`` if specified, otherwise it will be saved in \ + ``exp_name/demo_data``. + Arguments: + - env_num (:obj:`int`): The number of collector environments. Default to 8. + - save_data_path (:obj:`str`): The path to save the collected data. Default to None. \ + If not specified, the data will be saved in ``exp_name/demo_data``. + - n_sample (:obj:`int`): The number of samples to collect. Default to None. \ + If not specified, ``n_episode`` must be specified. + - n_episode (:obj:`int`): The number of episodes to collect. Default to None. \ + If not specified, ``n_sample`` must be specified. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + if n_episode is not None: + raise NotImplementedError + # define env and policy + env_num = env_num if env_num else self.cfg.env.collector_env_num + env = setup_ding_env_manager(self.env, env_num, context, debug, 'collector') + + if save_data_path is None: + save_data_path = os.path.join(self.exp_name, 'demo_data') + + # main execution task + with task.start(ctx=OnlineRLContext()): + task.use( + StepCollector( + self.cfg, self.policy.collect_mode, env, random_collect_size=self.cfg.policy.random_collect_size + ) + ) + task.use(offline_data_saver(save_data_path, data_type='hdf5')) + task.run(max_step=1) + logging.info( + f'SQL collecting is finished, more than {n_sample} samples are collected and saved in `{save_data_path}`' + ) + + def batch_evaluate( + self, + env_num: int = 4, + n_evaluator_episode: int = 4, + context: Optional[str] = None, + debug: bool = False + ) -> EvalReturn: + """ + Overview: + Evaluate the agent with SQL algorithm for ``n_evaluator_episode`` episodes with ``env_num`` evaluator \ + environments. The evaluation result will be returned. + The difference between methods ``batch_evaluate`` and ``deploy`` is that ``batch_evaluate`` will create \ + multiple evaluator environments to evaluate the agent to get an average performance, while ``deploy`` \ + will only create one evaluator environment to evaluate the agent and save the replay video. + Arguments: + - env_num (:obj:`int`): The number of evaluator environments. Default to 4. + - n_evaluator_episode (:obj:`int`): The number of episodes to evaluate. Default to 4. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + Returns: + - (:obj:`EvalReturn`): The evaluation result, of which the attributions are: + - eval_value (:obj:`np.float32`): The mean of evaluation return. + - eval_value_std (:obj:`np.float32`): The standard deviation of evaluation return. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + # define env and policy + env_num = env_num if env_num else self.cfg.env.evaluator_env_num + env = setup_ding_env_manager(self.env, env_num, context, debug, 'evaluator') + + # reset first to make sure the env is in the initial state + # env will be reset again in the main loop + env.launch() + env.reset() + + evaluate_cfg = self.cfg + evaluate_cfg.env.n_evaluator_episode = n_evaluator_episode + + # main execution task + with task.start(ctx=OnlineRLContext()): + task.use(interaction_evaluator(self.cfg, self.policy.eval_mode, env)) + task.run(max_step=1) + + return EvalReturn(eval_value=task.ctx.eval_value, eval_value_std=task.ctx.eval_value_std) + + @property + def best(self) -> 'SQLAgent': + """ + Overview: + Load the best model from the checkpoint directory, \ + which is by default in folder ``exp_name/ckpt/eval.pth.tar``. \ + The return value is the agent with the best model. + Returns: + - (:obj:`SQLAgent`): The agent with the best model. + Examples: + >>> agent = SQLAgent(env_id='LunarLander-v2') + >>> agent.train() + >>> agent = agent.best + + .. note:: + The best model is the model with the highest evaluation return. If this method is called, the current \ + model will be replaced by the best model. + """ + + best_model_file_path = os.path.join(self.checkpoint_save_dir, "eval.pth.tar") + # Load best model if it exists + if os.path.exists(best_model_file_path): + policy_state_dict = torch.load(best_model_file_path, map_location=torch.device("cpu")) + self.policy.learn_mode.load_state_dict(policy_state_dict) + return self diff --git a/ding/bonus/td3.py b/ding/bonus/td3.py new file mode 100644 index 0000000000..a2889a370d --- /dev/null +++ b/ding/bonus/td3.py @@ -0,0 +1,455 @@ +from typing import Optional, Union, List +from ditk import logging +from easydict import EasyDict +import os +import numpy as np +import torch +import treetensor.torch as ttorch +from ding.framework import task, OnlineRLContext +from ding.framework.middleware import CkptSaver, \ + wandb_online_logger, offline_data_saver, termination_checker, interaction_evaluator, StepCollector, data_pusher, \ + OffPolicyLearner, final_ctx_saver +from ding.envs import BaseEnv +from ding.envs import setup_ding_env_manager +from ding.policy import TD3Policy +from ding.utils import set_pkg_seed +from ding.utils import get_env_fps, render +from ding.config import save_config_py, compile_config +from ding.model import ContinuousQAC +from ding.data import DequeBuffer +from ding.bonus.common import TrainingReturn, EvalReturn +from ding.config.example.TD3 import supported_env_cfg +from ding.config.example.TD3 import supported_env + + +class TD3Agent: + """ + Overview: + Class of agent for training, evaluation and deployment of Reinforcement learning algorithm \ + Twin Delayed Deep Deterministic Policy Gradient(TD3). + For more information about the system design of RL agent, please refer to \ + . + Interface: + ``__init__``, ``train``, ``deploy``, ``collect_data``, ``batch_evaluate``, ``best`` + """ + supported_env_list = list(supported_env_cfg.keys()) + """ + Overview: + List of supported envs. + Examples: + >>> from ding.bonus.td3 import TD3Agent + >>> print(TD3Agent.supported_env_list) + """ + + def __init__( + self, + env_id: str = None, + env: BaseEnv = None, + seed: int = 0, + exp_name: str = None, + model: Optional[torch.nn.Module] = None, + cfg: Optional[Union[EasyDict, dict]] = None, + policy_state_dict: str = None, + ) -> None: + """ + Overview: + Initialize agent for TD3 algorithm. + Arguments: + - env_id (:obj:`str`): The environment id, which is a registered environment name in gym or gymnasium. \ + If ``env_id`` is not specified, ``env_id`` in ``cfg.env`` must be specified. \ + If ``env_id`` is specified, ``env_id`` in ``cfg.env`` will be ignored. \ + ``env_id`` should be one of the supported envs, which can be found in ``supported_env_list``. + - env (:obj:`BaseEnv`): The environment instance for training and evaluation. \ + If ``env`` is not specified, `env_id`` or ``cfg.env.env_id`` must be specified. \ + ``env_id`` or ``cfg.env.env_id`` will be used to create environment instance. \ + If ``env`` is specified, ``env_id`` and ``cfg.env.env_id`` will be ignored. + - seed (:obj:`int`): The random seed, which is set before running the program. \ + Default to 0. + - exp_name (:obj:`str`): The name of this experiment, which will be used to create the folder to save \ + log data. Default to None. If not specified, the folder name will be ``env_id``-``algorithm``. + - model (:obj:`torch.nn.Module`): The model of TD3 algorithm, which should be an instance of class \ + :class:`ding.model.ContinuousQAC`. \ + If not specified, a default model will be generated according to the configuration. + - cfg (:obj:Union[EasyDict, dict]): The configuration of TD3 algorithm, which is a dict. \ + Default to None. If not specified, the default configuration will be used. \ + The default configuration can be found in ``ding/config/example/TD3/gym_lunarlander_v2.py``. + - policy_state_dict (:obj:`str`): The path of policy state dict saved by PyTorch a in local file. \ + If specified, the policy will be loaded from this file. Default to None. + + .. note:: + An RL Agent Instance can be initialized in two basic ways. \ + For example, we have an environment with id ``LunarLanderContinuous-v2`` registered in gym, \ + and we want to train an agent with TD3 algorithm with default configuration. \ + Then we can initialize the agent in the following ways: + >>> agent = TD3Agent(env_id='LunarLanderContinuous-v2') + or, if we want can specify the env_id in the configuration: + >>> cfg = {'env': {'env_id': 'LunarLanderContinuous-v2'}, 'policy': ...... } + >>> agent = TD3Agent(cfg=cfg) + There are also other arguments to specify the agent when initializing. + For example, if we want to specify the environment instance: + >>> env = CustomizedEnv('LunarLanderContinuous-v2') + >>> agent = TD3Agent(cfg=cfg, env=env) + or, if we want to specify the model: + >>> model = ContinuousQAC(**cfg.policy.model) + >>> agent = TD3Agent(cfg=cfg, model=model) + or, if we want to reload the policy from a saved policy state dict: + >>> agent = TD3Agent(cfg=cfg, policy_state_dict='LunarLanderContinuous-v2.pth.tar') + Make sure that the configuration is consistent with the saved policy state dict. + """ + + assert env_id is not None or cfg is not None, "Please specify env_id or cfg." + + if cfg is not None and not isinstance(cfg, EasyDict): + cfg = EasyDict(cfg) + + if env_id is not None: + assert env_id in TD3Agent.supported_env_list, "Please use supported envs: {}".format( + TD3Agent.supported_env_list + ) + if cfg is None: + cfg = supported_env_cfg[env_id] + else: + assert cfg.env.env_id == env_id, "env_id in cfg should be the same as env_id in args." + else: + assert hasattr(cfg.env, "env_id"), "Please specify env_id in cfg." + assert cfg.env.env_id in TD3Agent.supported_env_list, "Please use supported envs: {}".format( + TD3Agent.supported_env_list + ) + default_policy_config = EasyDict({"policy": TD3Policy.default_config()}) + default_policy_config.update(cfg) + cfg = default_policy_config + + if exp_name is not None: + cfg.exp_name = exp_name + self.cfg = compile_config(cfg, policy=TD3Policy) + self.exp_name = self.cfg.exp_name + if env is None: + self.env = supported_env[cfg.env.env_id](cfg=cfg.env) + else: + assert isinstance(env, BaseEnv), "Please use BaseEnv as env data type." + self.env = env + + logging.getLogger().setLevel(logging.INFO) + self.seed = seed + set_pkg_seed(self.seed, use_cuda=self.cfg.policy.cuda) + if not os.path.exists(self.exp_name): + os.makedirs(self.exp_name) + save_config_py(self.cfg, os.path.join(self.exp_name, 'policy_config.py')) + if model is None: + model = ContinuousQAC(**self.cfg.policy.model) + self.buffer_ = DequeBuffer(size=self.cfg.policy.other.replay_buffer.replay_buffer_size) + self.policy = TD3Policy(self.cfg.policy, model=model) + if policy_state_dict is not None: + self.policy.learn_mode.load_state_dict(policy_state_dict) + self.checkpoint_save_dir = os.path.join(self.exp_name, "ckpt") + + def train( + self, + step: int = int(1e7), + collector_env_num: int = None, + evaluator_env_num: int = None, + n_iter_save_ckpt: int = 1000, + context: Optional[str] = None, + debug: bool = False, + wandb_sweep: bool = False, + ) -> TrainingReturn: + """ + Overview: + Train the agent with TD3 algorithm for ``step`` iterations with ``collector_env_num`` collector \ + environments and ``evaluator_env_num`` evaluator environments. Information during training will be \ + recorded and saved by wandb. + Arguments: + - step (:obj:`int`): The total training environment steps of all collector environments. Default to 1e7. + - collector_env_num (:obj:`int`): The collector environment number. Default to None. \ + If not specified, it will be set according to the configuration. + - evaluator_env_num (:obj:`int`): The evaluator environment number. Default to None. \ + If not specified, it will be set according to the configuration. + - n_iter_save_ckpt (:obj:`int`): The frequency of saving checkpoint every training iteration. \ + Default to 1000. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + - wandb_sweep (:obj:`bool`): Whether to use wandb sweep, \ + which is a hyper-parameter optimization process for seeking the best configurations. \ + Default to False. If True, the wandb sweep id will be used as the experiment name. + Returns: + - (:obj:`TrainingReturn`): The training result, of which the attributions are: + - wandb_url (:obj:`str`): The weight & biases (wandb) project url of the trainning experiment. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + logging.debug(self.policy._model) + # define env and policy + collector_env_num = collector_env_num if collector_env_num else self.cfg.env.collector_env_num + evaluator_env_num = evaluator_env_num if evaluator_env_num else self.cfg.env.evaluator_env_num + collector_env = setup_ding_env_manager(self.env, collector_env_num, context, debug, 'collector') + evaluator_env = setup_ding_env_manager(self.env, evaluator_env_num, context, debug, 'evaluator') + + with task.start(ctx=OnlineRLContext()): + task.use( + interaction_evaluator( + self.cfg, + self.policy.eval_mode, + evaluator_env, + render=self.cfg.policy.eval.render if hasattr(self.cfg.policy.eval, "render") else False + ) + ) + task.use(CkptSaver(policy=self.policy, save_dir=self.checkpoint_save_dir, train_freq=n_iter_save_ckpt)) + task.use( + StepCollector( + self.cfg, + self.policy.collect_mode, + collector_env, + random_collect_size=self.cfg.policy.random_collect_size + if hasattr(self.cfg.policy, 'random_collect_size') else 0, + ) + ) + task.use(data_pusher(self.cfg, self.buffer_)) + task.use(OffPolicyLearner(self.cfg, self.policy.learn_mode, self.buffer_)) + task.use( + wandb_online_logger( + metric_list=self.policy._monitor_vars_learn(), + model=self.policy._model, + anonymous=True, + project_name=self.exp_name, + wandb_sweep=wandb_sweep, + ) + ) + task.use(termination_checker(max_env_step=step)) + task.use(final_ctx_saver(name=self.exp_name)) + task.run() + + return TrainingReturn(wandb_url=task.ctx.wandb_url) + + def deploy( + self, + enable_save_replay: bool = False, + concatenate_all_replay: bool = False, + replay_save_path: str = None, + seed: Optional[Union[int, List]] = None, + debug: bool = False + ) -> EvalReturn: + """ + Overview: + Deploy the agent with TD3 algorithm by interacting with the environment, during which the replay video \ + can be saved if ``enable_save_replay`` is True. The evaluation result will be returned. + Arguments: + - enable_save_replay (:obj:`bool`): Whether to save the replay video. Default to False. + - concatenate_all_replay (:obj:`bool`): Whether to concatenate all replay videos into one video. \ + Default to False. If ``enable_save_replay`` is False, this argument will be ignored. \ + If ``enable_save_replay`` is True and ``concatenate_all_replay`` is False, \ + the replay video of each episode will be saved separately. + - replay_save_path (:obj:`str`): The path to save the replay video. Default to None. \ + If not specified, the video will be saved in ``exp_name/videos``. + - seed (:obj:`Union[int, List]`): The random seed, which is set before running the program. \ + Default to None. If not specified, ``self.seed`` will be used. \ + If ``seed`` is an integer, the agent will be deployed once. \ + If ``seed`` is a list of integers, the agent will be deployed once for each seed in the list. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + Returns: + - (:obj:`EvalReturn`): The evaluation result, of which the attributions are: + - eval_value (:obj:`np.float32`): The mean of evaluation return. + - eval_value_std (:obj:`np.float32`): The standard deviation of evaluation return. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + # define env and policy + env = self.env.clone(caller='evaluator') + + if seed is not None and isinstance(seed, int): + seeds = [seed] + elif seed is not None and isinstance(seed, list): + seeds = seed + else: + seeds = [self.seed] + + returns = [] + images = [] + if enable_save_replay: + replay_save_path = os.path.join(self.exp_name, 'videos') if replay_save_path is None else replay_save_path + env.enable_save_replay(replay_path=replay_save_path) + else: + logging.warning('No video would be generated during the deploy.') + if concatenate_all_replay: + logging.warning('concatenate_all_replay is set to False because enable_save_replay is False.') + concatenate_all_replay = False + + def single_env_forward_wrapper(forward_fn, cuda=True): + + def _forward(obs): + # unsqueeze means add batch dim, i.e. (O, ) -> (1, O) + obs = ttorch.as_tensor(obs).unsqueeze(0) + if cuda and torch.cuda.is_available(): + obs = obs.cuda() + action = forward_fn(obs, mode='compute_actor')["action"] + # squeeze means delete batch dim, i.e. (1, A) -> (A, ) + action = action.squeeze(0).detach().cpu().numpy() + return action + + return _forward + + forward_fn = single_env_forward_wrapper(self.policy._model, self.cfg.policy.cuda) + + # reset first to make sure the env is in the initial state + # env will be reset again in the main loop + env.reset() + + for seed in seeds: + env.seed(seed, dynamic_seed=False) + return_ = 0. + step = 0 + obs = env.reset() + images.append(render(env)[None]) if concatenate_all_replay else None + while True: + action = forward_fn(obs) + obs, rew, done, info = env.step(action) + images.append(render(env)[None]) if concatenate_all_replay else None + return_ += rew + step += 1 + if done: + break + logging.info(f'DQN deploy is finished, final episode return with {step} steps is: {return_}') + returns.append(return_) + + env.close() + + if concatenate_all_replay: + images = np.concatenate(images, axis=0) + import imageio + imageio.mimwrite(os.path.join(replay_save_path, 'deploy.mp4'), images, fps=get_env_fps(env)) + + return EvalReturn(eval_value=np.mean(returns), eval_value_std=np.std(returns)) + + def collect_data( + self, + env_num: int = 8, + save_data_path: Optional[str] = None, + n_sample: Optional[int] = None, + n_episode: Optional[int] = None, + context: Optional[str] = None, + debug: bool = False + ) -> None: + """ + Overview: + Collect data with TD3 algorithm for ``n_episode`` episodes with ``env_num`` collector environments. \ + The collected data will be saved in ``save_data_path`` if specified, otherwise it will be saved in \ + ``exp_name/demo_data``. + Arguments: + - env_num (:obj:`int`): The number of collector environments. Default to 8. + - save_data_path (:obj:`str`): The path to save the collected data. Default to None. \ + If not specified, the data will be saved in ``exp_name/demo_data``. + - n_sample (:obj:`int`): The number of samples to collect. Default to None. \ + If not specified, ``n_episode`` must be specified. + - n_episode (:obj:`int`): The number of episodes to collect. Default to None. \ + If not specified, ``n_sample`` must be specified. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + if n_episode is not None: + raise NotImplementedError + # define env and policy + env_num = env_num if env_num else self.cfg.env.collector_env_num + env = setup_ding_env_manager(self.env, env_num, context, debug, 'collector') + + if save_data_path is None: + save_data_path = os.path.join(self.exp_name, 'demo_data') + + # main execution task + with task.start(ctx=OnlineRLContext()): + task.use( + StepCollector( + self.cfg, self.policy.collect_mode, env, random_collect_size=self.cfg.policy.random_collect_size + ) + ) + task.use(offline_data_saver(save_data_path, data_type='hdf5')) + task.run(max_step=1) + logging.info( + f'TD3 collecting is finished, more than {n_sample} samples are collected and saved in `{save_data_path}`' + ) + + def batch_evaluate( + self, + env_num: int = 4, + n_evaluator_episode: int = 4, + context: Optional[str] = None, + debug: bool = False + ) -> EvalReturn: + """ + Overview: + Evaluate the agent with TD3 algorithm for ``n_evaluator_episode`` episodes with ``env_num`` evaluator \ + environments. The evaluation result will be returned. + The difference between methods ``batch_evaluate`` and ``deploy`` is that ``batch_evaluate`` will create \ + multiple evaluator environments to evaluate the agent to get an average performance, while ``deploy`` \ + will only create one evaluator environment to evaluate the agent and save the replay video. + Arguments: + - env_num (:obj:`int`): The number of evaluator environments. Default to 4. + - n_evaluator_episode (:obj:`int`): The number of episodes to evaluate. Default to 4. + - context (:obj:`str`): The multi-process context of the environment manager. Default to None. \ + It can be specified as ``spawn``, ``fork`` or ``forkserver``. + - debug (:obj:`bool`): Whether to use debug mode in the environment manager. Default to False. \ + If set True, base environment manager will be used for easy debugging. Otherwise, \ + subprocess environment manager will be used. + Returns: + - (:obj:`EvalReturn`): The evaluation result, of which the attributions are: + - eval_value (:obj:`np.float32`): The mean of evaluation return. + - eval_value_std (:obj:`np.float32`): The standard deviation of evaluation return. + """ + + if debug: + logging.getLogger().setLevel(logging.DEBUG) + # define env and policy + env_num = env_num if env_num else self.cfg.env.evaluator_env_num + env = setup_ding_env_manager(self.env, env_num, context, debug, 'evaluator') + + # reset first to make sure the env is in the initial state + # env will be reset again in the main loop + env.launch() + env.reset() + + evaluate_cfg = self.cfg + evaluate_cfg.env.n_evaluator_episode = n_evaluator_episode + + # main execution task + with task.start(ctx=OnlineRLContext()): + task.use(interaction_evaluator(self.cfg, self.policy.eval_mode, env)) + task.run(max_step=1) + + return EvalReturn(eval_value=task.ctx.eval_value, eval_value_std=task.ctx.eval_value_std) + + @property + def best(self) -> 'TD3Agent': + """ + Overview: + Load the best model from the checkpoint directory, \ + which is by default in folder ``exp_name/ckpt/eval.pth.tar``. \ + The return value is the agent with the best model. + Returns: + - (:obj:`TD3Agent`): The agent with the best model. + Examples: + >>> agent = TD3Agent(env_id='LunarLanderContinuous-v2') + >>> agent.train() + >>> agent.best + + .. note:: + The best model is the model with the highest evaluation return. If this method is called, the current \ + model will be replaced by the best model. + """ + + best_model_file_path = os.path.join(self.checkpoint_save_dir, "eval.pth.tar") + # Load best model if it exists + if os.path.exists(best_model_file_path): + policy_state_dict = torch.load(best_model_file_path, map_location=torch.device("cpu")) + self.policy.learn_mode.load_state_dict(policy_state_dict) + return self diff --git a/ding/config/__init__.py b/ding/config/__init__.py index 87edefa77c..162fc86c86 100644 --- a/ding/config/__init__.py +++ b/ding/config/__init__.py @@ -1,3 +1,4 @@ from .config import Config, read_config, save_config, compile_config, compile_config_parallel, read_config_directly, \ - read_config_with_system + read_config_with_system, save_config_py from .utils import parallel_transform, parallel_transform_slurm +from .example import A2C, C51, DDPG, DQN, PG, PPOF, PPOOffPolicy, SAC, SQL, TD3 diff --git a/ding/config/config.py b/ding/config/config.py index 78f04a6d5a..1cf462d6a8 100644 --- a/ding/config/config.py +++ b/ding/config/config.py @@ -9,11 +9,11 @@ import subprocess import datetime from importlib import import_module -from typing import Optional, Tuple, NoReturn +from typing import Optional, Tuple from easydict import EasyDict from copy import deepcopy -from ding.utils import deep_merge_dicts +from ding.utils import deep_merge_dicts, get_rank from ding.envs import get_env_cls, get_env_manager_cls, BaseEnvManager from ding.policy import get_policy_cls from ding.worker import BaseLearner, InteractionSerialEvaluator, BaseSerialCommander, Coordinator, \ @@ -131,7 +131,7 @@ def read_config_yaml(path: str) -> EasyDict: return EasyDict(config_) -def save_config_yaml(config_: dict, path: str) -> NoReturn: +def save_config_yaml(config_: dict, path: str) -> None: """ Overview: save configuration to path @@ -144,7 +144,7 @@ def save_config_yaml(config_: dict, path: str) -> NoReturn: yaml.safe_dump(json.loads(config_string), f) -def save_config_py(config_: dict, path: str) -> NoReturn: +def save_config_py(config_: dict, path: str) -> None: """ Overview: save configuration to python file @@ -180,7 +180,7 @@ def read_config_directly(path: str) -> dict: def read_config(path: str) -> Tuple[dict, dict]: """ Overview: - Read configuration from a file path(now only suport python file). And select some proper parts. + Read configuration from a file path(now only support python file). And select some proper parts. Arguments: - path (:obj:`str`): Path of configuration file Returns: @@ -200,7 +200,7 @@ def read_config(path: str) -> Tuple[dict, dict]: def read_config_with_system(path: str) -> Tuple[dict, dict, dict]: """ Overview: - Read configuration from a file path(now only suport python file). And select some proper parts + Read configuration from a file path(now only support python file). And select some proper parts Arguments: - path (:obj:`str`): Path of configuration file Returns: @@ -218,7 +218,7 @@ def read_config_with_system(path: str) -> Tuple[dict, dict, dict]: raise KeyError("invalid config file suffix: {}".format(suffix)) -def save_config(config_: dict, path: str, type_: str = 'py', save_formatted: bool = False) -> NoReturn: +def save_config(config_: dict, path: str, type_: str = 'py', save_formatted: bool = False) -> None: """ Overview: save configuration to python file or yaml file @@ -306,7 +306,7 @@ def compile_collector_config( other=dict(replay_buffer=dict()), ) policy_config_template = EasyDict(policy_config_template) -env_config_template = dict(manager=dict(), ) +env_config_template = dict(manager=dict(), stop_value=int(1e10), n_evaluator_episode=4) env_config_template = EasyDict(env_config_template) @@ -315,13 +315,13 @@ def save_project_state(exp_name: str) -> None: def _fn(cmd: str): return subprocess.run(cmd, shell=True, stdout=subprocess.PIPE).stdout.strip().decode("utf-8") - if subprocess.run("git status", shell=True, stderr=subprocess.PIPE).returncode == 0: + if subprocess.run("git status", shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE).returncode == 0: short_sha = _fn("git describe --always") log = _fn("git log --stat -n 5") diff = _fn("git diff") - with open(os.path.join(exp_name, "git_log.txt"), "w") as f: + with open(os.path.join(exp_name, "git_log.txt"), "w", encoding='utf-8') as f: f.write(short_sha + '\n\n' + log) - with open(os.path.join(exp_name, "git_diff.txt"), "w") as f: + with open(os.path.join(exp_name, "git_diff.txt"), "w", encoding='utf-8') as f: f.write(diff) @@ -341,6 +341,7 @@ def compile_config( create_cfg: dict = None, save_cfg: bool = True, save_path: str = 'total_config.py', + renew_dir: bool = True, ) -> EasyDict: """ Overview: @@ -361,6 +362,7 @@ def compile_config( - create_cfg (:obj:`dict`): Input create config dict - save_cfg (:obj:`bool`): Save config or not - save_path (:obj:`str`): Path of saving file + - renew_dir (:obj:`bool`): Whether to new a directory for saving config. Returns: - cfg (:obj:`EasyDict`): Config after compiling """ @@ -450,16 +452,17 @@ def compile_config( if len(world_model_config) > 0: default_config['world_model'] = world_model_config cfg = deep_merge_dicts(default_config, cfg) + if 'unroll_len' in cfg.policy: + cfg.policy.collect.unroll_len = cfg.policy.unroll_len cfg.seed = seed # check important key in config if evaluator in [InteractionSerialEvaluator, BattleInteractionSerialEvaluator]: # env interaction evaluation - if 'stop_value' in cfg.env: # data generation task doesn't need these fields - cfg.policy.eval.evaluator.n_episode = cfg.env.n_evaluator_episode - cfg.policy.eval.evaluator.stop_value = cfg.env.stop_value + cfg.policy.eval.evaluator.stop_value = cfg.env.stop_value + cfg.policy.eval.evaluator.n_episode = cfg.env.n_evaluator_episode if 'exp_name' not in cfg: cfg.exp_name = 'default_experiment' - if save_cfg: - if os.path.exists(cfg.exp_name): + if save_cfg and get_rank() == 0: + if os.path.exists(cfg.exp_name) and renew_dir: cfg.exp_name += datetime.datetime.now().strftime("_%y%m%d_%H%M%S") try: os.makedirs(cfg.exp_name) diff --git a/ding/config/example/A2C/__init__.py b/ding/config/example/A2C/__init__.py new file mode 100644 index 0000000000..37beafc248 --- /dev/null +++ b/ding/config/example/A2C/__init__.py @@ -0,0 +1,17 @@ +from easydict import EasyDict +from . import gym_bipedalwalker_v3 +from . import gym_lunarlander_v2 + +supported_env_cfg = { + gym_bipedalwalker_v3.cfg.env.env_id: gym_bipedalwalker_v3.cfg, + gym_lunarlander_v2.cfg.env.env_id: gym_lunarlander_v2.cfg, +} + +supported_env_cfg = EasyDict(supported_env_cfg) + +supported_env = { + gym_bipedalwalker_v3.cfg.env.env_id: gym_bipedalwalker_v3.env, + gym_lunarlander_v2.cfg.env.env_id: gym_lunarlander_v2.env, +} + +supported_env = EasyDict(supported_env) diff --git a/ding/config/example/A2C/gym_bipedalwalker_v3.py b/ding/config/example/A2C/gym_bipedalwalker_v3.py new file mode 100644 index 0000000000..a0f9637cec --- /dev/null +++ b/ding/config/example/A2C/gym_bipedalwalker_v3.py @@ -0,0 +1,43 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='Bipedalwalker-v3-A2C', + seed=0, + env=dict( + env_id='BipedalWalker-v3', + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + act_scale=True, + rew_clip=True, + ), + policy=dict( + cuda=True, + action_space='continuous', + model=dict( + action_space='continuous', + obs_shape=24, + action_shape=4, + ), + learn=dict( + batch_size=64, + learning_rate=0.0003, + value_weight=0.7, + entropy_weight=0.0005, + discount_factor=0.99, + adv_norm=True, + ), + collect=dict( + n_sample=64, + discount_factor=0.99, + ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/A2C/gym_lunarlander_v2.py b/ding/config/example/A2C/gym_lunarlander_v2.py new file mode 100644 index 0000000000..6092bb4126 --- /dev/null +++ b/ding/config/example/A2C/gym_lunarlander_v2.py @@ -0,0 +1,38 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='LunarLander-v2-A2C', + env=dict( + collector_env_num=8, + evaluator_env_num=8, + env_id='LunarLander-v2', + n_evaluator_episode=8, + stop_value=260, + ), + policy=dict( + cuda=True, + model=dict( + obs_shape=8, + action_shape=4, + ), + learn=dict( + batch_size=64, + learning_rate=3e-4, + entropy_weight=0.001, + adv_norm=True, + ), + collect=dict( + n_sample=64, + discount_factor=0.99, + gae_lambda=0.95, + ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/C51/__init__.py b/ding/config/example/C51/__init__.py new file mode 100644 index 0000000000..2704b04c53 --- /dev/null +++ b/ding/config/example/C51/__init__.py @@ -0,0 +1,23 @@ +from easydict import EasyDict +from . import gym_lunarlander_v2 +from . import gym_pongnoframeskip_v4 +from . import gym_qbertnoframeskip_v4 +from . import gym_spaceInvadersnoframeskip_v4 + +supported_env_cfg = { + gym_lunarlander_v2.cfg.env.env_id: gym_lunarlander_v2.cfg, + gym_pongnoframeskip_v4.cfg.env.env_id: gym_pongnoframeskip_v4.cfg, + gym_qbertnoframeskip_v4.cfg.env.env_id: gym_qbertnoframeskip_v4.cfg, + gym_spaceInvadersnoframeskip_v4.cfg.env.env_id: gym_spaceInvadersnoframeskip_v4.cfg, +} + +supported_env_cfg = EasyDict(supported_env_cfg) + +supported_env = { + gym_lunarlander_v2.cfg.env.env_id: gym_lunarlander_v2.env, + gym_pongnoframeskip_v4.cfg.env.env_id: gym_pongnoframeskip_v4.env, + gym_qbertnoframeskip_v4.cfg.env.env_id: gym_qbertnoframeskip_v4.env, + gym_spaceInvadersnoframeskip_v4.cfg.env.env_id: gym_spaceInvadersnoframeskip_v4.env, +} + +supported_env = EasyDict(supported_env) diff --git a/ding/config/example/C51/gym_lunarlander_v2.py b/ding/config/example/C51/gym_lunarlander_v2.py new file mode 100644 index 0000000000..8f929a964f --- /dev/null +++ b/ding/config/example/C51/gym_lunarlander_v2.py @@ -0,0 +1,52 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='lunarlander_c51', + seed=0, + env=dict( + collector_env_num=8, + evaluator_env_num=8, + env_id='LunarLander-v2', + n_evaluator_episode=8, + stop_value=260, + ), + policy=dict( + cuda=False, + model=dict( + obs_shape=8, + action_shape=4, + encoder_hidden_size_list=[512, 64], + v_min=-30, + v_max=30, + n_atom=51, + ), + discount_factor=0.99, + nstep=3, + learn=dict( + update_per_collect=10, + batch_size=64, + learning_rate=0.001, + target_update_freq=100, + ), + collect=dict( + n_sample=64, + unroll_len=1, + ), + other=dict( + eps=dict( + type='exp', + start=0.95, + end=0.1, + decay=50000, + ), replay_buffer=dict(replay_buffer_size=100000, ) + ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/C51/gym_pongnoframeskip_v4.py b/ding/config/example/C51/gym_pongnoframeskip_v4.py new file mode 100644 index 0000000000..c7b62b5fa8 --- /dev/null +++ b/ding/config/example/C51/gym_pongnoframeskip_v4.py @@ -0,0 +1,54 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='PongNoFrameskip-v4-C51', + seed=0, + env=dict( + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=30, + env_id='PongNoFrameskip-v4', + frame_stack=4, + env_wrapper='atari_default', + ), + policy=dict( + cuda=True, + priority=False, + model=dict( + obs_shape=[4, 84, 84], + action_shape=6, + encoder_hidden_size_list=[128, 128, 512], + v_min=-10, + v_max=10, + n_atom=51, + ), + nstep=3, + discount_factor=0.99, + learn=dict( + update_per_collect=10, + batch_size=32, + learning_rate=0.0001, + target_update_freq=500, + ), + collect=dict(n_sample=100, ), + eval=dict(evaluator=dict(eval_freq=4000, )), + other=dict( + eps=dict( + type='exp', + start=1., + end=0.05, + decay=250000, + ), + replay_buffer=dict(replay_buffer_size=100000, ), + ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/C51/gym_qbertnoframeskip_v4.py b/ding/config/example/C51/gym_qbertnoframeskip_v4.py new file mode 100644 index 0000000000..21c3bcf7ea --- /dev/null +++ b/ding/config/example/C51/gym_qbertnoframeskip_v4.py @@ -0,0 +1,54 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='QbertNoFrameskip-v4-C51', + seed=0, + env=dict( + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=30000, + env_id='QbertNoFrameskip-v4', + frame_stack=4, + env_wrapper='atari_default', + ), + policy=dict( + cuda=True, + priority=True, + model=dict( + obs_shape=[4, 84, 84], + action_shape=6, + encoder_hidden_size_list=[128, 128, 512], + v_min=-10, + v_max=10, + n_atom=51, + ), + nstep=3, + discount_factor=0.99, + learn=dict( + update_per_collect=10, + batch_size=32, + learning_rate=0.0001, + target_update_freq=500, + ), + collect=dict(n_sample=100, ), + eval=dict(evaluator=dict(eval_freq=4000, )), + other=dict( + eps=dict( + type='exp', + start=1., + end=0.05, + decay=1000000, + ), + replay_buffer=dict(replay_buffer_size=400000, ), + ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/C51/gym_spaceInvadersnoframeskip_v4.py b/ding/config/example/C51/gym_spaceInvadersnoframeskip_v4.py new file mode 100644 index 0000000000..d2fcf431c3 --- /dev/null +++ b/ding/config/example/C51/gym_spaceInvadersnoframeskip_v4.py @@ -0,0 +1,54 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='SpaceInvadersNoFrameskip-v4-C51', + seed=0, + env=dict( + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=10000000000, + env_id='SpaceInvadersNoFrameskip-v4', + frame_stack=4, + env_wrapper='atari_default', + ), + policy=dict( + cuda=True, + priority=False, + model=dict( + obs_shape=[4, 84, 84], + action_shape=6, + encoder_hidden_size_list=[128, 128, 512], + v_min=-10, + v_max=10, + n_atom=51, + ), + nstep=3, + discount_factor=0.99, + learn=dict( + update_per_collect=10, + batch_size=32, + learning_rate=0.0001, + target_update_freq=500, + ), + collect=dict(n_sample=100, ), + eval=dict(evaluator=dict(eval_freq=4000, )), + other=dict( + eps=dict( + type='exp', + start=1., + end=0.05, + decay=1000000, + ), + replay_buffer=dict(replay_buffer_size=400000, ), + ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/DDPG/__init__.py b/ding/config/example/DDPG/__init__.py new file mode 100644 index 0000000000..6e01f29d74 --- /dev/null +++ b/ding/config/example/DDPG/__init__.py @@ -0,0 +1,29 @@ +from easydict import EasyDict +from . import gym_bipedalwalker_v3 +from . import gym_halfcheetah_v3 +from . import gym_hopper_v3 +from . import gym_lunarlandercontinuous_v2 +from . import gym_pendulum_v1 +from . import gym_walker2d_v3 + +supported_env_cfg = { + gym_bipedalwalker_v3.cfg.env.env_id: gym_bipedalwalker_v3.cfg, + gym_halfcheetah_v3.cfg.env.env_id: gym_halfcheetah_v3.cfg, + gym_hopper_v3.cfg.env.env_id: gym_hopper_v3.cfg, + gym_lunarlandercontinuous_v2.cfg.env.env_id: gym_lunarlandercontinuous_v2.cfg, + gym_pendulum_v1.cfg.env.env_id: gym_pendulum_v1.cfg, + gym_walker2d_v3.cfg.env.env_id: gym_walker2d_v3.cfg, +} + +supported_env_cfg = EasyDict(supported_env_cfg) + +supported_env = { + gym_bipedalwalker_v3.cfg.env.env_id: gym_bipedalwalker_v3.env, + gym_halfcheetah_v3.cfg.env.env_id: gym_halfcheetah_v3.env, + gym_hopper_v3.cfg.env.env_id: gym_hopper_v3.env, + gym_lunarlandercontinuous_v2.cfg.env.env_id: gym_lunarlandercontinuous_v2.env, + gym_pendulum_v1.cfg.env.env_id: gym_pendulum_v1.env, + gym_walker2d_v3.cfg.env.env_id: gym_walker2d_v3.env, +} + +supported_env = EasyDict(supported_env) diff --git a/ding/config/example/DDPG/gym_bipedalwalker_v3.py b/ding/config/example/DDPG/gym_bipedalwalker_v3.py new file mode 100644 index 0000000000..cd26d46f68 --- /dev/null +++ b/ding/config/example/DDPG/gym_bipedalwalker_v3.py @@ -0,0 +1,45 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='Bipedalwalker-v3-DDPG', + seed=0, + env=dict( + env_id='BipedalWalker-v3', + collector_env_num=8, + evaluator_env_num=5, + n_evaluator_episode=5, + act_scale=True, + rew_clip=True, + ), + policy=dict( + cuda=True, + random_collect_size=10000, + model=dict( + obs_shape=24, + action_shape=4, + twin_critic=False, + action_space='regression', + actor_head_hidden_size=400, + critic_head_hidden_size=400, + ), + learn=dict( + update_per_collect=64, + batch_size=256, + learning_rate_actor=0.0003, + learning_rate_critic=0.0003, + target_theta=0.005, + discount_factor=0.99, + learner=dict(hook=dict(log_show_after_iter=1000, )) + ), + collect=dict(n_sample=64, ), + other=dict(replay_buffer=dict(replay_buffer_size=300000, ), ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/DDPG/gym_halfcheetah_v3.py b/ding/config/example/DDPG/gym_halfcheetah_v3.py new file mode 100644 index 0000000000..2bbf075a03 --- /dev/null +++ b/ding/config/example/DDPG/gym_halfcheetah_v3.py @@ -0,0 +1,53 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='HalfCheetah-v3-DDPG', + seed=0, + env=dict( + env_id='HalfCheetah-v3', + norm_obs=dict(use_norm=False, ), + norm_reward=dict(use_norm=False, ), + collector_env_num=1, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=11000, + env_wrapper='mujoco_default', + ), + policy=dict( + cuda=True, + random_collect_size=25000, + model=dict( + obs_shape=17, + action_shape=6, + twin_critic=False, + actor_head_hidden_size=256, + critic_head_hidden_size=256, + action_space='regression', + ), + learn=dict( + update_per_collect=1, + batch_size=256, + learning_rate_actor=1e-3, + learning_rate_critic=1e-3, + ignore_done=True, + target_theta=0.005, + discount_factor=0.99, + actor_update_freq=1, + noise=False, + ), + collect=dict( + n_sample=1, + unroll_len=1, + noise_sigma=0.1, + ), + other=dict(replay_buffer=dict(replay_buffer_size=1000000, ), ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/DDPG/gym_hopper_v3.py b/ding/config/example/DDPG/gym_hopper_v3.py new file mode 100644 index 0000000000..bd8d253807 --- /dev/null +++ b/ding/config/example/DDPG/gym_hopper_v3.py @@ -0,0 +1,53 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='Hopper-v3-DDPG', + seed=0, + env=dict( + env_id='Hopper-v3', + norm_obs=dict(use_norm=False, ), + norm_reward=dict(use_norm=False, ), + collector_env_num=1, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=6000, + env_wrapper='mujoco_default', + ), + policy=dict( + cuda=True, + random_collect_size=25000, + model=dict( + obs_shape=11, + action_shape=3, + twin_critic=False, + actor_head_hidden_size=256, + critic_head_hidden_size=256, + action_space='regression', + ), + learn=dict( + update_per_collect=1, + batch_size=256, + learning_rate_actor=1e-3, + learning_rate_critic=1e-3, + ignore_done=False, + target_theta=0.005, + discount_factor=0.99, + actor_update_freq=1, + noise=False, + ), + collect=dict( + n_sample=1, + unroll_len=1, + noise_sigma=0.1, + ), + other=dict(replay_buffer=dict(replay_buffer_size=1000000, ), ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/DDPG/gym_lunarlandercontinuous_v2.py b/ding/config/example/DDPG/gym_lunarlandercontinuous_v2.py new file mode 100644 index 0000000000..89b92e8de1 --- /dev/null +++ b/ding/config/example/DDPG/gym_lunarlandercontinuous_v2.py @@ -0,0 +1,60 @@ +from easydict import EasyDict +from functools import partial +import ding.envs.gym_env + +cfg = dict( + exp_name='LunarLanderContinuous-V2-DDPG', + seed=0, + env=dict( + env_id='LunarLanderContinuous-v2', + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=260, + act_scale=True, + ), + policy=dict( + cuda=True, + random_collect_size=0, + model=dict( + obs_shape=8, + action_shape=2, + twin_critic=True, + action_space='regression', + ), + learn=dict( + update_per_collect=2, + batch_size=128, + learning_rate_actor=0.001, + learning_rate_critic=0.001, + ignore_done=False, # TODO(pu) + # (int) When critic network updates once, how many times will actor network update. + # Delayed Policy Updates in original TD3 paper(https://arxiv.org/pdf/1802.09477.pdf). + # Default 1 for DDPG, 2 for TD3. + actor_update_freq=1, + # (bool) Whether to add noise on target network's action. + # Target Policy Smoothing Regularization in original TD3 paper(https://arxiv.org/pdf/1802.09477.pdf). + # Default True for TD3, False for DDPG. + noise=False, + noise_sigma=0.1, + noise_range=dict( + min=-0.5, + max=0.5, + ), + ), + collect=dict( + n_sample=48, + noise_sigma=0.1, + collector=dict(collect_print_freq=1000, ), + ), + eval=dict(evaluator=dict(eval_freq=100, ), ), + other=dict(replay_buffer=dict(replay_buffer_size=20000, ), ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = partial(ding.envs.gym_env.env, continuous=True) diff --git a/ding/config/example/DDPG/gym_pendulum_v1.py b/ding/config/example/DDPG/gym_pendulum_v1.py new file mode 100644 index 0000000000..63e85869e7 --- /dev/null +++ b/ding/config/example/DDPG/gym_pendulum_v1.py @@ -0,0 +1,52 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='Pendulum-v1-DDPG', + seed=0, + env=dict( + env_id='Pendulum-v1', + collector_env_num=8, + evaluator_env_num=5, + n_evaluator_episode=5, + stop_value=-250, + act_scale=True, + ), + policy=dict( + cuda=False, + priority=False, + random_collect_size=800, + model=dict( + obs_shape=3, + action_shape=1, + twin_critic=False, + action_space='regression', + ), + learn=dict( + update_per_collect=2, + batch_size=128, + learning_rate_actor=0.001, + learning_rate_critic=0.001, + ignore_done=True, + actor_update_freq=1, + noise=False, + ), + collect=dict( + n_sample=48, + noise_sigma=0.1, + collector=dict(collect_print_freq=1000, ), + ), + eval=dict(evaluator=dict(eval_freq=100, )), + other=dict(replay_buffer=dict( + replay_buffer_size=20000, + max_use=16, + ), ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/DDPG/gym_walker2d_v3.py b/ding/config/example/DDPG/gym_walker2d_v3.py new file mode 100644 index 0000000000..84e6407de4 --- /dev/null +++ b/ding/config/example/DDPG/gym_walker2d_v3.py @@ -0,0 +1,53 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='Walker2d-v3-DDPG', + seed=0, + env=dict( + env_id='Walker2d-v3', + norm_obs=dict(use_norm=False, ), + norm_reward=dict(use_norm=False, ), + collector_env_num=1, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=6000, + env_wrapper='mujoco_default', + ), + policy=dict( + cuda=True, + random_collect_size=25000, + model=dict( + obs_shape=17, + action_shape=6, + twin_critic=False, + actor_head_hidden_size=256, + critic_head_hidden_size=256, + action_space='regression', + ), + learn=dict( + update_per_collect=1, + batch_size=256, + learning_rate_actor=1e-3, + learning_rate_critic=1e-3, + ignore_done=False, + target_theta=0.005, + discount_factor=0.99, + actor_update_freq=1, + noise=False, + ), + collect=dict( + n_sample=1, + unroll_len=1, + noise_sigma=0.1, + ), + other=dict(replay_buffer=dict(replay_buffer_size=1000000, ), ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/DQN/__init__.py b/ding/config/example/DQN/__init__.py new file mode 100644 index 0000000000..2704b04c53 --- /dev/null +++ b/ding/config/example/DQN/__init__.py @@ -0,0 +1,23 @@ +from easydict import EasyDict +from . import gym_lunarlander_v2 +from . import gym_pongnoframeskip_v4 +from . import gym_qbertnoframeskip_v4 +from . import gym_spaceInvadersnoframeskip_v4 + +supported_env_cfg = { + gym_lunarlander_v2.cfg.env.env_id: gym_lunarlander_v2.cfg, + gym_pongnoframeskip_v4.cfg.env.env_id: gym_pongnoframeskip_v4.cfg, + gym_qbertnoframeskip_v4.cfg.env.env_id: gym_qbertnoframeskip_v4.cfg, + gym_spaceInvadersnoframeskip_v4.cfg.env.env_id: gym_spaceInvadersnoframeskip_v4.cfg, +} + +supported_env_cfg = EasyDict(supported_env_cfg) + +supported_env = { + gym_lunarlander_v2.cfg.env.env_id: gym_lunarlander_v2.env, + gym_pongnoframeskip_v4.cfg.env.env_id: gym_pongnoframeskip_v4.env, + gym_qbertnoframeskip_v4.cfg.env.env_id: gym_qbertnoframeskip_v4.env, + gym_spaceInvadersnoframeskip_v4.cfg.env.env_id: gym_spaceInvadersnoframeskip_v4.env, +} + +supported_env = EasyDict(supported_env) diff --git a/ding/config/example/DQN/gym_lunarlander_v2.py b/ding/config/example/DQN/gym_lunarlander_v2.py new file mode 100644 index 0000000000..6b79a4eeaa --- /dev/null +++ b/ding/config/example/DQN/gym_lunarlander_v2.py @@ -0,0 +1,53 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='LunarLander-v2-DQN', + seed=0, + env=dict( + env_id='LunarLander-v2', + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=260, + ), + policy=dict( + cuda=True, + random_collect_size=25000, + discount_factor=0.99, + nstep=3, + learn=dict( + update_per_collect=10, + batch_size=64, + learning_rate=0.001, + # Frequency of target network update. + target_update_freq=100, + ), + model=dict( + obs_shape=8, + action_shape=4, + encoder_hidden_size_list=[512, 64], + # Whether to use dueling head. + dueling=True, + ), + collect=dict( + n_sample=64, + unroll_len=1, + ), + other=dict( + eps=dict( + type='exp', + start=0.95, + end=0.1, + decay=50000, + ), replay_buffer=dict(replay_buffer_size=100000, ) + ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/DQN/gym_pongnoframeskip_v4.py b/ding/config/example/DQN/gym_pongnoframeskip_v4.py new file mode 100644 index 0000000000..0266783ac5 --- /dev/null +++ b/ding/config/example/DQN/gym_pongnoframeskip_v4.py @@ -0,0 +1,50 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='PongNoFrameskip-v4-DQN', + seed=0, + env=dict( + env_id='PongNoFrameskip-v4', + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=30, + fram_stack=4, + env_wrapper='atari_default', + ), + policy=dict( + cuda=True, + priority=False, + discount_factor=0.99, + nstep=3, + learn=dict( + update_per_collect=10, + batch_size=32, + learning_rate=0.0001, + # Frequency of target network update. + target_update_freq=500, + ), + model=dict( + obs_shape=[4, 84, 84], + action_shape=6, + encoder_hidden_size_list=[128, 128, 512], + ), + collect=dict(n_sample=96, ), + other=dict( + eps=dict( + type='exp', + start=1., + end=0.05, + decay=250000, + ), replay_buffer=dict(replay_buffer_size=100000, ) + ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/DQN/gym_qbertnoframeskip_v4.py b/ding/config/example/DQN/gym_qbertnoframeskip_v4.py new file mode 100644 index 0000000000..e782a12e9c --- /dev/null +++ b/ding/config/example/DQN/gym_qbertnoframeskip_v4.py @@ -0,0 +1,50 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='QbertNoFrameskip-v4-DQN', + seed=0, + env=dict( + env_id='QbertNoFrameskip-v4', + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + fram_stack=4, + stop_value=30000, + env_wrapper='atari_default', + ), + policy=dict( + cuda=True, + priority=False, + discount_factor=0.99, + nstep=3, + learn=dict( + update_per_collect=10, + batch_size=32, + learning_rate=0.0001, + # Frequency of target network update. + target_update_freq=500, + ), + model=dict( + obs_shape=[4, 84, 84], + action_shape=6, + encoder_hidden_size_list=[128, 128, 512], + ), + collect=dict(n_sample=100, ), + other=dict( + eps=dict( + type='exp', + start=1., + end=0.05, + decay=1000000, + ), replay_buffer=dict(replay_buffer_size=400000, ) + ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/DQN/gym_spaceInvadersnoframeskip_v4.py b/ding/config/example/DQN/gym_spaceInvadersnoframeskip_v4.py new file mode 100644 index 0000000000..f69a61e6dd --- /dev/null +++ b/ding/config/example/DQN/gym_spaceInvadersnoframeskip_v4.py @@ -0,0 +1,51 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='SpaceInvadersNoFrameskip-v4-DQN', + seed=0, + env=dict( + env_id='SpaceInvadersNoFrameskip-v4', + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + fram_stack=4, + stop_value=2000, + env_wrapper='atari_default', + ), + policy=dict( + cuda=True, + priority=False, + discount_factor=0.99, + nstep=3, + learn=dict( + update_per_collect=10, + batch_size=32, + learning_rate=0.0001, + # Frequency of target network update. + target_update_freq=500, + hook=dict(save_ckpt_after_iter=1000000, ) + ), + model=dict( + obs_shape=[4, 84, 84], + action_shape=6, + encoder_hidden_size_list=[128, 128, 512], + ), + collect=dict(n_sample=100, ), + other=dict( + eps=dict( + type='exp', + start=1., + end=0.05, + decay=1000000, + ), replay_buffer=dict(replay_buffer_size=400000, ) + ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/PG/__init__.py b/ding/config/example/PG/__init__.py new file mode 100644 index 0000000000..518449884f --- /dev/null +++ b/ding/config/example/PG/__init__.py @@ -0,0 +1,14 @@ +from easydict import EasyDict +from . import gym_pendulum_v1 + +supported_env_cfg = { + gym_pendulum_v1.cfg.env.env_id: gym_pendulum_v1.cfg, +} + +supported_env_cfg = EasyDict(supported_env_cfg) + +supported_env = { + gym_pendulum_v1.cfg.env.env_id: gym_pendulum_v1.env, +} + +supported_env = EasyDict(supported_env) diff --git a/ding/config/example/PG/gym_pendulum_v1.py b/ding/config/example/PG/gym_pendulum_v1.py new file mode 100644 index 0000000000..59b3e31eb8 --- /dev/null +++ b/ding/config/example/PG/gym_pendulum_v1.py @@ -0,0 +1,42 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='Pendulum-v1-PG', + seed=0, + env=dict( + env_id='Pendulum-v1', + collector_env_num=8, + evaluator_env_num=5, + n_evaluator_episode=5, + stop_value=-200, + act_scale=True, + ), + policy=dict( + cuda=False, + action_space='continuous', + model=dict( + action_space='continuous', + obs_shape=3, + action_shape=1, + ), + learn=dict( + batch_size=4000, + learning_rate=0.001, + entropy_weight=0.001, + ), + collect=dict( + n_episode=20, + unroll_len=1, + discount_factor=0.99, + ), + eval=dict(evaluator=dict(eval_freq=1, )) + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/PPOF/__init__.py b/ding/config/example/PPOF/__init__.py new file mode 100644 index 0000000000..2adaaf4df2 --- /dev/null +++ b/ding/config/example/PPOF/__init__.py @@ -0,0 +1,17 @@ +from easydict import EasyDict +from . import gym_lunarlander_v2 +from . import gym_lunarlandercontinuous_v2 + +supported_env_cfg = { + gym_lunarlander_v2.cfg.env_id: gym_lunarlander_v2.cfg, + gym_lunarlandercontinuous_v2.cfg.env_id: gym_lunarlandercontinuous_v2.cfg, +} + +supported_env_cfg = EasyDict(supported_env_cfg) + +supported_env = { + gym_lunarlander_v2.cfg.env_id: gym_lunarlander_v2.env, + gym_lunarlandercontinuous_v2.cfg.env_id: gym_lunarlandercontinuous_v2.env, +} + +supported_env = EasyDict(supported_env) diff --git a/ding/config/example/PPOF/gym_lunarlander_v2.py b/ding/config/example/PPOF/gym_lunarlander_v2.py new file mode 100644 index 0000000000..5a05266277 --- /dev/null +++ b/ding/config/example/PPOF/gym_lunarlander_v2.py @@ -0,0 +1,13 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='LunarLander-v2-PPO', + env_id='LunarLander-v2', + n_sample=400, + value_norm='popart', +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/PPOF/gym_lunarlandercontinuous_v2.py b/ding/config/example/PPOF/gym_lunarlandercontinuous_v2.py new file mode 100644 index 0000000000..c12c88fd42 --- /dev/null +++ b/ding/config/example/PPOF/gym_lunarlandercontinuous_v2.py @@ -0,0 +1,15 @@ +from easydict import EasyDict +from functools import partial +import ding.envs.gym_env + +cfg = dict( + exp_name='LunarLanderContinuous-V2-PPO', + env_id='LunarLanderContinuous-v2', + action_space='continuous', + n_sample=400, + act_scale=True, +) + +cfg = EasyDict(cfg) + +env = partial(ding.envs.gym_env.env, continuous=True) diff --git a/ding/config/example/PPOOffPolicy/__init__.py b/ding/config/example/PPOOffPolicy/__init__.py new file mode 100644 index 0000000000..2704b04c53 --- /dev/null +++ b/ding/config/example/PPOOffPolicy/__init__.py @@ -0,0 +1,23 @@ +from easydict import EasyDict +from . import gym_lunarlander_v2 +from . import gym_pongnoframeskip_v4 +from . import gym_qbertnoframeskip_v4 +from . import gym_spaceInvadersnoframeskip_v4 + +supported_env_cfg = { + gym_lunarlander_v2.cfg.env.env_id: gym_lunarlander_v2.cfg, + gym_pongnoframeskip_v4.cfg.env.env_id: gym_pongnoframeskip_v4.cfg, + gym_qbertnoframeskip_v4.cfg.env.env_id: gym_qbertnoframeskip_v4.cfg, + gym_spaceInvadersnoframeskip_v4.cfg.env.env_id: gym_spaceInvadersnoframeskip_v4.cfg, +} + +supported_env_cfg = EasyDict(supported_env_cfg) + +supported_env = { + gym_lunarlander_v2.cfg.env.env_id: gym_lunarlander_v2.env, + gym_pongnoframeskip_v4.cfg.env.env_id: gym_pongnoframeskip_v4.env, + gym_qbertnoframeskip_v4.cfg.env.env_id: gym_qbertnoframeskip_v4.env, + gym_spaceInvadersnoframeskip_v4.cfg.env.env_id: gym_spaceInvadersnoframeskip_v4.env, +} + +supported_env = EasyDict(supported_env) diff --git a/ding/config/example/PPOOffPolicy/gym_lunarlander_v2.py b/ding/config/example/PPOOffPolicy/gym_lunarlander_v2.py new file mode 100644 index 0000000000..1db3551b50 --- /dev/null +++ b/ding/config/example/PPOOffPolicy/gym_lunarlander_v2.py @@ -0,0 +1,44 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='LunarLander-v2-PPOOffPolicy', + env=dict( + collector_env_num=8, + evaluator_env_num=8, + env_id='LunarLander-v2', + n_evaluator_episode=8, + stop_value=260, + ), + policy=dict( + cuda=True, + model=dict( + obs_shape=8, + action_shape=4, + ), + learn=dict( + update_per_collect=4, + batch_size=64, + learning_rate=0.001, + value_weight=0.5, + entropy_weight=0.01, + clip_ratio=0.2, + nstep=1, + nstep_return=False, + adv_norm=True, + ), + collect=dict( + n_sample=128, + unroll_len=1, + discount_factor=0.99, + gae_lambda=0.95, + ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/PPOOffPolicy/gym_pongnoframeskip_v4.py b/ding/config/example/PPOOffPolicy/gym_pongnoframeskip_v4.py new file mode 100644 index 0000000000..0f376e2651 --- /dev/null +++ b/ding/config/example/PPOOffPolicy/gym_pongnoframeskip_v4.py @@ -0,0 +1,54 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='PongNoFrameskip-v4-PPOOffPolicy', + env=dict( + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=30, + env_id='PongNoFrameskip-v4', + frame_stack=4, + env_wrapper='atari_default', + ), + policy=dict( + cuda=True, + recompute_adv=True, + action_space='discrete', + model=dict( + obs_shape=[4, 84, 84], + action_shape=6, + action_space='discrete', + encoder_hidden_size_list=[64, 64, 128], + actor_head_hidden_size=128, + critic_head_hidden_size=128, + ), + learn=dict( + update_per_collect=10, + batch_size=320, + learning_rate=3e-4, + value_weight=0.5, + entropy_weight=0.001, + clip_ratio=0.2, + adv_norm=True, + # value_norm=True, + ignore_done=False, + grad_clip_type='clip_norm', + grad_clip_value=0.5, + ), + collect=dict( + n_sample=3200, + unroll_len=1, + discount_factor=0.99, + gae_lambda=0.95, + ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/PPOOffPolicy/gym_qbertnoframeskip_v4.py b/ding/config/example/PPOOffPolicy/gym_qbertnoframeskip_v4.py new file mode 100644 index 0000000000..7272ffebd6 --- /dev/null +++ b/ding/config/example/PPOOffPolicy/gym_qbertnoframeskip_v4.py @@ -0,0 +1,48 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='QbertNoFrameskip-v4-PPOOffPolicy', + env=dict( + collector_env_num=16, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=10000000000, + env_id='QbertNoFrameskip-v4', + frame_stack=4, + env_wrapper='atari_default', + ), + policy=dict( + cuda=True, + model=dict( + obs_shape=[4, 84, 84], + action_shape=6, + encoder_hidden_size_list=[32, 64, 64, 128], + actor_head_hidden_size=128, + critic_head_hidden_size=128, + critic_head_layer_num=2, + ), + learn=dict( + update_per_collect=18, + batch_size=128, + learning_rate=0.0001, + value_weight=1.0, + entropy_weight=0.005, + clip_ratio=0.1, + adv_norm=False, + ), + collect=dict( + n_sample=1024, + unroll_len=1, + discount_factor=0.99, + gae_lambda=0.95, + ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/PPOOffPolicy/gym_spaceInvadersnoframeskip_v4.py b/ding/config/example/PPOOffPolicy/gym_spaceInvadersnoframeskip_v4.py new file mode 100644 index 0000000000..18558553ac --- /dev/null +++ b/ding/config/example/PPOOffPolicy/gym_spaceInvadersnoframeskip_v4.py @@ -0,0 +1,48 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='SpaceInvadersNoFrameskip-v4-PPOOffPolicy', + env=dict( + collector_env_num=16, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=10000000000, + env_id='SpaceInvadersNoFrameskip-v4', + frame_stack=4, + env_wrapper='atari_default', + ), + policy=dict( + cuda=True, + model=dict( + obs_shape=[4, 84, 84], + action_shape=6, + encoder_hidden_size_list=[32, 64, 64, 128], + actor_head_hidden_size=128, + critic_head_hidden_size=128, + critic_head_layer_num=2, + ), + learn=dict( + update_per_collect=24, + batch_size=128, + learning_rate=0.0001, + value_weight=1.0, + entropy_weight=0.03, + clip_ratio=0.1, + adv_norm=False, + ), + collect=dict( + n_sample=1024, + unroll_len=1, + discount_factor=0.99, + gae_lambda=0.95, + ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/SAC/__init__.py b/ding/config/example/SAC/__init__.py new file mode 100644 index 0000000000..6e01f29d74 --- /dev/null +++ b/ding/config/example/SAC/__init__.py @@ -0,0 +1,29 @@ +from easydict import EasyDict +from . import gym_bipedalwalker_v3 +from . import gym_halfcheetah_v3 +from . import gym_hopper_v3 +from . import gym_lunarlandercontinuous_v2 +from . import gym_pendulum_v1 +from . import gym_walker2d_v3 + +supported_env_cfg = { + gym_bipedalwalker_v3.cfg.env.env_id: gym_bipedalwalker_v3.cfg, + gym_halfcheetah_v3.cfg.env.env_id: gym_halfcheetah_v3.cfg, + gym_hopper_v3.cfg.env.env_id: gym_hopper_v3.cfg, + gym_lunarlandercontinuous_v2.cfg.env.env_id: gym_lunarlandercontinuous_v2.cfg, + gym_pendulum_v1.cfg.env.env_id: gym_pendulum_v1.cfg, + gym_walker2d_v3.cfg.env.env_id: gym_walker2d_v3.cfg, +} + +supported_env_cfg = EasyDict(supported_env_cfg) + +supported_env = { + gym_bipedalwalker_v3.cfg.env.env_id: gym_bipedalwalker_v3.env, + gym_halfcheetah_v3.cfg.env.env_id: gym_halfcheetah_v3.env, + gym_hopper_v3.cfg.env.env_id: gym_hopper_v3.env, + gym_lunarlandercontinuous_v2.cfg.env.env_id: gym_lunarlandercontinuous_v2.env, + gym_pendulum_v1.cfg.env.env_id: gym_pendulum_v1.env, + gym_walker2d_v3.cfg.env.env_id: gym_walker2d_v3.env, +} + +supported_env = EasyDict(supported_env) diff --git a/ding/config/example/SAC/gym_bipedalwalker_v3.py b/ding/config/example/SAC/gym_bipedalwalker_v3.py new file mode 100644 index 0000000000..d97f131fd8 --- /dev/null +++ b/ding/config/example/SAC/gym_bipedalwalker_v3.py @@ -0,0 +1,47 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='BipedalWalker-v3-SAC', + seed=0, + env=dict( + env_id='BipedalWalker-v3', + collector_env_num=8, + evaluator_env_num=5, + n_evaluator_episode=5, + act_scale=True, + rew_clip=True, + ), + policy=dict( + cuda=True, + random_collect_size=10000, + model=dict( + obs_shape=24, + action_shape=4, + twin_critic=True, + action_space='reparameterization', + actor_head_hidden_size=128, + critic_head_hidden_size=128, + ), + learn=dict( + update_per_collect=64, + batch_size=256, + learning_rate_q=0.0003, + learning_rate_policy=0.0003, + learning_rate_alpha=0.0003, + target_theta=0.005, + discount_factor=0.99, + auto_alpha=True, + learner=dict(hook=dict(log_show_after_iter=1000, )) + ), + collect=dict(n_sample=64, ), + other=dict(replay_buffer=dict(replay_buffer_size=300000, ), ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/SAC/gym_halfcheetah_v3.py b/ding/config/example/SAC/gym_halfcheetah_v3.py new file mode 100644 index 0000000000..f2f0f8cc21 --- /dev/null +++ b/ding/config/example/SAC/gym_halfcheetah_v3.py @@ -0,0 +1,54 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='HalfCheetah-v3-SAC', + seed=0, + env=dict( + env_id='HalfCheetah-v3', + collector_env_num=1, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=12000, + env_wrapper='mujoco_default', + ), + policy=dict( + cuda=True, + random_collect_size=10000, + model=dict( + obs_shape=17, + action_shape=6, + twin_critic=True, + action_space='reparameterization', + actor_head_hidden_size=256, + critic_head_hidden_size=256, + ), + learn=dict( + update_per_collect=1, + batch_size=256, + learning_rate_q=1e-3, + learning_rate_policy=1e-3, + learning_rate_alpha=3e-4, + ignore_done=True, + target_theta=0.005, + discount_factor=0.99, + alpha=0.2, + reparameterization=True, + auto_alpha=False, + ), + collect=dict( + n_sample=1, + unroll_len=1, + ), + command=dict(), + eval=dict(), + other=dict(replay_buffer=dict(replay_buffer_size=1000000, ), ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/SAC/gym_hopper_v3.py b/ding/config/example/SAC/gym_hopper_v3.py new file mode 100644 index 0000000000..9a609e8c18 --- /dev/null +++ b/ding/config/example/SAC/gym_hopper_v3.py @@ -0,0 +1,41 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='Hopper-v3-SAC', + seed=0, + env=dict( + env_id='Hopper-v3', + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=6000, + env_wrapper='mujoco_default', + ), + policy=dict( + cuda=True, + random_collect_size=10000, + model=dict( + obs_shape=11, + action_shape=3, + action_space='reparameterization', + actor_head_hidden_size=256, + critic_head_hidden_size=256, + ), + learn=dict( + update_per_collect=1, + batch_size=256, + learning_rate_q=1e-3, + learning_rate_policy=1e-3, + reparameterization=True, + auto_alpha=False, + ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/SAC/gym_lunarlandercontinuous_v2.py b/ding/config/example/SAC/gym_lunarlandercontinuous_v2.py new file mode 100644 index 0000000000..37201da601 --- /dev/null +++ b/ding/config/example/SAC/gym_lunarlandercontinuous_v2.py @@ -0,0 +1,44 @@ +from easydict import EasyDict +from functools import partial +import ding.envs.gym_env + +cfg = dict( + exp_name='LunarLanderContinuous-v2-SAC', + seed=0, + env=dict( + env_id='LunarLanderContinuous-v2', + collector_env_num=4, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=260, + act_scale=True, + ), + policy=dict( + cuda=True, + random_collect_size=10000, + model=dict( + obs_shape=8, + action_shape=2, + action_space='reparameterization', + twin_critic=True, + ), + learn=dict( + update_per_collect=256, + batch_size=128, + learning_rate_q=1e-3, + learning_rate_policy=3e-4, + learning_rate_alpha=3e-4, + auto_alpha=True, + ), + collect=dict(n_sample=256, ), + eval=dict(evaluator=dict(eval_freq=1000, ), ), + other=dict(replay_buffer=dict(replay_buffer_size=int(1e5), ), ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = partial(ding.envs.gym_env.env, continuous=True) diff --git a/ding/config/example/SAC/gym_pendulum_v1.py b/ding/config/example/SAC/gym_pendulum_v1.py new file mode 100644 index 0000000000..45f2e9ddf3 --- /dev/null +++ b/ding/config/example/SAC/gym_pendulum_v1.py @@ -0,0 +1,49 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='Pendulum-v1-SAC', + seed=0, + env=dict( + env_id='Pendulum-v1', + collector_env_num=10, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=-250, + act_scale=True, + ), + policy=dict( + cuda=True, + priority=False, + random_collect_size=1000, + model=dict( + obs_shape=3, + action_shape=1, + twin_critic=True, + action_space='reparameterization', + actor_head_hidden_size=128, + critic_head_hidden_size=128, + ), + learn=dict( + update_per_collect=1, + batch_size=128, + learning_rate_q=0.001, + learning_rate_policy=0.001, + learning_rate_alpha=0.0003, + ignore_done=True, + target_theta=0.005, + discount_factor=0.99, + auto_alpha=True, + ), + collect=dict(n_sample=10, ), + eval=dict(evaluator=dict(eval_freq=100, )), + other=dict(replay_buffer=dict(replay_buffer_size=100000, ), ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/SAC/gym_walker2d_v3.py b/ding/config/example/SAC/gym_walker2d_v3.py new file mode 100644 index 0000000000..f40d68d117 --- /dev/null +++ b/ding/config/example/SAC/gym_walker2d_v3.py @@ -0,0 +1,54 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='Walker2d-v3-SAC', + seed=0, + env=dict( + env_id='Walker2d-v3', + collector_env_num=1, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=6000, + env_wrapper='mujoco_default', + ), + policy=dict( + cuda=True, + random_collect_size=10000, + model=dict( + obs_shape=17, + action_shape=6, + twin_critic=True, + action_space='reparameterization', + actor_head_hidden_size=256, + critic_head_hidden_size=256, + ), + learn=dict( + update_per_collect=1, + batch_size=256, + learning_rate_q=1e-3, + learning_rate_policy=1e-3, + learning_rate_alpha=3e-4, + ignore_done=False, + target_theta=0.005, + discount_factor=0.99, + alpha=0.2, + reparameterization=True, + auto_alpha=False, + ), + collect=dict( + n_sample=1, + unroll_len=1, + ), + command=dict(), + eval=dict(), + other=dict(replay_buffer=dict(replay_buffer_size=1000000, ), ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/SQL/__init__.py b/ding/config/example/SQL/__init__.py new file mode 100644 index 0000000000..9637366fb4 --- /dev/null +++ b/ding/config/example/SQL/__init__.py @@ -0,0 +1,14 @@ +from easydict import EasyDict +from . import gym_lunarlander_v2 + +supported_env_cfg = { + gym_lunarlander_v2.cfg.env.env_id: gym_lunarlander_v2.cfg, +} + +supported_env_cfg = EasyDict(supported_env_cfg) + +supported_env = { + gym_lunarlander_v2.cfg.env.env_id: gym_lunarlander_v2.env, +} + +supported_env = EasyDict(supported_env) diff --git a/ding/config/example/SQL/gym_lunarlander_v2.py b/ding/config/example/SQL/gym_lunarlander_v2.py new file mode 100644 index 0000000000..648564dadb --- /dev/null +++ b/ding/config/example/SQL/gym_lunarlander_v2.py @@ -0,0 +1,43 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='LunarLander-v2-SQL', + env=dict( + collector_env_num=8, + evaluator_env_num=8, + env_id='LunarLander-v2', + n_evaluator_episode=8, + stop_value=260, + ), + policy=dict( + cuda=True, + model=dict( + obs_shape=8, + action_shape=4, + encoder_hidden_size_list=[128, 128, 64], + dueling=True, + ), + nstep=1, + discount_factor=0.97, + learn=dict(batch_size=64, learning_rate=0.001, alpha=0.08), + collect=dict(n_sample=64), + eval=dict(evaluator=dict(eval_freq=50, )), # note: this is the times after which you learns to evaluate + other=dict( + eps=dict( + type='exp', + start=0.95, + end=0.1, + decay=10000, + ), + replay_buffer=dict(replay_buffer_size=20000, ), + ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/TD3/__init__.py b/ding/config/example/TD3/__init__.py new file mode 100644 index 0000000000..6e01f29d74 --- /dev/null +++ b/ding/config/example/TD3/__init__.py @@ -0,0 +1,29 @@ +from easydict import EasyDict +from . import gym_bipedalwalker_v3 +from . import gym_halfcheetah_v3 +from . import gym_hopper_v3 +from . import gym_lunarlandercontinuous_v2 +from . import gym_pendulum_v1 +from . import gym_walker2d_v3 + +supported_env_cfg = { + gym_bipedalwalker_v3.cfg.env.env_id: gym_bipedalwalker_v3.cfg, + gym_halfcheetah_v3.cfg.env.env_id: gym_halfcheetah_v3.cfg, + gym_hopper_v3.cfg.env.env_id: gym_hopper_v3.cfg, + gym_lunarlandercontinuous_v2.cfg.env.env_id: gym_lunarlandercontinuous_v2.cfg, + gym_pendulum_v1.cfg.env.env_id: gym_pendulum_v1.cfg, + gym_walker2d_v3.cfg.env.env_id: gym_walker2d_v3.cfg, +} + +supported_env_cfg = EasyDict(supported_env_cfg) + +supported_env = { + gym_bipedalwalker_v3.cfg.env.env_id: gym_bipedalwalker_v3.env, + gym_halfcheetah_v3.cfg.env.env_id: gym_halfcheetah_v3.env, + gym_hopper_v3.cfg.env.env_id: gym_hopper_v3.env, + gym_lunarlandercontinuous_v2.cfg.env.env_id: gym_lunarlandercontinuous_v2.env, + gym_pendulum_v1.cfg.env.env_id: gym_pendulum_v1.env, + gym_walker2d_v3.cfg.env.env_id: gym_walker2d_v3.env, +} + +supported_env = EasyDict(supported_env) diff --git a/ding/config/example/TD3/gym_bipedalwalker_v3.py b/ding/config/example/TD3/gym_bipedalwalker_v3.py new file mode 100644 index 0000000000..9e4cc24321 --- /dev/null +++ b/ding/config/example/TD3/gym_bipedalwalker_v3.py @@ -0,0 +1,52 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='Bipedalwalker-v3-TD3', + seed=0, + env=dict( + env_id='BipedalWalker-v3', + collector_env_num=8, + evaluator_env_num=5, + n_evaluator_episode=5, + act_scale=True, + rew_clip=True, + ), + policy=dict( + cuda=True, + random_collect_size=10000, + model=dict( + obs_shape=24, + action_shape=4, + twin_critic=True, + action_space='regression', + actor_head_hidden_size=400, + critic_head_hidden_size=400, + ), + learn=dict( + update_per_collect=64, + batch_size=256, + learning_rate_actor=0.0003, + learning_rate_critic=0.0003, + target_theta=0.005, + discount_factor=0.99, + actor_update_freq=2, + noise=True, + noise_sigma=0.2, + noise_range=dict( + min=-0.5, + max=0.5, + ), + learner=dict(hook=dict(log_show_after_iter=1000, )) + ), + collect=dict(n_sample=64, ), + other=dict(replay_buffer=dict(replay_buffer_size=300000, ), ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/TD3/gym_halfcheetah_v3.py b/ding/config/example/TD3/gym_halfcheetah_v3.py new file mode 100644 index 0000000000..ddd2f1a68e --- /dev/null +++ b/ding/config/example/TD3/gym_halfcheetah_v3.py @@ -0,0 +1,56 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='HalfCheetah-v3-TD3', + seed=0, + env=dict( + env_id='HalfCheetah-v3', + collector_env_num=1, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=11000, + env_wrapper='mujoco_default', + ), + policy=dict( + cuda=True, + random_collect_size=25000, + model=dict( + obs_shape=17, + action_shape=6, + twin_critic=True, + actor_head_hidden_size=256, + critic_head_hidden_size=256, + action_space='regression', + ), + learn=dict( + update_per_collect=1, + batch_size=256, + learning_rate_actor=1e-3, + learning_rate_critic=1e-3, + ignore_done=True, + target_theta=0.005, + discount_factor=0.99, + actor_update_freq=2, + noise=True, + noise_sigma=0.2, + noise_range=dict( + min=-0.5, + max=0.5, + ), + ), + collect=dict( + n_sample=1, + unroll_len=1, + noise_sigma=0.1, + ), + other=dict(replay_buffer=dict(replay_buffer_size=1000000, ), ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/TD3/gym_hopper_v3.py b/ding/config/example/TD3/gym_hopper_v3.py new file mode 100644 index 0000000000..e213323fd4 --- /dev/null +++ b/ding/config/example/TD3/gym_hopper_v3.py @@ -0,0 +1,35 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='Hopper-v3-TD3', + seed=0, + env=dict( + env_id='Hopper-v3', + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=6000, + env_wrapper='mujoco_default', + ), + policy=dict( + cuda=True, + random_collect_size=25000, + model=dict( + obs_shape=11, + action_shape=3, + actor_head_hidden_size=256, + critic_head_hidden_size=256, + action_space='regression', + ), + collect=dict(n_sample=1, ), + other=dict(replay_buffer=dict(replay_buffer_size=1000000, ), ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/TD3/gym_lunarlandercontinuous_v2.py b/ding/config/example/TD3/gym_lunarlandercontinuous_v2.py new file mode 100644 index 0000000000..9798705ff2 --- /dev/null +++ b/ding/config/example/TD3/gym_lunarlandercontinuous_v2.py @@ -0,0 +1,50 @@ +from easydict import EasyDict +from functools import partial +import ding.envs.gym_env + +cfg = dict( + exp_name='LunarLanderContinuous-V2-TD3', + seed=0, + env=dict( + env_id='LunarLanderContinuous-v2', + collector_env_num=4, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=240, + act_scale=True, + ), + policy=dict( + cuda=True, + random_collect_size=10000, + model=dict( + obs_shape=8, + action_shape=2, + action_space='regression', + ), + learn=dict( + update_per_collect=256, + batch_size=256, + learning_rate_actor=3e-4, + learning_rate_critic=1e-3, + noise=True, + noise_sigma=0.1, + noise_range=dict( + min=-0.5, + max=0.5, + ), + ), + collect=dict( + n_sample=256, + noise_sigma=0.1, + ), + eval=dict(evaluator=dict(eval_freq=1000, ), ), + other=dict(replay_buffer=dict(replay_buffer_size=100000, ), ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = partial(ding.envs.gym_env.env, continuous=True) diff --git a/ding/config/example/TD3/gym_pendulum_v1.py b/ding/config/example/TD3/gym_pendulum_v1.py new file mode 100644 index 0000000000..57ebeeae6c --- /dev/null +++ b/ding/config/example/TD3/gym_pendulum_v1.py @@ -0,0 +1,54 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='Pendulum-v1-TD3', + seed=0, + env=dict( + env_id='Pendulum-v1', + collector_env_num=8, + evaluator_env_num=5, + n_evaluator_episode=5, + stop_value=-250, + act_scale=True, + ), + policy=dict( + cuda=False, + priority=False, + random_collect_size=800, + model=dict( + obs_shape=3, + action_shape=1, + twin_critic=True, + action_space='regression', + ), + learn=dict( + update_per_collect=2, + batch_size=128, + learning_rate_actor=0.001, + learning_rate_critic=0.001, + ignore_done=True, + actor_update_freq=2, + noise=True, + noise_sigma=0.1, + noise_range=dict( + min=-0.5, + max=0.5, + ), + ), + collect=dict( + n_sample=48, + noise_sigma=0.1, + collector=dict(collect_print_freq=1000, ), + ), + eval=dict(evaluator=dict(eval_freq=100, ), ), + other=dict(replay_buffer=dict(replay_buffer_size=20000, ), ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/TD3/gym_walker2d_v3.py b/ding/config/example/TD3/gym_walker2d_v3.py new file mode 100644 index 0000000000..92b88e1e08 --- /dev/null +++ b/ding/config/example/TD3/gym_walker2d_v3.py @@ -0,0 +1,58 @@ +from easydict import EasyDict +import ding.envs.gym_env + +cfg = dict( + exp_name='Walker2d-v3-TD3', + seed=0, + env=dict( + env_id='Walker2d-v3', + norm_obs=dict(use_norm=False, ), + norm_reward=dict(use_norm=False, ), + collector_env_num=1, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=6000, + env_wrapper='mujoco_default', + ), + policy=dict( + cuda=True, + random_collect_size=25000, + model=dict( + obs_shape=17, + action_shape=6, + twin_critic=True, + actor_head_hidden_size=256, + critic_head_hidden_size=256, + action_space='regression', + ), + learn=dict( + update_per_collect=1, + batch_size=256, + learning_rate_actor=1e-3, + learning_rate_critic=1e-3, + ignore_done=False, + target_theta=0.005, + discount_factor=0.99, + actor_update_freq=2, + noise=True, + noise_sigma=0.2, + noise_range=dict( + min=-0.5, + max=0.5, + ), + ), + collect=dict( + n_sample=1, + unroll_len=1, + noise_sigma=0.1, + ), + other=dict(replay_buffer=dict(replay_buffer_size=1000000, ), ), + ), + wandb_logger=dict( + gradient_logger=True, video_logger=True, plot_logger=True, action_logger=True, return_logger=False + ), +) + +cfg = EasyDict(cfg) + +env = ding.envs.gym_env.env diff --git a/ding/config/example/__init__.py b/ding/config/example/__init__.py new file mode 100644 index 0000000000..7b2c8e750c --- /dev/null +++ b/ding/config/example/__init__.py @@ -0,0 +1,10 @@ +from . import A2C +from . import C51 +from . import DDPG +from . import DQN +from . import PG +from . import PPOF +from . import PPOOffPolicy +from . import SAC +from . import SQL +from . import TD3 diff --git a/ding/config/utils.py b/ding/config/utils.py index 8b11028034..5a9a2d6664 100644 --- a/ding/config/utils.py +++ b/ding/config/utils.py @@ -1,4 +1,4 @@ -from typing import Optional, List, NoReturn +from typing import Optional, List import copy from easydict import EasyDict @@ -240,7 +240,7 @@ def parallel_transform_k8s( return cfg -def save_config_formatted(config_: dict, path: str = 'formatted_total_config.py') -> NoReturn: +def save_config_formatted(config_: dict, path: str = 'formatted_total_config.py') -> None: """ Overview: save formatted configuration to python file that can be read by serial_pipeline directly. diff --git a/ding/data/__init__.py b/ding/data/__init__.py index 79ac868c86..b72987cac9 100644 --- a/ding/data/__init__.py +++ b/ding/data/__init__.py @@ -1,3 +1,7 @@ from torch.utils.data import Dataset, DataLoader from ding.utils.data import create_dataset, offline_data_save_type # for compatibility from .buffer import * +from .storage import * +from .storage_loader import StorageLoader, FileStorageLoader +from .shm_buffer import ShmBufferContainer, ShmBuffer +from .model_loader import ModelLoader, FileModelLoader diff --git a/ding/data/buffer/buffer.py b/ding/data/buffer/buffer.py index 86bfe0335d..53b3f39bd1 100644 --- a/ding/data/buffer/buffer.py +++ b/ding/data/buffer/buffer.py @@ -129,6 +129,26 @@ def delete(self, index: str): """ raise NotImplementedError + @abstractmethod + def save_data(self, file_name: str): + """ + Overview: + Save buffer data into a file. + Arguments: + - file_name (:obj:`str`): file name of buffer data + """ + raise NotImplementedError + + @abstractmethod + def load_data(self, file_name: str): + """ + Overview: + Load buffer data from a file. + Arguments: + - file_name (:obj:`str`): file name of buffer data + """ + raise NotImplementedError + @abstractmethod def count(self) -> int: raise NotImplementedError diff --git a/ding/data/buffer/deque_buffer.py b/ding/data/buffer/deque_buffer.py index ab9ba5444f..6c4af9cac5 100644 --- a/ding/data/buffer/deque_buffer.py +++ b/ding/data/buffer/deque_buffer.py @@ -1,7 +1,9 @@ +import os import itertools import random import uuid from ditk import logging +import hickle from typing import Any, Iterable, List, Optional, Tuple, Union from collections import Counter from collections import defaultdict, deque, OrderedDict @@ -52,16 +54,18 @@ class DequeBuffer(Buffer): A buffer implementation based on the deque structure. """ - def __init__(self, size: int) -> None: + def __init__(self, size: int, sliced: bool = False) -> None: """ Overview: The initialization method of DequeBuffer. Arguments: - size (:obj:`int`): The maximum number of objects that the buffer can hold. + - sliced (:obj:`bool`): The flag whether slice data by unroll_len when sample by group """ super().__init__(size=size) self.storage = deque(maxlen=size) self.indices = BufferIndex(maxlen=size) + self.sliced = sliced # Meta index is a dict which uses deque as values self.meta_index = {} @@ -140,7 +144,7 @@ def sample( sampled_data = [hashed_data[index] for index in indices] elif groupby: sampled_data = self._sample_by_group( - size=size, groupby=groupby, replace=replace, unroll_len=unroll_len, storage=storage + size=size, groupby=groupby, replace=replace, unroll_len=unroll_len, storage=storage, sliced=self.sliced ) else: if replace: @@ -209,6 +213,22 @@ def delete(self, indices: Union[str, Iterable[str]]) -> None: key_value_pairs = zip(remain_indices, range(len(indices))) self.indices = BufferIndex(self.storage.maxlen, key_value_pairs) + def save_data(self, file_name: str): + if not os.path.exists(os.path.dirname(file_name)): + # If the folder for the specified file does not exist, it will be created. + if os.path.dirname(file_name) != "": + os.makedirs(os.path.dirname(file_name)) + hickle.dump( + py_obj=( + self.storage, + self.indices, + self.meta_index, + ), file_obj=file_name + ) + + def load_data(self, file_name: str): + self.storage, self.indices, self.meta_index = hickle.load(file_name) + def count(self) -> int: """ Overview: @@ -219,10 +239,17 @@ def count(self) -> int: def get(self, idx: int) -> BufferedData: """ Overview: - The method that returns the BufferedData object given a specific index. + The method that returns the BufferedData object by subscript idx (int). """ return self.storage[idx] + def get_by_index(self, index: str) -> BufferedData: + """ + Overview: + The method that returns the BufferedData object given a specific index (str). + """ + return self.storage[self.indices.get(index)] + @apply_middleware("clear") def clear(self) -> None: """ @@ -233,25 +260,6 @@ def clear(self) -> None: self.indices.clear() self.meta_index = {} - def import_data(self, data_with_meta: List[Tuple[Any, dict]]) -> None: - """ - Overview: - The method that push data by sequence. - Arguments: - data_with_meta (List[Tuple[Any, dict]]): The sequence of (data, meta) tuples. - """ - for data, meta in data_with_meta: - self._push(data, meta) - - def export_data(self) -> List[BufferedData]: - """ - Overview: - The method that export all data in the buffer by sequence. - Returns: - storage (List[BufferedData]): All ``BufferedData`` objects stored in the buffer. - """ - return list(self.storage) - def _push(self, data: Any, meta: Optional[dict] = None) -> BufferedData: index = uuid.uuid1().hex if meta is None: @@ -302,7 +310,8 @@ def _sample_by_group( groupby: str, replace: bool = False, unroll_len: Optional[int] = None, - storage: deque = None + storage: deque = None, + sliced: bool = False ) -> List[List[BufferedData]]: """ Overview: @@ -325,6 +334,8 @@ def filter_by_unroll_len(): if unroll_len and unroll_len > 1: group_names = filter_by_unroll_len() + if len(group_names) == 0: + return [] else: group_names = list(set(self.meta_index[groupby])) @@ -349,8 +360,19 @@ def filter_by_unroll_len(): seq_data = sampled_data[group] # Filter records by unroll_len if unroll_len: - start_indice = random.choice(range(max(1, len(seq_data) - unroll_len))) - seq_data = seq_data[start_indice:start_indice + unroll_len] + # slice b unroll_len. If don’t do this, more likely obtain duplicate data, \ + # and the training will easily crash. + if sliced: + start_indice = random.choice(range(max(1, len(seq_data)))) + start_indice = start_indice // unroll_len + if start_indice == (len(seq_data) - 1) // unroll_len: + seq_data = seq_data[-unroll_len:] + else: + seq_data = seq_data[start_indice * unroll_len:start_indice * unroll_len + unroll_len] + else: + start_indice = random.choice(range(max(1, len(seq_data) - unroll_len))) + seq_data = seq_data[start_indice:start_indice + unroll_len] + final_sampled_data.append(seq_data) return final_sampled_data diff --git a/ding/data/buffer/deque_buffer_wrapper.py b/ding/data/buffer/deque_buffer_wrapper.py index 84fe7cc5de..f2e4945a9f 100644 --- a/ding/data/buffer/deque_buffer_wrapper.py +++ b/ding/data/buffer/deque_buffer_wrapper.py @@ -1,7 +1,9 @@ +import os from typing import Optional import copy from easydict import EasyDict import numpy as np +import hickle from ding.data.buffer import DequeBuffer from ding.data.buffer.middleware import use_time_check, PriorityExperienceReplay @@ -51,7 +53,6 @@ def __init__( self.buffer.use( PriorityExperienceReplay( self.buffer, - self.cfg.replay_buffer_size, IS_weight=self.cfg.priority_IS_weight, priority_power_factor=self.cfg.priority_power_factor, IS_weight_power_factor=self.cfg.IS_weight_power_factor, @@ -112,3 +113,9 @@ def update(self, meta: dict) -> None: def count(self) -> int: return self.buffer.count() + + def save_data(self, file_name): + self.buffer.save_data(file_name) + + def load_data(self, file_name: str): + self.buffer.load_data(file_name) diff --git a/ding/data/buffer/middleware/priority.py b/ding/data/buffer/middleware/priority.py index f6df6c1be3..b890bac78f 100644 --- a/ding/data/buffer/middleware/priority.py +++ b/ding/data/buffer/middleware/priority.py @@ -54,7 +54,10 @@ def __init__( def push(self, chain: Callable, data: Any, meta: Optional[dict] = None, *args, **kwargs) -> BufferedData: if meta is None: - meta = {'priority': self.max_priority} + if 'priority' in data: + meta = {'priority': data.pop('priority')} + else: + meta = {'priority': self.max_priority} else: if 'priority' not in meta: meta['priority'] = self.max_priority @@ -105,12 +108,12 @@ def update(self, chain: Callable, index: str, data: Any, meta: Any, *args, **kwa self.max_priority = max(self.max_priority, new_priority) def delete(self, chain: Callable, index: str, *args, **kwargs) -> None: - for item in self.buffer.storage: - meta = item.meta - priority_idx = meta['priority_idx'] - self.sum_tree[priority_idx] = self.sum_tree.neutral_element - self.min_tree[priority_idx] = self.min_tree.neutral_element - self.buffer_idx.pop(priority_idx) + item = self.buffer.get_by_index(index) + meta = item.meta + priority_idx = meta['priority_idx'] + self.sum_tree[priority_idx] = self.sum_tree.neutral_element + self.min_tree[priority_idx] = self.min_tree.neutral_element + self.buffer_idx.pop(priority_idx) return chain(index, *args, **kwargs) def clear(self, chain: Callable) -> None: diff --git a/ding/data/buffer/tests/test_buffer.py b/ding/data/buffer/tests/test_buffer.py index e50ab18c5a..647816d36d 100644 --- a/ding/data/buffer/tests/test_buffer.py +++ b/ding/data/buffer/tests/test_buffer.py @@ -1,7 +1,9 @@ +import os import pytest import time import random import functools +import tempfile from typing import Callable from ding.data.buffer import DequeBuffer from ding.data.buffer.buffer import BufferedData @@ -87,6 +89,25 @@ def test_rate_limit_push_sample(): assert 5 not in buffer.sample(5) +@pytest.mark.unittest +def test_load_and_save(): + buffer = DequeBuffer(size=10).use(RateLimit(max_rate=5)) + buffer.meta_index = {"label": []} + for i in range(10): + buffer.push(i, meta={"label": i}) + assert buffer.count() == 5 + assert 5 not in buffer.sample(5) + with tempfile.TemporaryDirectory() as tmpdirname: + test_file = os.path.join(tmpdirname, "data.hkl") + buffer.save_data(test_file) + buffer_new = DequeBuffer(size=10).use(RateLimit(max_rate=5)) + buffer_new.load_data(test_file) + assert buffer_new.count() == 5 + assert 5 not in buffer_new.sample(5) + assert len(buffer.meta_index["label"]) == 5 + assert all([index < 5 for index in buffer.meta_index["label"]]) + + @pytest.mark.unittest def test_buffer_view(): buf1 = DequeBuffer(size=10) @@ -250,17 +271,6 @@ def test_groupby(): assert len(sampled_data) == 2 -@pytest.mark.unittest -def test_import_export(): - buffer = DequeBuffer(size=10) - data_with_meta = [(i, {}) for i in range(10)] - buffer.import_data(data_with_meta) - assert buffer.count() == 10 - - sampled_data = buffer.export_data() - assert len(sampled_data) == 10 - - @pytest.mark.unittest def test_dataset(): buffer = DequeBuffer(size=10) @@ -316,3 +326,27 @@ def test_insufficient_unroll_len_in_group(): # Ensure samples in each group is continuous result = functools.reduce(lambda a, b: a and a.data + 1 == b.data and b, grouped_data) assert isinstance(result, BufferedData), "Not continuous" + + +@pytest.mark.unittest +def test_slice_unroll_len_in_group(): + buffer = DequeBuffer(size=100, sliced=True) + data_len = 10 + unroll_len = 4 + start_index = list(range(0, data_len, unroll_len)) + [data_len - unroll_len] + for i in range(data_len): + for env_id in list("ABC"): + buffer.push(i, {"env": env_id}) + + sampled_data = buffer.sample(3, groupby="env", unroll_len=unroll_len) + assert len(sampled_data) == 3 + for grouped_data in sampled_data: + assert len(grouped_data) == 4 + # Ensure each group has the same env + env_ids = set(map(lambda sample: sample.meta["env"], grouped_data)) + assert len(env_ids) == 1 + # Ensure samples in each group is continuous + result = functools.reduce(lambda a, b: a and a.data + 1 == b.data and b, grouped_data) + assert isinstance(result, BufferedData), "Not continuous" + # Ensure data after sliced start from correct index + assert grouped_data[0].data in start_index diff --git a/ding/data/buffer/tests/test_benchmark.py b/ding/data/buffer/tests/test_buffer_benchmark.py similarity index 100% rename from ding/data/buffer/tests/test_benchmark.py rename to ding/data/buffer/tests/test_buffer_benchmark.py diff --git a/ding/data/buffer/tests/test_middleware.py b/ding/data/buffer/tests/test_middleware.py index 9e7ca64fe0..94c51ba740 100644 --- a/ding/data/buffer/tests/test_middleware.py +++ b/ding/data/buffer/tests/test_middleware.py @@ -106,6 +106,39 @@ def test_priority(): assert data[0].meta['priority'] == 3.0 buffer.delete(data[0].index) assert buffer.count() == N + N - 1 + assert len(buffer._middleware[0].buffer_idx) == N + N - 1 + buffer.clear() + assert buffer.count() == 0 + + +@pytest.mark.unittest +def test_priority_from_collector(): + N = 5 + buffer = DequeBuffer(size=10) + buffer.use(PriorityExperienceReplay(buffer, IS_weight=True)) + for _ in range(N): + tmp_data = get_data() + tmp_data['priority'] = 2.0 + buffer.push(get_data()) + assert buffer.count() == N + for _ in range(N): + tmp_data = get_data() + tmp_data['priority'] = 2.0 + buffer.push(get_data()) + assert buffer.count() == N + N + data = buffer.sample(size=N + N, replace=False) + assert len(data) == N + N + for item in data: + meta = item.meta + assert set(meta.keys()).issuperset(set(['priority', 'priority_idx', 'priority_IS'])) + meta['priority'] = 3.0 + for item in data: + data, index, meta = item.data, item.index, item.meta + buffer.update(index, data, meta) + data = buffer.sample(size=1) + assert data[0].meta['priority'] == 3.0 + buffer.delete(data[0].index) + assert buffer.count() == N + N - 1 buffer.clear() assert buffer.count() == 0 diff --git a/dizoo/board_games/__init__.py b/ding/data/level_replay/__init__.py similarity index 100% rename from dizoo/board_games/__init__.py rename to ding/data/level_replay/__init__.py diff --git a/ding/data/level_replay/level_sampler.py b/ding/data/level_replay/level_sampler.py index b0930d9680..ac51fc4d4d 100644 --- a/ding/data/level_replay/level_sampler.py +++ b/ding/data/level_replay/level_sampler.py @@ -65,7 +65,7 @@ def __init__( self.unseen_seed_weights = np.ones(len(seeds)) self.seed_scores = np.zeros(len(seeds)) - self.partial_seed_scores = np.zeros((num_actors, len(seeds)), dtype=np.float) + self.partial_seed_scores = np.zeros((num_actors, len(seeds)), dtype=np.float32) self.partial_seed_steps = np.zeros((num_actors, len(seeds)), dtype=np.int64) self.seed_staleness = np.zeros(len(seeds)) @@ -183,6 +183,7 @@ def _update_with_rollouts(self, train_data: dict, num_actors: int, all_total_ste continue seed_t = level_seeds[start_t, actor_index].item() + seed_t = int(seed_t) seed_idx_t = self.seed2index[seed_t] score_function_kwargs = {} @@ -234,7 +235,7 @@ def _sample_replay_level(self): sample_weights = self._sample_weights() if np.isclose(np.sum(sample_weights), 0): - sample_weights = np.ones_like(sample_weights, dtype=np.float) / len(sample_weights) + sample_weights = np.ones_like(sample_weights, dtype=np.float32) / len(sample_weights) seed_idx = np.random.choice(range(len(self.seeds)), 1, p=sample_weights)[0] seed = self.seeds[seed_idx] diff --git a/ding/data/model_loader.py b/ding/data/model_loader.py new file mode 100644 index 0000000000..cd3182897b --- /dev/null +++ b/ding/data/model_loader.py @@ -0,0 +1,155 @@ +from abc import ABC, abstractmethod +import logging +from os import path +import os +from threading import Thread +from time import sleep, time +from typing import Callable, Optional +import uuid +import torch.multiprocessing as mp + +import torch +from ding.data.storage.file import FileModelStorage +from ding.data.storage.storage import Storage +from ding.framework import Supervisor +from ding.framework.supervisor import ChildType, SendPayload + + +class ModelWorker(): + + def __init__(self, model: torch.nn.Module) -> None: + self._model = model + + def save(self, storage: Storage) -> Storage: + storage.save(self._model.state_dict()) + return storage + + +class ModelLoader(Supervisor, ABC): + + def __init__(self, model: torch.nn.Module) -> None: + """ + Overview: + Save and send models asynchronously and load them synchronously. + Arguments: + - model (:obj:`torch.nn.Module`): Torch module. + """ + if next(model.parameters()).is_cuda: + super().__init__(type_=ChildType.PROCESS, mp_ctx=mp.get_context("spawn")) + else: + super().__init__(type_=ChildType.PROCESS) + self._model = model + self._send_callback_loop = None + self._send_callbacks = {} + self._model_worker = ModelWorker(self._model) + + def start(self): + if not self._running: + self._model.share_memory() + self.register(self._model_worker) + self.start_link() + self._send_callback_loop = Thread(target=self._loop_send_callback, daemon=True) + self._send_callback_loop.start() + + def shutdown(self, timeout: Optional[float] = None) -> None: + super().shutdown(timeout) + self._send_callback_loop = None + self._send_callbacks = {} + + def _loop_send_callback(self): + while True: + payload = self.recv(ignore_err=True) + if payload.err: + logging.warning("Got error when loading data: {}".format(payload.err)) + if payload.req_id in self._send_callbacks: + del self._send_callbacks[payload.req_id] + else: + if payload.req_id in self._send_callbacks: + callback = self._send_callbacks.pop(payload.req_id) + callback(payload.data) + + def load(self, storage: Storage) -> object: + """ + Overview: + Load model synchronously. + Arguments: + - storage (:obj:`Stroage`): The model should be wrapped in a storage object, e.g. FileModelStorage. + Returns: + - object (:obj:): The loaded model. + """ + return storage.load() + + @abstractmethod + def save(self, callback: Callable) -> Storage: + """ + Overview: + Save model asynchronously. + Arguments: + - callback (:obj:`Callable`): The callback function after saving model. + Returns: + - storage (:obj:`Storage`): The storage object is created synchronously, so it can be returned. + """ + raise NotImplementedError + + +class FileModelLoader(ModelLoader): + + def __init__(self, model: torch.nn.Module, dirname: str, ttl: int = 20) -> None: + """ + Overview: + Model loader using files as storage media. + Arguments: + - model (:obj:`torch.nn.Module`): Torch module. + - dirname (:obj:`str`): The directory for saving files. + - ttl (:obj:`int`): Files will be automatically cleaned after ttl. Note that \ + files that do not time out when the process is stopped are not cleaned up \ + (to avoid errors when other processes read the file), so you may need to \ + clean up the remaining files manually + """ + super().__init__(model) + self._dirname = dirname + self._ttl = ttl + self._files = [] + self._cleanup_thread = None + + def _start_cleanup(self): + """ + Overview: + Start a cleanup thread to clean up files that are taking up too much time on the disk. + """ + if self._cleanup_thread is None: + self._cleanup_thread = Thread(target=self._loop_cleanup, daemon=True) + self._cleanup_thread.start() + + def shutdown(self, timeout: Optional[float] = None) -> None: + super().shutdown(timeout) + self._cleanup_thread = None + + def _loop_cleanup(self): + while True: + if len(self._files) == 0 or time() - self._files[0][0] < self._ttl: + sleep(1) + continue + _, file_path = self._files.pop(0) + if path.exists(file_path): + os.remove(file_path) + + def save(self, callback: Callable) -> FileModelStorage: + if not self._running: + logging.warning("Please start model loader before saving model.") + return + if not path.exists(self._dirname): + os.mkdir(self._dirname) + file_path = "model_{}.pth.tar".format(uuid.uuid1()) + file_path = path.join(self._dirname, file_path) + model_storage = FileModelStorage(file_path) + payload = SendPayload(proc_id=0, method="save", args=[model_storage]) + self.send(payload) + + def clean_callback(storage: Storage): + self._files.append([time(), file_path]) + callback(storage) + + self._send_callbacks[payload.req_id] = clean_callback + self._start_cleanup() + return model_storage diff --git a/ding/data/shm_buffer.py b/ding/data/shm_buffer.py new file mode 100644 index 0000000000..b76f5d56e9 --- /dev/null +++ b/ding/data/shm_buffer.py @@ -0,0 +1,133 @@ +from typing import Any, Optional, Union, Tuple, Dict +from multiprocessing import Array +import ctypes +import numpy as np +import torch + +_NTYPE_TO_CTYPE = { + np.bool_: ctypes.c_bool, + np.uint8: ctypes.c_uint8, + np.uint16: ctypes.c_uint16, + np.uint32: ctypes.c_uint32, + np.uint64: ctypes.c_uint64, + np.int8: ctypes.c_int8, + np.int16: ctypes.c_int16, + np.int32: ctypes.c_int32, + np.int64: ctypes.c_int64, + np.float32: ctypes.c_float, + np.float64: ctypes.c_double, +} + + +class ShmBuffer(): + """ + Overview: + Shared memory buffer to store numpy array. + """ + + def __init__( + self, + dtype: Union[type, np.dtype], + shape: Tuple[int], + copy_on_get: bool = True, + ctype: Optional[type] = None + ) -> None: + """ + Overview: + Initialize the buffer. + Arguments: + - dtype (:obj:`Union[type, np.dtype]`): The dtype of the data to limit the size of the buffer. + - shape (:obj:`Tuple[int]`): The shape of the data to limit the size of the buffer. + - copy_on_get (:obj:`bool`): Whether to copy data when calling get method. + - ctype (:obj:`Optional[type]`): Origin class type, e.g. np.ndarray, torch.Tensor. + """ + if isinstance(dtype, np.dtype): # it is type of gym.spaces.dtype + dtype = dtype.type + self.buffer = Array(_NTYPE_TO_CTYPE[dtype], int(np.prod(shape))) + self.dtype = dtype + self.shape = shape + self.copy_on_get = copy_on_get + self.ctype = ctype + + def fill(self, src_arr: np.ndarray) -> None: + """ + Overview: + Fill the shared memory buffer with a numpy array. (Replace the original one.) + Arguments: + - src_arr (:obj:`np.ndarray`): array to fill the buffer. + """ + assert isinstance(src_arr, np.ndarray), type(src_arr) + # for np.array with shape (4, 84, 84) and float32 dtype, reshape is 15~20x faster than flatten + # for np.array with shape (4, 84, 84) and uint8 dtype, reshape is 5~7x faster than flatten + # so we reshape dst_arr rather than flatten src_arr + dst_arr = np.frombuffer(self.buffer.get_obj(), dtype=self.dtype).reshape(self.shape) + np.copyto(dst_arr, src_arr) + + def get(self) -> np.ndarray: + """ + Overview: + Get the array stored in the buffer. + Return: + - data (:obj:`np.ndarray`): A copy of the data stored in the buffer. + """ + data = np.frombuffer(self.buffer.get_obj(), dtype=self.dtype).reshape(self.shape) + if self.copy_on_get: + data = data.copy() # must use np.copy, torch.from_numpy and torch.as_tensor still use the same memory + if self.ctype is torch.Tensor: + data = torch.from_numpy(data) + return data + + +class ShmBufferContainer(object): + """ + Overview: + Support multiple shared memory buffers. Each key-value is name-buffer. + """ + + def __init__( + self, + dtype: Union[Dict[Any, type], type, np.dtype], + shape: Union[Dict[Any, tuple], tuple], + copy_on_get: bool = True + ) -> None: + """ + Overview: + Initialize the buffer container. + Arguments: + - dtype (:obj:`Union[type, np.dtype]`): The dtype of the data to limit the size of the buffer. + - shape (:obj:`Union[Dict[Any, tuple], tuple]`): If `Dict[Any, tuple]`, use a dict to manage \ + multiple buffers; If `tuple`, use single buffer. + - copy_on_get (:obj:`bool`): Whether to copy data when calling get method. + """ + if isinstance(shape, dict): + self._data = {k: ShmBufferContainer(dtype[k], v, copy_on_get) for k, v in shape.items()} + elif isinstance(shape, (tuple, list)): + self._data = ShmBuffer(dtype, shape, copy_on_get) + else: + raise RuntimeError("not support shape: {}".format(shape)) + self._shape = shape + + def fill(self, src_arr: Union[Dict[Any, np.ndarray], np.ndarray]) -> None: + """ + Overview: + Fill the one or many shared memory buffer. + Arguments: + - src_arr (:obj:`Union[Dict[Any, np.ndarray], np.ndarray]`): array to fill the buffer. + """ + if isinstance(self._shape, dict): + for k in self._shape.keys(): + self._data[k].fill(src_arr[k]) + elif isinstance(self._shape, (tuple, list)): + self._data.fill(src_arr) + + def get(self) -> Union[Dict[Any, np.ndarray], np.ndarray]: + """ + Overview: + Get the one or many arrays stored in the buffer. + Return: + - data (:obj:`np.ndarray`): The array(s) stored in the buffer. + """ + if isinstance(self._shape, dict): + return {k: self._data[k].get() for k in self._shape.keys()} + elif isinstance(self._shape, (tuple, list)): + return self._data.get() diff --git a/ding/data/storage/__init__.py b/ding/data/storage/__init__.py new file mode 100644 index 0000000000..962fbbbf18 --- /dev/null +++ b/ding/data/storage/__init__.py @@ -0,0 +1,2 @@ +from .storage import Storage +from .file import FileStorage, FileModelStorage diff --git a/ding/data/storage/file.py b/ding/data/storage/file.py new file mode 100644 index 0000000000..e6a89910b8 --- /dev/null +++ b/ding/data/storage/file.py @@ -0,0 +1,25 @@ +from typing import Any +from ding.data.storage import Storage +import pickle + +from ding.utils.file_helper import read_file, save_file + + +class FileStorage(Storage): + + def save(self, data: Any) -> None: + with open(self.path, "wb") as f: + pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL) + + def load(self) -> Any: + with open(self.path, "rb") as f: + return pickle.load(f) + + +class FileModelStorage(Storage): + + def save(self, state_dict: object) -> None: + save_file(self.path, state_dict) + + def load(self) -> object: + return read_file(self.path) diff --git a/ding/data/storage/storage.py b/ding/data/storage/storage.py new file mode 100644 index 0000000000..e6a0dae679 --- /dev/null +++ b/ding/data/storage/storage.py @@ -0,0 +1,16 @@ +from abc import ABC, abstractmethod +from typing import Any + + +class Storage(ABC): + + def __init__(self, path: str) -> None: + self.path = path + + @abstractmethod + def save(self, data: Any) -> None: + raise NotImplementedError + + @abstractmethod + def load(self) -> Any: + raise NotImplementedError diff --git a/ding/data/storage/tests/test_storage.py b/ding/data/storage/tests/test_storage.py new file mode 100644 index 0000000000..8f6f1d2c47 --- /dev/null +++ b/ding/data/storage/tests/test_storage.py @@ -0,0 +1,18 @@ +import tempfile +import pytest +import os +from os import path +from ding.data.storage import FileStorage + + +@pytest.mark.unittest +def test_file_storage(): + path_ = path.join(tempfile.gettempdir(), "test_storage.txt") + try: + storage = FileStorage(path=path_) + storage.save("test") + content = storage.load() + assert content == "test" + finally: + if path.exists(path_): + os.remove(path_) diff --git a/ding/data/storage_loader.py b/ding/data/storage_loader.py new file mode 100644 index 0000000000..daf18e2d82 --- /dev/null +++ b/ding/data/storage_loader.py @@ -0,0 +1,305 @@ +from dataclasses import dataclass +import os +import torch +import numpy as np +import uuid +import treetensor.torch as ttorch +from abc import ABC, abstractmethod +from ditk import logging +from time import sleep, time +from threading import Lock, Thread +from typing import Any, Callable, Dict, List, Optional, Union +from ding.data import FileStorage, Storage +from os import path +from ding.data.shm_buffer import ShmBuffer +from ding.framework.supervisor import RecvPayload, Supervisor, ChildType, SendPayload + + +@dataclass +class ShmObject: + id_: ShmBuffer + buf: Any + + +class StorageWorker: + + def load(self, storage: Storage) -> Any: + return storage.load() + + +class StorageLoader(Supervisor, ABC): + + def __init__(self, worker_num: int = 3) -> None: + """ + Overview: + Save and send data synchronously and load them asynchronously. + Arguments: + - worker_num (:obj:`int`): Subprocess worker number. + """ + super().__init__(type_=ChildType.PROCESS) + self._load_lock = Lock() # Load (first meet) should be called one by one. + self._callback_map: Dict[str, Callable] = {} + self._shm_obj_map: Dict[int, ShmObject] = {} + self._worker_num = worker_num + self._req_count = 0 + + def shutdown(self, timeout: Optional[float] = None) -> None: + super().shutdown(timeout) + self._recv_loop = None + self._callback_map = {} + self._shm_obj_map = {} + self._req_count = 0 + + def start_link(self) -> None: + if not self._running: + super().start_link() + self._recv_loop = Thread(target=self._loop_recv, daemon=True) + self._recv_loop.start() + + @property + def _next_proc_id(self): + return self._req_count % self._worker_num + + @abstractmethod + def save(self, obj: Union[Dict, List]) -> Storage: + """ + Overview: + Save data with a storage object synchronously. + Arguments: + - obj (:obj:`Union[Dict, List]`): The data (traj or episodes), can be numpy, tensor or treetensor. + Returns: + - storage (:obj:`Storage`): The storage object. + """ + raise NotImplementedError + + def load(self, storage: Storage, callback: Callable): + """ + Overview: + Load data from a storage object asynchronously. \ + This function will analysis the data structure when first meet a new data, \ + then alloc a shared memory buffer for each subprocess, these shared memory buffer \ + will be responsible for asynchronously loading data into memory. + Arguments: + - storage (:obj:`Storage`): The storage object. + - callback (:obj:`Callable`): Callback function after data loaded. + """ + with self._load_lock: + if not self._running: + self._first_meet(storage, callback) + return + + payload = SendPayload(proc_id=self._next_proc_id, method="load", args=[storage]) + self._callback_map[payload.req_id] = callback + self.send(payload) + self._req_count += 1 + + def _first_meet(self, storage: Storage, callback: Callable): + """ + Overview: + When first meet an object type, we'll load this object directly and analysis the structure, + to allocate the shared memory object and create subprocess workers. + Arguments: + - storage (:obj:`Storage`): The storage object. + - callback (:obj:`Callable`): Callback function after data loaded. + """ + obj = storage.load() + # Create three workers for each usage type. + for i in range(self._worker_num): + shm_obj = self._create_shm_buffer(obj) + self._shm_obj_map[i] = shm_obj + self.register(StorageWorker, shm_buffer=shm_obj, shm_callback=self._shm_callback) + self.start_link() + callback(obj) + + def _loop_recv(self): + while True: + payload = self.recv(ignore_err=True) + if payload.err: + logging.warning("Got error when loading data: {}".format(payload.err)) + if payload.req_id in self._callback_map: + del self._callback_map[payload.req_id] + else: + self._shm_putback(payload, self._shm_obj_map[payload.proc_id]) + if payload.req_id in self._callback_map: + callback = self._callback_map.pop(payload.req_id) + callback(payload.data) + + def _create_shm_buffer(self, obj: Union[Dict, List]) -> Optional[ShmObject]: + """ + Overview: + Create shared object (buf and callback) by walk through the data structure. + Arguments: + - obj (:obj:`Union[Dict, List]`): The data (traj or episodes), can be numpy, tensor or treetensor. + Returns: + - shm_buf (:obj:`Optional[ShmObject]`): The shared memory buffer. + """ + max_level = 2 + + def to_shm(obj: Dict, level: int): + if level > max_level: + return + shm_buf = None + if isinstance(obj, Dict) or isinstance(obj, ttorch.Tensor): + shm_buf = {} + for key, val in obj.items(): + # Only numpy array can fill into shm buffer + if isinstance(val, np.ndarray): + shm_buf[key] = ShmBuffer(val.dtype, val.shape, copy_on_get=False) + elif isinstance(val, torch.Tensor): + shm_buf[key] = ShmBuffer( + val.numpy().dtype, val.numpy().shape, copy_on_get=False, ctype=torch.Tensor + ) + # Recursive parsing structure + elif isinstance(val, Dict) or isinstance(val, ttorch.Tensor) or isinstance(val, List): + buf = to_shm(val, level=level + 1) + if buf: + shm_buf[key] = buf + elif isinstance(obj, List): + # Double the size of buffer + shm_buf = [to_shm(o, level=level) for o in obj] * 2 + if all(s is None for s in shm_buf): + shm_buf = [] + return shm_buf + + shm_buf = to_shm(obj, level=0) + if shm_buf is not None: + random_id = self._random_id() + shm_buf = ShmObject(id_=ShmBuffer(random_id.dtype, random_id.shape, copy_on_get=False), buf=shm_buf) + return shm_buf + + def _random_id(self) -> np.ndarray: + return np.random.randint(1, 9e6, size=(1)) + + def _shm_callback(self, payload: RecvPayload, shm_obj: ShmObject): + """ + Overview: + Called in subprocess, put payload.data into buf. + Arguments: + - payload (:obj:`RecvPayload`): The recv payload with meta info of the data. + - shm_obj (:obj:`ShmObject`): The shm buffer. + """ + assert isinstance(payload.data, type( + shm_obj.buf + )), "Data type ({}) and buf type ({}) are not match!".format(type(payload.data), type(shm_obj.buf)) + + # Sleep while shm object is not ready. + while shm_obj.id_.get()[0] != 0: + sleep(0.001) + + max_level = 2 + + def shm_callback(data: Union[Dict, List, ttorch.Tensor], buf: Union[Dict, List], level: int): + if level > max_level: + return + + if isinstance(buf, List): + assert isinstance(data, List), "Data ({}) and buf ({}) type not match".format(type(data), type(buf)) + elif isinstance(buf, Dict): + assert isinstance(data, ttorch.Tensor) or isinstance( + data, Dict + ), "Data ({}) and buf ({}) type not match".format(type(data), type(buf)) + + if isinstance(data, Dict) or isinstance(data, ttorch.Tensor): + for key, val in data.items(): + if isinstance(val, torch.Tensor): + val = val.numpy() + buf_val = buf.get(key) + if buf_val is None: + continue + if isinstance(buf_val, ShmBuffer) and isinstance(val, np.ndarray): + buf_val.fill(val) + data[key] = None + else: + shm_callback(val, buf_val, level=level + 1) + elif isinstance(data, List): + for i, data_ in enumerate(data): + shm_callback(data_, buf[i], level=level) + + shm_callback(payload.data, buf=shm_obj.buf, level=0) + id_ = self._random_id() + shm_obj.id_.fill(id_) + payload.extra = id_ + + def _shm_putback(self, payload: RecvPayload, shm_obj: ShmObject): + """ + Overview: + Called in main process, put buf back into payload.data. + Arguments: + - payload (:obj:`RecvPayload`): The recv payload with meta info of the data. + - shm_obj (:obj:`ShmObject`): The shm buffer. + """ + assert isinstance(payload.data, type( + shm_obj.buf + )), "Data type ({}) and buf type ({}) are not match!".format(type(payload.data), type(shm_obj.buf)) + + assert shm_obj.id_.get()[0] == payload.extra[0], "Shm object and payload do not match ({} - {}).".format( + shm_obj.id_.get()[0], payload.extra[0] + ) + + def shm_putback(data: Union[Dict, List], buf: Union[Dict, List]): + if isinstance(data, Dict) or isinstance(data, ttorch.Tensor): + for key, val in data.items(): + buf_val = buf.get(key) + if buf_val is None: + continue + if val is None and isinstance(buf_val, ShmBuffer): + data[key] = buf[key].get() + else: + shm_putback(val, buf_val) + elif isinstance(data, List): + for i, data_ in enumerate(data): + shm_putback(data_, buf[i]) + + shm_putback(payload.data, buf=shm_obj.buf) + shm_obj.id_.fill(np.array([0])) + + +class FileStorageLoader(StorageLoader): + + def __init__(self, dirname: str, ttl: int = 20, worker_num: int = 3) -> None: + """ + Overview: + Dump and load object with file storage. + Arguments: + - dirname (:obj:`str`): The directory to save files. + - ttl (:obj:`str`): Maximum time to keep a file, after which it will be deleted. + - worker_num (:obj:`int`): Number of subprocess worker loaders. + """ + super().__init__(worker_num) + self._dirname = dirname + self._files = [] + self._cleanup_thread = None + self._ttl = ttl # # Delete files created 10 minutes ago. + + def save(self, obj: Union[Dict, List]) -> FileStorage: + if not path.exists(self._dirname): + os.mkdir(self._dirname) + filename = "{}.pkl".format(uuid.uuid1()) + full_path = path.join(self._dirname, filename) + f = FileStorage(full_path) + f.save(obj) + self._files.append([time(), f.path]) + self._start_cleanup() + return f + + def _start_cleanup(self): + """ + Overview: + Start a cleanup thread to clean up files that are taking up too much time on the disk. + """ + if self._cleanup_thread is None: + self._cleanup_thread = Thread(target=self._loop_cleanup, daemon=True) + self._cleanup_thread.start() + + def shutdown(self, timeout: Optional[float] = None) -> None: + super().shutdown(timeout) + self._cleanup_thread = None + + def _loop_cleanup(self): + while True: + if len(self._files) == 0 or time() - self._files[0][0] < self._ttl: + sleep(1) + continue + _, file_path = self._files.pop(0) + if path.exists(file_path): + os.remove(file_path) diff --git a/ding/data/tests/test_model_loader.py b/ding/data/tests/test_model_loader.py new file mode 100644 index 0000000000..caf8c07186 --- /dev/null +++ b/ding/data/tests/test_model_loader.py @@ -0,0 +1,74 @@ +import shutil +import tempfile +from time import sleep, time +import pytest +from ding.data.model_loader import FileModelLoader +from ding.data.storage.file import FileModelStorage +from ding.model import DQN +from ding.config import compile_config +from dizoo.atari.config.serial.pong.pong_dqn_config import main_config, create_config +from os import path +import torch + + +@pytest.mark.tmp # gitlab ci and local test pass, github always fail +def test_model_loader(): + tempdir = path.join(tempfile.gettempdir(), "test_model_loader") + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + model = DQN(**cfg.policy.model) + loader = FileModelLoader(model=model, dirname=tempdir, ttl=1) + try: + loader.start() + model_storage = None + + def save_model(storage): + nonlocal model_storage + model_storage = storage + + start = time() + loader.save(save_model) + save_time = time() - start + print("Save time: {:.4f}s".format(save_time)) + assert save_time < 0.1 + sleep(0.5) + assert isinstance(model_storage, FileModelStorage) + assert len(loader._files) > 0 + + state_dict = loader.load(model_storage) + model.load_state_dict(state_dict) + + sleep(2) + assert not path.exists(model_storage.path) + assert len(loader._files) == 0 + finally: + if path.exists(tempdir): + shutil.rmtree(tempdir) + + +@pytest.mark.benchmark +def test_model_loader_benchmark(): + model = torch.nn.Sequential(torch.nn.Linear(1024, 1024), torch.nn.Linear(1024, 100)) # 40MB + tempdir = path.join(tempfile.gettempdir(), "test_model_loader") + loader = FileModelLoader(model=model, dirname=tempdir) + + try: + loader.start() + count = 0 + + def send_callback(_): + nonlocal count + count += 1 + + start = time() + for _ in range(5): + loader.save(send_callback) + sleep(0.2) + + while count < 5: + sleep(0.001) + + assert time() - start < 1.2 + finally: + if path.exists(tempdir): + shutil.rmtree(tempdir) + loader.shutdown() diff --git a/ding/data/tests/test_shm_buffer.py b/ding/data/tests/test_shm_buffer.py new file mode 100644 index 0000000000..04334b4799 --- /dev/null +++ b/ding/data/tests/test_shm_buffer.py @@ -0,0 +1,20 @@ +import pytest +import numpy as np +import timeit +from ding.data.shm_buffer import ShmBuffer +import multiprocessing as mp + + +def subprocess(shm_buf): + data = np.random.rand(1024, 1024).astype(np.float32) + res = timeit.repeat(lambda: shm_buf.fill(data), repeat=5, number=1000) + print("Mean: {:.4f}s, STD: {:.4f}s, Mean each call: {:.4f}ms".format(np.mean(res), np.std(res), np.mean(res))) + + +@pytest.mark.benchmark +def test_shm_buffer(): + data = np.random.rand(1024, 1024).astype(np.float32) + shm_buf = ShmBuffer(data.dtype, data.shape, copy_on_get=False) + proc = mp.Process(target=subprocess, args=[shm_buf]) + proc.start() + proc.join() diff --git a/ding/data/tests/test_storage_loader.py b/ding/data/tests/test_storage_loader.py new file mode 100644 index 0000000000..5ab07acd73 --- /dev/null +++ b/ding/data/tests/test_storage_loader.py @@ -0,0 +1,176 @@ +import os +import timeit +import pytest +import tempfile +import shutil +import numpy as np +import torch +import treetensor.torch as ttorch +from ding.data.shm_buffer import ShmBuffer +from ding.data.storage_loader import FileStorageLoader +from time import sleep, time +from os import path +from ding.framework.supervisor import RecvPayload + + +@pytest.mark.tmp # gitlab ci and local test pass, github always fail +def test_file_storage_loader(): + tempdir = path.join(tempfile.gettempdir(), "test_storage_loader") + loader = FileStorageLoader(dirname=tempdir) + try: + total_num = 200 + storages = [] + for i in range(10): + # 21MB + data = [ + { + "s": "abc", + "obs": np.random.rand(4, 84, 84).astype(np.float32), + # "next_obs": np.random.rand(4, 84, 84).astype(np.float32), + # "obs": torch.rand(4, 84, 84, dtype=torch.float32), + "next_obs": torch.rand(4, 84, 84, dtype=torch.float32) + } for _ in range(96) + ] + storage = loader.save(data) + storages.append(storage) + + start = time() + for i in range(total_num): + storage = storages[i % 10] + data = storage.load() + origin_time_cost = time() - start + print("Load time cost: {:.4f}s".format(origin_time_cost)) + + call_times = 0 + + def callback(data): + assert data[0]['obs'] is not None + nonlocal call_times + call_times += 1 + + # First initialize shared memory is very slow, discard this time cost. + start = time() + loader._first_meet(storage=storages[0], callback=callback) + print("Initialize shared memory time: {:.4f}s".format(time() - start)) + + start = time() + for i in range(1, total_num): + storage = storages[i % 10] + loader.load(storage, callback) + + while True: + if call_times == total_num: + break + sleep(0.01) + new_time_cost = time() - start + print("Loader time cost: {:.4f}s".format(new_time_cost)) + + assert new_time_cost < origin_time_cost + finally: + if path.exists(tempdir): + shutil.rmtree(tempdir) + loader.shutdown() + + +@pytest.mark.unittest +def test_file_storage_loader_cleanup(): + tempdir = path.join(tempfile.gettempdir(), "test_storage_loader") + loader = FileStorageLoader(dirname=tempdir, ttl=1) + try: + storages = [] + for _ in range(4): + data = np.random.rand(4, 84, 84).astype(np.float32) + storage = loader.save(data) + storages.append(storage) + sleep(0.5) + assert len(os.listdir(tempdir)) < 4 + finally: + if path.exists(tempdir): + shutil.rmtree(tempdir) + loader.shutdown() + + +@pytest.mark.unittest +def test_shared_object(): + loader = FileStorageLoader(dirname="") + + # ========== Test array ========== + obj = [{"obs": np.random.rand(100, 100)} for _ in range(10)] + shm_obj = loader._create_shm_buffer(obj) + assert len(shm_obj.buf) == len(obj) * 2 + assert isinstance(shm_obj.buf[0]["obs"], ShmBuffer) + + # Callback + payload = RecvPayload(proc_id=0, data=obj) + loader._shm_callback(payload=payload, shm_obj=shm_obj) + assert len(payload.data) == 10 + assert [d["obs"] is None for d in payload.data] + + # ========== Putback ========== + loader._shm_putback(payload=payload, shm_obj=shm_obj) + obj = payload.data + assert len(obj) == 10 + for o in obj: + assert isinstance(o["obs"], np.ndarray) + assert o["obs"].shape == (100, 100) + + # ========== Test dict ========== + obj = {"obs": torch.rand(100, 100, dtype=torch.float32)} + shm_obj = loader._create_shm_buffer(obj) + assert isinstance(shm_obj.buf["obs"], ShmBuffer) + + payload = RecvPayload(proc_id=0, data=obj) + loader._shm_callback(payload=payload, shm_obj=shm_obj) + assert payload.data["obs"] is None + + loader._shm_putback(payload=payload, shm_obj=shm_obj) + assert isinstance(payload.data["obs"], torch.Tensor) + assert payload.data["obs"].shape == (100, 100) + + # ========== Test treetensor ========== + obj = {"trajectories": [ttorch.as_tensor({"obs": torch.rand(10, 10, dtype=torch.float32)}) for _ in range(10)]} + shm_obj = loader._create_shm_buffer(obj) + + payload = RecvPayload(proc_id=0, data=obj) + loader._shm_callback(payload=payload, shm_obj=shm_obj) + assert len(payload.data["trajectories"]) == 10 + for traj in payload.data["trajectories"]: + assert traj["obs"] is None + + loader._shm_putback(payload=payload, shm_obj=shm_obj) + for traj in payload.data["trajectories"]: + assert isinstance(traj["obs"], torch.Tensor) + assert traj["obs"].shape == (10, 10) + + +@pytest.mark.benchmark +def test_shared_object_benchmark(): + loader = FileStorageLoader(dirname="") + # ========== Test treetensor ========== + obj = { + "env_step": 0, + "trajectories": [ + ttorch.as_tensor( + { + "done": False, + "reward": torch.tensor([1, 0], dtype=torch.int32), + "obs": torch.rand(4, 84, 84, dtype=torch.float32), + "next_obs": torch.rand(4, 84, 84, dtype=torch.float32), + "action": torch.tensor([1], dtype=torch.int32), + "collect_train_iter": torch.tensor([1], dtype=torch.int32), + "env_data_id": torch.tensor([1], dtype=torch.int32), + } + ) for _ in range(10) + ] + } + buf = loader._create_shm_buffer(obj) + payload = RecvPayload(proc_id=0, data=obj) + loader._shm_callback(payload=payload, shm_obj=buf) + + def stmt(): + payload.extra = buf.id_.get() + loader._shm_putback(payload=payload, shm_obj=buf) + + res = timeit.repeat(stmt, repeat=5, number=1000) + print("Mean: {:.4f}s, STD: {:.4f}s, Mean each call: {:.4f}ms".format(np.mean(res), np.std(res), np.mean(res))) + assert np.mean(res) < 1 diff --git a/ding/design/serial_evaluator-activity.puml b/ding/design/serial_evaluator-activity.puml index 1dea306034..aa84c4eef5 100644 --- a/ding/design/serial_evaluator-activity.puml +++ b/ding/design/serial_evaluator-activity.puml @@ -26,6 +26,6 @@ repeat endif repeat while (evaluate episodes are not enough?) |#FFCCCC|evaluator| -:return eval_episode_reward; +:return eval_episode_return; stop @enduml diff --git a/ding/design/serial_main.puml b/ding/design/serial_main.puml index 75e00d2bd1..f710639dd8 100644 --- a/ding/design/serial_main.puml +++ b/ding/design/serial_main.puml @@ -31,7 +31,7 @@ loop evaluator -> evaluator: eval_performance alt reach eval stop_value learner -> learner: save checkpoint and exit - else eval_reward is new highest + else episode_return is new highest learner -> learner: save checkpoint end end diff --git a/ding/entry/__init__.py b/ding/entry/__init__.py index 1e90351ee4..11cccf0e13 100644 --- a/ding/entry/__init__.py +++ b/ding/entry/__init__.py @@ -6,7 +6,6 @@ from .serial_entry_onpolicy_ppg import serial_pipeline_onpolicy_ppg from .serial_entry_offline import serial_pipeline_offline from .serial_entry_ngu import serial_pipeline_ngu -from .serial_entry_decision_transformer import serial_pipeline_dt from .serial_entry_reward_model_offpolicy import serial_pipeline_reward_model_offpolicy from .serial_entry_reward_model_onpolicy import serial_pipeline_reward_model_onpolicy from .serial_entry_bc import serial_pipeline_bc @@ -24,6 +23,6 @@ import serial_pipeline_preference_based_irl from .serial_entry_preference_based_irl_onpolicy \ import serial_pipeline_preference_based_irl_onpolicy -from .application_entry_drex_collect_data import drex_collecting_data -from .serial_entry_mbrl import serial_pipeline_dyna, serial_pipeline_dream +from .serial_entry_mbrl import serial_pipeline_dyna, serial_pipeline_dream, serial_pipeline_dreamer from .serial_entry_bco import serial_pipeline_bco +from .serial_entry_pc import serial_pipeline_pc diff --git a/ding/entry/application_entry.py b/ding/entry/application_entry.py index a9fd1ae41b..bb8fe882df 100644 --- a/ding/entry/application_entry.py +++ b/ding/entry/application_entry.py @@ -26,9 +26,9 @@ def eval( load_path: Optional[str] = None, replay_path: Optional[str] = None, ) -> float: - r""" + """ Overview: - Pure evaluation entry. + Pure policy evaluation entry. Evaluate mean episode return and save replay videos. Arguments: - input_cfg (:obj:`Union[str, Tuple[dict, dict]]`): Config in dict type. \ ``str`` type means config file path. \ @@ -72,10 +72,9 @@ def eval( # Evaluate _, episode_info = evaluator.eval() - reward = [e['final_eval_reward'] for e in episode_info] - eval_reward = np.mean(to_ndarray(reward)) - print('Eval is over! The performance of your RL policy is {}'.format(eval_reward)) - return eval_reward + episode_return = np.mean(episode_info['eval_episode_return']) + print('Eval is over! The performance of your RL policy is {}'.format(episode_return)) + return episode_return def collect_demo_data( @@ -271,8 +270,8 @@ def episode_to_transitions_filter(data_path: str, expert_data_path: str, nstep: _dict = pickle.load(f) # class is list; length is cfg.reward_model.collect_count post_process_data = [] for i in range(len(_dict)): - episode_rewards = torch.stack([_dict[i][j]['reward'] for j in range(_dict[i].__len__())], axis=0) - if episode_rewards.sum() < min_episode_return: + episode_returns = torch.stack([_dict[i][j]['reward'] for j in range(_dict[i].__len__())], axis=0) + if episode_returns.sum() < min_episode_return: continue data = get_nstep_return_data(_dict[i], nstep) post_process_data.extend(data) diff --git a/ding/entry/application_entry_drex_collect_data.py b/ding/entry/application_entry_drex_collect_data.py deleted file mode 100644 index 6dafbe1b9c..0000000000 --- a/ding/entry/application_entry_drex_collect_data.py +++ /dev/null @@ -1,257 +0,0 @@ -import argparse -import copy -import pickle -import random -import easydict -import torch -import os -from typing import Optional, List, Any -from functools import partial -from copy import deepcopy - -from ding.config import compile_config, read_config -from ding.worker import EpisodeSerialCollector, create_buffer, BaseLearner -from ding.envs import create_env_manager, get_vec_env_setting -from ding.policy import create_policy, bc -from ding.torch_utils import to_device -from ding.utils import set_pkg_seed -from ding.utils.data import default_collate, offline_data_save_type -from functools import reduce - - -def collect_episodic_demo_data_for_drex( - cfg: easydict, - seed: int, - collect_count: int, - rank: int, - save_cfg_path: str, - noise: float, - env_setting: Optional[List[Any]] = None, - model: Optional[torch.nn.Module] = None, - state_dict: Optional[dict] = None, - state_dict_path: Optional[str] = None, -): - r""" - Overview: - Collect episodic demonstration data by the trained policy for trex specifically. - Arguments: - - cfg (:obj:`easydict`): Config in dict type. - - seed (:obj:`int`): Random seed. - - collect_count (:obj:`int`): The count of collected data. - - rank (:obj:`int`) the episode ranking. - - save_cfg_path(:obj:'str') where to save the collector config - - env_setting (:obj:`Optional[List[Any]]`): A list with 3 elements: \ - ``BaseEnv`` subclass, collector env config, and evaluator env config. - - model (:obj:`Optional[torch.nn.Module]`): Instance of torch.nn.Module. - - state_dict (:obj:`Optional[dict]`): The state_dict of policy or model. - - state_dict_path (:obj:'str') the abs path of the state dict - """ - cfg.env.collector_env_num = 1 - if not os.path.exists(save_cfg_path): - os.makedirs(save_cfg_path) - - # Create components: env, policy, collector - if env_setting is None: - env_fn, collector_env_cfg, _ = get_vec_env_setting(cfg.env) - else: - env_fn, collector_env_cfg, _ = env_setting - collector_env = create_env_manager(cfg.env.manager, [partial(env_fn, cfg=c) for c in collector_env_cfg]) - collector_env.seed(seed) - - set_pkg_seed(seed, use_cuda=cfg.policy.cuda) - policy = create_policy(cfg.policy, model=model, enable_field=['collect', 'eval']) - collect_demo_policy = policy.collect_mode - if state_dict is None: - assert state_dict_path is not None - state_dict = torch.load(state_dict_path, map_location='cpu') - policy.collect_mode.load_state_dict(state_dict) - collector = EpisodeSerialCollector( - cfg.policy.collect.collector, collector_env, collect_demo_policy, exp_name=cfg.exp_name - ) - - policy_kwargs = {'eps': noise} - - # Let's collect some sub-optimal demonstrations - exp_data = collector.collect(n_episode=collect_count, policy_kwargs=policy_kwargs) - - if cfg.policy.cuda: - exp_data = to_device(exp_data, 'cpu') - # Save data transitions. - print('Collect {}th episodic demo data successfully'.format(rank)) - return exp_data - - -def drex_get_args(): - parser = argparse.ArgumentParser() - parser.add_argument('--cfg', type=str, default='abs path for a config') - parser.add_argument('--seed', type=int, default=0) - parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu') - args = parser.parse_known_args()[0] - return args - - -def eval_bc(validation_set, policy, use_cuda): - tot = 0 - tot_acc = 0 - device = 'cuda' if use_cuda else 'cpu' - for _, data in enumerate(validation_set): - x, y = {'obs': data['obs'].to(device).squeeze(0)}, data['action'].squeeze(-1) - y_pred = policy.forward(x, eps=-1)['obs']['action'] - tot += y_pred.shape[0] - tot_acc += (y_pred == y).sum().item() - acc = tot_acc / tot - return acc - - -def train_bc(cfg, pre_expert_data=None, max_iterations=6000): - cfg_new = copy.deepcopy(cfg).policy - cfg_new.continuous = False - bc_policy = bc.BehaviourCloningPolicy(cfg_new) - - if pre_expert_data is None: - with open(cfg.reward_model.offline_data_path + '/suboptimal_data.pkl', 'rb') as f: - pre_expert_data = pickle.load(f) - - set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) - replay_buffer = create_buffer(cfg.policy.other.replay_buffer) - push_data = [] - for i in range(len(pre_expert_data)): - push_data += pre_expert_data[i] - - random.shuffle(push_data) - validation_set = push_data[-len(push_data) // 10:] - push_data = push_data[:-len(push_data) // 10] - replay_buffer.push(push_data, cur_collector_envstep=0) - learner = BaseLearner(cfg.policy.learn.learner, bc_policy.learn_mode) - - best_acc = 0 - cnt = 0 - for i in range(max_iterations): - train_data = replay_buffer.sample(learner.policy.get_attribute('batch_size'), learner.train_iter) - if i % 100 == 0: - acc = eval_bc(validation_set, bc_policy.collect_mode, cfg.policy.cuda) - if acc < best_acc: - cnt += 1 - if cnt > 100: - break - else: - cnt = 0 - best_acc = acc - torch.save(bc_policy.collect_mode.state_dict(), cfg.reward_model.offline_data_path + '/bc_best.pth.tar') - if train_data is None: - replay_buffer.push(push_data, cur_collector_envstep=0) - train_data = replay_buffer.sample(learner.policy.get_attribute('batch_size'), learner.train_iter) - learner.train(train_data) - - ckpt = torch.load(cfg.reward_model.offline_data_path + '/bc_best.pth.tar') - bc_policy.collect_mode.load_state_dict(ckpt) - return bc_policy - - -def load_bc(load_path, cfg): - cfg_new = copy.deepcopy(cfg).policy - cfg_new.continuous = False - bc_policy = bc.BehaviourCloningPolicy(cfg_new) - state_dict = torch.load(load_path, map_location='cpu') - bc_policy.collect_mode.load_state_dict(state_dict) - print('Load bc from {}'.format(load_path)) - return bc_policy - - -def cal_mean(lis): - return reduce(lambda x, y: x + y, lis) / len(lis) - - -def create_data_drex(bc_policy, cfg): - env_fn, collector_env_cfg, evaluator_env_cfg = get_vec_env_setting(copy.deepcopy(cfg.env)) - collector_env = create_env_manager( - copy.deepcopy(cfg.env.manager), [partial(env_fn, cfg=c) for c in collector_env_cfg] - ) - collector = EpisodeSerialCollector( - cfg.policy.collect.collector, - env=collector_env, - policy=bc_policy.collect_mode, - ) - eps_list = cfg.reward_model.eps_list - eps_list.sort(reverse=True) - created_data = [] - created_data_returns = [] - for eps in eps_list: - policy_kwargs = {'eps': eps} - # Let's collect some sub-optimal demonstrations - exp_data = collector.collect(n_episode=cfg.reward_model.num_trajs_per_bin, policy_kwargs=policy_kwargs) - episodes = [default_collate(data)['obs'].numpy() for data in exp_data] - returns = [torch.sum(default_collate(data)['reward']).item() for data in exp_data] - - created_data.append(episodes) - created_data_returns.append(returns) - print('noise: {}, returns: {}, avg: {}'.format(eps, returns, cal_mean(returns))) - return created_data, created_data_returns - - -def drex_generating_data(compiled_cfg): - offline_data_path = compiled_cfg.reward_model.offline_data_path - expert_model_path = compiled_cfg.reward_model.expert_model_path - data_for_save = {} - learning_returns = [] - learning_rewards = [] - episodes_data = [] - for i in range(10): - model_path = expert_model_path - seed = compiled_cfg.seed + i - exp_data = collect_episodic_demo_data_for_drex( - deepcopy(compiled_cfg), - seed, - noise=-1, - state_dict_path=model_path, - save_cfg_path=offline_data_path, - collect_count=1, - rank=i - ) - data_for_save[i] = exp_data[0] - obs = list(default_collate(exp_data[0])['obs'].numpy()) - learning_rewards.append(default_collate(exp_data[0])['reward'].tolist()) - sum_reward = torch.sum(default_collate(exp_data[0])['reward']).item() - learning_returns.append(sum_reward) - episodes_data.append(obs) - - offline_data_save_type( - data_for_save, - offline_data_path + '/suboptimal_data.pkl', - data_type=compiled_cfg.policy.collect.get('data_type', 'naive') - ) - - return data_for_save - - -def drex_collecting_data(args=drex_get_args(), seed=0): - if isinstance(args.cfg, str): - cfg, create_cfg = read_config(args.cfg) - else: - cfg, create_cfg = deepcopy(args.cfg) - bc_iteration = getattr(cfg, 'bc_iteration', 50000) - cfg = compile_config(cfg, seed=seed, env=None, auto=True, create_cfg=create_cfg, save_cfg=True) - pre_expert_data = drex_generating_data(cfg) - returns = [torch.sum(default_collate(pre_expert_data[i])['reward']).item() for i in range(len(pre_expert_data))] - print('demonstrations rewards: ' + str(returns)) - - if 'bc_path' in cfg.reward_model and os.path.exists(cfg.reward_model.bc_path): - bc_policy = load_bc(cfg.reward_model.bc_path, cfg) - else: - bc_policy = train_bc(cfg=cfg, pre_expert_data=pre_expert_data, max_iterations=bc_iteration) - created_data, created_data_returns = create_data_drex(bc_policy, cfg) - offline_data_path = cfg.reward_model.offline_data_path - - offline_data_save_type( - created_data, offline_data_path + '/episodes_data.pkl', data_type=cfg.policy.collect.get('data_type', 'naive') - ) - - offline_data_save_type( - created_data_returns, - offline_data_path + '/learning_returns.pkl', - data_type=cfg.policy.collect.get('data_type', 'naive') - ) - - -if __name__ == '__main__': - drex_collecting_data() diff --git a/ding/entry/cli_ditask.py b/ding/entry/cli_ditask.py index 68ec836fe6..443fe1a6b6 100644 --- a/ding/entry/cli_ditask.py +++ b/ding/entry/cli_ditask.py @@ -43,8 +43,7 @@ def print_version(ctx: Context, param: Option, value: bool) -> None: @click.option( "--ports", type=str, - default="50515", - help="The port addresses that the tasks listen to, e.g. 50515,50516, default: 50515" + help="The port addresses that the tasks listen to, e.g. 50515,50516, default: k8s, local: 50515, slurm: 15151" ) @click.option("--attach-to", type=str, help="The addresses to connect to.") @click.option("--address", type=str, help="The address to listen to (without port).") @@ -62,6 +61,8 @@ def print_version(ctx: Context, param: Option, value: bool) -> None: @click.option("--redis-host", type=str, help="Redis host.") @click.option("--redis-port", type=int, help="Redis port.") @click.option("-m", "--main", type=str, help="Main function of entry module.") +@click.option("--startup-interval", type=int, default=1, help="Start up interval between each task.") +@click.option("--local_rank", type=int, default=0, help="Compatibility with PyTorch DDP") def cli_ditask(*args, **kwargs): return _cli_ditask(*args, **kwargs) @@ -86,7 +87,7 @@ def _parse_platform_args(platform: str, platform_spec: str, all_args: dict): parsed_args = PLATFORM_PARSERS[platform](platform_spec, **all_args) except Exception as e: click.echo("error when parse platform spec configure: {}".format(e)) - exit(1) + raise e return parsed_args @@ -105,6 +106,8 @@ def _cli_ditask( mq_type: str, redis_host: str, redis_port: int, + startup_interval: int, + local_rank: int = 0, platform: str = None, platform_spec: str = None, ): @@ -128,6 +131,7 @@ def _cli_ditask( mod = importlib.import_module(mod_name) main_func = getattr(mod, func_name) # Parse arguments + ports = ports or 50515 if not isinstance(ports, int): ports = ports.split(",") ports = list(map(lambda i: int(i), ports)) @@ -152,5 +156,6 @@ def _cli_ditask( node_ids=node_ids, mq_type=mq_type, redis_host=redis_host, - redis_port=redis_port + redis_port=redis_port, + startup_interval=startup_interval )(main_func) diff --git a/ding/entry/cli_parsers/k8s_parser.py b/ding/entry/cli_parsers/k8s_parser.py index 3767ef2879..6f2b0aebe7 100644 --- a/ding/entry/cli_parsers/k8s_parser.py +++ b/ding/entry/cli_parsers/k8s_parser.py @@ -1,11 +1,12 @@ import os import numpy as np -from typing import List, Optional +from time import sleep +from typing import Dict, List, Optional class K8SParser(): - def __init__(self, platform_spec: Optional[str] = None, **kwargs) -> None: + def __init__(self, platform_spec: Optional[Dict] = None, **kwargs) -> None: """ Overview: Should only set global cluster properties @@ -14,9 +15,9 @@ def __init__(self, platform_spec: Optional[str] = None, **kwargs) -> None: self.nodelist = self._parse_node_list() self.ntasks = len(self.nodelist) self.platform_spec = platform_spec - self.parallel_workers = kwargs.get("parallel_workers", 1) - self.topology = kwargs.get("topology", "alone") - self.ports = kwargs.get("ports", 50515) + self.parallel_workers = kwargs.get("parallel_workers") or 1 + self.topology = kwargs.get("topology") or "alone" + self.ports = int(kwargs.get("ports") or 50515) self.tasks = {} def parse(self) -> dict: @@ -49,13 +50,13 @@ def _get_task(self, procid: int) -> dict: else: task = {} if "ports" not in task: - task["ports"] = self._get_ports() + task["ports"] = self.kwargs.get("ports") or self._get_ports() if "address" not in task: - task["address"] = self._get_address(procid) + task["address"] = self.kwargs.get("address") or self._get_address(procid) if "node_ids" not in task: - task["node_ids"] = self._get_node_id(procid) + task["node_ids"] = self.kwargs.get("node_ids") or self._get_node_id(procid) - task["attach_to"] = self._get_attach_to(procid, task.get("attach_to")) + task["attach_to"] = self.kwargs.get("attach_to") or self._get_attach_to(procid, task.get("attach_to")) task["topology"] = self.topology task["parallel_workers"] = self.parallel_workers diff --git a/ding/entry/cli_parsers/slurm_parser.py b/ding/entry/cli_parsers/slurm_parser.py index 3a335eb758..c46716438b 100644 --- a/ding/entry/cli_parsers/slurm_parser.py +++ b/ding/entry/cli_parsers/slurm_parser.py @@ -1,30 +1,55 @@ import os import re -from typing import List +from time import sleep +import numpy as np +from typing import Any, Dict, List, Optional class SlurmParser(): - def __init__(self, platform_spec: str, **kwargs) -> None: + def __init__(self, platform_spec: Optional[Dict] = None, **kwargs) -> None: """ Overview: Should only set global cluster properties """ self.kwargs = kwargs self.ntasks = int(os.environ["SLURM_NTASKS"]) - self.tasks = platform_spec["tasks"] + self.platform_spec = platform_spec + self.tasks = {} self.ntasks_per_node = int(os.environ["SLURM_NTASKS_PER_NODE"]) self.nodelist = self._parse_node_list() + self.ports = int(kwargs.get("ports") or 15151) + self.parallel_workers = kwargs.get("parallel_workers") or 1 + self.topology = kwargs.get("topology") or "alone" def parse(self) -> dict: - assert len(self.tasks) == self.ntasks procid = int(os.environ["SLURM_PROCID"]) - nodename = os.environ["SLURMD_NODENAME"] - task = self._get_node_args(procid) + task = self._get_task(procid) # Validation - assert task["address"] == nodename + assert task["address"] == os.environ["SLURMD_NODENAME"] return {**self.kwargs, **task} + def _get_task(self, procid: int) -> Dict[str, Any]: + if procid in self.tasks: + return self.tasks.get(procid) + if self.platform_spec: + task = self.platform_spec["tasks"][procid] + else: + task = {} + if "ports" not in task: + task["ports"] = self._get_ports(procid) + if "address" not in task: + task["address"] = self._get_address(procid) + if "node_ids" not in task: + task["node_ids"] = self._get_node_id(procid) + + task["attach_to"] = self._get_attach_to(procid, task.get("attach_to")) + task["topology"] = self.topology + task["parallel_workers"] = self.parallel_workers + + self.tasks[procid] = task + return task + def _parse_node_list(self) -> List[str]: nodelist = os.environ["SLURM_NODELIST"] result = re.match(r"(.*)?\[(.*)\]$", nodelist) @@ -40,58 +65,86 @@ def _parse_node_list(self) -> List[str]: nodelist.append(prefix + tail) elif isinstance(nodelist, str): nodelist = [nodelist] + if self.ntasks_per_node > 1: + expand_nodelist = [] # Expand node for each task + for node in nodelist: + for _ in range(self.ntasks_per_node): + expand_nodelist.append(node) + nodelist = expand_nodelist return nodelist - def _get_node_args(self, procid: int) -> dict: - """ - Overview: - Complete node properties, use environment vars in list instead of on current node. - For example, if you want to set nodename in this function, please derive it from SLURM_NODELIST, - the variable from SLURMD_NODENAME should only be used in validation. - """ - task = self.tasks[procid] - if "attach_to" in task: - task["attach_to"] = self._get_attach_to(task["attach_to"]) - if "address" not in task: - task["address"] = self._get_address(procid) - if "ports" not in task: - task["ports"] = self._get_ports(procid) - if "node_ids" not in task: - task["node_ids"] = procid - return task + def _get_attach_to(self, procid: int, attach_to: Optional[str] = None) -> str: + if attach_to: + attach_to = [self._get_attach_to_part(part) for part in attach_to.split(",")] + elif procid == 0: + attach_to = [] + else: + if self.topology == "mesh": + prev_tasks = [self._get_task(i) for i in range(procid)] + attach_to = [self._get_attach_to_from_task(task) for task in prev_tasks] + attach_to = list(np.concatenate(attach_to)) + elif self.topology == "star": + head_task = self._get_task(0) + attach_to = self._get_attach_to_from_task(head_task) + else: + attach_to = [] - def _get_attach_to(self, attach_to: str) -> str: - attach_to = [self._get_attach_to_part(part) for part in attach_to.split(",")] return ",".join(attach_to) def _get_attach_to_part(self, attach_part: str) -> str: + """ + Overview: + Parse each part of attach_to. + Arguments: + - attach_part (:obj:`str`): The attach_to field with specific pattern, e.g. $node:0 + Returns + - attach_to (:obj:`str`): The real address, e.g. tcp://SH-0:50000 + """ if not attach_part.startswith("$node."): return attach_part attach_node_id = int(attach_part[6:]) - attach_node = self._get_node_args(self._get_procid_from_nodeid(attach_node_id)) - return "tcp://{}:{}".format(attach_node["address"], attach_node["ports"]) + attach_task = self._get_task(self._get_procid_from_nodeid(attach_node_id)) + return self._get_tcp_link(attach_task["address"], attach_task["ports"]) + + def _get_attach_to_from_task(self, task: dict) -> List[str]: + """ + Overview: + Get attach_to list from task, note that parallel_workers will affact the connected processes. + Arguments: + - task (:obj:`dict`): The task object. + Returns + - attach_to (:obj:`str`): The real address, e.g. tcp://SH-0:50000 + """ + port = task.get("ports") + address = task.get("address") + ports = [int(port) + i for i in range(self.parallel_workers)] + attach_to = [self._get_tcp_link(address, port) for port in ports] + return attach_to def _get_procid_from_nodeid(self, nodeid: int) -> int: procid = None - for i, task in enumerate(self.tasks): - if task.get("node_ids") == nodeid: - procid = i - break - elif nodeid == i: + for i in range(self.ntasks): + task = self._get_task(i) + if task["node_ids"] == nodeid: procid = i break if procid is None: raise Exception("Can not find procid from nodeid: {}".format(nodeid)) return procid - def _get_ports(self, procid: int) -> List[int]: - ports = 15151 + procid % self.ntasks_per_node - return ports + def _get_ports(self, procid) -> int: + return self.ports + (procid % self.ntasks_per_node) * self.parallel_workers def _get_address(self, procid: int) -> str: - address = self.nodelist[procid // self.ntasks_per_node] + address = self.nodelist[procid] return address + def _get_node_id(self, procid: int) -> int: + return procid * self.parallel_workers + + def _get_tcp_link(self, address: str, port: int) -> str: + return "tcp://{}:{}".format(address, port) + def slurm_parser(platform_spec: str, **kwargs) -> dict: return SlurmParser(platform_spec, **kwargs).parse() diff --git a/ding/entry/cli_parsers/tests/test_slurm_parser.py b/ding/entry/cli_parsers/tests/test_slurm_parser.py index f56efdd663..9b817ba48a 100644 --- a/ding/entry/cli_parsers/tests/test_slurm_parser.py +++ b/ding/entry/cli_parsers/tests/test_slurm_parser.py @@ -10,12 +10,7 @@ def set_slurm_env(): os.environ["SLURM_NTASKS"] = '6' # Parameter n，Process count / Task count os.environ["SLURM_NTASKS_PER_NODE"] = '3' # Parameter ntasks-per-node，process count of each node os.environ["SLURM_NODELIST"] = 'SH-IDC1-10-5-38-[190,215]' # All the nodes - os.environ["SLURM_SRUN_COMM_PORT"] = '42932' # Available ports - os.environ["SLURM_TOPOLOGY_ADDR"] = 'SH-IDC1-10-5-38-215' # Name of current node - os.environ["SLURM_NODEID"] = '1' # Node order，start from 0 os.environ["SLURM_PROCID"] = '3' # Proc order，start from 0，the read proc order may be different from nominal order - os.environ["SLURM_LOCALID"] = '0' # Proc order on current node, smaller or equal than ntasks-per-node - 1 - os.environ["SLURM_GTIDS"] = '2,3' # All the proc ids on current node os.environ["SLURMD_NODENAME"] = 'SH-IDC1-10-5-38-215' # Name of current node yield @@ -23,12 +18,7 @@ def set_slurm_env(): del os.environ["SLURM_NTASKS"] del os.environ["SLURM_NTASKS_PER_NODE"] del os.environ["SLURM_NODELIST"] - del os.environ["SLURM_SRUN_COMM_PORT"] - del os.environ["SLURM_TOPOLOGY_ADDR"] - del os.environ["SLURM_NODEID"] del os.environ["SLURM_PROCID"] - del os.environ["SLURM_LOCALID"] - del os.environ["SLURM_GTIDS"] del os.environ["SLURMD_NODENAME"] @@ -73,8 +63,22 @@ def test_slurm_parser(): "tcp://SH-IDC1-10-5-38-190:15152," +\ "tcp://SH-IDC1-10-5-38-190:15153" + # Test without platform_spec + all_args = slurm_parser(None, topology="mesh", mq_type="nng") + assert all_args["address"] == "SH-IDC1-10-5-38-215" + assert all_args["node_ids"] == 3 + assert all_args["parallel_workers"] == 1 + assert all_args[ + "attach_to" + ] == "tcp://SH-IDC1-10-5-38-190:15151," +\ + "tcp://SH-IDC1-10-5-38-190:15152," +\ + "tcp://SH-IDC1-10-5-38-190:15153" + # Test _parse_node_list sp = SlurmParser(platform_spec) os.environ["SLURM_NODELIST"] = 'SH-IDC1-10-5-[38-40]' - nodelist = sp._parse_node_list() - assert nodelist == ['SH-IDC1-10-5-38', 'SH-IDC1-10-5-39', 'SH-IDC1-10-5-40'] + nodelist = sp._parse_node_list() # Nodes * parallel_workers + assert nodelist == [ + 'SH-IDC1-10-5-38', 'SH-IDC1-10-5-38', 'SH-IDC1-10-5-38', 'SH-IDC1-10-5-39', 'SH-IDC1-10-5-39', + 'SH-IDC1-10-5-39', 'SH-IDC1-10-5-40', 'SH-IDC1-10-5-40', 'SH-IDC1-10-5-40' + ] diff --git a/ding/entry/serial_entry.py b/ding/entry/serial_entry.py index 2161704191..f7c039e494 100644 --- a/ding/entry/serial_entry.py +++ b/ding/entry/serial_entry.py @@ -11,7 +11,7 @@ create_serial_collector, create_serial_evaluator from ding.config import read_config, compile_config from ding.policy import create_policy -from ding.utils import set_pkg_seed +from ding.utils import set_pkg_seed, get_rank from .utils import random_collect @@ -22,6 +22,7 @@ def serial_pipeline( model: Optional[torch.nn.Module] = None, max_train_iter: Optional[int] = int(1e10), max_env_step: Optional[int] = int(1e10), + dynamic_seed: Optional[bool] = True, ) -> 'Policy': # noqa """ Overview: @@ -36,6 +37,7 @@ def serial_pipeline( - model (:obj:`Optional[torch.nn.Module]`): Instance of torch.nn.Module. - max_train_iter (:obj:`Optional[int]`): Maximum policy update iterations in training. - max_env_step (:obj:`Optional[int]`): Maximum collected environment interaction steps. + - dynamic_seed(:obj:`Optional[bool]`): set dynamic seed for collector. Returns: - policy (:obj:`Policy`): Converged policy. """ @@ -45,7 +47,15 @@ def serial_pipeline( cfg, create_cfg = deepcopy(input_cfg) create_cfg.policy.type = create_cfg.policy.type + '_command' env_fn = None if env_setting is None else env_setting[0] - cfg = compile_config(cfg, seed=seed, env=env_fn, auto=True, create_cfg=create_cfg, save_cfg=True) + cfg = compile_config( + cfg, + seed=seed, + env=env_fn, + auto=True, + create_cfg=create_cfg, + save_cfg=True, + renew_dir=not cfg.policy.learn.get('resume_training', False) + ) # Create main components: env, policy if env_setting is None: env_fn, collector_env_cfg, evaluator_env_cfg = get_vec_env_setting(cfg.env) @@ -53,13 +63,13 @@ def serial_pipeline( env_fn, collector_env_cfg, evaluator_env_cfg = env_setting collector_env = create_env_manager(cfg.env.manager, [partial(env_fn, cfg=c) for c in collector_env_cfg]) evaluator_env = create_env_manager(cfg.env.manager, [partial(env_fn, cfg=c) for c in evaluator_env_cfg]) - collector_env.seed(cfg.seed) + collector_env.seed(cfg.seed, dynamic_seed=dynamic_seed) evaluator_env.seed(cfg.seed, dynamic_seed=False) set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) policy = create_policy(cfg.policy, model=model, enable_field=['learn', 'collect', 'eval', 'command']) # Create worker components: learner, collector, evaluator, replay buffer, commander. - tb_logger = SummaryWriter(os.path.join('./{}/log/'.format(cfg.exp_name), 'serial')) + tb_logger = SummaryWriter(os.path.join('./{}/log/'.format(cfg.exp_name), 'serial')) if get_rank() == 0 else None learner = BaseLearner(cfg.policy.learn.learner, policy.learn_mode, tb_logger, exp_name=cfg.exp_name) collector = create_serial_collector( cfg.policy.collect.collector, @@ -84,6 +94,8 @@ def serial_pipeline( # ========== # Learner's before_run hook. learner.call_hook('before_run') + if cfg.policy.learn.get('resume_training', False): + collector.envstep = learner.collector_envstep # Accumulate plenty of data at the beginning of training. if cfg.policy.get('random_collect_size', 0) > 0: @@ -117,18 +129,19 @@ def serial_pipeline( # Learner's after_run hook. learner.call_hook('after_run') - import time - import pickle - import numpy as np - with open(os.path.join(cfg.exp_name, 'result.pkl'), 'wb') as f: - eval_value_raw = [d['final_eval_reward'] for d in eval_info] - final_data = { - 'stop': stop, - 'env_step': collector.envstep, - 'train_iter': learner.train_iter, - 'eval_value': np.mean(eval_value_raw), - 'eval_value_raw': eval_value_raw, - 'finish_time': time.ctime(), - } - pickle.dump(final_data, f) + if get_rank() == 0: + import time + import pickle + import numpy as np + with open(os.path.join(cfg.exp_name, 'result.pkl'), 'wb') as f: + eval_value_raw = eval_info['eval_episode_return'] + final_data = { + 'stop': stop, + 'env_step': collector.envstep, + 'train_iter': learner.train_iter, + 'eval_value': np.mean(eval_value_raw), + 'eval_value_raw': eval_value_raw, + 'finish_time': time.ctime(), + } + pickle.dump(final_data, f) return policy diff --git a/ding/entry/serial_entry_decision_transformer.py b/ding/entry/serial_entry_decision_transformer.py deleted file mode 100644 index d0431839c4..0000000000 --- a/ding/entry/serial_entry_decision_transformer.py +++ /dev/null @@ -1,88 +0,0 @@ -""" -The code is adapted from https://github.com/nikhilbarhate99/min-decision-transformer -""" -from typing import Union, Optional, List, Any, Tuple -from functools import partial -from tensorboardX import SummaryWriter -from copy import deepcopy -from torch.utils.data import DataLoader - -from ding.envs import get_vec_env_setting, create_env_manager -from ding.worker import BaseLearner, InteractionSerialEvaluator, BaseSerialCommander, create_buffer -from ding.config import read_config, compile_config -from ding.policy import create_policy -from ding.utils import set_pkg_seed -from ding.utils.data import create_dataset -import os -import logging -import random -import time -import pickle -import torch -import numpy as np -from torch.utils.data import Dataset -from ding.rl_utils import discount_cumsum -from ding.utils.data.dataset import D4RLTrajectoryDataset - - -def serial_pipeline_dt( - input_cfg: Union[str, Tuple[dict, dict]], - seed: int = 0, - env_setting: Optional[List[Any]] = None, - model: Optional[torch.nn.Module] = None, - max_train_iter: Optional[int] = int(5e2), -) -> 'Policy': # noqa - """ - Overview: - Serial pipeline entry. - Arguments: - - input_cfg (:obj:`Union[str, Tuple[dict, dict]]`): Config in dict type. \ - ``str`` type means config file path. \ - ``Tuple[dict, dict]`` type means [user_config, create_cfg]. - - seed (:obj:`int`): Random seed. - - env_setting (:obj:`Optional[List[Any]]`): A list with 3 elements: \ - ``BaseEnv`` subclass, collector env config, and evaluator env config. - - model (:obj:`Optional[torch.nn.Module]`): Instance of torch.nn.Module. - - max_train_iter (:obj:`Optional[int]`): Maximum policy update iterations in training. - Returns: - - policy (:obj:`Policy`): Converged policy. - """ - if isinstance(input_cfg, str): - cfg, create_cfg = read_config(input_cfg) - else: - cfg, create_cfg = deepcopy(input_cfg) - create_cfg.policy.type = create_cfg.policy.type + '_command' - cfg = compile_config(cfg, seed=seed, auto=True, create_cfg=create_cfg) - - # Dataset - traj_dataset = D4RLTrajectoryDataset(cfg.policy.learn.dataset_path, cfg.policy.context_len, cfg.policy.rtg_scale) - traj_data_loader = DataLoader( - traj_dataset, batch_size=cfg.policy.batch_size, shuffle=True, pin_memory=True, drop_last=True - ) - # get state stats from dataset - state_mean, state_std = traj_dataset.get_state_stats() - - policy = create_policy(cfg.policy, model=model, enable_field=['learn', 'eval']) - - tb_logger = SummaryWriter(os.path.join('./{}/log/'.format(cfg.exp_name), 'serial')) - learner = BaseLearner(cfg.policy.learn.learner, policy.learn_mode, tb_logger, exp_name=cfg.exp_name) - - # ========== - # Main loop - # ========== - # Learner's before_run hook. - learner.call_hook('before_run') - stop = False - total_update_times = 0 - for i in range(max_train_iter): - for j, train_data in enumerate(traj_data_loader): - learner.train(train_data) - total_update_times += 1 - if total_update_times != 0 and total_update_times % 1000 == 0: - stop = policy.evaluate(total_update_times, state_mean, state_std) - if stop: - break - if stop: - break - learner.call_hook('after_run') - return policy, stop diff --git a/ding/entry/serial_entry_gail.py b/ding/entry/serial_entry_gail.py index 407dbaaa21..4060291fac 100644 --- a/ding/entry/serial_entry_gail.py +++ b/ding/entry/serial_entry_gail.py @@ -127,7 +127,7 @@ def serial_pipeline_gail( # Evaluate policy performance if evaluator.should_eval(learner.train_iter): stop, reward = evaluator.eval(learner.save_checkpoint, learner.train_iter, collector.envstep) - reward_mean = np.array([r['final_eval_reward'] for r in reward]).mean() + reward_mean = np.array([r['eval_episode_return'] for r in reward]).mean() if reward_mean >= best_reward: save_reward_model(cfg.exp_name, reward_model, 'best') best_reward = reward_mean diff --git a/ding/entry/serial_entry_mbrl.py b/ding/entry/serial_entry_mbrl.py index 27e7be6c4d..edb97e0c62 100644 --- a/ding/entry/serial_entry_mbrl.py +++ b/ding/entry/serial_entry_mbrl.py @@ -30,7 +30,15 @@ def mbrl_entry_setup( cfg, create_cfg = deepcopy(input_cfg) create_cfg.policy.type = create_cfg.policy.type + '_command' env_fn = None if env_setting is None else env_setting[0] - cfg = compile_config(cfg, seed=seed, env=env_fn, auto=True, create_cfg=create_cfg, save_cfg=True) + cfg = compile_config( + cfg, + seed=seed, + env=env_fn, + auto=True, + create_cfg=create_cfg, + save_cfg=True, + renew_dir=not cfg.policy.learn.get('resume_training', False) + ) if env_setting is None: env_fn, collector_env_cfg, evaluator_env_cfg = get_vec_env_setting(cfg.env) @@ -70,18 +78,7 @@ def mbrl_entry_setup( cfg.policy.other.commander, learner, collector, evaluator, env_buffer, policy.command_mode ) - return ( - cfg, - policy, - world_model, - env_buffer, - learner, - collector, - collector_env, - evaluator, - commander, - tb_logger, - ) + return (cfg, policy, world_model, env_buffer, learner, collector, collector_env, evaluator, commander, tb_logger) def create_img_buffer( @@ -131,6 +128,8 @@ def serial_pipeline_dyna( img_buffer = create_img_buffer(cfg, input_cfg, world_model, tb_logger) learner.call_hook('before_run') + if cfg.policy.learn.get('resume_training', False): + collector.envstep = learner.collector_envstep if cfg.policy.get('random_collect_size', 0) > 0: random_collect(cfg.policy, policy, collector, collector_env, commander, env_buffer) @@ -202,6 +201,8 @@ def serial_pipeline_dream( mbrl_entry_setup(input_cfg, seed, env_setting, model) learner.call_hook('before_run') + if cfg.policy.learn.get('resume_training', False): + collector.envstep = learner.collector_envstep if cfg.policy.get('random_collect_size', 0) > 0: random_collect(cfg.policy, policy, collector, collector_env, commander, env_buffer) @@ -243,3 +244,83 @@ def serial_pipeline_dream( learner.call_hook('after_run') return policy + + +def serial_pipeline_dreamer( + input_cfg: Union[str, Tuple[dict, dict]], + seed: int = 0, + env_setting: Optional[List[Any]] = None, + model: Optional[torch.nn.Module] = None, + max_train_iter: Optional[int] = int(1e10), + max_env_step: Optional[int] = int(1e10), +) -> 'Policy': # noqa + """ + Overview: + Serial pipeline entry for dreamerv3. + Arguments: + - input_cfg (:obj:`Union[str, Tuple[dict, dict]]`): Config in dict type. \ + ``str`` type means config file path. \ + ``Tuple[dict, dict]`` type means [user_config, create_cfg]. + - seed (:obj:`int`): Random seed. + - env_setting (:obj:`Optional[List[Any]]`): A list with 3 elements: \ + ``BaseEnv`` subclass, collector env config, and evaluator env config. + - model (:obj:`Optional[torch.nn.Module]`): Instance of torch.nn.Module. + - max_train_iter (:obj:`Optional[int]`): Maximum policy update iterations in training. + - max_env_step (:obj:`Optional[int]`): Maximum collected environment interaction steps. + Returns: + - policy (:obj:`Policy`): Converged policy. + """ + cfg, policy, world_model, env_buffer, learner, collector, collector_env, evaluator, commander, tb_logger = \ + mbrl_entry_setup(input_cfg, seed, env_setting, model) + + learner.call_hook('before_run') + + # prefill environment buffer + if cfg.policy.get('random_collect_size', 0) > 0: + cfg.policy.random_collect_size = cfg.policy.random_collect_size // cfg.policy.collect.unroll_len + random_collect(cfg.policy, policy, collector, collector_env, commander, env_buffer) + + while True: + collect_kwargs = commander.step() + # eval the policy + if evaluator.should_eval(collector.envstep): + stop, reward = evaluator.eval( + learner.save_checkpoint, + learner.train_iter, + collector.envstep, + policy_kwargs=dict(world_model=world_model) + ) + if stop: + break + + # train world model and fill imagination buffer + steps = ( + cfg.world_model.pretrain + if world_model.should_pretrain() else int(world_model.should_train(collector.envstep)) + ) + for _ in range(steps): + batch_size = learner.policy.get_attribute('batch_size') + batch_length = cfg.policy.learn.batch_length + post, context = world_model.train( + env_buffer, collector.envstep, learner.train_iter, batch_size, batch_length + ) + + start = post + + learner.train( + start, collector.envstep, policy_kwargs=dict(world_model=world_model, envstep=collector.envstep) + ) + + # fill environment buffer + data = collector.collect( + train_iter=learner.train_iter, + policy_kwargs=dict(world_model=world_model, envstep=collector.envstep, **collect_kwargs) + ) + env_buffer.push(data, cur_collector_envstep=collector.envstep) + + if collector.envstep >= max_env_step or learner.train_iter >= max_train_iter: + break + + learner.call_hook('after_run') + + return policy diff --git a/ding/entry/serial_entry_ngu.py b/ding/entry/serial_entry_ngu.py index 176f5558cd..1fcce53dc7 100644 --- a/ding/entry/serial_entry_ngu.py +++ b/ding/entry/serial_entry_ngu.py @@ -47,7 +47,15 @@ def serial_pipeline_ngu( cfg, create_cfg = deepcopy(input_cfg) create_cfg.policy.type = create_cfg.policy.type + '_command' env_fn = None if env_setting is None else env_setting[0] - cfg = compile_config(cfg, seed=seed, env=env_fn, auto=True, create_cfg=create_cfg, save_cfg=True) + cfg = compile_config( + cfg, + seed=seed, + env=env_fn, + auto=True, + create_cfg=create_cfg, + save_cfg=True, + renew_dir=not cfg.policy.learn.get('resume_training', False) + ) # Create main components: env, policy if env_setting is None: env_fn, collector_env_cfg, evaluator_env_cfg = get_vec_env_setting(cfg.env) @@ -89,6 +97,8 @@ def serial_pipeline_ngu( # ========== # Learner's before_run hook. learner.call_hook('before_run') + if cfg.policy.learn.get('resume_training', False): + collector.envstep = learner.collector_envstep # Accumulate plenty of data at the beginning of training. if cfg.policy.get('random_collect_size', 0) > 0: diff --git a/ding/entry/serial_entry_offline.py b/ding/entry/serial_entry_offline.py old mode 100644 new mode 100755 index 9243294865..b92b5c7dda --- a/ding/entry/serial_entry_offline.py +++ b/ding/entry/serial_entry_offline.py @@ -75,6 +75,8 @@ def serial_pipeline_offline( evaluator_env.seed(cfg.seed, dynamic_seed=False) set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) policy = create_policy(cfg.policy, model=model, enable_field=['learn', 'eval']) + if cfg.policy.collect.data_type == 'diffuser_traj': + policy.init_data_normalizer(dataset.normalizer) if hasattr(policy, 'set_statistic'): # useful for setting action bounds for ibc diff --git a/ding/entry/serial_entry_onpolicy.py b/ding/entry/serial_entry_onpolicy.py index 40631257f6..a044be05f3 100644 --- a/ding/entry/serial_entry_onpolicy.py +++ b/ding/entry/serial_entry_onpolicy.py @@ -12,7 +12,7 @@ from ding.config import read_config, compile_config from ding.policy import create_policy, PolicyFactory from ding.reward_model import create_reward_model -from ding.utils import set_pkg_seed +from ding.utils import set_pkg_seed, get_rank def serial_pipeline_onpolicy( @@ -45,7 +45,16 @@ def serial_pipeline_onpolicy( cfg, create_cfg = deepcopy(input_cfg) create_cfg.policy.type = create_cfg.policy.type + '_command' env_fn = None if env_setting is None else env_setting[0] - cfg = compile_config(cfg, seed=seed, env=env_fn, auto=True, create_cfg=create_cfg, save_cfg=True) + cfg = compile_config( + cfg, + seed=seed, + env=env_fn, + auto=True, + create_cfg=create_cfg, + save_cfg=True, + renew_dir=not cfg.policy.learn.get('resume_training', False) + ) + # Create main components: env, policy if env_setting is None: env_fn, collector_env_cfg, evaluator_env_cfg = get_vec_env_setting(cfg.env) @@ -59,7 +68,7 @@ def serial_pipeline_onpolicy( policy = create_policy(cfg.policy, model=model, enable_field=['learn', 'collect', 'eval', 'command']) # Create worker components: learner, collector, evaluator, replay buffer, commander. - tb_logger = SummaryWriter(os.path.join('./{}/log/'.format(cfg.exp_name), 'serial')) + tb_logger = SummaryWriter(os.path.join('./{}/log/'.format(cfg.exp_name), 'serial')) if get_rank() == 0 else None learner = BaseLearner(cfg.policy.learn.learner, policy.learn_mode, tb_logger, exp_name=cfg.exp_name) collector = create_serial_collector( cfg.policy.collect.collector, @@ -80,12 +89,14 @@ def serial_pipeline_onpolicy( # ========== # Learner's before_run hook. learner.call_hook('before_run') + if cfg.policy.learn.get('resume_training', False): + collector.envstep = learner.collector_envstep while True: collect_kwargs = commander.step() # Evaluate policy performance if evaluator.should_eval(learner.train_iter): - stop, reward = evaluator.eval(learner.save_checkpoint, learner.train_iter, collector.envstep) + stop, eval_info = evaluator.eval(learner.save_checkpoint, learner.train_iter, collector.envstep) if stop: break # Collect data by default config n_sample/n_episode @@ -98,4 +109,18 @@ def serial_pipeline_onpolicy( # Learner's after_run hook. learner.call_hook('after_run') + import time + import pickle + import numpy as np + with open(os.path.join(cfg.exp_name, 'result.pkl'), 'wb') as f: + eval_value_raw = eval_info['eval_episode_return'] + final_data = { + 'stop': stop, + 'env_step': collector.envstep, + 'train_iter': learner.train_iter, + 'eval_value': np.mean(eval_value_raw), + 'eval_value_raw': eval_value_raw, + 'finish_time': time.ctime(), + } + pickle.dump(final_data, f) return policy diff --git a/ding/entry/serial_entry_onpolicy_ppg.py b/ding/entry/serial_entry_onpolicy_ppg.py index 02c6dee307..90f31891ab 100644 --- a/ding/entry/serial_entry_onpolicy_ppg.py +++ b/ding/entry/serial_entry_onpolicy_ppg.py @@ -45,7 +45,15 @@ def serial_pipeline_onpolicy_ppg( cfg, create_cfg = deepcopy(input_cfg) create_cfg.policy.type = create_cfg.policy.type + '_command' env_fn = None if env_setting is None else env_setting[0] - cfg = compile_config(cfg, seed=seed, env=env_fn, auto=True, create_cfg=create_cfg, save_cfg=True) + cfg = compile_config( + cfg, + seed=seed, + env=env_fn, + auto=True, + create_cfg=create_cfg, + save_cfg=True, + renew_dir=not cfg.policy.learn.get('resume_training', False) + ) # Create main components: env, policy if env_setting is None: env_fn, collector_env_cfg, evaluator_env_cfg = get_vec_env_setting(cfg.env) @@ -80,6 +88,8 @@ def serial_pipeline_onpolicy_ppg( # ========== # Learner's before_run hook. learner.call_hook('before_run') + if cfg.policy.learn.get('resume_training', False): + collector.envstep = learner.collector_envstep while True: collect_kwargs = commander.step() diff --git a/ding/entry/serial_entry_pc.py b/ding/entry/serial_entry_pc.py new file mode 100644 index 0000000000..386d6f0ec9 --- /dev/null +++ b/ding/entry/serial_entry_pc.py @@ -0,0 +1,108 @@ +from typing import Union, Optional, Tuple +import os +from functools import partial +from copy import deepcopy + +import torch +from tensorboardX import SummaryWriter +from torch.utils.data import DataLoader + +from ding.envs import get_vec_env_setting, create_env_manager +from ding.worker import BaseLearner, InteractionSerialEvaluator +from ding.config import read_config, compile_config +from ding.policy import create_policy +from ding.utils import set_pkg_seed +from ding.utils.data.dataset import load_bfs_datasets + + +def serial_pipeline_pc( + input_cfg: Union[str, Tuple[dict, dict]], + seed: int = 0, + model: Optional[torch.nn.Module] = None, + max_iter=int(1e6), +) -> Union['Policy', bool]: # noqa + r""" + Overview: + Serial pipeline entry of procedure cloning using BFS as expert policy. + Arguments: + - input_cfg (:obj:`Union[str, Tuple[dict, dict]]`): Config in dict type. \ + ``str`` type means config file path. \ + ``Tuple[dict, dict]`` type means [user_config, create_cfg]. + - seed (:obj:`int`): Random seed. + - model (:obj:`Optional[torch.nn.Module]`): Instance of torch.nn.Module. + - max_iter (:obj:`Optional[int]`): Max iteration for executing PC training. + Returns: + - policy (:obj:`Policy`): Converged policy. + - convergence (:obj:`bool`): whether the training is converged + """ + if isinstance(input_cfg, str): + cfg, create_cfg = read_config(input_cfg) + else: + cfg, create_cfg = deepcopy(input_cfg) + cfg = compile_config(cfg, seed=seed, auto=True, create_cfg=create_cfg) + + # Env, Policy + env_fn, _, evaluator_env_cfg = get_vec_env_setting(cfg.env) + evaluator_env = create_env_manager(cfg.env.manager, [partial(env_fn, cfg=c) for c in evaluator_env_cfg]) + # Random seed + evaluator_env.seed(cfg.seed, dynamic_seed=False) + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + policy = create_policy(cfg.policy, model=model, enable_field=['learn', 'eval']) + + # Main components + tb_logger = SummaryWriter(os.path.join('./{}/log/'.format(cfg.exp_name), 'serial')) + train_data, test_data = load_bfs_datasets(train_seeds=cfg.train_seeds) + dataloader = DataLoader(train_data, batch_size=cfg.policy.learn.batch_size, shuffle=True) + test_dataloader = DataLoader(test_data, batch_size=cfg.policy.learn.batch_size, shuffle=True) + learner = BaseLearner(cfg.policy.learn.learner, policy.learn_mode, tb_logger, exp_name=cfg.exp_name) + evaluator = InteractionSerialEvaluator( + cfg.policy.eval.evaluator, evaluator_env, policy.eval_mode, tb_logger, exp_name=cfg.exp_name + ) + + # ========== + # Main loop + # ========== + learner.call_hook('before_run') + stop = False + iter_cnt = 0 + for epoch in range(cfg.policy.learn.train_epoch): + # train + criterion = torch.nn.CrossEntropyLoss() + for i, train_data in enumerate(dataloader): + learner.train(train_data) + iter_cnt += 1 + if iter_cnt >= max_iter: + stop = True + break + if epoch % 69 == 0: + policy._optimizer.param_groups[0]['lr'] /= 10 + if stop: + break + losses = [] + acces = [] + # Evaluation + for _, test_data in enumerate(test_dataloader): + observations, bfs_input_maps, bfs_output_maps = test_data['obs'], test_data['bfs_in'].long(), \ + test_data['bfs_out'].long() + states = observations + bfs_input_onehot = torch.nn.functional.one_hot(bfs_input_maps, 5).float() + + bfs_states = torch.cat([ + states, + bfs_input_onehot, + ], dim=-1).cuda() + logits = policy._model(bfs_states)['logit'] + logits = logits.flatten(0, -2) + labels = bfs_output_maps.flatten(0, -1).cuda() + + loss = criterion(logits, labels).item() + preds = torch.argmax(logits, dim=-1) + acc = torch.sum((preds == labels)) / preds.shape[0] + + losses.append(loss) + acces.append(acc) + print('Test Finished! Loss: {} acc: {}'.format(sum(losses) / len(losses), sum(acces) / len(acces))) + stop, reward = evaluator.eval(learner.save_checkpoint, learner.train_iter) + learner.call_hook('after_run') + print('final reward is: {}'.format(reward)) + return policy, stop diff --git a/ding/entry/serial_entry_preference_based_irl.py b/ding/entry/serial_entry_preference_based_irl.py index 8ceaf52977..682e662baa 100644 --- a/ding/entry/serial_entry_preference_based_irl.py +++ b/ding/entry/serial_entry_preference_based_irl.py @@ -47,7 +47,7 @@ def serial_pipeline_preference_based_irl( create_cfg.policy.type = create_cfg.policy.type + '_command' create_cfg.reward_model = dict(type=cfg.reward_model.type) env_fn = None if env_setting is None else env_setting[0] - cfg = compile_config(cfg, seed=seed, env=env_fn, auto=True, create_cfg=create_cfg, save_cfg=True) + cfg = compile_config(cfg, seed=seed, env=env_fn, auto=True, create_cfg=create_cfg, save_cfg=True, renew_dir=False) cfg_bak = copy.deepcopy(cfg) # Create main components: env, policy if env_setting is None: diff --git a/ding/entry/serial_entry_preference_based_irl_onpolicy.py b/ding/entry/serial_entry_preference_based_irl_onpolicy.py index 3bc1a8875b..3941f3337e 100644 --- a/ding/entry/serial_entry_preference_based_irl_onpolicy.py +++ b/ding/entry/serial_entry_preference_based_irl_onpolicy.py @@ -46,7 +46,7 @@ def serial_pipeline_preference_based_irl_onpolicy( create_cfg.policy.type = create_cfg.policy.type + '_command' create_cfg.reward_model = dict(type=cfg.reward_model.type) env_fn = None if env_setting is None else env_setting[0] - cfg = compile_config(cfg, seed=seed, env=env_fn, auto=True, create_cfg=create_cfg, save_cfg=True) + cfg = compile_config(cfg, seed=seed, env=env_fn, auto=True, create_cfg=create_cfg, save_cfg=True, renew_dir=False) # Create main components: env, policy if env_setting is None: env_fn, collector_env_cfg, evaluator_env_cfg = get_vec_env_setting(cfg.env) diff --git a/ding/entry/serial_entry_reward_model_offpolicy.py b/ding/entry/serial_entry_reward_model_offpolicy.py index 7269fbd29b..f1b4c004b3 100644 --- a/ding/entry/serial_entry_reward_model_offpolicy.py +++ b/ding/entry/serial_entry_reward_model_offpolicy.py @@ -1,5 +1,6 @@ from typing import Union, Optional, List, Any, Tuple import os +import numpy as np import torch from ditk import logging from functools import partial @@ -87,11 +88,17 @@ def serial_pipeline_reward_model_offpolicy( # Accumulate plenty of data at the beginning of training. if cfg.policy.get('random_collect_size', 0) > 0: random_collect(cfg.policy, policy, collector, collector_env, commander, replay_buffer) + count = 0 + best_return = -np.inf while True: collect_kwargs = commander.step() # Evaluate policy performance if evaluator.should_eval(learner.train_iter): - stop, reward = evaluator.eval(learner.save_checkpoint, learner.train_iter, collector.envstep) + stop, eval_info = evaluator.eval(learner.save_checkpoint, learner.train_iter, collector.envstep) + eval_return_mean = np.mean(eval_info['eval_episode_return']) + if eval_return_mean >= best_return: + reward_model.save(path=cfg.exp_name, name='best') + best_return = eval_return_mean if stop: break new_data_count, target_new_data_count = 0, cfg.reward_model.get('target_new_data_count', 1) @@ -103,7 +110,9 @@ def serial_pipeline_reward_model_offpolicy( replay_buffer.push(new_data, cur_collector_envstep=collector.envstep) # update reward_model reward_model.train() - reward_model.clear_data() + # clear buffer per fix iters to make sure replay buffer's data count isn't too few. + if count % cfg.reward_model.clear_buffer_per_iters == 0: + reward_model.clear_data() # Learn policy from collected data for i in range(cfg.policy.learn.update_per_collect): # Learner will train ``update_per_collect`` times in one iteration. @@ -122,7 +131,9 @@ def serial_pipeline_reward_model_offpolicy( replay_buffer.update(learner.priority_info) if collector.envstep >= max_env_step or learner.train_iter >= max_train_iter: break + count += 1 # Learner's after_run hook. learner.call_hook('after_run') + reward_model.save(path=cfg.exp_name, name='last') return policy diff --git a/ding/entry/serial_entry_reward_model_onpolicy.py b/ding/entry/serial_entry_reward_model_onpolicy.py index b648c86cfb..b01864f98f 100644 --- a/ding/entry/serial_entry_reward_model_onpolicy.py +++ b/ding/entry/serial_entry_reward_model_onpolicy.py @@ -1,5 +1,6 @@ from typing import Union, Optional, List, Any, Tuple import os +import numpy as np import torch from ditk import logging from functools import partial @@ -88,11 +89,16 @@ def serial_pipeline_reward_model_onpolicy( if cfg.policy.get('random_collect_size', 0) > 0: random_collect(cfg.policy, policy, collector, collector_env, commander, replay_buffer) count = 0 + best_return = -np.inf while True: collect_kwargs = commander.step() # Evaluate policy performance if evaluator.should_eval(learner.train_iter): - stop, reward = evaluator.eval(learner.save_checkpoint, learner.train_iter, collector.envstep) + stop, eval_info = evaluator.eval(learner.save_checkpoint, learner.train_iter, collector.envstep) + eval_return_mean = np.mean(eval_info['eval_episode_return']) + if eval_return_mean >= best_return: + reward_model.save(path=cfg.exp_name, name='best') + best_return = eval_return_mean if stop: break new_data_count, target_new_data_count = 0, cfg.reward_model.get('target_new_data_count', 1) @@ -127,4 +133,5 @@ def serial_pipeline_reward_model_onpolicy( # Learner's after_run hook. learner.call_hook('after_run') + reward_model.save(path=cfg.exp_name, name='last') return policy diff --git a/ding/entry/tests/test_application_entry.py b/ding/entry/tests/test_application_entry.py index 706f50efbf..9276d9e6e5 100644 --- a/ding/entry/tests/test_application_entry.py +++ b/ding/entry/tests/test_application_entry.py @@ -3,8 +3,8 @@ import os import pickle -from dizoo.classic_control.cartpole.config.cartpole_offppo_config import cartpole_offppo_config, \ - cartpole_offppo_create_config # noqa +from dizoo.classic_control.cartpole.config.cartpole_ppo_offpolicy_config import cartpole_ppo_offpolicy_config, \ + cartpole_ppo_offpolicy_create_config # noqa from dizoo.classic_control.cartpole.config.cartpole_trex_offppo_config import cartpole_trex_offppo_config,\ cartpole_trex_offppo_create_config from dizoo.classic_control.cartpole.envs import CartPoleEnv @@ -15,7 +15,7 @@ @pytest.fixture(scope='module') def setup_state_dict(): - config = deepcopy(cartpole_offppo_config), deepcopy(cartpole_offppo_create_config) + config = deepcopy(cartpole_ppo_offpolicy_config), deepcopy(cartpole_ppo_offpolicy_create_config) try: policy = serial_pipeline(config, seed=0) except Exception: @@ -31,22 +31,24 @@ def setup_state_dict(): class TestApplication: def test_eval(self, setup_state_dict): - cfg_for_stop_value = compile_config(cartpole_offppo_config, auto=True, create_cfg=cartpole_offppo_create_config) + cfg_for_stop_value = compile_config( + cartpole_ppo_offpolicy_config, auto=True, create_cfg=cartpole_ppo_offpolicy_create_config + ) stop_value = cfg_for_stop_value.env.stop_value - config = deepcopy(cartpole_offppo_config), deepcopy(cartpole_offppo_create_config) - eval_reward = eval(config, seed=0, state_dict=setup_state_dict['eval']) - assert eval_reward >= stop_value - config = deepcopy(cartpole_offppo_config), deepcopy(cartpole_offppo_create_config) - eval_reward = eval( + config = deepcopy(cartpole_ppo_offpolicy_config), deepcopy(cartpole_ppo_offpolicy_create_config) + episode_return = eval(config, seed=0, state_dict=setup_state_dict['eval']) + assert episode_return >= stop_value + config = deepcopy(cartpole_ppo_offpolicy_config), deepcopy(cartpole_ppo_offpolicy_create_config) + episode_return = eval( config, seed=0, env_setting=[CartPoleEnv, None, [{} for _ in range(5)]], state_dict=setup_state_dict['eval'] ) - assert eval_reward >= stop_value + assert episode_return >= stop_value def test_collect_demo_data(self, setup_state_dict): - config = deepcopy(cartpole_offppo_config), deepcopy(cartpole_offppo_create_config) + config = deepcopy(cartpole_ppo_offpolicy_config), deepcopy(cartpole_ppo_offpolicy_create_config) collect_count = 16 expert_data_path = './expert.data' collect_demo_data( diff --git a/ding/entry/tests/test_application_entry_drex_collect_data.py b/ding/entry/tests/test_application_entry_drex_collect_data.py deleted file mode 100644 index de1ded03ee..0000000000 --- a/ding/entry/tests/test_application_entry_drex_collect_data.py +++ /dev/null @@ -1,74 +0,0 @@ -from easydict import EasyDict -from copy import deepcopy -from itertools import product - -import pytest -import os -import torch - -from ding.entry import serial_pipeline -from ding.entry.application_entry_drex_collect_data import collect_episodic_demo_data_for_drex, drex_collecting_data -from ding.config import compile_config -from dizoo.classic_control.cartpole.config.cartpole_drex_dqn_config import cartpole_drex_dqn_config,\ - cartpole_drex_dqn_create_config -from dizoo.classic_control.cartpole.config.cartpole_dqn_config import cartpole_dqn_config,\ - cartpole_dqn_create_config - - -@pytest.mark.unittest -def test_collect_episodic_demo_data_for_drex(): - expert_policy_state_dict_path = './expert_policy.pth' - expert_policy_state_dict_path = os.path.abspath('./expert_policy.pth') - config = [deepcopy(cartpole_dqn_config), deepcopy(cartpole_dqn_create_config)] - expert_policy = serial_pipeline(config, seed=0, max_train_iter=1) - torch.save(expert_policy.collect_mode.state_dict(), expert_policy_state_dict_path) - - config = deepcopy(cartpole_drex_dqn_config) - config = compile_config( - config, seed=0, env=None, auto=True, create_cfg=cartpole_drex_dqn_create_config, save_cfg=True - ) - collect_count = 1 - save_cfg_path = './cartpole_drex_dqn' - save_cfg_path = os.path.abspath(save_cfg_path) - exp_data = collect_episodic_demo_data_for_drex( - config, - seed=0, - state_dict_path=expert_policy_state_dict_path, - save_cfg_path=save_cfg_path, - collect_count=collect_count, - rank=1, - noise=-1, - ) - assert isinstance(exp_data, list) - assert isinstance(exp_data[0][0], dict) - os.popen('rm -rf {}'.format(save_cfg_path)) - os.popen('rm -rf {}'.format(expert_policy_state_dict_path)) - - -@pytest.mark.unittest -def test_drex_collecting_data(): - expert_policy_state_dict_path = './cartpole_dqn_seed0' - expert_policy_state_dict_path = os.path.abspath(expert_policy_state_dict_path) - config = [deepcopy(cartpole_dqn_config), deepcopy(cartpole_dqn_create_config)] - expert_policy = serial_pipeline(config, seed=0, max_train_iter=1) - - args = EasyDict( - { - 'cfg': [deepcopy(cartpole_drex_dqn_config), - deepcopy(cartpole_drex_dqn_create_config)], - 'seed': 0, - 'device': 'cpu' - } - ) - args.cfg[0].reward_model.offline_data_path = './cartpole_drex_dqn_seed0' - args.cfg[0].reward_model.offline_data_path = os.path.abspath(args.cfg[0].reward_model.offline_data_path) - args.cfg[0].reward_model.reward_model_path = args.cfg[0].reward_model.offline_data_path + '/cartpole.params' - args.cfg[0].reward_model.expert_model_path = './cartpole_dqn_seed0/ckpt/ckpt_best.pth.tar' - args.cfg[0].reward_model.expert_model_path = os.path.abspath(args.cfg[0].reward_model.expert_model_path) - args.cfg[0].reward_model.bc_iterations = 6 - args.cfg[0].reward_model.num_trajs_per_bin = 8 - args.cfg[0].bc_iteration = 1000 # for unittest - args.cfg[1].policy.type = 'bc' - drex_collecting_data(args=args) - os.popen('rm -rf {}'.format(expert_policy_state_dict_path)) - os.popen('rm -rf {}'.format(args.cfg[0].reward_model.offline_data_path)) diff --git a/ding/entry/tests/test_application_entry_trex_collect_data.py b/ding/entry/tests/test_application_entry_trex_collect_data.py index 13e418b5f0..f5cb3c16b9 100644 --- a/ding/entry/tests/test_application_entry_trex_collect_data.py +++ b/ding/entry/tests/test_application_entry_trex_collect_data.py @@ -8,8 +8,8 @@ from dizoo.classic_control.cartpole.config.cartpole_trex_offppo_config import cartpole_trex_offppo_config,\ cartpole_trex_offppo_create_config -from dizoo.classic_control.cartpole.config.cartpole_offppo_config import cartpole_offppo_config,\ - cartpole_offppo_create_config +from dizoo.classic_control.cartpole.config.cartpole_ppo_offpolicy_config import cartpole_ppo_offpolicy_config,\ + cartpole_ppo_offpolicy_create_config from ding.entry.application_entry_trex_collect_data import collect_episodic_demo_data_for_trex, trex_collecting_data from ding.entry import serial_pipeline @@ -18,7 +18,7 @@ def test_collect_episodic_demo_data_for_trex(): exp_name = "test_collect_episodic_demo_data_for_trex_expert" expert_policy_state_dict_path = os.path.join(exp_name, 'expert_policy.pth.tar') - config = [deepcopy(cartpole_offppo_config), deepcopy(cartpole_offppo_create_config)] + config = [deepcopy(cartpole_ppo_offpolicy_config), deepcopy(cartpole_ppo_offpolicy_create_config)] config[0].exp_name = exp_name expert_policy = serial_pipeline(config, seed=0) torch.save(expert_policy.collect_mode.state_dict(), expert_policy_state_dict_path) @@ -41,7 +41,7 @@ def test_collect_episodic_demo_data_for_trex(): @pytest.mark.unittest def test_trex_collecting_data(): expert_policy_dir = 'test_trex_collecting_data_expert' - config = [deepcopy(cartpole_offppo_config), deepcopy(cartpole_offppo_create_config)] + config = [deepcopy(cartpole_ppo_offpolicy_config), deepcopy(cartpole_ppo_offpolicy_create_config)] config[0].exp_name = expert_policy_dir config[0].policy.learn.learner.hook.save_ckpt_after_iter = 100 serial_pipeline(config, seed=0) diff --git a/ding/entry/tests/test_cli_ditask.py b/ding/entry/tests/test_cli_ditask.py index 66cd37e906..6bb64e5e6e 100644 --- a/ding/entry/tests/test_cli_ditask.py +++ b/ding/entry/tests/test_cli_ditask.py @@ -25,7 +25,8 @@ def test_cli_ditask(): "node_ids": 0, "mq_type": "nng", "redis_host": "", - "redis_port": "" + "redis_port": "", + "startup_interval": 1 } os.environ["DI_NODES"] = '127.0.0.1' os.environ["DI_RANK"] = '0' diff --git a/ding/entry/tests/test_serial_entry.py b/ding/entry/tests/test_serial_entry.py index 917f73554d..d36f6bc717 100644 --- a/ding/entry/tests/test_serial_entry.py +++ b/ding/entry/tests/test_serial_entry.py @@ -9,8 +9,8 @@ from dizoo.classic_control.cartpole.config.cartpole_dqn_stdim_config import cartpole_dqn_stdim_config, \ cartpole_dqn_stdim_create_config from dizoo.classic_control.cartpole.config.cartpole_ppo_config import cartpole_ppo_config, cartpole_ppo_create_config -from dizoo.classic_control.cartpole.config.cartpole_offppo_config import cartpole_offppo_config, \ - cartpole_offppo_create_config +from dizoo.classic_control.cartpole.config.cartpole_ppo_offpolicy_config import cartpole_ppo_offpolicy_config, \ + cartpole_ppo_offpolicy_create_config from dizoo.classic_control.cartpole.config.cartpole_impala_config import cartpole_impala_config, cartpole_impala_create_config # noqa from dizoo.classic_control.cartpole.config.cartpole_rainbow_config import cartpole_rainbow_config, cartpole_rainbow_create_config # noqa from dizoo.classic_control.cartpole.config.cartpole_iqn_config import cartpole_iqn_config, cartpole_iqn_create_config # noqa @@ -37,6 +37,7 @@ from dizoo.petting_zoo.config import ptz_simple_spread_qtran_config, ptz_simple_spread_qtran_create_config # noqa from dizoo.petting_zoo.config import ptz_simple_spread_vdn_config, ptz_simple_spread_vdn_create_config # noqa from dizoo.petting_zoo.config import ptz_simple_spread_wqmix_config, ptz_simple_spread_wqmix_create_config # noqa +from dizoo.petting_zoo.config import ptz_simple_spread_madqn_config, ptz_simple_spread_madqn_create_config # noqa from dizoo.league_demo.league_demo_ppo_config import league_demo_ppo_config from dizoo.league_demo.selfplay_demo_ppo_main import main as selfplay_main from dizoo.league_demo.league_demo_ppo_main import main as league_main @@ -44,12 +45,15 @@ from dizoo.classic_control.pendulum.config.pendulum_cql_config import pendulum_cql_config, pendulum_cql_create_config # noqa from dizoo.classic_control.cartpole.config.cartpole_qrdqn_generation_data_config import cartpole_qrdqn_generation_data_config, cartpole_qrdqn_generation_data_create_config # noqa from dizoo.classic_control.cartpole.config.cartpole_cql_config import cartpole_discrete_cql_config, cartpole_discrete_cql_create_config # noqa +from dizoo.classic_control.cartpole.config.cartpole_dt_config import cartpole_discrete_dt_config, cartpole_discrete_dt_create_config # noqa from dizoo.classic_control.pendulum.config.pendulum_td3_data_generation_config import pendulum_td3_generation_config, pendulum_td3_generation_create_config # noqa from dizoo.classic_control.pendulum.config.pendulum_td3_bc_config import pendulum_td3_bc_config, pendulum_td3_bc_create_config # noqa from dizoo.classic_control.pendulum.config.pendulum_ibc_config import pendulum_ibc_config, pendulum_ibc_create_config from dizoo.gym_hybrid.config.gym_hybrid_ddpg_config import gym_hybrid_ddpg_config, gym_hybrid_ddpg_create_config from dizoo.gym_hybrid.config.gym_hybrid_pdqn_config import gym_hybrid_pdqn_config, gym_hybrid_pdqn_create_config from dizoo.gym_hybrid.config.gym_hybrid_mpdqn_config import gym_hybrid_mpdqn_config, gym_hybrid_mpdqn_create_config +from dizoo.classic_control.pendulum.config.pendulum_bdq_config import pendulum_bdq_config, pendulum_bdq_create_config # noqa +from dizoo.classic_control.cartpole.config.cartpole_mdqn_config import cartpole_mdqn_config, cartpole_mdqn_create_config @pytest.mark.platformtest @@ -66,6 +70,34 @@ def test_dqn(): os.popen('rm -rf cartpole_dqn_unittest') +@pytest.mark.platformtest +@pytest.mark.unittest +def test_mdqn(): + config = [deepcopy(cartpole_mdqn_config), deepcopy(cartpole_mdqn_create_config)] + config[0].policy.learn.update_per_collect = 1 + config[0].exp_name = 'cartpole_mdqn_unittest' + try: + serial_pipeline(config, seed=0, max_train_iter=1, dynamic_seed=False) + except Exception: + assert False, "pipeline fail" + finally: + os.popen('rm -rf cartpole_mdqn_unittest') + + +@pytest.mark.platformtest +@pytest.mark.unittest +def test_bdq(): + config = [deepcopy(pendulum_bdq_config), deepcopy(pendulum_bdq_create_config)] + config[0].policy.learn.update_per_collect = 1 + config[0].exp_name = 'pendulum_bdq_unittest' + try: + serial_pipeline(config, seed=0, max_train_iter=1) + except Exception: + assert False, "pipeline fail" + finally: + os.popen('rm -rf pendulum_bdq_unittest') + + @pytest.mark.platformtest @pytest.mark.unittest def test_ddpg(): @@ -193,7 +225,7 @@ def test_qrdqn(): @pytest.mark.platformtest @pytest.mark.unittest def test_ppo(): - config = [deepcopy(cartpole_offppo_config), deepcopy(cartpole_offppo_create_config)] + config = [deepcopy(cartpole_ppo_offpolicy_config), deepcopy(cartpole_ppo_offpolicy_create_config)] config[0].policy.learn.update_per_collect = 1 config[0].exp_name = 'ppo_offpolicy_unittest' try: @@ -205,7 +237,7 @@ def test_ppo(): @pytest.mark.platformtest @pytest.mark.unittest def test_ppo_nstep_return(): - config = [deepcopy(cartpole_offppo_config), deepcopy(cartpole_offppo_create_config)] + config = [deepcopy(cartpole_ppo_offpolicy_config), deepcopy(cartpole_ppo_offpolicy_create_config)] config[0].policy.learn.update_per_collect = 1 config[0].policy.nstep_return = True try: @@ -252,15 +284,10 @@ def test_sac_log_space(): assert False, "pipeline fail" -auto_alpha = [True, False] -log_space = [True, False] -args = [item for item in product(*[auto_alpha, log_space])] - - @pytest.mark.platformtest @pytest.mark.unittest -@pytest.mark.parametrize('auto_alpha, log_space', args) -def test_discrete_sac(auto_alpha, log_space): +def test_discrete_sac(): + auto_alpha, log_space = True, False config = [deepcopy(cartpole_sac_config), deepcopy(cartpole_sac_create_config)] config[0].policy.learn.update_per_collect = 1 config[0].policy.learn.auto_alpha = auto_alpha @@ -378,6 +405,20 @@ def test_wqmix(): os.popen('rm -rf log ckpt*') +@pytest.mark.platformtest +@pytest.mark.unittest +def test_madqn(): + config = [deepcopy(ptz_simple_spread_madqn_config), deepcopy(ptz_simple_spread_madqn_create_config)] + config[0].policy.cuda = False + config[0].policy.learn.update_per_collect = 1 + try: + serial_pipeline(config, seed=0, max_train_iter=1) + except Exception: + assert False, "pipeline fail" + finally: + os.popen('rm -rf log ckpt*') + + @pytest.mark.platformtest @pytest.mark.unittest def test_qtran(): @@ -576,6 +617,71 @@ def test_discrete_cql(): os.popen('rm -rf cartpole cartpole_cql') +@pytest.mark.platformtest +@pytest.mark.unittest +def test_discrete_dt(): + # train expert + config = [deepcopy(cartpole_qrdqn_config), deepcopy(cartpole_qrdqn_create_config)] + config[0].policy.learn.update_per_collect = 1 + config[0].exp_name = 'dt_cartpole' + try: + serial_pipeline(config, seed=0, max_train_iter=1) + except Exception: + assert False, "pipeline fail" + # collect expert data + import torch + config = [deepcopy(cartpole_qrdqn_generation_data_config), deepcopy(cartpole_qrdqn_generation_data_create_config)] + state_dict = torch.load('./dt_cartpole/ckpt/iteration_0.pth.tar', map_location='cpu') + try: + collect_demo_data(config, seed=0, collect_count=1000, state_dict=state_dict) + except Exception as e: + assert False, "pipeline fail" + print(repr(e)) + + # train dt + config = [deepcopy(cartpole_discrete_dt_config), deepcopy(cartpole_discrete_dt_create_config)] + config[0].policy.eval.evaluator.eval_freq = 5 + try: + from ding.framework import task, ding_init + from ding.framework.context import OfflineRLContext + from ding.envs import SubprocessEnvManagerV2, BaseEnvManagerV2 + from ding.envs.env_wrappers.env_wrappers import AllinObsWrapper + from dizoo.classic_control.cartpole.envs import CartPoleEnv + from ding.utils import set_pkg_seed + from ding.data import create_dataset + from ding.config import compile_config + from ding.model import DecisionTransformer + from ding.policy import DTPolicy + from ding.framework.middleware import interaction_evaluator, trainer, CkptSaver, \ + OfflineMemoryDataFetcher, offline_logger, termination_checker + ding_init(config[0]) + config = compile_config(config[0], create_cfg=config[1], auto=True) + with task.start(async_mode=False, ctx=OfflineRLContext()): + evaluator_env = BaseEnvManagerV2( + env_fn=[lambda: AllinObsWrapper(CartPoleEnv(config.env)) for _ in range(config.env.evaluator_env_num)], + cfg=config.env.manager + ) + + set_pkg_seed(config.seed, use_cuda=config.policy.cuda) + + dataset = create_dataset(config) + + model = DecisionTransformer(**config.policy.model) + policy = DTPolicy(config.policy, model=model) + + task.use(termination_checker(max_train_iter=1)) + task.use(interaction_evaluator(config, policy.eval_mode, evaluator_env)) + task.use(OfflineMemoryDataFetcher(config, dataset)) + task.use(trainer(config, policy.learn_mode)) + task.use(CkptSaver(policy, config.exp_name, train_freq=100)) + task.use(offline_logger()) + task.run() + except Exception: + assert False, "pipeline fail" + finally: + os.popen('rm -rf cartpole cartpole_dt') + + @pytest.mark.platformtest @pytest.mark.unittest def test_td3_bc(): diff --git a/ding/entry/tests/test_serial_entry_algo.py b/ding/entry/tests/test_serial_entry_algo.py index 4057ce2fe8..640b1e800c 100644 --- a/ding/entry/tests/test_serial_entry_algo.py +++ b/ding/entry/tests/test_serial_entry_algo.py @@ -11,6 +11,7 @@ from dizoo.classic_control.cartpole.config.cartpole_sqil_config import cartpole_sqil_config, cartpole_sqil_create_config from dizoo.classic_control.cartpole.config.cartpole_dqn_config import cartpole_dqn_config, cartpole_dqn_create_config from dizoo.classic_control.cartpole.config.cartpole_ppo_config import cartpole_ppo_config, cartpole_ppo_create_config +from dizoo.classic_control.cartpole.config.cartpole_pg_config import cartpole_pg_config, cartpole_pg_create_config from dizoo.classic_control.cartpole.config.cartpole_a2c_config import cartpole_a2c_config, cartpole_a2c_create_config from dizoo.classic_control.cartpole.config.cartpole_impala_config import cartpole_impala_config, cartpole_impala_create_config # noqa from dizoo.classic_control.cartpole.config.cartpole_rainbow_config import cartpole_rainbow_config, cartpole_rainbow_create_config # noqa @@ -44,6 +45,7 @@ from dizoo.petting_zoo.config import ptz_simple_spread_qtran_config, ptz_simple_spread_qtran_create_config # noqa from dizoo.petting_zoo.config import ptz_simple_spread_vdn_config, ptz_simple_spread_vdn_create_config # noqa from dizoo.petting_zoo.config import ptz_simple_spread_wqmix_config, ptz_simple_spread_wqmix_create_config # noqa +from dizoo.classic_control.cartpole.config import cartpole_mdqn_config, cartpole_mdqn_create_config with open("./algo_record.log", "w+") as f: f.write("ALGO TEST STARTS\n") @@ -170,6 +172,17 @@ def test_r2d2(): f.write("11. r2d2\n") +@pytest.mark.algotest +def test_pg(): + config = [deepcopy(cartpole_pg_config), deepcopy(cartpole_pg_create_config)] + try: + serial_pipeline_onpolicy(config, seed=0) + except Exception: + assert False, "pipeline fail" + with open("./algo_record.log", "a+") as f: + f.write("12. pg\n") + + # @pytest.mark.algotest def test_atoc(): config = [deepcopy(ptz_simple_spread_atoc_config), deepcopy(ptz_simple_spread_atoc_create_config)] @@ -393,6 +406,17 @@ def test_wqmix(): f.write("28. wqmix\n") +@pytest.mark.algotest +def test_mdqn(): + config = [deepcopy(cartpole_mdqn_config), deepcopy(cartpole_mdqn_create_config)] + try: + serial_pipeline(config, seed=0) + except Exception: + assert False, "pipeline fail" + with open("./algo_record.log", "a+") as f: + f.write("29. mdqn\n") + + # @pytest.mark.algotest def test_td3_bc(): # train expert diff --git a/ding/entry/tests/test_serial_entry_bc.py b/ding/entry/tests/test_serial_entry_bc.py index 4c2b1ae9cb..f2c0923ad2 100644 --- a/ding/entry/tests/test_serial_entry_bc.py +++ b/ding/entry/tests/test_serial_entry_bc.py @@ -14,7 +14,7 @@ from ding.utils import POLICY_REGISTRY from ding.utils.data import default_collate, default_decollate from dizoo.classic_control.cartpole.config import cartpole_dqn_config, cartpole_dqn_create_config, \ - cartpole_offppo_config, cartpole_offppo_create_config + cartpole_ppo_offpolicy_config, cartpole_ppo_offpolicy_create_config from dizoo.classic_control.pendulum.config import pendulum_sac_config, pendulum_sac_create_config @@ -53,7 +53,7 @@ def _monitor_vars_learn(self) -> list: @pytest.mark.unittest def test_serial_pipeline_bc_ppo(): # train expert policy - train_config = [deepcopy(cartpole_offppo_config), deepcopy(cartpole_offppo_create_config)] + train_config = [deepcopy(cartpole_ppo_offpolicy_config), deepcopy(cartpole_ppo_offpolicy_create_config)] train_config[0].exp_name = 'test_serial_pipeline_bc_ppo' expert_policy = serial_pipeline(train_config, seed=0) @@ -61,14 +61,14 @@ def test_serial_pipeline_bc_ppo(): collect_count = 10000 expert_data_path = 'expert_data_ppo_bc.pkl' state_dict = expert_policy.collect_mode.state_dict() - collect_config = [deepcopy(cartpole_offppo_config), deepcopy(cartpole_offppo_create_config)] + collect_config = [deepcopy(cartpole_ppo_offpolicy_config), deepcopy(cartpole_ppo_offpolicy_create_config)] collect_config[0].exp_name = 'test_serial_pipeline_bc_ppo_collect' collect_demo_data( collect_config, seed=0, state_dict=state_dict, expert_data_path=expert_data_path, collect_count=collect_count ) # il training 1 - il_config = [deepcopy(cartpole_offppo_config), deepcopy(cartpole_offppo_create_config)] + il_config = [deepcopy(cartpole_ppo_offpolicy_config), deepcopy(cartpole_ppo_offpolicy_create_config)] il_config[0].policy.eval.evaluator.multi_gpu = False il_config[0].policy.learn.train_epoch = 20 il_config[1].policy.type = 'ppo_bc' diff --git a/ding/entry/tests/test_serial_entry_onpolicy.py b/ding/entry/tests/test_serial_entry_onpolicy.py index b1cb13ac44..5b43f4068d 100644 --- a/ding/entry/tests/test_serial_entry_onpolicy.py +++ b/ding/entry/tests/test_serial_entry_onpolicy.py @@ -4,6 +4,7 @@ from copy import deepcopy from ding.entry import serial_pipeline_onpolicy +from dizoo.classic_control.cartpole.config.cartpole_pg_config import cartpole_pg_config, cartpole_pg_create_config from dizoo.classic_control.cartpole.config.cartpole_ppo_config import cartpole_ppo_config, cartpole_ppo_create_config from dizoo.classic_control.cartpole.config.cartpole_ppopg_config import cartpole_ppopg_config, cartpole_ppopg_create_config # noqa from dizoo.classic_control.cartpole.config.cartpole_a2c_config import cartpole_a2c_config, cartpole_a2c_create_config @@ -12,6 +13,16 @@ from dizoo.classic_control.cartpole.config.cartpole_ppo_stdim_config import cartpole_ppo_stdim_config, cartpole_ppo_stdim_create_config # noqa +@pytest.mark.platformtest +@pytest.mark.unittest +def test_pg(): + config = [deepcopy(cartpole_pg_config), deepcopy(cartpole_pg_create_config)] + try: + serial_pipeline_onpolicy(config, seed=0, max_train_iter=1) + except Exception: + assert False, "pipeline fail" + + @pytest.mark.platformtest @pytest.mark.unittest def test_a2c(): diff --git a/ding/entry/tests/test_serial_entry_preference_based_irl.py b/ding/entry/tests/test_serial_entry_preference_based_irl.py index 9718c01f0f..7e9198f929 100644 --- a/ding/entry/tests/test_serial_entry_preference_based_irl.py +++ b/ding/entry/tests/test_serial_entry_preference_based_irl.py @@ -9,8 +9,8 @@ from ding.entry import serial_pipeline_preference_based_irl from dizoo.classic_control.cartpole.config.cartpole_trex_offppo_config import cartpole_trex_offppo_config,\ cartpole_trex_offppo_create_config -from dizoo.classic_control.cartpole.config.cartpole_offppo_config import cartpole_offppo_config,\ - cartpole_offppo_create_config +from dizoo.classic_control.cartpole.config.cartpole_ppo_offpolicy_config import cartpole_ppo_offpolicy_config,\ + cartpole_ppo_offpolicy_create_config from ding.entry.application_entry_trex_collect_data import trex_collecting_data from ding.reward_model.trex_reward_model import TrexConvEncoder from ding.torch_utils import is_differentiable @@ -19,7 +19,7 @@ @pytest.mark.unittest def test_serial_pipeline_trex(): exp_name = 'test_serial_pipeline_trex_expert' - config = [deepcopy(cartpole_offppo_config), deepcopy(cartpole_offppo_create_config)] + config = [deepcopy(cartpole_ppo_offpolicy_config), deepcopy(cartpole_ppo_offpolicy_create_config)] config[0].policy.learn.learner.hook.save_ckpt_after_iter = 100 config[0].exp_name = exp_name expert_policy = serial_pipeline(config, seed=0) @@ -27,8 +27,6 @@ def test_serial_pipeline_trex(): exp_name = 'test_serial_pipeline_trex_collect' config = [deepcopy(cartpole_trex_offppo_config), deepcopy(cartpole_trex_offppo_create_config)] config[0].exp_name = exp_name - config[0].reward_model.data_path = exp_name - config[0].reward_model.reward_model_path = exp_name + '/cartpole.params' config[0].reward_model.expert_model_path = 'test_serial_pipeline_trex_expert' config[0].reward_model.checkpoint_max = 100 config[0].reward_model.checkpoint_step = 100 diff --git a/ding/entry/tests/test_serial_entry_preference_based_irl_onpolicy.py b/ding/entry/tests/test_serial_entry_preference_based_irl_onpolicy.py index dfeee9b4e4..f50ade18ba 100644 --- a/ding/entry/tests/test_serial_entry_preference_based_irl_onpolicy.py +++ b/ding/entry/tests/test_serial_entry_preference_based_irl_onpolicy.py @@ -15,18 +15,18 @@ @pytest.mark.unittest def test_serial_pipeline_trex_onpolicy(): - exp_name = 'test_serial_pipeline_trex_onpolicy_expert' + exp_name = 'trex_onpolicy_test_serial_pipeline_trex_onpolicy_expert' config = [deepcopy(cartpole_ppo_config), deepcopy(cartpole_ppo_create_config)] + config[0].policy.learn.learner = dict() + config[0].policy.learn.learner.hook = dict() config[0].policy.learn.learner.hook.save_ckpt_after_iter = 100 config[0].exp_name = exp_name expert_policy = serial_pipeline_onpolicy(config, seed=0) - exp_name = 'test_serial_pipeline_trex_onpolicy_collect' + exp_name = 'trex_onpolicy_test_serial_pipeline_trex_onpolicy_collect' config = [deepcopy(cartpole_trex_ppo_onpolicy_config), deepcopy(cartpole_trex_ppo_onpolicy_create_config)] config[0].exp_name = exp_name - config[0].reward_model.data_path = exp_name - config[0].reward_model.reward_model_path = exp_name + '/cartpole.params' - config[0].reward_model.expert_model_path = 'test_serial_pipeline_trex_onpolicy_expert' + config[0].reward_model.expert_model_path = 'trex_onpolicy_test_serial_pipeline_trex_onpolicy_expert' config[0].reward_model.checkpoint_max = 100 config[0].reward_model.checkpoint_step = 100 config[0].reward_model.num_snippets = 100 diff --git a/ding/entry/tests/test_serial_entry_reward_model.py b/ding/entry/tests/test_serial_entry_reward_model.py index a81c5a4b6e..404cb6d78c 100644 --- a/ding/entry/tests/test_serial_entry_reward_model.py +++ b/ding/entry/tests/test_serial_entry_reward_model.py @@ -5,7 +5,7 @@ from copy import deepcopy from dizoo.classic_control.cartpole.config.cartpole_dqn_config import cartpole_dqn_config, cartpole_dqn_create_config -from dizoo.classic_control.cartpole.config.cartpole_offppo_config import cartpole_offppo_config, cartpole_offppo_create_config # noqa +from dizoo.classic_control.cartpole.config.cartpole_ppo_offpolicy_config import cartpole_ppo_offpolicy_config, cartpole_ppo_offpolicy_create_config # noqa from dizoo.classic_control.cartpole.config.cartpole_rnd_onppo_config import cartpole_ppo_rnd_config, cartpole_ppo_rnd_create_config # noqa from dizoo.classic_control.cartpole.config.cartpole_ppo_icm_config import cartpole_ppo_icm_config, cartpole_ppo_icm_create_config # noqa from ding.entry import serial_pipeline, collect_demo_data, serial_pipeline_reward_model_offpolicy, \ @@ -44,13 +44,13 @@ @pytest.mark.parametrize('reward_model_config', cfg) def test_irl(reward_model_config): reward_model_config = EasyDict(reward_model_config) - config = deepcopy(cartpole_offppo_config), deepcopy(cartpole_offppo_create_config) + config = deepcopy(cartpole_ppo_offpolicy_config), deepcopy(cartpole_ppo_offpolicy_create_config) expert_policy = serial_pipeline(config, seed=0, max_train_iter=2) # collect expert demo data collect_count = 10000 expert_data_path = 'expert_data.pkl' state_dict = expert_policy.collect_mode.state_dict() - config = deepcopy(cartpole_offppo_config), deepcopy(cartpole_offppo_create_config) + config = deepcopy(cartpole_ppo_offpolicy_config), deepcopy(cartpole_ppo_offpolicy_create_config) collect_demo_data( config, seed=0, state_dict=state_dict, expert_data_path=expert_data_path, collect_count=collect_count ) diff --git a/ding/entry/utils.py b/ding/entry/utils.py index dcc80d9833..bbfbaa83bd 100644 --- a/ding/entry/utils.py +++ b/ding/entry/utils.py @@ -59,7 +59,12 @@ def random_collect( if policy_cfg.collect.collector.type == 'episode': new_data = collector.collect(n_episode=policy_cfg.random_collect_size, policy_kwargs=collect_kwargs) else: - new_data = collector.collect(n_sample=policy_cfg.random_collect_size, policy_kwargs=collect_kwargs) + new_data = collector.collect( + n_sample=policy_cfg.random_collect_size, + random_collect=True, + record_random_collect=False, + policy_kwargs=collect_kwargs + ) # 'record_random_collect=False' means random collect without output log if postprocess_data_fn is not None: new_data = postprocess_data_fn(new_data) replay_buffer.push(new_data, cur_collector_envstep=0) diff --git a/ding/envs/__init__.py b/ding/envs/__init__.py index 5c9ecbf1b4..fc1a3fed7b 100644 --- a/ding/envs/__init__.py +++ b/ding/envs/__init__.py @@ -1,3 +1,5 @@ from .env import * from .env_wrappers import * from .env_manager import * +from .env_manager.ding_env_manager import setup_ding_env_manager +from . import gym_env diff --git a/ding/envs/common/__init__.py b/ding/envs/common/__init__.py index 2938adb909..c4b3a2e011 100644 --- a/ding/envs/common/__init__.py +++ b/ding/envs/common/__init__.py @@ -1,5 +1,5 @@ from .common_function import num_first_one_hot, sqrt_one_hot, div_one_hot, div_func, clip_one_hot, \ reorder_one_hot, reorder_one_hot_array, reorder_boolean_vector, affine_transform, \ - batch_binary_encode, get_postion_vector + batch_binary_encode, get_postion_vector, save_frames_as_gif from .env_element import EnvElement, EnvElementInfo from .env_element_runner import EnvElementRunner diff --git a/ding/envs/common/common_function.py b/ding/envs/common/common_function.py index 7b35efad5f..71db317280 100644 --- a/ding/envs/common/common_function.py +++ b/ding/envs/common/common_function.py @@ -242,6 +242,7 @@ def get_postion_vector(x: list) -> torch.Tensor: def affine_transform( data: Any, + action_clip: Optional[bool] = True, alpha: Optional[float] = None, beta: Optional[float] = None, min_val: Optional[float] = None, @@ -252,6 +253,7 @@ def affine_transform( do affine transform for data in range [-1, 1], :math:`\alpha \times data + \beta` Arguments: - data (:obj:`Any`): the input data + - action_clip (:obj:`bool`): whether to do action clip operation ([-1, 1]) - alpha (:obj:`float`): affine transform weight - beta (:obj:`float`): affine transform bias - min_val (:obj:`float`): min value, if `min_val` and `max_val` are indicated, scale input data\ @@ -260,7 +262,8 @@ def affine_transform( Returns: - transformed_data (:obj:`Any`): affine transformed data """ - data = np.clip(data, -1, 1) + if action_clip: + data = np.clip(data, -1, 1) if min_val is not None: assert max_val is not None alpha = (max_val - min_val) / 2 @@ -268,3 +271,21 @@ def affine_transform( assert alpha is not None beta = beta if beta is not None else 0. return data * alpha + beta + + +def save_frames_as_gif(frames: list, path: str) -> None: + """ + Overview: + save frames as gif to a specified path. + Arguments: + - frames (:obj:`List`): list of frames + - path (:obj:`str`): the path to save gif + """ + try: + import imageio + except ImportError: + from ditk import logging + import sys + logging.warning("Please install imageio first.") + sys.exit(1) + imageio.mimsave(path, frames, fps=20) diff --git a/ding/envs/common/tests/test_common_function.py b/ding/envs/common/tests/test_common_function.py index ca5b421de5..11712331bb 100644 --- a/ding/envs/common/tests/test_common_function.py +++ b/ding/envs/common/tests/test_common_function.py @@ -1,13 +1,14 @@ +import os import random +import shutil import numpy as np import pytest import torch - -from ding.envs.common.common_function import num_first_one_hot, sqrt_one_hot, div_one_hot, div_func, clip_one_hot, \ +from ding.envs.common.common_function import sqrt_one_hot, div_one_hot, div_func, clip_one_hot, \ reorder_one_hot, reorder_one_hot_array, reorder_boolean_vector, \ - get_to_and, batch_binary_encode, compute_denominator, get_postion_vector, \ - affine_transform + batch_binary_encode, get_postion_vector, \ + affine_transform, save_frames_as_gif VALUES = [2, 3, 5, 7, 11] @@ -113,3 +114,16 @@ def test_affine_transform(self): ans = affine_transform(a, alpha=4, beta=1) assert ans.shape == (3, 5) assert ans.min() == -3 and ans.max() == 5 + + +@pytest.mark.other +def test_save_frames_as_gif(): + frames = [np.random.randint(0, 255, [84, 84, 3]) for _ in range(100)] + replay_path_gif = './replay_path_gif' + env_id = 'test' + save_replay_count = 1 + if not os.path.exists(replay_path_gif): + os.makedirs(replay_path_gif) + path = os.path.join(replay_path_gif, '{}_episode_{}.gif'.format(env_id, save_replay_count)) + save_frames_as_gif(frames, path) + shutil.rmtree(replay_path_gif) diff --git a/ding/envs/env/default_wrapper.py b/ding/envs/env/default_wrapper.py index f44649583a..d0e1401b46 100644 --- a/ding/envs/env/default_wrapper.py +++ b/ding/envs/env/default_wrapper.py @@ -2,14 +2,28 @@ from typing import Optional, List import copy -final_eval_reward_wrapper = EasyDict(type='final_eval_reward') +eval_episode_return_wrapper = EasyDict(type='eval_episode_return') -def get_default_wrappers(env_wrapper_name: str, env_id: Optional[str] = None) -> List[dict]: +def get_default_wrappers(env_wrapper_name: str, env_id: Optional[str] = None, caller: str = 'collector') -> List[dict]: + """ + Overview: + Get default wrappers for different environments used in ``DingEnvWrapper``. + Arguments: + - env_wrapper_name (:obj:`str`): The name of the environment wrapper. + - env_id (:obj:`Optional[str]`): The id of the specific environment, such as ``PongNoFrameskip-v4``. + - caller (:obj:`str`): The caller of the environment, including ``collector`` or ``evaluator``. Different \ + caller may need different wrappers. + Returns: + - wrapper_list (:obj:`List[dict]`): The list of wrappers, each element is a config of the concrete wrapper. + Raises: + - NotImplementedError: ``env_wrapper_name`` is not in ``['mujoco_default', 'atari_default', \ + 'gym_hybrid_default', 'default']`` + """ + assert caller == 'collector' or 'evaluator', caller if env_wrapper_name == 'mujoco_default': return [ - EasyDict(type='delay_reward', kwargs=dict(delay_reward_step=3)), - copy.deepcopy(final_eval_reward_wrapper), + copy.deepcopy(eval_episode_return_wrapper), ] elif env_wrapper_name == 'atari_default': wrapper_list = [] @@ -21,16 +35,17 @@ def get_default_wrappers(env_wrapper_name: str, env_id: Optional[str] = None) -> wrapper_list.append(EasyDict(type='fire_reset')) wrapper_list.append(EasyDict(type='warp_frame')) wrapper_list.append(EasyDict(type='scaled_float_frame')) - wrapper_list.append(EasyDict(type='clip_reward')) + if caller == 'collector': + wrapper_list.append(EasyDict(type='clip_reward')) wrapper_list.append(EasyDict(type='frame_stack', kwargs=dict(n_frames=4))) - wrapper_list.append(copy.deepcopy(final_eval_reward_wrapper)) + wrapper_list.append(copy.deepcopy(eval_episode_return_wrapper)) return wrapper_list elif env_wrapper_name == 'gym_hybrid_default': return [ EasyDict(type='gym_hybrid_dict_action'), - copy.deepcopy(final_eval_reward_wrapper), + copy.deepcopy(eval_episode_return_wrapper), ] elif env_wrapper_name == 'default': - return [copy.deepcopy(final_eval_reward_wrapper)] + return [copy.deepcopy(eval_episode_return_wrapper)] else: - raise NotImplementedError() + raise NotImplementedError("not supported env_wrapper_name: {}".format(env_wrapper_name)) diff --git a/ding/envs/env/ding_env_wrapper.py b/ding/envs/env/ding_env_wrapper.py index 8a6db4f0a9..7bd440fd70 100644 --- a/ding/envs/env/ding_env_wrapper.py +++ b/ding/envs/env/ding_env_wrapper.py @@ -1,39 +1,85 @@ from typing import List, Optional, Union, Dict +from easydict import EasyDict import gym +import gymnasium import copy import numpy as np +import treetensor.numpy as tnp from ding.envs.common.common_function import affine_transform -from ding.envs.env_wrappers import create_env_wrapper +from ding.envs.env_wrappers import create_env_wrapper, GymToGymnasiumWrapper from ding.torch_utils import to_ndarray +from ding.utils import CloudPickleWrapper from .base_env import BaseEnv, BaseEnvTimestep from .default_wrapper import get_default_wrappers class DingEnvWrapper(BaseEnv): + """ + Overview: + This is a wrapper for the BaseEnv class, used to provide a consistent environment interface. + Interfaces: + __init__, reset, step, close, seed, random_action, _wrap_env, __repr__, create_collector_env_cfg, + create_evaluator_env_cfg, enable_save_replay, observation_space, action_space, reward_space, clone + """ - def __init__(self, env: gym.Env = None, cfg: dict = None) -> None: - ''' - You can pass in either an env instance, or a config to create an env instance: - - An env instance: Parameter `env` must not be `None`, but should be the instance. - Do not support subprocess env manager; Thus usually used in simple env. - - A config to create an env instance: Parameter `cfg` dict must contain `env_id`. - ''' + def __init__( + self, + env: Union[gym.Env, gymnasium.Env] = None, + cfg: dict = None, + seed_api: bool = True, + caller: str = 'collector', + is_gymnasium: bool = False + ) -> None: + """ + Overview: + Initialize the DingEnvWrapper. Either an environment instance or a config to create the environment \ + instance should be passed in. For the former, i.e., an environment instance: The `env` parameter must not \ + be `None`, but should be the instance. It does not support subprocess environment manager. Thus, it is \ + usually used in simple environments. For the latter, i.e., a config to create an environment instance: \ + The `cfg` parameter must contain `env_id`. + Arguments: + - env (:obj:`Union[gym.Env, gymnasium.Env]`): An environment instance to be wrapped. + - cfg (:obj:`dict`): The configuration dictionary to create an environment instance. + - seed_api (:obj:`bool`): Whether to use seed API. Defaults to True. + - caller (:obj:`str`): A string representing the caller of this method, including ``collector`` or \ + ``evaluator``. Different caller may need different wrappers. Default is 'collector'. + - is_gymnasium (:obj:`bool`): Whether the environment is a gymnasium environment. Defaults to False, i.e., \ + the environment is a gym environment. + """ + self._env = None + self._raw_env = env self._cfg = cfg + self._seed_api = seed_api # some env may disable `env.seed` api + self._caller = caller + if self._cfg is None: - self._cfg = dict() + self._cfg = {} + self._cfg = EasyDict(self._cfg) + if 'act_scale' not in self._cfg: + self._cfg.act_scale = False + if 'rew_clip' not in self._cfg: + self._cfg.rew_clip = False + if 'env_wrapper' not in self._cfg: + self._cfg.env_wrapper = 'default' + if 'env_id' not in self._cfg: + self._cfg.env_id = None if env is not None: - self._init_flag = True + self._is_gymnasium = isinstance(env, gymnasium.Env) self._env = env - self._wrap_env() + self._wrap_env(caller) self._observation_space = self._env.observation_space self._action_space = self._env.action_space self._action_space.seed(0) # default seed - self._reward_space = gym.spaces.Box( - low=self._env.reward_range[0], high=self._env.reward_range[1], shape=(1, ), dtype=np.float32 - ) + try: + low, high = self._env.reward_range + except: # for compatibility with gymnasium high-version API + low, high = -1, 1 + self._reward_space = gym.spaces.Box(low=low, high=high, shape=(1, ), dtype=np.float32) + self._init_flag = True else: assert 'env_id' in self._cfg + self._is_gymnasium = is_gymnasium self._init_flag = False self._observation_space = None self._action_space = None @@ -42,10 +88,17 @@ def __init__(self, env: gym.Env = None, cfg: dict = None) -> None: self._replay_path = None # override - def reset(self) -> None: + def reset(self) -> np.ndarray: + """ + Overview: + Resets the state of the environment. If the environment is not initialized, it will be created first. + Returns: + - obs (:obj:`Dict`): The new observation after reset. + """ if not self._init_flag: - self._env = gym.make(self._cfg.env_id) - self._wrap_env() + gym_proxy = gymnasium if self._is_gymnasium else gym + self._env = gym_proxy.make(self._cfg.env_id) + self._wrap_env(self._caller) self._observation_space = self._env.observation_space self._action_space = self._env.action_space self._reward_space = gym.spaces.Box( @@ -62,46 +115,97 @@ def reset(self) -> None: self._replay_path = None if hasattr(self, '_seed') and hasattr(self, '_dynamic_seed') and self._dynamic_seed: np_seed = 100 * np.random.randint(1, 1000) - self._env.seed(self._seed + np_seed) + if self._seed_api: + self._env.seed(self._seed + np_seed) self._action_space.seed(self._seed + np_seed) elif hasattr(self, '_seed'): - self._env.seed(self._seed) + if self._seed_api: + self._env.seed(self._seed) self._action_space.seed(self._seed) obs = self._env.reset() - obs = to_ndarray(obs, np.float32) + if self.observation_space.dtype == np.float32: + obs = to_ndarray(obs, dtype=np.float32) + else: + obs = to_ndarray(obs) return obs # override def close(self) -> None: - self._env.close() + """ + Overview: + Clean up the environment by closing and deleting it. + This method should be called when the environment is no longer needed. + Failing to call this method can lead to memory leaks. + """ + try: + self._env.close() + del self._env + except: # noqa + pass # override def seed(self, seed: int, dynamic_seed: bool = True) -> None: + """ + Overview: + Set the seed for the environment. + Arguments: + - seed (:obj:`int`): The seed to set. + - dynamic_seed (:obj:`bool`): Whether to use dynamic seed, default is True. + """ self._seed = seed self._dynamic_seed = dynamic_seed np.random.seed(self._seed) # override def step(self, action: Union[np.int64, np.ndarray]) -> BaseEnvTimestep: + """ + Overview: + Execute the given action in the environment, and return the timestep (observation, reward, done, info). + Arguments: + - action (:obj:`Union[np.int64, np.ndarray]`): The action to execute in the environment. + Returns: + - timestep (:obj:`BaseEnvTimestep`): The timestep after the action execution. + """ action = self._judge_action_type(action) - if self._cfg.get('act_scale', False): + if self._cfg.act_scale: action = affine_transform(action, min_val=self._env.action_space.low, max_val=self._env.action_space.high) obs, rew, done, info = self._env.step(action) - obs = to_ndarray(obs, np.float32) + if self._cfg.rew_clip: + rew = max(-10, rew) + rew = np.float32(rew) + if self.observation_space.dtype == np.float32: + obs = to_ndarray(obs, dtype=np.float32) + else: + obs = to_ndarray(obs) rew = to_ndarray([rew], np.float32) return BaseEnvTimestep(obs, rew, done, info) def _judge_action_type(self, action: Union[np.ndarray, dict]) -> Union[np.ndarray, dict]: + """ + Overview: + Ensure the action taken by the agent is of the correct type. + This method is used to standardize different action types to a common format. + Arguments: + - action (Union[np.ndarray, dict]): The action taken by the agent. + Returns: + - action (Union[np.ndarray, dict]): The formatted action. + """ if isinstance(action, int): return action - if isinstance(action, np.ndarray): - if action.shape == (1, ) and action.dtype == np.int64: + elif isinstance(action, np.int64): + return int(action) + elif isinstance(action, np.ndarray): + if action.shape == (): + action = action.item() + elif action.shape == (1, ) and action.dtype == np.int64: action = action.item() return action elif isinstance(action, dict): for k, v in action.items(): action[k] = self._judge_action_type(v) return action + elif isinstance(action, tnp.ndarray): + return self._judge_action_type(action.json()) else: raise TypeError( '`action` should be either int/np.ndarray or dict of int/np.ndarray, but get {}: {}'.format( @@ -110,10 +214,16 @@ def _judge_action_type(self, action: Union[np.ndarray, dict]) -> Union[np.ndarra ) def random_action(self) -> np.ndarray: + """ + Overview: + Return a random action from the action space of the environment. + Returns: + - action (:obj:`np.ndarray`): The random action. + """ random_action = self.action_space.sample() if isinstance(random_action, np.ndarray): pass - elif isinstance(random_action, int): + elif isinstance(random_action, (int, np.int64)): random_action = to_ndarray([random_action], dtype=np.int64) elif isinstance(random_action, dict): random_action = to_ndarray(random_action) @@ -125,11 +235,20 @@ def random_action(self) -> np.ndarray: ) return random_action - def _wrap_env(self) -> None: + def _wrap_env(self, caller: str = 'collector') -> None: + """ + Overview: + Wrap the environment according to the configuration. + Arguments: + - caller (:obj:`str`): The caller of the environment, including ``collector`` or ``evaluator``. \ + Different caller may need different wrappers. Default is 'collector'. + """ + if self._is_gymnasium: + self._env = GymToGymnasiumWrapper(self._env) # wrapper_cfgs: Union[str, List] - wrapper_cfgs = self._cfg.get('env_wrapper', 'default') + wrapper_cfgs = self._cfg.env_wrapper if isinstance(wrapper_cfgs, str): - wrapper_cfgs = get_default_wrappers(wrapper_cfgs, self._cfg.get('env_id', None)) + wrapper_cfgs = get_default_wrappers(wrapper_cfgs, self._cfg.env_id, caller) # self._wrapper_cfgs: List[Union[Callable, Dict]] self._wrapper_cfgs = wrapper_cfgs for wrapper in self._wrapper_cfgs: @@ -140,10 +259,24 @@ def _wrap_env(self) -> None: self._env = wrapper(self._env) def __repr__(self) -> str: + """ + Overview: + Return the string representation of the instance. + Returns: + - str (:obj:`str`): The string representation of the instance. + """ return "DI-engine Env({}), generated by DingEnvWrapper".format(self._cfg.env_id) @staticmethod def create_collector_env_cfg(cfg: dict) -> List[dict]: + """ + Overview: + Create a list of environment configuration for collectors based on the input configuration. + Arguments: + - cfg (:obj:`dict`): The input configuration dictionary. + Returns: + - env_cfgs (:obj:`List[dict]`): The list of environment configurations for collectors. + """ actor_env_num = cfg.pop('collector_env_num') cfg = copy.deepcopy(cfg) cfg.is_train = True @@ -151,24 +284,86 @@ def create_collector_env_cfg(cfg: dict) -> List[dict]: @staticmethod def create_evaluator_env_cfg(cfg: dict) -> List[dict]: + """ + Overview: + Create a list of environment configuration for evaluators based on the input configuration. + Arguments: + - cfg (:obj:`dict`): The input configuration dictionary. + Returns: + - env_cfgs (:obj:`List[dict]`): The list of environment configurations for evaluators. + """ evaluator_env_num = cfg.pop('evaluator_env_num') cfg = copy.deepcopy(cfg) cfg.is_train = False return [cfg for _ in range(evaluator_env_num)] def enable_save_replay(self, replay_path: Optional[str] = None) -> None: + """ + Overview: + Enable the save replay functionality. The replay will be saved at the specified path. + Arguments: + - replay_path (:obj:`Optional[str]`): The path to save the replay, default is None. + """ if replay_path is None: replay_path = './video' self._replay_path = replay_path @property def observation_space(self) -> gym.spaces.Space: + """ + Overview: + Return the observation space of the wrapped environment. + The observation space represents the range and shape of possible observations + that the environment can provide to the agent. + Note: + If the data type of the observation space is float64, it's converted to float32 + for better compatibility with most machine learning libraries. + Returns: + - observation_space (gym.spaces.Space): The observation space of the environment. + """ + if self._observation_space.dtype == np.float64: + self._observation_space.dtype = np.float32 return self._observation_space @property def action_space(self) -> gym.spaces.Space: + """ + Overview: + Return the action space of the wrapped environment. + The action space represents the range and shape of possible actions + that the agent can take in the environment. + Returns: + - action_space (gym.spaces.Space): The action space of the environment. + """ return self._action_space @property def reward_space(self) -> gym.spaces.Space: + """ + Overview: + Return the reward space of the wrapped environment. + The reward space represents the range and shape of possible rewards + that the agent can receive as a result of its actions. + Returns: + - reward_space (gym.spaces.Space): The reward space of the environment. + """ return self._reward_space + + def clone(self, caller: str = 'collector') -> BaseEnv: + """ + Overview: + Clone the current environment wrapper, creating a new environment with the same settings. + Arguments: + - caller (str): A string representing the caller of this method, including ``collector`` or ``evaluator``. \ + Different caller may need different wrappers. Default is 'collector'. + Returns: + - DingEnvWrapper: A new instance of the environment with the same settings. + """ + try: + spec = copy.deepcopy(self._raw_env.spec) + raw_env = CloudPickleWrapper(self._raw_env) + raw_env = copy.deepcopy(raw_env).data + raw_env.__setattr__('spec', spec) + except Exception: + raw_env = self._raw_env + return DingEnvWrapper(raw_env, self._cfg, self._seed_api, caller, self._is_gymnasium) diff --git a/ding/envs/env/env_implementation_check.py b/ding/envs/env/env_implementation_check.py index f972446d74..4bf8c2976c 100644 --- a/ding/envs/env/env_implementation_check.py +++ b/ding/envs/env/env_implementation_check.py @@ -1,15 +1,15 @@ -from tabnanny import check -from typing import Any, Callable, List, Tuple +from typing import Any, Callable, List, Tuple, Union, Dict, TYPE_CHECKING import numpy as np from collections.abc import Sequence -from easydict import EasyDict - -from ding.envs.env import BaseEnv, BaseEnvTimestep +from gym.spaces import Space, Box, Discrete, MultiDiscrete, MultiBinary from ding.envs.env.tests import DemoEnv # from dizoo.atari.envs import AtariEnv +if TYPE_CHECKING: + from ding.envs.env import BaseEnv + -def check_space_dtype(env: BaseEnv) -> None: +def check_space_dtype(env: 'BaseEnv') -> None: print("== 0. Test obs/act/rew space's dtype") env.reset() for name, space in zip(['obs', 'act', 'rew'], [env.observation_space, env.action_space, env.reward_space]): @@ -24,44 +24,47 @@ def check_space_dtype(env: BaseEnv) -> None: # Util function -def check_array_space(ndarray, space, name) -> bool: - if isinstance(ndarray, np.ndarray): +def check_array_space(data: Union[np.ndarray, Sequence, Dict], space: Union['Space', Dict], name: str) -> None: + if isinstance(data, np.ndarray): # print("{}'s type should be np.ndarray".format(name)) - assert ndarray.dtype == space.dtype, "{}'s dtype is {}, but requires {}".format( - name, ndarray.dtype, space.dtype - ) - assert ndarray.shape == space.shape, "{}'s shape is {}, but requires {}".format( - name, ndarray.shape, space.shape - ) - assert (space.low <= ndarray).all() and (ndarray <= space.high).all( - ), "{}'s value is {}, but requires in range ({},{})".format(name, ndarray, space.low, space.high) - elif isinstance(ndarray, Sequence): - for i in range(len(ndarray)): + assert data.dtype == space.dtype, "{}'s dtype is {}, but requires {}".format(name, data.dtype, space.dtype) + assert data.shape == space.shape, "{}'s shape is {}, but requires {}".format(name, data.shape, space.shape) + if isinstance(space, Box): + assert (space.low <= data).all() and (data <= space.high).all( + ), "{}'s value is {}, but requires in range ({},{})".format(name, data, space.low, space.high) + elif isinstance(space, (Discrete, MultiDiscrete, MultiBinary)): + if isinstance(space, Discrete): + assert (data >= space.start) and (data <= space.n) + else: + assert (data >= 0).all() + assert all([d < n for d, n in zip(data, space.nvec)]) + elif isinstance(data, Sequence): + for i in range(len(data)): try: - check_array_space(ndarray[i], space[i], name) + check_array_space(data[i], space[i], name) except AssertionError as e: - print("The following error happens at {}-th index".format(i)) + print("The following error happens at {}-th index".format(i)) raise e - elif isinstance(ndarray, dict): - for k in ndarray.keys(): + elif isinstance(data, dict): + for k in data.keys(): try: - check_array_space(ndarray[k], space[k], name) + check_array_space(data[k], space[k], name) except AssertionError as e: print("The following error happens at key {}".format(k)) raise e else: raise TypeError( - "Input array should be np.ndarray or sequence/dict of np.ndarray, but found {}".format(type(ndarray)) + "Input array should be np.ndarray or sequence/dict of np.ndarray, but found {}".format(type(data)) ) -def check_reset(env: BaseEnv) -> None: +def check_reset(env: 'BaseEnv') -> None: print('== 1. Test reset method') obs = env.reset() check_array_space(obs, env.observation_space, 'obs') -def check_step(env: BaseEnv) -> None: +def check_step(env: 'BaseEnv') -> None: done_times = 0 print('== 2. Test step method') _ = env.reset() @@ -74,7 +77,7 @@ def check_step(env: BaseEnv) -> None: for ndarray, space, name in zip([obs, rew], [env.observation_space, env.reward_space], ['obs', 'rew']): check_array_space(ndarray, space, name) if done: - assert 'final_eval_reward' in info, "info dict should have 'final_eval_reward' key." + assert 'eval_episode_return' in info, "info dict should have 'eval_episode_return' key." done_times += 1 _ = env.reset() if done_times == 3: @@ -82,7 +85,9 @@ def check_step(env: BaseEnv) -> None: # Util function -def check_different_memory(array1, array2, step_times) -> None: +def check_different_memory( + array1: Union[np.ndarray, Sequence, Dict], array2: Union[np.ndarray, Sequence, Dict], step_times: int +) -> None: assert type(array1) == type( array2 ), "In step times {}, obs_last_frame({}) and obs_this_frame({}) are not of the same type".format( @@ -121,7 +126,7 @@ def check_different_memory(array1, array2, step_times) -> None: ) -def check_obs_deepcopy(env: BaseEnv) -> None: +def check_obs_deepcopy(env: 'BaseEnv') -> None: step_times = 0 print('== 3. Test observation deepcopy') @@ -139,14 +144,14 @@ def check_obs_deepcopy(env: BaseEnv) -> None: break -def check_all(env: BaseEnv) -> None: +def check_all(env: 'BaseEnv') -> None: check_space_dtype(env) check_reset(env) check_step(env) check_obs_deepcopy(env) -def demonstrate_correct_procedure(env_fn: Callable) -> None: +def demonstrate_correct_procedure(env_fn: Callable[[Dict], 'BaseEnv']) -> None: print('== 4. Demonstrate the correct procudures') done_times = 0 # Init the env. @@ -163,7 +168,7 @@ def demonstrate_correct_procedure(env_fn: Callable) -> None: action = env.random_action() obs, rew, done, info = env.step(action) if done: - assert 'final_eval_reward' in info + assert 'eval_episode_return' in info done_times += 1 obs = env.reset() # Seed will not change unless `seed` method is called again. @@ -174,7 +179,7 @@ def demonstrate_correct_procedure(env_fn: Callable) -> None: if __name__ == "__main__": ''' - # Moethods `check_*` are for user to check whether their implemented env obeys DI-engine's rules. + # Methods `check_*` are for user to check whether their implemented env obeys DI-engine's rules. # You can replace `AtariEnv` with your own env. atari_env = AtariEnv(EasyDict(env_id='PongNoFrameskip-v4', frame_stack=4, is_train=False)) check_reset(atari_env) diff --git a/ding/envs/env/tests/demo_env.py b/ding/envs/env/tests/demo_env.py index 93d8b65507..4867667f10 100644 --- a/ding/envs/env/tests/demo_env.py +++ b/ding/envs/env/tests/demo_env.py @@ -59,7 +59,7 @@ def step(self, action: Any) -> 'BaseEnv.timestep': done = False info = {} if done: - info['final_eval_reward'] = self.reward_space.sample() * 30 + info['eval_episode_return'] = self.reward_space.sample() * 30 return BaseEnvTimestep(obs, rew, done, info) def seed(self, seed: int) -> None: diff --git a/ding/envs/env/tests/test_ding_env_wrapper.py b/ding/envs/env/tests/test_ding_env_wrapper.py index ec0ca04ffc..0ac2a44d04 100644 --- a/ding/envs/env/tests/test_ding_env_wrapper.py +++ b/ding/envs/env/tests/test_ding_env_wrapper.py @@ -1,4 +1,5 @@ import gym +import gymnasium import numpy as np import pytest from easydict import EasyDict @@ -68,6 +69,26 @@ def test_cartpole_pendulum(self, env_id): # assert isinstance(action, np.ndarray) print('random_action: {}, action_space: {}'.format(action.shape, ding_env.action_space)) + @pytest.mark.unittest + @pytest.mark.parametrize('env_id', ['CartPole-v0', 'Pendulum-v1']) + def test_cartpole_pendulum_gymnasium(self, env_id): + env = gymnasium.make(env_id) + ding_env = DingEnvWrapper(env=env) + cfg = EasyDict(dict( + collector_env_num=16, + evaluator_env_num=3, + is_train=True, + )) + l1 = ding_env.create_collector_env_cfg(cfg) + assert isinstance(l1, list) + l1 = ding_env.create_evaluator_env_cfg(cfg) + assert isinstance(l1, list) + obs = ding_env.reset() + assert isinstance(obs, np.ndarray) + action = ding_env.random_action() + # assert isinstance(action, np.ndarray) + print('random_action: {}, action_space: {}'.format(action.shape, ding_env.action_space)) + @pytest.mark.envtest def test_mujoco(self): env_cfg = EasyDict( @@ -85,7 +106,7 @@ def test_mujoco(self): # print(_, timestep.reward) assert timestep.reward.shape == (1, ), timestep.reward.shape if timestep.done: - assert 'final_eval_reward' in timestep.info, timestep.info + assert 'eval_episode_return' in timestep.info, timestep.info break print(ding_env_mujoco.observation_space, ding_env_mujoco.action_space, ding_env_mujoco.reward_space) action = ding_env_mujoco.random_action() @@ -113,14 +134,14 @@ def test_atari(self, atari_env_id): # print(timestep.reward) assert timestep.reward.shape == ding_env_atari.reward_space.shape, timestep.reward.shape # (1, ) if timestep.done: - assert 'final_eval_reward' in timestep.info, timestep.info + assert 'eval_episode_return' in timestep.info, timestep.info break print(ding_env_atari.observation_space, ding_env_atari.action_space, ding_env_atari.reward_space) action = ding_env_atari.random_action() # assert isinstance(action, np.ndarray) assert action.shape == (1, ) - @pytest.mark.unittest + @pytest.mark.envtest @pytest.mark.parametrize('lun_bip_env_id', ['LunarLander-v2', 'LunarLanderContinuous-v2', 'BipedalWalker-v3']) def test_lunarlander_bipedalwalker(self, lun_bip_env_id): env_cfg = EasyDict( @@ -147,7 +168,7 @@ def test_lunarlander_bipedalwalker(self, lun_bip_env_id): # print(timestep.reward) assert timestep.reward.shape == ding_env_lun_bip.reward_space.shape, timestep.reward.shape # (1, ) if timestep.done: - assert 'final_eval_reward' in timestep.info, timestep.info + assert 'eval_episode_return' in timestep.info, timestep.info break print(ding_env_lun_bip.observation_space, ding_env_lun_bip.action_space, ding_env_lun_bip.reward_space) action = ding_env_lun_bip.random_action() @@ -174,9 +195,28 @@ def test_hybrid(self): # print(timestep.reward) assert timestep.reward.shape == ding_env_hybrid.reward_space.shape, timestep.reward.shape # (1, ) if timestep.done: - assert 'final_eval_reward' in timestep.info, timestep.info + assert 'eval_episode_return' in timestep.info, timestep.info break print(ding_env_hybrid.observation_space, ding_env_hybrid.action_space, ding_env_hybrid.reward_space) action = ding_env_hybrid.random_action() print('random_action', action) assert isinstance(action, dict) + + @pytest.mark.envtest + def test_AllinObsWrapper(self): + env_cfg = EasyDict(env_id='PongNoFrameskip-v4', env_wrapper='reward_in_obs') + ding_env_aio = DingEnvWrapper(cfg=env_cfg) + + data = ding_env_aio.reset() + assert isinstance(data, dict) + assert 'obs' in data.keys() and 'reward' in data.keys() + assert data['obs'].shape == ding_env_aio.observation_space + while True: + action = ding_env_aio.random_action() + timestep = ding_env_aio.step(action) + # print(timestep.reward) + assert isinstance(timestep.obs, dict) + if timestep.done: + assert 'eval_episode_return' in timestep.info, timestep.info + break + print(ding_env_aio.observation_space, ding_env_aio.action_space, ding_env_aio.reward_space) diff --git a/ding/envs/env/tests/test_env_implementation_check.py b/ding/envs/env/tests/test_env_implementation_check.py index fb413304ce..41e66b62b2 100644 --- a/ding/envs/env/tests/test_env_implementation_check.py +++ b/ding/envs/env/tests/test_env_implementation_check.py @@ -1,5 +1,4 @@ import pytest -from easydict import EasyDict import numpy as np import gym from copy import deepcopy @@ -17,6 +16,17 @@ def test_an_implemented_env(): @pytest.mark.unittest def test_check_array_space(): + discrete_space = gym.spaces.Discrete(10) + discrete_array = np.array(2, dtype=np.int64) + check_array_space(discrete_array, discrete_space, 'test_discrete') + discrete_array = np.array(11, dtype=np.int64) + with pytest.raises(AssertionError): + check_array_space(discrete_array, discrete_space, 'test_discrete') + + multi_discrete_space = gym.spaces.MultiDiscrete([2, 3]) + multi_discrete_array = np.array([1, 2], dtype=np.int64) + check_array_space(multi_discrete_array, multi_discrete_space, 'test_multi_discrete') + seq_array = (np.array([1, 2, 3], dtype=np.int64), np.array([4., 5., 6.], dtype=np.float32)) seq_space = [gym.spaces.Box(low=0, high=10, shape=(3, ), dtype=np.int64) for _ in range(2)] with pytest.raises(AssertionError): diff --git a/ding/envs/env_manager/__init__.py b/ding/envs/env_manager/__init__.py index e5947902f9..62d45baf27 100644 --- a/ding/envs/env_manager/__init__.py +++ b/ding/envs/env_manager/__init__.py @@ -1,5 +1,5 @@ from .base_env_manager import BaseEnvManager, BaseEnvManagerV2, create_env_manager, get_env_manager_cls from .subprocess_env_manager import AsyncSubprocessEnvManager, SyncSubprocessEnvManager, SubprocessEnvManagerV2 from .gym_vector_env_manager import GymVectorEnvManager -# Do not import PoolEnvManager, because it depends on installation of `envpool` +# Do not import PoolEnvManager here, because it depends on installation of `envpool` from .env_supervisor import EnvSupervisor diff --git a/ding/envs/env_manager/base_env_manager.py b/ding/envs/env_manager/base_env_manager.py index e00fc7649c..6ef4bad28c 100644 --- a/ding/envs/env_manager/base_env_manager.py +++ b/ding/envs/env_manager/base_env_manager.py @@ -11,7 +11,7 @@ import treetensor.numpy as tnp from ding.utils import ENV_MANAGER_REGISTRY, import_module, one_time_warning, make_key_as_identifier, WatchDog, \ remove_illegal_item -from ding.envs.env import BaseEnvTimestep +from ding.envs import BaseEnv, BaseEnvTimestep global space_log_flag space_log_flag = True @@ -64,27 +64,45 @@ def wrapper(*args, **kwargs): class BaseEnvManager(object): """ Overview: - Create a BaseEnvManager to manage multiple environments. + The basic class of env manager to manage multiple vectorized environments. BaseEnvManager define all the + necessary interfaces and derived class must extend this basic class. + + The class is implemented by the pseudo-parallelism (i.e. serial) mechanism, therefore, this class is only + used in some tiny environments and for debug purpose. Interfaces: - reset, step, seed, close, enable_save_replay, launch, default_config, env_state_done + reset, step, seed, close, enable_save_replay, launch, default_config, reward_shaping, enable_save_figure Properties: - env_num, ready_obs, done, method_name_list - observation_space, action_space, reward_space + env_num, env_ref, ready_obs, ready_obs_id, ready_imgs, done, closed, method_name_list, observation_space, \ + action_space, reward_space """ @classmethod def default_config(cls: type) -> EasyDict: + """ + Overview: + Return the deepcopyed default config of env manager. + Returns: + - cfg (:obj:`EasyDict`): The default config of env manager. + """ cfg = EasyDict(copy.deepcopy(cls.config)) cfg.cfg_type = cls.__name__ + 'Dict' return cfg config = dict( + # (int) The total episode number to be executed, defaults to inf, which means no episode limits. episode_num=float("inf"), + # (int) The maximum retry times when the env is in error state, defaults to 1, i.e. no retry. max_retry=1, + # (str) The retry type when the env is in error state, including ['reset', 'renew'], defaults to 'reset'. + # The former is to reset the env to the last reset state, while the latter is to create a new env. retry_type='reset', + # (bool) Whether to automatically reset sub-environments when they are done, defaults to True. auto_reset=True, + # (float) WatchDog timeout (second) for ``step`` method, defaults to None, which means no timeout. step_timeout=None, + # (float) WatchDog timeout (second) for ``reset`` method, defaults to None, which means no timeout. reset_timeout=None, + # (float) The interval waiting time for automatically retry mechanism, defaults to 0.1. retry_waiting_time=0.1, ) @@ -95,10 +113,18 @@ def __init__( ) -> None: """ Overview: - Initialize the BaseEnvManager. + Initialize the base env manager with callable the env function and the EasyDict-type config. Here we use + ``env_fn`` to ensure the lazy initialization of sub-environments, which is benetificial to resource + allocation and parallelism. ``cfg`` is the merged result between the default config of this class + and user's config. + This construction function is in lazy-initialization mode, the actual initialization is in ``launch``. Arguments: - - env_fn (:obj:`List[Callable]`): The function to create environment - - cfg (:obj:`EasyDict`): Config + - env_fn (:obj:`List[Callable]`): A list of functions to create ``env_num`` sub-environments. + - cfg (:obj:`EasyDict`): Final merged config. + + .. note:: + For more details about how to merge config, please refer to the system document of DI-engine \ + (`en link1 <../03_system/config.html>`_). """ self._cfg = cfg self._env_fn = env_fn @@ -122,16 +148,6 @@ def __init__( self._action_space = self._env_ref.action_space self._reward_space = self._env_ref.reward_space self._env_ref.close() - try: - global space_log_flag - if space_log_flag: - logging.info("Env Space Information:") - logging.info("\tObservation Space: {}".format(self._observation_space)) - logging.info("\tAction Space: {}".format(self._action_space)) - logging.info("\tReward Space: {}".format(self._reward_space)) - space_log_flag = False - except: - pass self._env_states = {i: EnvState.VOID for i in range(self._env_num)} self._env_seed = {i: None for i in range(self._env_num)} self._episode_num = self._cfg.episode_num @@ -145,35 +161,67 @@ def __init__( @property def env_num(self) -> int: + """ + Overview: + ``env_num`` is the number of sub-environments in env manager. + Returns: + - env_num (:obj:`int`): The number of sub-environments. + """ return self._env_num @property - def env_ref(self) -> int: + def env_ref(self) -> 'BaseEnv': + """ + Overview: + ``env_ref`` is used to acquire some common attributes of env, like obs_shape and act_shape. + Returns: + - env_ref (:obj:`BaseEnv`): The reference of sub-environment. + """ return self._env_ref @property def observation_space(self) -> 'gym.spaces.Space': # noqa + """ + Overview: + ``observation_space`` is the observation space of sub-environment, following the format of gym.spaces. + Returns: + - observation_space (:obj:`gym.spaces.Space`): The observation space of sub-environment. + """ return self._observation_space @property def action_space(self) -> 'gym.spaces.Space': # noqa + """ + Overview: + ``action_space`` is the action space of sub-environment, following the format of gym.spaces. + Returns: + - action_space (:obj:`gym.spaces.Space`): The action space of sub-environment. + """ return self._action_space @property def reward_space(self) -> 'gym.spaces.Space': # noqa + """ + Overview: + ``reward_space`` is the reward space of sub-environment, following the format of gym.spaces. + Returns: + - reward_space (:obj:`gym.spaces.Space`): The reward space of sub-environment. + """ return self._reward_space @property def ready_obs(self) -> Dict[int, Any]: """ Overview: - Get the ready (next) observation, which is uniform for both aysnc/sync scenarios. - Return: - - ready_obs (:obj:`Dict[int, Any]:`): Dict with env_id keys and observation values. + Get the ready (next) observation, which is a special design to unify both aysnc/sync env manager. + For each interaction between policy and env, the policy will input the ready_obs and output the action. + Then the env_manager will ``step`` with the action and prepare the next ready_obs. + Returns: + - ready_obs (:obj:`Dict[int, Any]`): A dict with env_id keys and observation values. Example: >>> obs = env_manager.ready_obs >>> stacked_obs = np.concatenate(list(obs.values())) - >>> action = model(obs) # model input np obs and output np action + >>> action = policy(obs) # here policy inputs np obs and outputs np action >>> action = {env_id: a for env_id, a in zip(obs.keys(), action)} >>> timesteps = env_manager.step(action) """ @@ -182,6 +230,12 @@ def ready_obs(self) -> Dict[int, Any]: @property def ready_obs_id(self) -> List[int]: + """ + Overview: + Get the ready (next) observation id, which is a special design to unify both aysnc/sync env manager. + Returns: + - ready_obs_id (:obj:`List[int]`): A list of env_ids for ready observations. + """ # In BaseEnvManager, if env_episode_count equals episode_num, this env is done. return [i for i, s in self._env_states.items() if s == EnvState.RUN] @@ -189,21 +243,43 @@ def ready_obs_id(self) -> List[int]: def ready_imgs(self, render_mode: Optional[str] = 'rgb_array') -> Dict[int, Any]: """ Overview: - Get the next ready renderd frame and corresponding env id. - Return: - - ready_imgs (:obj:`Dict[int, np.ndarray]:`): Dict with env_id keys and rendered frames. + Sometimes, we need to render the envs, this function is used to get the next ready renderd frame and \ + corresponding env id. + Arguments: + - render_mode (:obj:`Optional[str]`): The render mode, can be 'rgb_array' or 'depth_array', which follows \ + the definition in the ``render`` function of ``ding.utils`` . + Returns: + - ready_imgs (:obj:`Dict[int, np.ndarray]`): A dict with env_id keys and rendered frames. """ from ding.utils import render - assert render_mode in ['rgb_array', 'depth_array'] - return {i: render(self._envs[i], render_mode) for i in self.ready_obs.keys()} + assert render_mode in ['rgb_array', 'depth_array'], render_mode + return {i: render(self._envs[i], render_mode) for i in self.ready_obs_id} @property def done(self) -> bool: + """ + Overview: + ``done`` is a flag to indicate whether env manager is done, i.e., whether all sub-environments have \ + executed enough episodes. + Returns: + - done (:obj:`bool`): Whether env manager is done. + """ return all([s == EnvState.DONE for s in self._env_states.values()]) @property def method_name_list(self) -> list: - return ['reset', 'step', 'seed', 'close', 'enable_save_replay', 'render'] + """ + Overview: + The public methods list of sub-environments that can be directly called from the env manager level. Other \ + methods and attributes will be accessed with the ``__getattr__`` method. + Methods defined in this list can be regarded as the vectorized extension of methods in sub-environments. + Sub-class of ``BaseEnvManager`` can override this method to add more methods. + Returns: + - method_name_list (:obj:`list`): The public methods list of sub-environments. + """ + return [ + 'reset', 'step', 'seed', 'close', 'enable_save_replay', 'render', 'reward_shaping', 'enable_save_figure' + ] def env_state_done(self, env_id: int) -> bool: return self._env_states[env_id] == EnvState.DONE @@ -232,12 +308,22 @@ def _check_closed(self): def launch(self, reset_param: Optional[Dict] = None) -> None: """ Overview: - Set up the environments and their parameters. + Launch the env manager, instantiate the sub-environments and set up the environments and their parameters. Arguments: - - reset_param (:obj:`Optional[Dict]`): Dict of reset parameters for each environment, key is the env_id, \ - value is the cooresponding reset parameters. + - reset_param (:obj:`Optional[Dict]`): A dict of reset parameters for each environment, key is the env_id, \ + value is the corresponding reset parameter, defaults to None. """ assert self._closed, "Please first close the env manager" + try: + global space_log_flag + if space_log_flag: + logging.info("Env Space Information:") + logging.info("\tObservation Space: {}".format(self._observation_space)) + logging.info("\tAction Space: {}".format(self._action_space)) + logging.info("\tReward Space: {}".format(self._reward_space)) + space_log_flag = False + except: + pass if reset_param is not None: assert len(reset_param) == len(self._env_fn) self._create_state() @@ -258,10 +344,12 @@ def _create_state(self) -> None: def reset(self, reset_param: Optional[Dict] = None) -> None: """ Overview: - Reset the environments their parameters. + Forcely reset the sub-environments their corresponding parameters. Because in env manager all the \ + sub-environments usually are reset automatically as soon as they are done, this method is only called when \ + the caller must forcely reset all the sub-environments, such as in evaluation. Arguments: - reset_param (:obj:`List`): Dict of reset parameters for each environment, key is the env_id, \ - value is the cooresponding reset parameters. + value is the corresponding reset parameters. """ self._check_closed() # set seed if necessary @@ -319,26 +407,32 @@ def reset_fn(): logging.error("Env {} reset has exceeded max retries({})".format(env_id, self._max_retry)) runtime_error = RuntimeError( "Env {} reset has exceeded max retries({}), and the latest exception is: {}".format( - env_id, self._max_retry, repr(exceptions[-1]) + env_id, self._max_retry, str(exceptions[-1]) ) ) runtime_error.__traceback__ = exceptions[-1].__traceback__ raise runtime_error - def step(self, actions: Dict[int, Any]) -> List[Dict[int, BaseEnvTimestep]]: + def step(self, actions: Dict[int, Any]) -> Dict[int, BaseEnvTimestep]: """ Overview: - Execute env step according to input actions. And reset an env if done. + Execute env step according to input actions. If some sub-environments are done after this execution, \ + they will be reset automatically when ``self._auto_reset`` is True, otherwise they need to be reset when \ + the caller use the ``reset`` method of env manager. Arguments: - - actions (:obj:`Dict[int, Any]`): actions came from outer caller like policy + - actions (:obj:`Dict[int, Any]`): A dict of actions, key is the env_id, value is corresponding action. \ + action can be any type, it depends on the env, and the env will handle it. Ususlly, the action is \ + a dict of numpy array, and the value is generated by the outer caller like ``policy``. Returns: - - timesteps (:obj:`List[Dict[int, BaseEnvTimestep]]`): Each timestep is a BaseEnvTimestep object, \ - usually including observation, reward, done, info. + - timesteps (:obj:`Dict[int, BaseEnvTimestep]`): Each timestep is a ``BaseEnvTimestep`` object, \ + usually including observation, reward, done, info. Some special customized environments will have \ + the special timestep definition. The length of timesteps is the same as the length of actions in \ + synchronous env manager. Example: >>> timesteps = env_manager.step(action) - >>> for i, timestep in enumerate(timesteps): + >>> for env_id, timestep in enumerate(timesteps): >>> if timestep.done: - >>> print('Env {} is done'.format(timestep.env_id)) + >>> print('Env {} is done'.format(env_id)) """ self._check_closed() timesteps = {} @@ -373,7 +467,7 @@ def step_fn(): logging.error("Env {} step has exceeded max retries({})".format(env_id, self._max_retry)) runtime_error = RuntimeError( "Env {} step has exceeded max retries({}), and the latest exception is: {}".format( - env_id, self._max_retry, repr(exceptions[-1]) + env_id, self._max_retry, str(exceptions[-1]) ) ) runtime_error.__traceback__ = exceptions[-1].__traceback__ @@ -382,10 +476,15 @@ def step_fn(): def seed(self, seed: Union[Dict[int, int], List[int], int], dynamic_seed: bool = None) -> None: """ Overview: - Set the seed for each environment. + Set the random seed for each environment. Arguments: - - seed (:obj:`Union[Dict[int, int], List[int], int]`): List of seeds for each environment; \ - Or one seed for the first environment and other seeds are generated automatically. + - seed (:obj:`Union[Dict[int, int], List[int], int]`): Dict or List of seeds for each environment; \ + If only one seed is provided, it will be used in the same way for all environments. + - dynamic_seed (:obj:`bool`): Whether to use dynamic seed. + + .. note:: + For more details about ``dynamic_seed``, please refer to the best practice document of DI-engine \ + (`en link2 <../04_best_practice/random_seed.html>`_). """ if isinstance(seed, numbers.Integral): seed = [seed + i for i in range(self.env_num)] @@ -409,7 +508,7 @@ def seed(self, seed: Union[Dict[int, int], List[int], int], dynamic_seed: bool = def enable_save_replay(self, replay_path: Union[List[str], str]) -> None: """ Overview: - Set each env's replay save path. + Enable all environments to save replay video after each episode terminates. Arguments: - replay_path (:obj:`Union[List[str], str]`): List of paths for each environment; \ Or one path for all environments. @@ -418,10 +517,20 @@ def enable_save_replay(self, replay_path: Union[List[str], str]) -> None: replay_path = [replay_path] * self.env_num self._env_replay_path = replay_path + def enable_save_figure(self, env_id: int, figure_path: str) -> None: + """ + Overview: + Enable a specific env to save figure (e.g. environment statistics or episode return curve). + Arguments: + - figure_path (:obj:`str`): The file directory path for all environments to save figures. + """ + assert figure_path is not None + self._envs[env_id].enable_save_figure(figure_path) + def close(self) -> None: """ Overview: - Release the environment resources. + Close the env manager and release all the environment resources. """ if self._closed: return @@ -431,43 +540,96 @@ def close(self) -> None: self._env_states[i] = EnvState.VOID self._closed = True + def reward_shaping(self, env_id: int, transitions: List[dict]) -> List[dict]: + """ + Overview: + Execute reward shaping for a specific environment, which is often called when a episode terminates. + Arguments: + - env_id (:obj:`int`): The id of the environment to be shaped. + - transitions (:obj:`List[dict]`): The transition data list of the environment to be shaped. + Returns: + - transitions (:obj:`List[dict]`): The shaped transition data list. + """ + return self._envs[env_id].reward_shaping(transitions) + @property def closed(self) -> bool: + """ + Overview: + ``closed`` is a property that returns whether the env manager is closed. + Returns: + - closed (:obj:`bool`): Whether the env manager is closed. + """ return self._closed + def random_action(self) -> Dict: + return {env_id: self._env_ref.action_space.sample() for env_id in self.ready_obs_id} + @ENV_MANAGER_REGISTRY.register('base_v2') class BaseEnvManagerV2(BaseEnvManager): """ Overview: - BaseEnvManager for new task pipeline and interfaces coupled with treetensor. + The basic class of env manager to manage multiple vectorized environments. BaseEnvManager define all the + necessary interfaces and derived class must extend this basic class. + + The class is implemented by the pseudo-parallelism (i.e. serial) mechanism, therefore, this class is only + used in some tiny environments and for debug purpose. + + ``V2`` means this env manager is designed for new task pipeline and interfaces coupled with treetensor.` + + .. note:: + For more details about new task pipeline, please refer to the system document of DI-engine \ + (`system en link3 <../03_system/index.html>`_). + + Interfaces: + reset, step, seed, close, enable_save_replay, launch, default_config, reward_shaping, enable_save_figure + Properties: + env_num, env_ref, ready_obs, ready_obs_id, ready_imgs, done, closed, method_name_list, observation_space, \ + action_space, reward_space """ @property def ready_obs(self) -> tnp.array: """ Overview: - Get the ready (next) observation in ``tnp.array`` type, which is uniform for both async/sync scenarios. + Get the ready (next) observation, which is a special design to unify both aysnc/sync env manager. + For each interaction between policy and env, the policy will input the ready_obs and output the action. + Then the env_manager will ``step`` with the action and prepare the next ready_obs. + For ``V2`` version, the observation is transformed and packed up into ``tnp.array`` type, which allows + more convenient operations. Return: - ready_obs (:obj:`tnp.array`): A stacked treenumpy-type observation data. Example: >>> obs = env_manager.ready_obs - >>> action = model(obs) # model input np obs and output np action + >>> action = policy(obs) # here policy inputs treenp obs and output np action >>> timesteps = env_manager.step(action) """ active_env = [i for i, s in self._env_states.items() if s == EnvState.RUN] obs = [self._ready_obs[i] for i in active_env] + if isinstance(obs[0], dict): # transform each element to treenumpy array + obs = [tnp.array(o) for o in obs] return tnp.stack(obs) def step(self, actions: List[tnp.ndarray]) -> List[tnp.ndarray]: """ Overview: - Execute env step according to input actions. And reset an env if done. + Execute env step according to input actions. If some sub-environments are done after this execution, \ + they will be reset automatically by default. Arguments: - - actions (:obj:`List[tnp.ndarray]`): actions came from outer caller like policy + - actions (:obj:`List[tnp.ndarray]`): A list of treenumpy-type actions, the value is generated by the \ + outer caller like ``policy``. Returns: - - timesteps (:obj:`List[tnp.ndarray]`): Each timestep is a tnp.array with observation, reward, done, \ - info, env_id. + - timesteps (:obj:`List[tnp.ndarray]`): A list of timestep, Each timestep is a ``tnp.ndarray`` object, \ + usually including observation, reward, done, info, env_id. Some special environments will have \ + the special timestep definition. The length of timesteps is the same as the length of actions in \ + synchronous env manager. For the compatibility of treenumpy, here we use ``make_key_as_identifier`` \ + and ``remove_illegal_item`` functions to modify the original timestep. + Example: + >>> timesteps = env_manager.step(action) + >>> for timestep in timesteps: + >>> if timestep.done: + >>> print('Env {} is done'.format(timestep.env_id)) """ actions = {env_id: a for env_id, a in zip(self.ready_obs_id, actions)} timesteps = super().step(actions) @@ -482,15 +644,22 @@ def step(self, actions: List[tnp.ndarray]) -> List[tnp.ndarray]: return new_data -def create_env_manager(manager_cfg: dict, env_fn: List[Callable]) -> BaseEnvManager: - r""" +def create_env_manager(manager_cfg: EasyDict, env_fn: List[Callable]) -> BaseEnvManager: + """ Overview: - Create an env manager according to manager cfg and env function. + Create an env manager according to ``manager_cfg`` and env functions. Arguments: - - manager_cfg (:obj:`EasyDict`): Env manager config. - - env_fn (:obj:` List[Callable]`): A list of envs' functions. + - manager_cfg (:obj:`EasyDict`): Final merged env manager config. + - env_fn (:obj:`List[Callable]`): A list of functions to create ``env_num`` sub-environments. ArgumentsKeys: - - `manager_cfg`'s necessary: `type` + - type (:obj:`str`): Env manager type set in ``ENV_MANAGER_REGISTRY.register`` , such as ``base`` . + - import_names (:obj:`List[str]`): A list of module names (paths) to import before creating env manager, such \ + as ``ding.envs.env_manager.base_env_manager`` . + Returns: + - env_manager (:obj:`BaseEnvManager`): The created env manager. + + .. tip:: + This method will not modify the ``manager_cfg`` , it will deepcopy the ``manager_cfg`` and then modify it. """ manager_cfg = copy.deepcopy(manager_cfg) if 'import_names' in manager_cfg: @@ -500,13 +669,17 @@ def create_env_manager(manager_cfg: dict, env_fn: List[Callable]) -> BaseEnvMana def get_env_manager_cls(cfg: EasyDict) -> type: - r""" + """ Overview: - Get an env manager class according to cfg. + Get the env manager class according to config, which is used to access related class variables/methods. Arguments: - - cfg (:obj:`EasyDict`): Env manager config. + - manager_cfg (:obj:`EasyDict`): Final merged env manager config. ArgumentsKeys: - - necessary: `type` + - type (:obj:`str`): Env manager type set in ``ENV_MANAGER_REGISTRY.register`` , such as ``base`` . + - import_names (:obj:`List[str]`): A list of module names (paths) to import before creating env manager, such \ + as ``ding.envs.env_manager.base_env_manager`` . + Returns: + - env_manager_cls (:obj:`type`): The corresponding env manager class. """ import_module(cfg.get('import_names', [])) return ENV_MANAGER_REGISTRY.get(cfg.type) diff --git a/ding/envs/env_manager/ding_env_manager.py b/ding/envs/env_manager/ding_env_manager.py new file mode 100644 index 0000000000..12a3ecf88f --- /dev/null +++ b/ding/envs/env_manager/ding_env_manager.py @@ -0,0 +1,23 @@ +from . import BaseEnvManagerV2, SubprocessEnvManagerV2 +from ..env import DingEnvWrapper +from typing import Optional +from functools import partial + + +def setup_ding_env_manager( + env: DingEnvWrapper, + env_num: int, + context: Optional[str] = None, + debug: bool = False, + caller: str = 'collector' +) -> BaseEnvManagerV2: + assert caller in ['evaluator', 'collector'] + if debug: + env_cls = BaseEnvManagerV2 + manager_cfg = env_cls.default_config() + else: + env_cls = SubprocessEnvManagerV2 + manager_cfg = env_cls.default_config() + if context is not None: + manager_cfg.context = context + return env_cls([partial(env.clone, caller) for _ in range(env_num)], manager_cfg) diff --git a/ding/envs/env_manager/env_supervisor.py b/ding/envs/env_manager/env_supervisor.py index b3e1c17fcf..ec5e29beab 100644 --- a/ding/envs/env_manager/env_supervisor.py +++ b/ding/envs/env_manager/env_supervisor.py @@ -5,10 +5,10 @@ import gym from ding.framework import Supervisor from typing import TYPE_CHECKING, Any, List, Union, Dict, Optional, Callable -from ding.framework.supervisor import ChildType, RecvPayload, SendPayload, SharedObject +from ding.framework.supervisor import ChildType, RecvPayload, SendPayload from ding.utils import make_key_as_identifier from ditk import logging -from ding.envs.env_manager.subprocess_env_manager import ShmBufferContainer +from ding.data import ShmBufferContainer import enum import treetensor.numpy as tnp import numbers @@ -106,9 +106,7 @@ def __init__( for env_id in range(len(self._env_fn)) } for env_init in env_fn: - self.register( - env_init, shared_object=SharedObject(buf=self._obs_buffers, callback=self._shm_callback) - ) + self.register(env_init, shm_buffer=self._obs_buffers, shm_callback=self._shm_callback) else: for env_init in env_fn: self.register(env_init) @@ -136,6 +134,11 @@ def _init_states(self): self._last_called = defaultdict(lambda: {"step": math.inf, "reset": math.inf}) def _shm_callback(self, payload: RecvPayload, obs_buffers: Any): + """ + Overview: + This method will be called in child worker, so we can put large data into shared memory + and replace the original payload data to none, then reduce the serialization/deserialization cost. + """ if payload.method == "reset" and payload.data is not None: obs_buffers[payload.proc_id].fill(payload.data) payload.data = None diff --git a/ding/envs/env_manager/envpool_env_manager.py b/ding/envs/env_manager/envpool_env_manager.py index 3defc8e6eb..a8d1a4ae03 100644 --- a/ding/envs/env_manager/envpool_env_manager.py +++ b/ding/envs/env_manager/envpool_env_manager.py @@ -80,7 +80,7 @@ def reset(self) -> None: self._ready_obs = deep_merge_dicts({i: o for i, o in zip(env_id, obs)}, self._ready_obs) if len(self._ready_obs) == self._env_num: break - self._final_eval_reward = [0. for _ in range(self._env_num)] + self._eval_episode_return = [0. for _ in range(self._env_num)] def step(self, action: dict) -> Dict[int, namedtuple]: env_id = np.array(list(action.keys())) @@ -98,11 +98,11 @@ def step(self, action: dict) -> Dict[int, namedtuple]: for i in range(len(env_id)): d = bool(done[i]) r = to_ndarray([rew[i]]) - self._final_eval_reward[env_id[i]] += r + self._eval_episode_return[env_id[i]] += r timesteps[env_id[i]] = BaseEnvTimestep(obs[i], r, d, info={'env_id': i}) if d: - timesteps[env_id[i]].info['final_eval_reward'] = self._final_eval_reward[env_id[i]] - self._final_eval_reward[env_id[i]] = 0. + timesteps[env_id[i]].info['eval_episode_return'] = self._eval_episode_return[env_id[i]] + self._eval_episode_return[env_id[i]] = 0. self._ready_obs[env_id[i]] = obs[i] return timesteps diff --git a/ding/envs/env_manager/gym_vector_env_manager.py b/ding/envs/env_manager/gym_vector_env_manager.py index ea87310151..46bd8c0769 100644 --- a/ding/envs/env_manager/gym_vector_env_manager.py +++ b/ding/envs/env_manager/gym_vector_env_manager.py @@ -48,7 +48,7 @@ def __init__(self, env_fn: List[Callable], cfg: EasyDict) -> None: shared_memory=cfg.shared_memory, ) self._env_states = {i: EnvState.INIT for i in range(self._env_num)} - self._final_eval_reward = [0. for _ in range(self._env_num)] + self._eval_episode_return = [0. for _ in range(self._env_num)] def reset(self, reset_param: Optional[Dict] = None) -> None: assert reset_param is None @@ -58,7 +58,7 @@ def reset(self, reset_param: Optional[Dict] = None) -> None: self._ready_obs = self._env_manager.reset() for env_id in range(self.env_num): self._env_states[env_id] = EnvState.RUN - self._final_eval_reward = [0. for _ in range(self._env_num)] + self._eval_episode_return = [0. for _ in range(self._env_num)] def step(self, actions: Dict[int, Any]) -> Dict[int, namedtuple]: assert isinstance(actions, Dict), type(actions) @@ -95,10 +95,10 @@ def step(self, actions: Dict[int, Any]) -> Dict[int, namedtuple]: timestep_collate_result[i] = BaseEnvTimestep( timestep[0][i], timestep[1][i], timestep[2][i], timestep[3][i] ) - self._final_eval_reward[i] += timestep_collate_result[i].reward + self._eval_episode_return[i] += timestep_collate_result[i].reward if timestep_collate_result[i].done: - timestep_collate_result[i].info['final_eval_reward'] = self._final_eval_reward[i] - self._final_eval_reward[i] = 0 + timestep_collate_result[i].info['eval_episode_return'] = self._eval_episode_return[i] + self._eval_episode_return[i] = 0 self._env_episode_count[i] += 1 if self._env_episode_count[i] >= self._episode_num: self._env_states[i] = EnvState.DONE diff --git a/ding/envs/env_manager/subprocess_env_manager.py b/ding/envs/env_manager/subprocess_env_manager.py index 38c01b906e..b30fe10394 100644 --- a/ding/envs/env_manager/subprocess_env_manager.py +++ b/ding/envs/env_manager/subprocess_env_manager.py @@ -1,40 +1,26 @@ from typing import Any, Union, List, Tuple, Dict, Callable, Optional -from multiprocessing import Pipe, connection, get_context, Array +from multiprocessing import connection, get_context from collections import namedtuple from ditk import logging import platform import time import copy +import gymnasium import gym import traceback import torch -import ctypes import pickle -import cloudpickle import numpy as np import treetensor.numpy as tnp from easydict import EasyDict from types import MethodType +from ding.data import ShmBufferContainer, ShmBuffer from ding.envs.env import BaseEnvTimestep from ding.utils import PropagatingThread, LockContextType, LockContext, ENV_MANAGER_REGISTRY, make_key_as_identifier, \ - remove_illegal_item + remove_illegal_item, CloudPickleWrapper from .base_env_manager import BaseEnvManager, EnvState, timeout_wrapper -_NTYPE_TO_CTYPE = { - np.bool_: ctypes.c_bool, - np.uint8: ctypes.c_uint8, - np.uint16: ctypes.c_uint16, - np.uint32: ctypes.c_uint32, - np.uint64: ctypes.c_uint64, - np.int8: ctypes.c_int8, - np.int16: ctypes.c_int16, - np.int32: ctypes.c_int32, - np.int64: ctypes.c_int64, - np.float32: ctypes.c_float, - np.float64: ctypes.c_double, -} - def is_abnormal_timestep(timestep: namedtuple) -> bool: if isinstance(timestep.info, dict): @@ -45,129 +31,6 @@ def is_abnormal_timestep(timestep: namedtuple) -> bool: raise TypeError("invalid env timestep type: {}".format(type(timestep.info))) -class ShmBuffer(): - """ - Overview: - Shared memory buffer to store numpy array. - """ - - def __init__(self, dtype: Union[type, np.dtype], shape: Tuple[int], copy_on_get: bool = True) -> None: - """ - Overview: - Initialize the buffer. - Arguments: - - dtype (:obj:`Union[type, np.dtype]`): The dtype of the data to limit the size of the buffer. - - shape (:obj:`Tuple[int]`): The shape of the data to limit the size of the buffer. - - copy_on_get (:obj:`bool`): Whether to copy data when calling get method. - """ - if isinstance(dtype, np.dtype): # it is type of gym.spaces.dtype - dtype = dtype.type - self.buffer = Array(_NTYPE_TO_CTYPE[dtype], int(np.prod(shape))) - self.dtype = dtype - self.shape = shape - self.copy_on_get = copy_on_get - - def fill(self, src_arr: np.ndarray) -> None: - """ - Overview: - Fill the shared memory buffer with a numpy array. (Replace the original one.) - Arguments: - - src_arr (:obj:`np.ndarray`): array to fill the buffer. - """ - assert isinstance(src_arr, np.ndarray), type(src_arr) - # for np.array with shape (4, 84, 84) and float32 dtype, reshape is 15~20x faster than flatten - # for np.array with shape (4, 84, 84) and uint8 dtype, reshape is 5~7x faster than flatten - # so we reshape dst_arr rather than flatten src_arr - dst_arr = np.frombuffer(self.buffer.get_obj(), dtype=self.dtype).reshape(self.shape) - np.copyto(dst_arr, src_arr) - - def get(self) -> np.ndarray: - """ - Overview: - Get the array stored in the buffer. - Return: - - data (:obj:`np.ndarray`): A copy of the data stored in the buffer. - """ - data = np.frombuffer(self.buffer.get_obj(), dtype=self.dtype).reshape(self.shape) - if self.copy_on_get: - data = data.copy() # must use np.copy, torch.from_numpy and torch.as_tensor still use the same memory - return data - - -class ShmBufferContainer(object): - """ - Overview: - Support multiple shared memory buffers. Each key-value is name-buffer. - """ - - def __init__( - self, - dtype: Union[Dict[Any, type], type, np.dtype], - shape: Union[Dict[Any, tuple], tuple], - copy_on_get: bool = True - ) -> None: - """ - Overview: - Initialize the buffer container. - Arguments: - - dtype (:obj:`Union[type, np.dtype]`): The dtype of the data to limit the size of the buffer. - - shape (:obj:`Union[Dict[Any, tuple], tuple]`): If `Dict[Any, tuple]`, use a dict to manage \ - multiple buffers; If `tuple`, use single buffer. - - copy_on_get (:obj:`bool`): Whether to copy data when calling get method. - """ - if isinstance(shape, dict): - self._data = {k: ShmBufferContainer(dtype[k], v, copy_on_get) for k, v in shape.items()} - elif isinstance(shape, (tuple, list)): - self._data = ShmBuffer(dtype, shape, copy_on_get) - else: - raise RuntimeError("not support shape: {}".format(shape)) - self._shape = shape - - def fill(self, src_arr: Union[Dict[Any, np.ndarray], np.ndarray]) -> None: - """ - Overview: - Fill the one or many shared memory buffer. - Arguments: - - src_arr (:obj:`Union[Dict[Any, np.ndarray], np.ndarray]`): array to fill the buffer. - """ - if isinstance(self._shape, dict): - for k in self._shape.keys(): - self._data[k].fill(src_arr[k]) - elif isinstance(self._shape, (tuple, list)): - self._data.fill(src_arr) - - def get(self) -> Union[Dict[Any, np.ndarray], np.ndarray]: - """ - Overview: - Get the one or many arrays stored in the buffer. - Return: - - data (:obj:`np.ndarray`): The array(s) stored in the buffer. - """ - if isinstance(self._shape, dict): - return {k: self._data[k].get() for k in self._shape.keys()} - elif isinstance(self._shape, (tuple, list)): - return self._data.get() - - -class CloudPickleWrapper: - """ - Overview: - CloudPickleWrapper can be able to pickle more python object(e.g: an object with lambda expression) - """ - - def __init__(self, data: Any) -> None: - self.data = data - - def __getstate__(self) -> bytes: - return cloudpickle.dumps(self.data) - - def __setstate__(self, data: bytes) -> None: - if isinstance(data, (tuple, list, np.ndarray)): # pickle is faster - self.data = pickle.loads(data) - else: - self.data = cloudpickle.loads(data) - - @ENV_MANAGER_REGISTRY.register('async_subprocess') class AsyncSubprocessEnvManager(BaseEnvManager): """ @@ -180,7 +43,7 @@ class AsyncSubprocessEnvManager(BaseEnvManager): config = dict( episode_num=float("inf"), - max_retry=5, + max_retry=1, step_timeout=None, auto_reset=True, retry_type='reset', @@ -242,7 +105,7 @@ def _create_state(self) -> None: self._reset_param = {i: {} for i in range(self.env_num)} if self._shared_memory: obs_space = self._observation_space - if isinstance(obs_space, gym.spaces.Dict): + if isinstance(obs_space, (gym.spaces.Dict, gymnasium.spaces.Dict)): # For multi_agent case, such as multiagent_mujoco and petting_zoo mpe. # Now only for the case that each agent in the team have the same obs structure # and corresponding shape. @@ -313,7 +176,7 @@ def ready_obs(self) -> Dict[int, Any]: no_done_env_idx = [i for i, s in self._env_states.items() if s != EnvState.DONE] sleep_count = 0 while not any([self._env_states[i] == EnvState.RUN for i in no_done_env_idx]): - if sleep_count % 1000 == 0: + if sleep_count != 0 and sleep_count % 10000 == 0: logging.warning( 'VEC_ENV_MANAGER: all the not done envs are resetting, sleep {} times'.format(sleep_count) ) @@ -376,7 +239,7 @@ def reset(self, reset_param: Optional[Dict] = None) -> None: sleep_count = 0 while any([self._env_states[i] == EnvState.RESET for i in reset_env_list]): - if sleep_count % 1000 == 0: + if sleep_count != 0 and sleep_count % 10000 == 0: logging.warning( 'VEC_ENV_MANAGER: not all the envs finish resetting, sleep {} times'.format(sleep_count) ) @@ -397,7 +260,10 @@ def reset(self, reset_param: Optional[Dict] = None) -> None: self._check_data({env_id: ret}) self._env_seed[env_id] = None # seed only use once except BaseException as e: - logging.warning("subprocess reset set seed failed, ignore and continue...") + logging.warning( + "subprocess reset set seed failed, ignore and continue... \n subprocess exception traceback: \n" + + traceback.format_exc() + ) self._env_states[env_id] = EnvState.RESET reset_thread = PropagatingThread(target=self._reset, args=(env_id, )) reset_thread.daemon = True @@ -439,6 +305,7 @@ def reset_fn(): reset_fn() return except BaseException as e: + logging.info("subprocess exception traceback: \n" + traceback.format_exc()) if self._retry_type == 'renew' or isinstance(e, pickle.UnpicklingError): self._pipe_parents[env_id].close() if self._subprocesses[env_id].is_alive(): @@ -450,7 +317,7 @@ def reset_fn(): logging.error("Env {} reset has exceeded max retries({})".format(env_id, self._max_retry)) runtime_error = RuntimeError( "Env {} reset has exceeded max retries({}), and the latest exception is: {}".format( - env_id, self._max_retry, repr(exceptions[-1]) + env_id, self._max_retry, str(exceptions[-1]) ) ) runtime_error.__traceback__ = exceptions[-1].__traceback__ @@ -616,6 +483,7 @@ def worker_fn( except Exception as e: # when there are some errors in env, worker_fn will send the errors to env manager # directly send error to another process will lose the stack trace, so we create a new Exception + logging.warning("subprocess exception traceback: \n" + traceback.format_exc()) c.send( e.__class__( '\nEnv Process Exception:\n' + ''.join(traceback.format_tb(e.__traceback__)) + repr(e) @@ -671,6 +539,7 @@ def reset_fn(*args, **kwargs): ret = None return ret except BaseException as e: + logging.warning("subprocess exception traceback: \n" + traceback.format_exc()) env.close() raise e @@ -704,6 +573,7 @@ def reset_fn(*args, **kwargs): logging.debug("Sub env '{}' error when executing {}".format(str(env), cmd)) # when there are some errors in env, worker_fn will send the errors to env manager # directly send error to another process will lose the stack trace, so we create a new Exception + logging.warning("subprocess exception traceback: \n" + traceback.format_exc()) child.send( e.__class__('\nEnv Process Exception:\n' + ''.join(traceback.format_tb(e.__traceback__)) + repr(e)) ) @@ -808,7 +678,7 @@ def wait(rest_conn: list, wait_num: int, timeout: Optional[float] = None) -> Tup class SyncSubprocessEnvManager(AsyncSubprocessEnvManager): config = dict( episode_num=float("inf"), - max_retry=5, + max_retry=1, step_timeout=None, auto_reset=True, reset_timeout=None, @@ -927,26 +797,31 @@ def ready_obs(self) -> tnp.array: no_done_env_idx = [i for i, s in self._env_states.items() if s != EnvState.DONE] sleep_count = 0 while not any([self._env_states[i] == EnvState.RUN for i in no_done_env_idx]): - if sleep_count % 1000 == 0: + if sleep_count != 0 and sleep_count % 10000 == 0: logging.warning( 'VEC_ENV_MANAGER: all the not done envs are resetting, sleep {} times'.format(sleep_count) ) time.sleep(0.001) sleep_count += 1 - obs = [self._ready_obs[i] for i in self.ready_env] - return tnp.stack(obs) + return tnp.stack([tnp.array(self._ready_obs[i]) for i in self.ready_env]) - def step(self, actions: List[tnp.ndarray]) -> List[tnp.ndarray]: + def step(self, actions: Union[List[tnp.ndarray], tnp.ndarray]) -> List[tnp.ndarray]: """ Overview: Execute env step according to input actions. And reset an env if done. Arguments: - - actions (:obj:`List[tnp.ndarray]`): actions came from outer caller like policy + - actions (:obj:`Union[List[tnp.ndarray], tnp.ndarray]`): actions came from outer caller like policy. Returns: - timesteps (:obj:`List[tnp.ndarray]`): Each timestep is a tnp.array with observation, reward, done, \ info, env_id. """ - actions = {env_id: a for env_id, a in zip(self.ready_obs_id, actions)} + if isinstance(actions, tnp.ndarray): + # zip operation will lead to wrong behaviour if not split data + split_action = tnp.split(actions, actions.shape[0]) + split_action = [s.squeeze(0) for s in split_action] + else: + split_action = actions + actions = {env_id: a for env_id, a in zip(self.ready_obs_id, split_action)} timesteps = super().step(actions) new_data = [] for env_id, timestep in timesteps.items(): diff --git a/ding/envs/env_manager/tests/conftest.py b/ding/envs/env_manager/tests/conftest.py index f74f051e59..f824899a0d 100644 --- a/ding/envs/env_manager/tests/conftest.py +++ b/ding/envs/env_manager/tests/conftest.py @@ -164,7 +164,7 @@ def setup_model_type(): return FakeModel -def get_base_manager_cfg(env_num=4): +def get_base_manager_cfg(env_num=3): manager_cfg = { 'env_cfg': [{ 'name': 'name{}'.format(i), @@ -178,7 +178,7 @@ def get_base_manager_cfg(env_num=4): return EasyDict(manager_cfg) -def get_subprecess_manager_cfg(env_num=4): +def get_subprecess_manager_cfg(env_num=3): manager_cfg = { 'env_cfg': [{ 'name': 'name{}'.format(i), @@ -194,7 +194,7 @@ def get_subprecess_manager_cfg(env_num=4): return EasyDict(manager_cfg) -def get_gym_vector_manager_cfg(env_num=4): +def get_gym_vector_manager_cfg(env_num=3): manager_cfg = { 'env_cfg': [{ 'name': 'name{}'.format(i), @@ -210,7 +210,7 @@ def get_gym_vector_manager_cfg(env_num=4): @pytest.fixture(scope='function') def setup_base_manager_cfg(): - manager_cfg = get_base_manager_cfg(4) + manager_cfg = get_base_manager_cfg(3) env_cfg = manager_cfg.pop('env_cfg') manager_cfg['env_fn'] = [partial(FakeEnv, cfg=c) for c in env_cfg] return deep_merge_dicts(BaseEnvManager.default_config(), EasyDict(manager_cfg)) @@ -218,7 +218,7 @@ def setup_base_manager_cfg(): @pytest.fixture(scope='function') def setup_fast_base_manager_cfg(): - manager_cfg = get_base_manager_cfg(4) + manager_cfg = get_base_manager_cfg(3) env_cfg = manager_cfg.pop('env_cfg') for e in env_cfg: e['scale'] = 0.1 @@ -228,7 +228,7 @@ def setup_fast_base_manager_cfg(): @pytest.fixture(scope='function') def setup_sync_manager_cfg(): - manager_cfg = get_subprecess_manager_cfg(4) + manager_cfg = get_subprecess_manager_cfg(3) env_cfg = manager_cfg.pop('env_cfg') # TODO(nyz) test fail when shared_memory = True manager_cfg['shared_memory'] = False @@ -238,7 +238,7 @@ def setup_sync_manager_cfg(): @pytest.fixture(scope='function') def setup_async_manager_cfg(): - manager_cfg = get_subprecess_manager_cfg(4) + manager_cfg = get_subprecess_manager_cfg(3) env_cfg = manager_cfg.pop('env_cfg') manager_cfg['env_fn'] = [partial(FakeAsyncEnv, cfg=c) for c in env_cfg] manager_cfg['shared_memory'] = False @@ -247,7 +247,7 @@ def setup_async_manager_cfg(): @pytest.fixture(scope='function') def setup_gym_vector_manager_cfg(): - manager_cfg = get_subprecess_manager_cfg(4) + manager_cfg = get_subprecess_manager_cfg(3) env_cfg = manager_cfg.pop('env_cfg') manager_cfg['env_fn'] = [partial(FakeGymEnv, cfg=c) for c in env_cfg] manager_cfg['shared_memory'] = False diff --git a/ding/envs/env_manager/tests/test_base_env_manager.py b/ding/envs/env_manager/tests/test_base_env_manager.py index 379589314e..973a0a1291 100644 --- a/ding/envs/env_manager/tests/test_base_env_manager.py +++ b/ding/envs/env_manager/tests/test_base_env_manager.py @@ -93,7 +93,7 @@ def test_error(self, setup_base_manager_cfg): assert timestep[0].info.abnormal assert all(['abnormal' not in timestep[i].info for i in range(1, env_manager.env_num)]) assert all([env_manager._env_states[i] == EnvState.RUN for i in range(env_manager.env_num)]) - assert len(env_manager.ready_obs) == 4 + assert len(env_manager.ready_obs) == 3 timestep = env_manager.step({i: np.random.randn(4) for i in range(env_manager.env_num)}) # Test step error action[0] = 'error' @@ -103,7 +103,7 @@ def test_error(self, setup_base_manager_cfg): assert all([env_manager._env_states[i] == EnvState.RUN for i in range(1, env_manager.env_num)]) obs = env_manager.reset(reset_param) assert all([env_manager._env_states[i] == EnvState.RUN for i in range(env_manager.env_num)]) - assert len(env_manager.ready_obs) == 4 + assert len(env_manager.ready_obs) == 3 timestep = env_manager.step({i: np.random.randn(4) for i in range(env_manager.env_num)}) env_manager.close() diff --git a/ding/envs/env_manager/tests/test_env_supervisor.py b/ding/envs/env_manager/tests/test_env_supervisor.py index 4deb220eb3..c16a7d1c6d 100644 --- a/ding/envs/env_manager/tests/test_env_supervisor.py +++ b/ding/envs/env_manager/tests/test_env_supervisor.py @@ -119,7 +119,7 @@ def test_renew_error(self, setup_base_manager_cfg, type_): assert not env_supervisor.closed # If retry type is renew, time id should not be equal assert env_supervisor.time_id[0] != env_id_0 - assert len(env_supervisor.ready_obs) == 4 + assert len(env_supervisor.ready_obs) == 3 for i, obs in enumerate(env_supervisor.ready_obs): assert all(x == y for x, y in zip(obs, env_supervisor._ready_obs.get(i))) @@ -132,7 +132,7 @@ def test_renew_error(self, setup_base_manager_cfg, type_): assert all(['abnormal' not in timestep[i].info for i in range(1, env_supervisor.env_num)]) # With auto_reset, abnormal timestep with done==True will be auto reset. assert all([env_supervisor.env_states[i] == EnvState.RUN for i in range(env_supervisor.env_num)]) - assert len(env_supervisor.ready_obs) == 4 + assert len(env_supervisor.ready_obs) == 3 env_supervisor.close() @pytest.mark.tmp # gitlab ci and local test pass, github always fail @@ -215,18 +215,18 @@ def test_auto_reset(self, setup_base_manager_cfg, type_): ) env_supervisor.launch(reset_param={i: {'stat': 'stat_test'} for i in range(env_supervisor.env_num)}) - assert len(env_supervisor.ready_obs) == 4 - assert len(env_supervisor.ready_obs_id) == 4 + assert len(env_supervisor.ready_obs) == 3 + assert len(env_supervisor.ready_obs_id) == 3 timesteps = [] for _ in range(10): action = {i: np.random.randn(4) for i in range(env_supervisor.env_num)} timesteps.append(env_supervisor.step(action)) - assert len(env_supervisor.ready_obs) == 4 + assert len(env_supervisor.ready_obs) == 3 time.sleep(1) timesteps = tnp.stack(timesteps).reshape(-1) - assert len(timesteps.done) == 40 + assert len(timesteps.done) == 30 assert any(done for done in timesteps.done) assert all([env_supervisor.env_states[env_id] == EnvState.RUN for env_id in range(env_supervisor.env_num)]) env_supervisor.close() @@ -297,7 +297,7 @@ def test_reset_error_once(self, setup_base_manager_cfg, type_): # Normal step env_supervisor.step({i: np.random.randn(4) for i in range(env_supervisor.env_num)}, block=False) timestep = [] - while len(timestep) != 4: + while len(timestep) != 3: payload = env_supervisor.recv() if payload.method == "step": timestep.append(payload.data) @@ -311,7 +311,7 @@ def test_reset_error_once(self, setup_base_manager_cfg, type_): env_supervisor.reset(reset_param, block=False) # Second try, error and recover reset_obs = [] - while len(reset_obs) != 8: + while len(reset_obs) != 6: reset_obs.append(env_supervisor.recv(ignore_err=True)) assert env_supervisor.time_id[0] == env_id_0 assert all([state == EnvState.RUN for state in env_supervisor.env_states.values()]) @@ -334,11 +334,11 @@ def test_renew_error_once(self, setup_base_manager_cfg, type_): env_supervisor.reset(reset_param, block=False) reset_obs = [] - while len(reset_obs) != 8: + while len(reset_obs) != 6: reset_obs.append(env_supervisor.recv(ignore_err=True)) assert env_supervisor.time_id[0] != env_id_0 - assert len(env_supervisor.ready_obs) == 4 + assert len(env_supervisor.ready_obs) == 3 # Test step catched error action = [np.random.randn(4) for i in range(env_supervisor.env_num)] @@ -346,7 +346,7 @@ def test_renew_error_once(self, setup_base_manager_cfg, type_): env_supervisor.step(action, block=False) timestep = {} - while len(timestep) != 4: + while len(timestep) != 3: payload = env_supervisor.recv() if payload.method == "step": timestep[payload.proc_id] = payload.data diff --git a/ding/envs/env_manager/tests/test_gym_vector_env_manager.py b/ding/envs/env_manager/tests/test_gym_vector_env_manager.py index 51d1c3a3ea..2ea79f2f47 100644 --- a/ding/envs/env_manager/tests/test_gym_vector_env_manager.py +++ b/ding/envs/env_manager/tests/test_gym_vector_env_manager.py @@ -9,7 +9,8 @@ from gym.vector.async_vector_env import AsyncState -@pytest.mark.unittest +@pytest.mark.tmp +# @pytest.mark.unittest class TestGymVectorEnvManager: def test_naive(self, setup_gym_vector_manager_cfg): @@ -31,7 +32,7 @@ def test_naive(self, setup_gym_vector_manager_cfg): while not env_manager.done: env_id = env_manager.ready_obs.keys() assert all(env_manager._env_episode_count[i] < env_manager._episode_num for i in env_id) - action = {i: np.random.randn(4) for i in env_id} + action = {i: np.random.randn(3) for i in env_id} timestep = env_manager.step(action) assert len(timestep) == len(env_id) print('Count {}'.format(count)) diff --git a/ding/envs/env_manager/tests/test_subprocess_env_manager.py b/ding/envs/env_manager/tests/test_subprocess_env_manager.py index 103b69f2ff..218b87e383 100644 --- a/ding/envs/env_manager/tests/test_subprocess_env_manager.py +++ b/ding/envs/env_manager/tests/test_subprocess_env_manager.py @@ -112,13 +112,13 @@ def test_error(self, setup_sync_manager_cfg): assert timestep[0].info['abnormal'] assert all(['abnormal' not in timestep[i].info for i in range(1, env_manager.env_num)]) assert env_manager._env_states[0] == EnvState.ERROR - assert len(env_manager.ready_obs) == 3 + assert len(env_manager.ready_obs) == 2 # wait for reset env_manager.reset({0: {'stat': 'stat_test'}}) while not len(env_manager.ready_obs) == env_manager.env_num: time.sleep(0.1) assert env_manager._env_states[0] == EnvState.RUN - assert len(env_manager.ready_obs) == 4 + assert len(env_manager.ready_obs) == 3 timestep = env_manager.step({i: np.random.randn(4) for i in range(env_manager.env_num)}) # # Test step error diff --git a/ding/envs/env_wrappers/env_wrappers.py b/ding/envs/env_wrappers/env_wrappers.py index fbc6092bf6..a33938d20a 100644 --- a/ding/envs/env_wrappers/env_wrappers.py +++ b/ding/envs/env_wrappers/env_wrappers.py @@ -1,68 +1,95 @@ -# Borrow a lot from openai baselines: -# https://github.com/openai/baselines/blob/master/baselines/common/atari_wrappers.py +""" +This code is adapted from OpenAI Baselines: + https://github.com/openai/baselines/blob/master/baselines/common/atari_wrappers.py + +List of Environment Wrappers: +- NoopResetWrapper: This wrapper facilitates the sampling of initial states by executing a random number of + no-operation actions upon environment reset. +- MaxAndSkipWrapper: Incorporates max pooling across time steps, a method that reduces the temporal dimension by taking + the maximum value over specified time intervals. +- WarpFrameWrapper: Implements frame warping by resizing the images to 84x84, a common preprocessing step in + reinforcement learning on visual data, as described in the DeepMind Nature paper and subsequent works. +- ScaledFloatFrameWrapper: Normalizes observations to a range of 0 to 1, which is a common requirement for neural + network inputs. +- ClipRewardWrapper: Clips the reward to {-1, 0, +1} based on its sign. This simplifies the reward structure and + can make learning more stable in environments with high variance in rewards. +- DelayRewardWrapper: Returns cumulative reward at defined intervals, and at all other times, returns a reward of 0. + This can be useful for sparse reward problems. +- FrameStackWrapper: Stacks the latest 'n' frames as a single observation. This allows the agent to have a sense of + dynamics and motion from the stacked frames. +- ObsTransposeWrapper: Transposes the observation to bring the channel to the first dimension, a common requirement + for convolutional neural networks. +- ObsNormWrapper: Normalizes observations based on a running mean and standard deviation. This can help to standardize + inputs for the agent and speed up learning. +- RewardNormWrapper: Normalizes reward based on a running standard deviation, which can stabilize learning in + environments with high variance in rewards. +- RamWrapper: Wraps a RAM-based environment into an image-like environment. This can be useful for applying + image-based algorithms to RAM-based Atari games. +- EpisodicLifeWrapper: Treats end of life as the end of an episode, but only resets on true game over. This can help + the agent better differentiate between losing a life and losing the game. +- FireResetWrapper: Executes the 'fire' action upon environment reset. This is specific to certain Atari games where + the 'fire' action starts the game. +- GymHybridDictActionWrapper: Transforms the original `gym.spaces.Tuple` action space into a `gym.spaces.Dict`. +- FlatObsWrapper: Flattens image and language observations into a single vector, which can be helpful for input into + certain types of models. +- StaticObsNormWrapper: Provides functionality for normalizing observations according to a static mean and + standard deviation. +- EvalEpisodeReturnWrapper: Evaluates the return over an episode during evaluation, providing a more comprehensive + view of the agent's performance. +- GymToGymnasiumWrapper: Adapts environments from the Gym library to be compatible with the Gymnasium library. +- AllinObsWrapper: Consolidates all information into the observation, useful for environments where the agent's + observation should include additional information such as the current score or time remaining. +- ObsPlusPrevActRewWrapper: This wrapper is used in policy NGU. It sets a dict as the new wrapped observation, + which includes the current observation, previous action and previous reward. +""" -from typing import Union, List, Tuple -from easydict import EasyDict -from collections import deque import copy +import operator +from collections import deque +from functools import reduce +from typing import Union, Any, Tuple, Dict, List + import gym +import gymnasium import numpy as np -from torch import float32 +from easydict import EasyDict from ding.torch_utils import to_ndarray from ding.utils import ENV_WRAPPER_REGISTRY, import_module -''' -Env Wrapper List: - - NoopResetWrapper: Sample initial states by taking random number of no-ops on reset. - - MaxAndSkipWrapper: Max pooling across time steps - - WarpFrameWrapper: Warp frames to 84x84 as done in the Nature paper and later work. - - ScaledFloatFrameWrapper: Normalize observations to 0~1. - - ClipRewardWrapper: Clip the reward to {+1, 0, -1} by its sign. - - DelayRewardWrapper: Return cumulative reward at intervals; At other time, return reward of 0. - - FrameStackWrapper: Stack latest n frames(usually 4 in Atari) as one observation. - - ObsTransposeWrapper: Transpose observation to put channel to first dim. - - ObsNormWrapper: Normalize observations according to running mean and std. - - RewardNormWrapper: Normalize reward according to running std. - - RamWrapper: Wrap ram env into image-like env - - EpisodicLifeWrapper: Make end-of-life == end-of-episode, but only reset on true game over. - - FireResetWrapper: Take fire action at environment reset. - - GymHybridDictActionWrapper: Transform Gym-Hybrid's original ``gym.spaces.Tuple`` action space - to ``gym.spaces.Dict``. -''' @ENV_WRAPPER_REGISTRY.register('noop_reset') class NoopResetWrapper(gym.Wrapper): """ Overview: - Sample initial states by taking random number of no-ops on reset. \ - No-op is assumed to be action 0. - Interface: - ``__init__``, ``reset``, ``new_shape`` + Sample initial states by taking random number of no-ops on reset. No-op is assumed to be action 0. + Interfaces: + __init__, reset Properties: - env (:obj:`gym.Env`): the environment to wrap. - noop_max (:obj:`int`): the maximum value of no-ops to run. """ - def __init__(self, env, noop_max=30): + def __init__(self, env: gym.Env, noop_max: int = 30): """ Overview: - Initialize ``self.`` See ``help(type(self))`` for accurate signature. + Initialize the NoopResetWrapper. Arguments: - env (:obj:`gym.Env`): the environment to wrap. - - noop_max (:obj:`int`): the maximum value of no-ops to run. + - noop_max (:obj:`int`): the maximum value of no-ops to run. Defaults to 30. """ super().__init__(env) self.noop_max = noop_max self.noop_action = 0 assert env.unwrapped.get_action_meanings()[0] == 'NOOP' - def reset(self): + def reset(self) -> np.ndarray: """ Overview: - Resets the state of the environment and returns an initial observation. + Resets the state of the environment and returns an initial observation, + after taking a random number of no-ops. Returns: - - observation (:obj:`Any`): the initial observation. + - observation (:obj:`Any`): The initial observation after no-ops. """ self.env.reset() noops = np.random.randint(1, self.noop_max + 1) @@ -77,41 +104,39 @@ def reset(self): class MaxAndSkipWrapper(gym.Wrapper): """ Overview: - Return only every `skip`-th frame (frameskipping) using most \ - recent raw observations (for max pooling across time steps) - Interface: - ``__init__``, ``step``, ``new_shape`` + Wraps the environment to return only every ``skip``-th frame (frameskipping) \ + using most recent raw observations (for max pooling across time steps). + Interfaces: + __init__, step Properties: - - env (:obj:`gym.Env`): the environment to wrap. - - skip (:obj:`int`): number of `skip`-th frame. + - env (:obj:`gym.Env`): The environment to wrap. + - skip (:obj:`int`): Number of ``skip``-th frame. Defaults to 4. """ - def __init__(self, env, skip=4): + def __init__(self, env: gym.Env, skip: int = 4): """ Overview: - Initialize ``self.`` See ``help(type(self))`` for accurate signature. + Initialize the MaxAndSkipWrapper. Arguments: - - env (:obj:`gym.Env`): the environment to wrap. - - skip (:obj:`int`): number of `skip`-th frame. + - env (:obj:`gym.Env`): The environment to wrap. + - skip (:obj:`int`): Number of ``skip``-th frame. Defaults to 4. """ super().__init__(env) self._skip = skip - def step(self, action): + def step(self, action: Union[int, np.ndarray]) -> tuple: """ Overview: - Step the environment with the given action. Repeat action, \ - sum reward, and max over last observations. + Take the given action and repeat it for a specified number of steps. \ + The rewards are summed up and the maximum frame over the last observations is returned. Arguments: - - action (:obj:`Any`): the given action to step with. + - action (:obj:`Any`): The action to repeat. Returns: - - max_frame (:obj:`np.array`) : max over last observations - - total_reward (:obj:`Any`) : amount of reward returned after previous action - - done (:obj:`Bool`) : whether the episode has ended, in which case further step() \ - calls will return undefined results - - info (:obj:`Dict`) : contains auxiliary diagnostic information (helpful for \ + - max_frame (:obj:`np.array`): Max over last observations + - total_reward (:obj:`Any`): Sum of rewards after previous action. + - done (:obj:`Bool`): Whether the episode has ended. + - info (:obj:`Dict`): Contains auxiliary diagnostic information (helpful for \ debugging, and sometimes learning) - """ obs_list, total_reward, done = [], 0., False for i in range(self._skip): @@ -128,21 +153,25 @@ def step(self, action): class WarpFrameWrapper(gym.ObservationWrapper): """ Overview: - Warp frames to 84x84 as done in the Nature paper and later work. - Interface: - ``__init__``, ``observation``, ``new_shape`` + The WarpFrameWrapper class is a gym observation wrapper that resizes + the frame of an environment observation to a specified size (default is 84x84). + This is often used in the preprocessing pipeline of observations in reinforcement learning, + especially for visual observations from Atari environments. + Interfaces: + __init__, observation Properties: - env (:obj:`gym.Env`): the environment to wrap. - - ``size=84``, ``obs_space``, ``self.observation_space`` - + - size (:obj:`int`): the size to which the frames are to be resized. + - observation_space (:obj:`gym.Space`): the observation space of the wrapped environment. """ - def __init__(self, env, size=84): + def __init__(self, env: gym.Env, size: int = 84): """ Overview: - Initialize ``self.`` See ``help(type(self))`` for accurate signature. + Constructor for WarpFrameWrapper class, initializes the environment and the size. Arguments: - env (:obj:`gym.Env`): the environment to wrap. + - size (:obj:`int`): the size to which the frames are to be resized. Default is 84. """ super().__init__(env) self.size = size @@ -162,14 +191,14 @@ def __init__(self, env, size=84): if len(self.observation_space) == 1: self.observation_space = self.observation_space[0] - def observation(self, frame): + def observation(self, frame: np.ndarray) -> np.ndarray: """ Overview: - Returns the current observation from a frame + Resize the frame (observation) to the desired size. Arguments: - - frame (:obj:`Any`): the frame to get observation from + - frame (:obj:`np.ndarray`): the frame to be resized. Returns: - - observation (:obj:`Any`): Framed observation + - frame (:obj:`np.ndarray`): the resized frame. """ try: import cv2 @@ -178,23 +207,31 @@ def observation(self, frame): import sys logging.warning("Please install opencv-python first.") sys.exit(1) - frame = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY) - return cv2.resize(frame, (self.size, self.size), interpolation=cv2.INTER_AREA) + # deal with the `channel_first` case + if frame.shape[0] < 10: + frame = frame.transpose(1, 2, 0) + frame = cv2.resize(frame, (self.size, self.size), interpolation=cv2.INTER_AREA) + frame = frame.transpose(2, 0, 1) + else: + frame = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY) + frame = cv2.resize(frame, (self.size, self.size), interpolation=cv2.INTER_AREA) + + return frame @ENV_WRAPPER_REGISTRY.register('scaled_float_frame') class ScaledFloatFrameWrapper(gym.ObservationWrapper): """ Overview: - Normalize observations to 0~1. - Interface: - ``__init__``, ``observation``, ``new_shape`` + The ScaledFloatFrameWrapper normalizes observations to between 0 and 1. + Interfaces: + __init__, observation """ - def __init__(self, env): + def __init__(self, env: gym.Env): """ Overview: - Initialize ``self.`` See ``help(type(self))`` for accurate signature; setup the properties. + Initialize the ScaledFloatFrameWrapper, setting the scale and bias for normalization. Arguments: - env (:obj:`gym.Env`): the environment to wrap. """ @@ -205,16 +242,15 @@ def __init__(self, env): self.scale = high - low self.observation_space = gym.spaces.Box(low=0., high=1., shape=env.observation_space.shape, dtype=np.float32) - def observation(self, observation): + def observation(self, observation: np.ndarray) -> np.ndarray: """ Overview: - Returns the scaled observation + Scale the observation to be within the range [0, 1]. Arguments: - - observation(:obj:`Float`): The original observation + - observation (:obj:`np.ndarray`): the original observation. Returns: - - observation (:obj:`Float`): The Scaled Float observation + - scaled_observation (:obj:`np.ndarray`): the scaled observation. """ - return ((observation - self.bias) / self.scale).astype('float32') @@ -222,66 +258,158 @@ def observation(self, observation): class ClipRewardWrapper(gym.RewardWrapper): """ Overview: - Clip the reward to {+1, 0, -1} by its sign. - Interface: - ``__init__``, ``reward``, ``new_shape`` + The ClipRewardWrapper class is a gym reward wrapper that clips the reward to {-1, 0, +1} based on its sign. + This can be used to normalize the scale of the rewards in reinforcement learning algorithms. + Interfaces: + __init__, reward Properties: - env (:obj:`gym.Env`): the environment to wrap. - - ``reward_range`` - + - reward_range (:obj:`Tuple[int, int]`): the range of the reward values after clipping. """ - def __init__(self, env): + def __init__(self, env: gym.Env): """ Overview: - Initialize ``self.`` See ``help(type(self))`` for accurate signature; setup the properties. + Initialize the ClipRewardWrapper class. Arguments: - env (:obj:`gym.Env`): the environment to wrap. """ super().__init__(env) self.reward_range = (-1, 1) - def reward(self, reward): + def reward(self, reward: float) -> float: """ Overview: - Bin reward to {+1, 0, -1} by its sign. Note: np.sign(0) == 0. + Clip the reward to {-1, 0, +1} based on its sign. Note: np.sign(0) == 0. Arguments: - - reward(:obj:`Float`): Raw Reward + - reward (:obj:`float`): the original reward. Returns: - - reward(:obj:`Float`): Clipped Reward + - reward (:obj:`float`): the clipped reward. """ return np.sign(reward) +@ENV_WRAPPER_REGISTRY.register('action_repeat') +class ActionRepeatWrapper(gym.Wrapper): + """ + Overview: + The ActionRepeatWrapper class is a gym wrapper that repeats the same action for a number of steps. + This wrapper is particularly useful in environments where the desired effect is achieved by maintaining + the same action across multiple time steps. For instance, some physical environments like motion control + tasks might require consistent force input to produce a significant state change. + + Using this wrapper can reduce the temporal complexity of the problem, as it allows the agent to perform + multiple actions within a single time step. This can speed up learning, as the agent has fewer decisions + to make within a time step. However, it may also sacrifice some level of decision-making precision, as the + agent cannot change its action across successive time steps. + + Note that the use of the ActionRepeatWrapper may not be suitable for all types of environments. Specifically, + it may not be the best choice for environments where new decisions must be made at each time step, or where + the time sequence of actions has a significant impact on the outcome. + Interfaces: + __init__, step + Properties: + - env (:obj:`gym.Env`): the environment to wrap. + - action_repeat (:obj:`int`): the number of times to repeat the action. + """ + + def __init__(self, env: gym.Env, action_repeat: int = 1): + """ + Overview: + Initialize the ActionRepeatWrapper class. + Arguments: + - env (:obj:`gym.Env`): the environment to wrap. + - action_repeat (:obj:`int`): the number of times to repeat the action. Default is 1. + """ + super().__init__(env) + self.action_repeat = action_repeat + + def step(self, action: Union[int, np.ndarray]) -> tuple: + """ + Overview: + Take the given action and repeat it for a specified number of steps. The rewards are summed up. + Arguments: + - action (:obj:`Union[int, np.ndarray]`): The action to repeat. + Returns: + - obs (:obj:`np.ndarray`): The observation after repeating the action. + - reward (:obj:`float`): The sum of rewards after repeating the action. + - done (:obj:`bool`): Whether the episode has ended. + - info (:obj:`Dict`): Contains auxiliary diagnostic information. + """ + reward = 0 + for _ in range(self.action_repeat): + obs, rew, done, info = self.env.step(action) + reward += rew or 0 + if done: + break + return obs, reward, done, info + + @ENV_WRAPPER_REGISTRY.register('delay_reward') class DelayRewardWrapper(gym.Wrapper): """ Overview: - Return cumulative reward at intervals; At other time, return reward of 0. - Interface: - ``__init__``, ``reset``, ``step``, ``new_shape`` + The DelayRewardWrapper class is a gym wrapper that delays the reward. It cumulates the reward over a + predefined number of steps and returns the cumulated reward only at the end of this interval. + At other times, it returns a reward of 0. + + This wrapper is particularly useful in environments where the impact of an action is not immediately + observable, but rather delayed over several steps. For instance, in strategic games or planning tasks, + the effect of an action may not be directly noticeable, but it contributes to a sequence of actions that + leads to a reward. In these cases, delaying the reward to match the action-effect delay can make the + learning process more consistent with the problem's nature. + + However, using this wrapper may increase the difficulty of learning, as the agent needs to associate its + actions with delayed outcomes. It also introduces a non-standard reward structure, which could limit the + applicability of certain reinforcement learning algorithms. + + Note that the use of the DelayRewardWrapper may not be suitable for all types of environments. Specifically, + it may not be the best choice for environments where the effect of actions is immediately observable and the + reward should be assigned accordingly. + Interfaces: + __init__, reset, step Properties: - env (:obj:`gym.Env`): the environment to wrap. - - ``reward_range`` + - delay_reward_step (:obj:`int`): the number of steps over which to delay and cumulate the reward. """ - def __init__(self, env, delay_reward_step=0): + def __init__(self, env: gym.Env, delay_reward_step: int = 0): """ Overview: - Initialize ``self.`` See ``help(type(self))`` for accurate signature; setup the properties. + Initialize the DelayRewardWrapper class. Arguments: - env (:obj:`gym.Env`): the environment to wrap. + - delay_reward_step (:obj:`int`): the number of steps over which to delay and cumulate the reward. """ super().__init__(env) self._delay_reward_step = delay_reward_step - def reset(self): + def reset(self) -> np.ndarray: + """ + Overview: + Resets the state of the environment and resets the delay reward duration and current delay reward. + Returns: + - obs (:obj:`np.ndarray`): the initial observation of the environment. + """ self._delay_reward_duration = 0 self._current_delay_reward = 0. obs = self.env.reset() return obs - def step(self, action): + def step(self, action: Union[int, np.ndarray]) -> tuple: + """ + Overview: + Take the given action and repeat it for a specified number of steps. The rewards are summed up. + If the number of steps equals the delay reward step, return the cumulated reward and reset the + delay reward duration and current delay reward. Otherwise, return a reward of 0. + Arguments: + - action (:obj:`Union[int, np.ndarray]`): the action to take in the step. + Returns: + - obs (:obj:`np.ndarray`): The observation after the step. + - reward (:obj:`float`): The cumulated reward after the delay reward step or 0. + - done (:obj:`bool`): Whether the episode has ended. + - info (:obj:`Dict`): Contains auxiliary diagnostic information. + """ obs, reward, done, info = self.env.step(action) self._current_delay_reward += reward self._delay_reward_duration += 1 @@ -294,58 +422,92 @@ def step(self, action): return obs, reward, done, info -@ENV_WRAPPER_REGISTRY.register('final_eval_reward') -class FinalEvalRewardEnv(gym.Wrapper): +@ENV_WRAPPER_REGISTRY.register('eval_episode_return') +class EvalEpisodeReturnWrapper(gym.Wrapper): """ Overview: - Accumulate rewards at every timestep, and return at the end of the episode in `info`. - Interface: - ``__init__``, ``reset``, ``step``, ``new_shape`` + A wrapper for a gym environment that accumulates rewards at every timestep, and returns the total reward at the + end of the episode in `info`. This is used for evaluation purposes. + Interfaces: + __init__, reset, step Properties: - env (:obj:`gym.Env`): the environment to wrap. """ - def __init__(self, env): + def __init__(self, env: gym.Env): """ Overview: - Initialize ``self.`` See ``help(type(self))`` for accurate signature; setup the properties. + Initialize the EvalEpisodeReturnWrapper. This involves setting up the environment to wrap. Arguments: - - env (:obj:`gym.Env`): the environment to wrap. + - env (:obj:`gym.Env`): The environment to wrap. """ super().__init__(env) - def reset(self): - self._final_eval_reward = 0. + def reset(self) -> np.ndarray: + """ + Overview: + Reset the environment and initialize the accumulated reward to zero. + Returns: + - obs (:obj:`np.ndarray`): The initial observation from the environment. + """ + self._eval_episode_return = 0. return self.env.reset() - def step(self, action): + def step(self, action: Any) -> tuple: + """ + Overview: + Step the environment with the provided action, accumulate the returned reward, and add the total reward to + `info` if the episode is done. + Arguments: + - action (:obj:`Any`): The action to take in the environment. + Returns: + - obs (:obj:`np.ndarray`): The next observation from the environment. + - reward (:obj:`float`): The reward from taking the action. + - done (:obj:`bool`): Whether the episode is done. + - info (:obj:`Dict[str, Any]`): A dictionary of extra information, which includes 'eval_episode_return' if + the episode is done. + Examples: + >>> env = gym.make("CartPole-v1") + >>> env = EvalEpisodeReturnWrapper(env) + >>> obs = env.reset() + >>> done = False + >>> while not done: + ... action = env.action_space.sample() # Replace with your own policy + ... obs, reward, done, info = env.step(action) + ... if done: + ... print("Total episode reward:", info['eval_episode_return']) + """ obs, reward, done, info = self.env.step(action) - self._final_eval_reward += reward + self._eval_episode_return += reward if done: - info['final_eval_reward'] = to_ndarray([self._final_eval_reward], dtype=np.float32) + info['eval_episode_return'] = to_ndarray([self._eval_episode_return], dtype=np.float32) return obs, reward, done, info @ENV_WRAPPER_REGISTRY.register('frame_stack') class FrameStackWrapper(gym.Wrapper): """ - Overview: - Stack latest n frames(usually 4 in Atari) as one observation. - Interface: - ``__init__``, ``reset``, ``step``, ``_get_ob``, ``new_shape`` - Properties: - - env (:obj:`gym.Env`): the environment to wrap. - - n_frame (:obj:`int`): the number of frames to stack. - - ``observation_space``, ``frames`` - """ - - def __init__(self, env, n_frames=4): + Overview: + FrameStackWrapper is a gym environment wrapper that stacks the latest n frames (generally 4 in Atari) + as a single observation. It is commonly used in environments where the observation is an image, + and consecutive frames provide useful temporal information for the agent. + Interfaces: + __init__, reset, step, _get_ob + Properties: + - env (:obj:`gym.Env`): The environment to wrap. + - n_frames (:obj:`int`): The number of frames to stack. + - frames (:obj:`collections.deque`): A queue that holds the most recent frames. + - observation_space (:obj:`gym.Space`): The space of the stacked observations. + """ + + def __init__(self, env: gym.Env, n_frames: int = 4) -> None: """ Overview: - Initialize ``self.`` See ``help(type(self))`` for accurate signature; setup the properties. + Initialize the FrameStackWrapper. This process includes setting up the environment to wrap, + the number of frames to stack, and the observation space. Arguments: - - env (:obj:`gym.Env`): the environment to wrap. - - n_frame (:obj:`int`): the number of frames to stack. + - env (:obj:`gym.Env`): The environment to wrap. + - n_frame (:obj:`int`): The number of frames to stack. """ super().__init__(env) self.n_frames = n_frames @@ -364,42 +526,43 @@ def __init__(self, env, n_frames=4): if len(self.observation_space) == 1: self.observation_space = self.observation_space[0] - def reset(self): + def reset(self) -> np.ndarray: """ Overview: - Resets the state of the environment and append new observation to frames + Reset the environment and initialize frames with the initial observation. Returns: - - ``self._get_ob()``: observation + - init_obs (:obj:`np.ndarray`): The stacked initial observations. """ obs = self.env.reset() for _ in range(self.n_frames): self.frames.append(obs) return self._get_ob() - def step(self, action): + def step(self, action: Any) -> Tuple[np.ndarray, float, bool, Dict[str, Any]]: """ Overview: - Step the environment with the given action. Repeat action, sum reward, \ - and max over last observations, and append new observation to frames + Perform a step in the environment with the given action, append the returned observation + to frames, and return the stacked observations. Arguments: - - action (:obj:`Any`): the given action to step with. + - action (:obj:`Any`): The action to perform a step with. Returns: - - ``self._get_ob()`` : observation - - reward (:obj:`Any`) : amount of reward returned after previous action - - done (:obj:`Bool`) : whether the episode has ended, in which case further \ - step() calls will return undefined results - - info (:obj:`Dict`) : contains auxiliary diagnostic information (helpful \ - for debugging, and sometimes learning) + - self._get_ob() (:obj:`np.ndarray`): The stacked observations. + - reward (:obj:`float`): The amount of reward returned after the previous action. + - done (:obj:`bool`): Whether the episode has ended, in which case further step() calls will return + undefined results. + - info (:obj:`Dict[str, Any]`): Contains auxiliary diagnostic information (helpful for debugging, + and sometimes learning). """ - obs, reward, done, info = self.env.step(action) self.frames.append(obs) return self._get_ob(), reward, done, info - def _get_ob(self): + def _get_ob(self) -> np.ndarray: """ Overview: - The original wrapper use `LazyFrames` but since we use np buffer, it has no effect + The original wrapper used `LazyFrames`, but since we use an np buffer, it has no effect. + Returns: + - stacked_frames (:obj:`np.ndarray`): The stacked frames. """ return np.stack(self.frames, axis=0) @@ -408,21 +571,23 @@ def _get_ob(self): class ObsTransposeWrapper(gym.ObservationWrapper): """ Overview: - Transpose observation to put channel to first dim. - Interface: - ``__init__``, ``observation``, ``new_shape`` + The ObsTransposeWrapper class is a gym wrapper that transposes the observation to put the channel dimension + first. This can be helpful for certain types of neural networks that expect the channel dimension to be + the first dimension. + Interfaces: + __init__, observation Properties: - - env (:obj:`gym.Env`): the environment to wrap. - - ``observation_space`` + - env (:obj:`gym.Env`): The environment to wrap. + - observation_space (:obj:`gym.spaces.Box`): The transformed observation space. """ - def __init__(self, env): + def __init__(self, env: gym.Env): """ Overview: - Initialize ``self.`` See ``help(type(self))`` for accurate signature; \ - setup the properties. + Initialize the ObsTransposeWrapper class and update the observation space according to the environment's + observation space. Arguments: - - env (:obj:`gym.Env`): the environment to wrap. + - env (:obj:`gym.Env`): The environment to wrap. """ super().__init__(env) obs_space = env.observation_space @@ -444,11 +609,12 @@ def __init__(self, env): def observation(self, obs: Union[tuple, np.ndarray]) -> Union[tuple, np.ndarray]: """ Overview: - Returns the transposed observation + Transpose the observation to put the channel dimension first. If the observation is a tuple, each element + in the tuple is transposed independently. Arguments: - - observation (:obj:`Union[tuple, np.ndarray]`): The original observation + - obs (:obj:`Union[tuple, np.ndarray]`): The original observation. Returns: - - observation (:obj:`Union[tuple, np.ndarray]`): The transposed observation + - obs (:obj:`Union[tuple, np.ndarray]`): The transposed observation. """ if isinstance(obs, tuple): new_obs = [] @@ -463,35 +629,40 @@ def observation(self, obs: Union[tuple, np.ndarray]) -> Union[tuple, np.ndarray] class RunningMeanStd(object): """ Overview: - Wrapper to update new variable, new mean, and new count - Interface: - ``__init__``, ``update``, ``reset``, ``new_shape`` + The RunningMeanStd class is a utility that maintains a running mean and standard deviation calculation over + a stream of data. + Interfaces: + __init__, update, reset, mean, std Properties: - - env (:obj:`gym.Env`): the environment to wrap. - - ``mean``, ``std``, ``_epsilon``, ``_shape``, ``_mean``, ``_var``, ``_count`` + - mean (:obj:`np.ndarray`): The running mean. + - std (:obj:`np.ndarray`): The running standard deviation. + - _epsilon (:obj:`float`): A small number to prevent division by zero when calculating standard deviation. + - _shape (:obj:`tuple`): The shape of the data stream. + - _mean (:obj:`np.ndarray`): The current mean of the data stream. + - _var (:obj:`np.ndarray`): The current variance of the data stream. + - _count (:obj:`float`): The number of data points processed. """ - def __init__(self, epsilon=1e-4, shape=()): + def __init__(self, epsilon: float = 1e-4, shape: tuple = ()): """ Overview: - Initialize ``self.`` See ``help(type(self))`` for accurate \ - signature; setup the properties. + Initialize the RunningMeanStd object. Arguments: - - env (:obj:`gym.Env`): the environment to wrap. - - epsilon (:obj:`Float`): the epsilon used for self for the std output - - shape (:obj: `np.array`): the np array shape used for the expression \ - of this wrapper on attibutes of mean and variance + - epsilon (:obj:`float`, optional): A small number to prevent division by zero when calculating standard + deviation. Default is 1e-4. + - shape (:obj:`tuple`, optional): The shape of the data stream. Default is an empty tuple, which + corresponds to scalars. """ self._epsilon = epsilon self._shape = shape self.reset() - def update(self, x): + def update(self, x: np.array): """ Overview: - Update mean, variable, and count + Update the running statistics with a new batch of data. Arguments: - - ``x``: the batch + - x (:obj:`np.array`): A batch of data. """ batch_mean = np.mean(x, axis=0) batch_var = np.var(x, axis=0) @@ -523,7 +694,9 @@ def reset(self): def mean(self) -> np.ndarray: """ Overview: - Property ``mean`` gotten from ``self._mean`` + Get the current running mean. + Returns: + The current running mean. """ return self._mean @@ -531,7 +704,9 @@ def mean(self) -> np.ndarray: def std(self) -> np.ndarray: """ Overview: - Property ``std`` calculated from ``self._var`` and the epsilon value of ``self._epsilon`` + Get the current running standard deviation. + Returns: + The current running mean. """ return np.sqrt(self._var) + self._epsilon @@ -540,20 +715,21 @@ def std(self) -> np.ndarray: class ObsNormWrapper(gym.ObservationWrapper): """ Overview: - Normalize observations according to running mean and std. - Interface: - ``__init__``, ``step``, ``reset``, ``observation``, ``new_shape`` + The ObsNormWrapper class is a gym observation wrapper that normalizes + observations according to running mean and standard deviation (std). + Interfaces: + __init__, step, reset, observation Properties: - env (:obj:`gym.Env`): the environment to wrap. - - - ``data_count``, ``clip_range``, ``rms`` + - data_count (:obj:`int`): the count of data points observed so far. + - clip_range (:obj:`Tuple[int, int]`): the range to clip the normalized observation. + - rms (:obj:`RunningMeanStd`): running mean and standard deviation of the observations. """ - def __init__(self, env): + def __init__(self, env: gym.Env): """ Overview: - Initialize ``self.`` See ``help(type(self))`` for accurate signature; \ - setup the properties according to running mean and std. + Initialize the ObsNormWrapper class. Arguments: - env (:obj:`gym.Env`): the environment to wrap. """ @@ -562,38 +738,33 @@ def __init__(self, env): self.clip_range = (-3, 3) self.rms = RunningMeanStd(shape=env.observation_space.shape) - def step(self, action): + def step(self, action: Union[int, np.ndarray]): """ Overview: - Step the environment with the given action. Repeat action, sum reward, \ - and update ``data_count``, and also update the ``self.rms`` property \ - once after integrating with the input ``action``. + Take an action in the environment, update the running mean and std, + and return the normalized observation. Arguments: - - action (:obj:`Any`): the given action to step with. + - action (:obj:`Union[int, np.ndarray]`): the action to take in the environment. Returns: - - ``self.observation(observation)`` : normalized observation after the \ - input action and updated ``self.rms`` - - reward (:obj:`Any`) : amount of reward returned after previous action - - done (:obj:`Bool`) : whether the episode has ended, in which case further \ - step() calls will return undefined results - - info (:obj:`Dict`) : contains auxiliary diagnostic information (helpful \ - for debugging, and sometimes learning) - + - obs (:obj:`np.ndarray`): the normalized observation after the action. + - reward (:obj:`float`): the reward after the action. + - done (:obj:`bool`): whether the episode has ended. + - info (:obj:`Dict`): contains auxiliary diagnostic information. """ self.data_count += 1 observation, reward, done, info = self.env.step(action) self.rms.update(observation) return self.observation(observation), reward, done, info - def observation(self, observation): + def observation(self, observation: np.ndarray) -> np.ndarray: """ Overview: - Get obeservation + Normalize the observation using the current running mean and std. + If less than 30 data points have been observed, return the original observation. Arguments: - - observation (:obj:`Any`): Original observation + - observation (:obj:`np.ndarray`): the original observation. Returns: - - observation (:obj:`Any`): Normalized new observation - + - observation (:obj:`np.ndarray`): the normalized observation. """ if self.data_count > 30: return np.clip((observation - self.rms.mean) / self.rms.std, self.clip_range[0], self.clip_range[1]) @@ -603,12 +774,11 @@ def observation(self, observation): def reset(self, **kwargs): """ Overview: - Resets the state of the environment and reset properties. + Reset the environment and the properties related to the running mean and std. Arguments: - - kwargs (:obj:`Dict`): Reset with this key argumets + - kwargs (:obj:`Dict`): keyword arguments to be passed to the environment's reset function. Returns: - - observation (:obj:`Any`): New observation after reset - + - observation (:obj:`np.ndarray`): the initial observation of the environment. """ self.data_count = 0 self.rms.reset() @@ -620,39 +790,40 @@ def reset(self, **kwargs): class StaticObsNormWrapper(gym.ObservationWrapper): """ Overview: - Normalize observations according to the mean and std in the fixed dataset. - Interface: - ``__init__``, ``observation`` + The StaticObsNormWrapper class is a gym observation wrapper that normalizes + observations according to a precomputed mean and standard deviation (std) from a fixed dataset. + Interfaces: + __init__, observation Properties: - env (:obj:`gym.Env`): the environment to wrap. - - - ``mean``, ``std``, ``clip_range`` + - mean (:obj:`numpy.ndarray`): the mean of the observations in the fixed dataset. + - std (:obj:`numpy.ndarray`): the standard deviation of the observations in the fixed dataset. + - clip_range (:obj:`Tuple[int, int]`): the range to clip the normalized observation. """ - def __init__(self, env, mean, std): + def __init__(self, env: gym.Env, mean: np.ndarray, std: np.ndarray): """ Overview: - Initialize ``self.`` See ``help(type(self))`` for accurate signature; \ - setup the properties according to dataset mean and std. + Initialize the StaticObsNormWrapper class. Arguments: - env (:obj:`gym.Env`): the environment to wrap. - - mean (:obj:`numpy.ndarray`): the mean of observation in the dataset. - - std (:obj:`numpy.ndarray`): the standard deviation of observation in the dataset. + - mean (:obj:`numpy.ndarray`): the mean of the observations in the fixed dataset. + - std (:obj:`numpy.ndarray`): the standard deviation of the observations in the fixed dataset. """ super().__init__(env) self.mean = mean self.std = std self.clip_range = (-3, 3) - def observation(self, observation): + def observation(self, observation: np.ndarray) -> np.ndarray: """ Overview: - Get obeservation + Normalize the given observation using the precomputed mean and std. + The normalized observation is then clipped within the specified range. Arguments: - - observation (:obj:`Any`): Original observation + - observation (:obj:`np.ndarray`): the original observation. Returns: - - observation (:obj:`Any`): Normalized new observation - + - observation (:obj:`np.ndarray`): the normalized and clipped observation. """ return np.clip((observation - self.mean) / self.std, self.clip_range[0], self.clip_range[1]) @@ -661,21 +832,24 @@ def observation(self, observation): class RewardNormWrapper(gym.RewardWrapper): """ Overview: - Normalize reward according to running std. - Interface: - ``__init__``, ``step``, ``reward``, ``reset``, ``new_shape`` + This wrapper class normalizes the reward according to running std. It extends the `gym.RewardWrapper`. + Interfaces: + __init__, step, reward, reset Properties: - - env (:obj:`gym.Env`): the environment to wrap. - - ``cum_reward``, ``reward_discount``, ``data_count``, ``rms`` + - env (:obj:`gym.Env`): The environment to wrap. + - cum_reward (:obj:`numpy.ndarray`): The cumulated reward, initialized as zero and updated in `step` method. + - reward_discount (:obj:`float`): The discount factor for reward. + - data_count (:obj:`int`): A counter for data, incremented in each `step` call. + - rms (:obj:`RunningMeanStd`): An instance of RunningMeanStd to compute the running mean and std of reward. """ - def __init__(self, env, reward_discount): + def __init__(self, env: gym.Env, reward_discount: float) -> None: """ Overview: - Initialize ``self.`` See ``help(type(self))`` for accurate signature; \ - setup the properties according to running mean and std. + Initialize the RewardNormWrapper, setup the properties according to running mean and std. Arguments: - - env (:obj:`gym.Env`): the environment to wrap. + - env (:obj:`gym.Env`): The environment to wrap. + - reward_discount (:obj:`float`): The discount factor for reward. """ super().__init__(env) self.cum_reward = np.zeros((1, ), 'float64') @@ -683,23 +857,21 @@ def __init__(self, env, reward_discount): self.data_count = 0 self.rms = RunningMeanStd(shape=(1, )) - def step(self, action): + def step(self, action: Any) -> Tuple[np.ndarray, float, bool, Dict]: """ Overview: - Step the environment with the given action. Repeat action, sum reward, \ - and update ``data_count``, and also update the ``self.rms`` and ``self.cum_reward`` \ - properties once after integrating with the input ``action``. + Step the environment with the given action, update properties and return the new observation, reward, + done status and info. Arguments: - - action (:obj:`Any`): the given action to step with. + - action (:obj:`Any`): The action to execute in the environment. Returns: - - observation : normalized observation after the input action and updated ``self.rms`` - - ``self.reward(reward)`` : amount of reward returned after previous action \ - (normalized) and update ``self.cum_reward`` - - done (:obj:`Bool`) : whether the episode has ended, in which case further \ - step() calls will return undefined results - - info (:obj:`Dict`) : contains auxiliary diagnostic information (helpful for \ - debugging, and sometimes learning) - + - observation (:obj:`np.ndarray`): Normalized observation after executing the action and updated `self.rms`. + - reward (:obj:`float`): Amount of reward returned after the action execution (normalized) and updated + `self.cum_reward`. + - done (:obj:`bool`): Whether the episode has ended, in which case further step() calls will return + undefined results. + - info (:obj:`Dict`): Contains auxiliary diagnostic information (helpful for debugging, and sometimes + learning). """ self.data_count += 1 observation, reward, done, info = self.env.step(action) @@ -708,14 +880,14 @@ def step(self, action): self.rms.update(self.cum_reward) return observation, self.reward(reward), done, info - def reward(self, reward): + def reward(self, reward: float) -> float: """ Overview: - Normalize reward if ``data_count`` is more than 30 + Normalize reward if `data_count` is more than 30. Arguments: - - reward(:obj:`Float`): Raw Reward + - reward (:obj:`float`): The raw reward. Returns: - - reward(:obj:`Float`): Normalized Reward + - reward (:obj:`float`): Normalized reward. """ if self.data_count > 30: return float(reward / self.rms.std) @@ -740,21 +912,21 @@ def reset(self, **kwargs): class RamWrapper(gym.Wrapper): """ Overview: - Wrap ram env into image-like env - Interface: - ``__init__``, ``reset``, ``step``, ``new_shape`` + This wrapper class wraps a RAM environment into an image-like environment. It extends the `gym.Wrapper`. + Interfaces: + __init__, reset, step Properties: - - env (:obj:`gym.Env`): the environment to wrap. - - n_frame (:obj:`int`): the number of frames to stack. - - ``observation_space`` + - env (:obj:`gym.Env`): The environment to wrap. + - observation_space (:obj:`gym.spaces.Box`): The observation space of the wrapped environment. """ - def __init__(self, env, render=False): + def __init__(self, env: gym.Env, render: bool = False) -> None: """ Overview: - Initialize ``self.`` See ``help(type(self))`` for accurate signature; + Initialize the RamWrapper and set up the observation space to wrap the RAM environment. Arguments: - - env (:obj:`gym.Env`): the environment to wrap. + - env (:obj:`gym.Env`): The environment to wrap. + - render (:obj:`bool`): Whether to render the environment, default is False. """ super().__init__(env) shape = env.observation_space.shape + (1, 1) @@ -765,34 +937,30 @@ def __init__(self, env, render=False): dtype=np.float32 ) - def reset(self): + def reset(self) -> np.ndarray: """ Overview: - Resets the state of the environment and reset properties. - + Resets the state of the environment and returns a reshaped observation. Returns: - - observation (:obj:`Any`): New observation after reset and reshaped - + - observation (:obj:`np.ndarray`): New observation after reset and reshaped. """ obs = self.env.reset() return obs.reshape(128, 1, 1).astype(np.float32) - def step(self, action): + def step(self, action: Any) -> Tuple[np.ndarray, Any, bool, Dict]: """ Overview: - Step the environment with the given action. Repeat action, sum reward and \ - reshape the observation. + Execute one step within the environment with the given action. Repeat action, sum reward and reshape the + observation. Arguments: - - action (:obj:`Any`): the given action to step with. + - action (:obj:`Any`): The action to take in the environment. Returns: - - ``obs.reshape(128, 1, 1).astype(np.float32)`` : reshaped observation after \ - step with type restriction. - - reward (:obj:`Any`) : amount of reward returned after previous action - - done (:obj:`Bool`) : whether the episode has ended, in which case further \ - step() calls will return undefined results - - info (:obj:`Dict`) : contains auxiliary diagnostic information (helpful for \ - debugging, and sometimes learning) - + - observation (:obj:`np.ndarray`): Reshaped observation after step with type restriction. + - reward (:obj:`Any`): Amount of reward returned after previous action. + - done (:obj:`bool`): Whether the episode has ended, in which case further step() calls will return + undefined results. + - info (:obj:`Dict`): Contains auxiliary diagnostic information (helpful for debugging, and sometimes + learning). """ obs, reward, done, info = self.env.step(action) return obs.reshape(128, 1, 1).astype(np.float32), reward, done, info @@ -802,45 +970,41 @@ def step(self, action): class EpisodicLifeWrapper(gym.Wrapper): """ Overview: - Make end-of-life == end-of-episode, but only reset on true game over. It helps \ - the value estimation. - Interface: - ``__init__``, ``step``, ``reset``, ``observation``, ``new_shape`` + This wrapper makes end-of-life equivalent to end-of-episode, but only resets on + true game over. This helps in better value estimation. + Interfaces: + __init__, step, reset Properties: - - env (:obj:`gym.Env`): the environment to wrap. - - - ``lives``, ``was_real_done`` + - env (:obj:`gym.Env`): The environment to wrap. + - lives (:obj:`int`): The current number of lives. + - was_real_done (:obj:`bool`): Whether the last episode was ended due to game over. """ - def __init__(self, env): + def __init__(self, env: gym.Env) -> None: """ Overview: - Initialize ``self.`` See ``help(type(self))`` for accurate signature; set \ - lives to 0 at set done. + Initialize the EpisodicLifeWrapper, setting lives to 0 and was_real_done to True. Arguments: - - env (:obj:`gym.Env`): the environment to wrap. + - env (:obj:`gym.Env`): The environment to wrap. """ super().__init__(env) self.lives = 0 self.was_real_done = True - def step(self, action): + def step(self, action: Any) -> Tuple[np.ndarray, float, bool, Dict]: """ Overview: - Step the environment with the given action. Repeat action, sum reward; set \ - ``self.was_real_done`` as done, and step according to lives i.e. check \ - current lives, make loss of life terminal, then update lives to \ - handle bonus lives. + Execute the given action in the environment, update properties based on the new + state and return the new observation, reward, done status and info. Arguments: - - action (:obj:`Any`): the given action to step with. + - action (:obj:`Any`): The action to execute in the environment. Returns: - - obs (:obj:`Any`): normalized observation after the input action and updated ``self.rms`` - - reward (:obj:`Any`) : amount of reward returned after previous action - - done (:obj:`Bool`) : whether the episode has ended, in which case further step() \ - calls will return undefined results - - info (:obj:`Dict`) : contains auxiliary diagnostic information (helpful for debugging,\ - and sometimes learning) - + - observation (:obj:`np.ndarray`): Normalized observation after the action execution and updated `self.rms`. + - reward (:obj:`float`): Amount of reward returned after the action execution. + - done (:obj:`bool`): Whether the episode has ended, in which case further step() calls will return + undefined results. + - info (:obj:`Dict`): Contains auxiliary diagnostic information (helpful for debugging, and + sometimes learning). """ obs, reward, done, info = self.env.step(action) self.was_real_done = done @@ -855,16 +1019,15 @@ def step(self, action): self.lives = lives return obs, reward, done, info - def reset(self): + def reset(self) -> np.ndarray: """ Overview: - Calls the Gym environment reset, only when lives are exhausted. This way all states are \ - still reachable even though lives are episodic, and the learner need not know about \ - any of this behind-the-scenes. + Resets the state of the environment and updates the number of lives, only when + lives are exhausted. This way all states are still reachable even though lives + are episodic, and the learner need not know about any of this behind-the-scenes. Returns: - - obs (:obj:`Any`): New observation after reset with no-op step to advance from terminal/lost \ - life state in case of not ``self.was_real_done``. - + - observation (:obj:`np.ndarray`): New observation after reset with no-op step to advance from + terminal/lost life state. """ if self.was_real_done: obs = self.env.reset() @@ -879,29 +1042,32 @@ def reset(self): class FireResetWrapper(gym.Wrapper): """ Overview: - Take fire action at environment reset. + This wrapper takes a fire action at environment reset. Related discussion: https://github.com/openai/baselines/issues/240 - Interface: - ``__init__``, ``reset``, ``new_shape`` + Interfaces: + __init__, reset Properties: - - env (:obj:`gym.Env`): the environment to wrap. + - env (:obj:`gym.Env`): The environment to wrap. """ - def __init__(self, env): + def __init__(self, env: gym.Env) -> None: """ Overview: - Initialize ``self.`` See ``help(type(self))`` for accurate signature. + Initialize the FireResetWrapper. Assume that the second action of the environment + is 'FIRE' and there are at least three actions. Arguments: - - env (:obj:`gym.Env`): the environment to wrap. + - env (:obj:`gym.Env`): The environment to wrap. """ super().__init__(env) assert env.unwrapped.get_action_meanings()[1] == 'FIRE' assert len(env.unwrapped.get_action_meanings()) >= 3 - def reset(self): + def reset(self) -> np.ndarray: """ Overview: - Resets the state of the environment and reset properties i.e. reset with action 1 + Resets the state of the environment and executes a fire action, i.e. reset with action 1. + Returns: + - observation (:obj:`np.ndarray`): New observation after reset and fire action. """ self.env.reset() return self.env.step(1)[0] @@ -911,20 +1077,20 @@ def reset(self): class GymHybridDictActionWrapper(gym.ActionWrapper): """ Overview: - Transform Gym-Hybrid's original ``gym.spaces.Tuple`` action space to ``gym.spaces.Dict``. - Interface: - ``__init__``, ``action`` + Transform Gym-Hybrid's original `gym.spaces.Tuple` action space to `gym.spaces.Dict`. + Interfaces: + __init__, action Properties: - - env (:obj:`gym.Env`): the environment to wrap. - - ``self.action_space`` + - env (:obj:`gym.Env`): The environment to wrap. + - action_space (:obj:`gym.spaces.Dict`): The new action space. """ - def __init__(self, env): + def __init__(self, env: gym.Env) -> None: """ Overview: - Initialize ``self.`` See ``help(type(self))`` for accurate signature. + Initialize the GymHybridDictActionWrapper, setting up the new action space. Arguments: - - env (:obj:`gym.Env`): the environment to wrap. + - env (:obj:`gym.Env`): The environment to wrap. """ super().__init__(env) self.action_space = gym.spaces.Dict( @@ -941,7 +1107,22 @@ def __init__(self, env): } ) - def step(self, action): + def step(self, action: Dict) -> Tuple[Dict, float, bool, Dict]: + """ + Overview: + Execute the given action in the environment, transform the action from Dict to Tuple, + and return the new observation, reward, done status and info. + Arguments: + - action (:obj:`Dict`): The action to execute in the environment, structured as a dictionary. + Returns: + - observation (:obj:`Dict`): The wrapped observation, which includes the current observation, + previous action and previous reward. + - reward (:obj:`float`): Amount of reward returned after the action execution. + - done (:obj:`bool`): Whether the episode has ended, in which case further step() calls will return + undefined results. + - info (:obj:`Dict`): Contains auxiliary diagnostic information (helpful for debugging, and + sometimes learning). + """ # # From Dict to Tuple # action_type = action[0] # if action_type == 0: @@ -958,40 +1139,27 @@ def step(self, action): action_type, action_mask, action_args = action['type'], action['mask'], action['args'] return self.env.step((action_type, action_args)) - # @staticmethod - # def new_shape(obs_shape, act_shape, rew_shape, size=84): - # """ - # Overview: - # Get new shape of observation, acton, and reward; in this case only \ - # observation space wrapped to (4, 84, 84); others unchanged. - # Arguments: - # obs_shape (:obj:`Any`), act_shape (:obj:`Any`), rew_shape (:obj:`Any`) - # Returns: - # obs_shape (:obj:`Any`), act_shape (:obj:`Any`), rew_shape (:obj:`Any`) - # """ - # return (size, size), act_shape, rew_shape - @ENV_WRAPPER_REGISTRY.register('obs_plus_prev_action_reward') class ObsPlusPrevActRewWrapper(gym.Wrapper): """ Overview: - This wrapper is used in policy NGU. - Set a dict {'obs': obs, 'prev_action': self.prev_action, 'prev_reward_extrinsic': self.prev_reward_extrinsic} - as the new wrapped observation, - which including the current obs, previous action and previous reward. - Interface: - ``__init__``, ``reset``, ``step`` + This wrapper is used in policy NGU. It sets a dict as the new wrapped observation, + which includes the current observation, previous action and previous reward. + Interfaces: + __init__, reset, step Properties: - - env (:obj:`gym.Env`): the environment to wrap. + - env (:obj:`gym.Env`): The environment to wrap. + - prev_action (:obj:`int`): The previous action. + - prev_reward_extrinsic (:obj:`float`): The previous reward. """ - def __init__(self, env): + def __init__(self, env: gym.Env) -> None: """ Overview: - Initialize ``self.`` See ``help(type(self))`` for accurate signature; setup the properties. + Initialize the ObsPlusPrevActRewWrapper, setting up the previous action and reward. Arguments: - - env (:obj:`gym.Env`): the environment to wrap. + - env (:obj:`gym.Env`): The environment to wrap. """ super().__init__(env) self.observation_space = gym.spaces.Dict( @@ -1006,35 +1174,35 @@ def __init__(self, env): self.prev_action = -1 # null action self.prev_reward_extrinsic = 0 # null reward - def reset(self): + def reset(self) -> Dict: """ Overview: - Resets the state of the environment. + Resets the state of the environment, and returns the wrapped observation. Returns: - - obs (:obj:`Dict`) : the wrapped observation, which including the current obs, \ + - observation (:obj:`Dict`): The wrapped observation, which includes the current observation, previous action and previous reward. """ obs = self.env.reset() obs = {'obs': obs, 'prev_action': self.prev_action, 'prev_reward_extrinsic': self.prev_reward_extrinsic} return obs - def step(self, action): + def step(self, action: Any) -> Tuple[Dict, float, bool, Dict]: """ Overview: - Step the environment with the given action. - Save the previous action and reward to be used in next new obs + Execute the given action in the environment, save the previous action and reward + to be used in the next observation, and return the new observation, reward, + done status and info. Arguments: - - action (:obj:`Any`): the given action to step with. + - action (:obj:`Any`): The action to execute in the environment. Returns: - - obs (:obj:`Dict`) : the wrapped observation, which including the current obs, \ + - observation (:obj:`Dict`): The wrapped observation, which includes the current observation, previous action and previous reward. - - reward (:obj:`Any`) : amount of reward returned after previous action - - done (:obj:`Bool`) : whether the episode has ended, in which case further \ - step() calls will return undefined results - - info (:obj:`Dict`) : contains auxiliary diagnostic information (helpful \ - for debugging, and sometimes learning) + - reward (:obj:`float`): Amount of reward returned after the action execution. + - done (:obj:`bool`): Whether the episode has ended, in which case further step() calls will return + undefined results. + - info (:obj:`Dict`): Contains auxiliary diagnostic information (helpful for debugging, and sometimes + learning). """ - obs, reward, done, info = self.env.step(action) obs = {'obs': obs, 'prev_action': self.prev_action, 'prev_reward_extrinsic': self.prev_reward_extrinsic} self.prev_action = action @@ -1042,14 +1210,359 @@ def step(self, action): return obs, reward, done, info -def update_shape(obs_shape, act_shape, rew_shape, wrapper_names): +class TransposeWrapper(gym.Wrapper): + """ + Overview: + This class is used to transpose the observation space of the environment. + + Interfaces: + __init__, _process_obs, step, reset + """ + + def __init__(self, env: gym.Env) -> None: + """ + Overview: + Initialize the TransposeWrapper, setting up the new observation space. + Arguments: + - env (:obj:`gym.Env`): The environment to wrap. + """ + super().__init__(env) + old_space = copy.deepcopy(env.observation_space) + new_shape = (old_space.shape[-1], *old_space.shape[:-1]) + self._observation_space = gym.spaces.Box( + low=old_space.low.min(), high=old_space.high.max(), shape=new_shape, dtype=old_space.dtype + ) + + def _process_obs(self, obs: np.ndarray) -> np.ndarray: + """ + Overview: + Transpose the observation into the format (channels, height, width). + Arguments: + - obs (:obj:`np.ndarray`): The observation to transform. + Returns: + - obs (:obj:`np.ndarray`): The transposed observation. + """ + obs = to_ndarray(obs) + obs = np.transpose(obs, (2, 0, 1)) + return obs + + def step(self, action: Any) -> Tuple[np.ndarray, float, bool, Dict]: + """ + Overview: + Execute the given action in the environment, process the observation and return + the new observation, reward, done status, and info. + Arguments: + - action (:obj:`Any`): The action to execute in the environment. + Returns: + - observation (:obj:`np.ndarray`): The processed observation after the action execution. + - reward (:obj:`float`): Amount of reward returned after the action execution. + - done (:obj:`bool`): Whether the episode has ended, in which case further step() calls will return + undefined results. + - info (:obj:`Dict`): Contains auxiliary diagnostic information (helpful for debugging, and sometimes + learning). + """ + obs, reward, done, info = self.env.step(action) + return self._process_obs(obs), reward, done, info + + def reset(self) -> np.ndarray: + """ + Overview: + Resets the state of the environment and returns the processed observation. + Returns: + - observation (:obj:`np.ndarray`): The processed observation after reset. + """ + obs = self.env.reset() + return self._process_obs(obs) + + +class TimeLimitWrapper(gym.Wrapper): + """ + Overview: + This class is used to enforce a time limit on the environment. + Interfaces: + __init__, reset, step + """ + + def __init__(self, env: gym.Env, max_limit: int) -> None: + """ + Overview: + Initialize the TimeLimitWrapper, setting up the maximum limit of time steps. + Arguments: + - env (:obj:`gym.Env`): The environment to wrap. + - max_limit (:obj:`int`): The maximum limit of time steps. + """ + super().__init__(env) + self.max_limit = max_limit + + def reset(self) -> np.ndarray: + """ + Overview: + Resets the state of the environment and the time counter. + Returns: + - observation (:obj:`np.ndarray`): The new observation after reset. + """ + self.time_count = 0 + return self.env.reset() + + def step(self, action: Any) -> Tuple[np.ndarray, float, bool, Dict]: + """ + Overview: + Execute the given action in the environment, update the time counter, and + return the new observation, reward, done status and info. + Arguments: + - action (:obj:`Any`): The action to execute in the environment. + Returns: + - observation (:obj:`np.ndarray`): The new observation after the action execution. + - reward (:obj:`float`): Amount of reward returned after the action execution. + - done (:obj:`bool`): Whether the episode has ended, in which case further step() calls will return + undefined results. + - info (:obj:`Dict`): Contains auxiliary diagnostic information (helpful for debugging, and sometimes + learning). + """ + obs, reward, done, info = self.env.step(action) + self.time_count += 1 + if self.time_count >= self.max_limit: + done = True + info['time_limit'] = True + else: + info['time_limit'] = False + info['time_count'] = self.time_count + return obs, reward, done, info + + +class FlatObsWrapper(gym.Wrapper): + """ + Overview: + This class is used to flatten the observation space of the environment. + Note: only suitable for environments like minigrid. + Interfaces: + __init__, observation, reset, step + """ + + def __init__(self, env: gym.Env, maxStrLen: int = 96) -> None: + """ + Overview: + Initialize the FlatObsWrapper, setup the new observation space. + Arguments: + - env (:obj:`gym.Env`): The environment to wrap. + - maxStrLen (:obj:`int`): The maximum length of mission string, default is 96. + """ + super().__init__(env) + + self.maxStrLen = maxStrLen + self.numCharCodes = 28 + + imgSpace = env.observation_space.spaces["image"] + imgSize = reduce(operator.mul, imgSpace.shape, 1) + + self.observation_space = gym.spaces.Box( + low=0, + high=255, + shape=(imgSize + self.numCharCodes * self.maxStrLen, ), + dtype="float32", + ) + + self.cachedStr: str = None + + def observation(self, obs: Union[np.ndarray, Tuple]) -> np.ndarray: + """ + Overview: + Process the observation, convert the mission into one-hot encoding and concatenate + it with the image data. + Arguments: + - obs (:obj:`Union[np.ndarray, Tuple]`): The raw observation to process. + Returns: + - obs (:obj:`np.ndarray`): The processed observation. + """ + if isinstance(obs, tuple): # for compatibility of gymnasium + obs = obs[0] + image = obs["image"] + mission = obs["mission"] + + # Cache the last-encoded mission string + if mission != self.cachedStr: + assert (len(mission) <= self.maxStrLen), f"mission string too long ({len(mission)} chars)" + mission = mission.lower() + + strArray = np.zeros(shape=(self.maxStrLen, self.numCharCodes), dtype="float32") + + for idx, ch in enumerate(mission): + if ch >= "a" and ch <= "z": + chNo = ord(ch) - ord("a") + elif ch == " ": + chNo = ord("z") - ord("a") + 1 + elif ch == ",": + chNo = ord("z") - ord("a") + 2 + else: + raise ValueError(f"Character {ch} is not available in mission string.") + assert chNo < self.numCharCodes, "%s : %d" % (ch, chNo) + strArray[idx, chNo] = 1 + + self.cachedStr = mission + self.cachedArray = strArray + + obs = np.concatenate((image.flatten(), self.cachedArray.flatten())) + + return obs + + def reset(self, *args, **kwargs) -> np.ndarray: + """ + Overview: + Resets the state of the environment and returns the processed observation. + Returns: + - observation (:obj:`np.ndarray`): The processed observation after reset. + """ + obs = self.env.reset(*args, **kwargs) + return self.observation(obs) + + def step(self, *args, **kwargs) -> Tuple[np.ndarray, float, bool, Dict]: + """ + Overview: + Execute the given action in the environment, and return the processed observation, + reward, done status, and info. + Returns: + - observation (:obj:`np.ndarray`): The processed observation after the action execution. + - reward (:obj:`float`): Amount of reward returned after the action execution. + - done (:obj:`bool`): Whether the episode has ended, in which case further step() calls will return + undefined results. + - info (:obj:`Dict`): Contains auxiliary diagnostic information (helpful for debugging, and sometimes + learning). + """ + o, r, d, i = self.env.step(*args, **kwargs) + o = self.observation(o) + return o, r, d, i + + +class GymToGymnasiumWrapper(gym.Wrapper): """ Overview: - Get new shape of observation, acton, and reward given the wrapper. + This class is used to wrap a gymnasium environment to a gym environment. + Interfaces: + __init__, seed, reset, step + """ + + def __init__(self, env: gymnasium.Env) -> None: + """ + Overview: + Initialize the GymToGymnasiumWrapper. + Arguments: + - env (:obj:`gymnasium.Env`): The gymnasium environment to wrap. + """ + assert isinstance(env, gymnasium.Env), type(env) + super().__init__(env) + self._seed = None + + def seed(self, seed: int) -> None: + """ + Overview: + Set the seed for the environment. + Arguments: + - seed (:obj:`int`): The seed to set. + """ + self._seed = seed + + def reset(self) -> np.ndarray: + """ + Overview: + Resets the state of the environment and returns the new observation. If a seed + was set, use it in the reset. + Returns: + - observation (:obj:`np.ndarray`): The new observation after reset. + """ + if self.seed is not None: + obs, info = self.env.reset(seed=self._seed) + else: + obs, info = self.env.reset() + return obs + + def step(self, *args, **kwargs): + """ + Overview: + Execute the given action in the environment, and return the new observation, + reward, done status, and info. To keep consistency with gym, the done status should be the either \ + terminated=True or truncated=True. + """ + obs, rew, terminated, truncated, info = self.env.step(*args, **kwargs) + return obs, rew, terminated or truncated, info + + +@ENV_WRAPPER_REGISTRY.register('reward_in_obs') +class AllinObsWrapper(gym.Wrapper): + """ + Overview: + This wrapper is used in policy ``Decision Transformer``, which is proposed in paper + https://arxiv.org/abs/2106.01345. It sets a dict {'obs': obs, 'reward': reward} + as the new wrapped observation, which includes the current observation and previous reward. + Interfaces: + __init__, reset, step, seed + Properties: + - env (:obj:`gym.Env`): The environment to wrap. + """ + + def __init__(self, env: gym.Env) -> None: + """ + Overview: + Initialize the AllinObsWrapper. + Arguments: + - env (:obj:`gym.Env`): The environment to wrap. + """ + super().__init__(env) + + def reset(self) -> Dict: + """ + Overview: + Resets the state of the environment and returns the new observation. + Returns: + - observation (:obj:`Dict`): The new observation after reset, includes the current observation and reward. + """ + ret = {'obs': self.env.reset(), 'reward': np.array([0])} + self._observation_space = gym.spaces.Dict( + { + 'obs': self.env.observation_space, + 'reward': gym.spaces.Box(low=-np.inf, high=np.inf, dtype=np.float32, shape=(1, )) + } + ) + return ret + + def step(self, action: Any): + """ + Overview: + Execute the given action in the environment, and return the new observation, + reward, done status, and info. + Arguments: + - action (:obj:`Any`): The action to execute in the environment. + Returns: + - timestep (:obj:`BaseEnvTimestep`): The timestep after the action execution. + """ + obs, reward, done, info = self.env.step(action) + obs = {'obs': obs, 'reward': reward} + from ding.envs import BaseEnvTimestep + return BaseEnvTimestep(obs, reward, done, info) + + def seed(self, seed: int, dynamic_seed: bool = True) -> None: + """ + Overview: + Set the seed for the environment. + Arguments: + - seed (:obj:`int`): The seed to set. + - dynamic_seed (:obj:`bool`): Whether to use dynamic seed, default is True. + """ + self.env.seed(seed, dynamic_seed) + + +def update_shape(obs_shape: Any, act_shape: Any, rew_shape: Any, wrapper_names: List[str]) -> Tuple[Any, Any, Any]: + """ + Overview: + Get new shapes of observation, action, and reward given the wrapper. Arguments: - obs_shape (:obj:`Any`), act_shape (:obj:`Any`), rew_shape (:obj:`Any`), wrapper_names (:obj:`Any`) + - obs_shape (:obj:`Any`): The original shape of observation. + - act_shape (:obj:`Any`): The original shape of action. + - rew_shape (:obj:`Any`): The original shape of reward. + - wrapper_names (:obj:`List[str]`): The names of the wrappers. Returns: - obs_shape (:obj:`Any`), act_shape (:obj:`Any`), rew_shape (:obj:`Any`) + - obs_shape (:obj:`Any`): The new shape of observation. + - act_shape (:obj:`Any`): The new shape of action. + - rew_shape (:obj:`Any`): The new shape of reward. """ for wrapper_name in wrapper_names: if wrapper_name: @@ -1060,32 +1573,18 @@ def update_shape(obs_shape, act_shape, rew_shape, wrapper_names): return obs_shape, act_shape, rew_shape -def create_env_wrapper(env: gym.Env, env_wrapper_cfg: dict) -> gym.Wrapper: - r""" +def create_env_wrapper(env: gym.Env, env_wrapper_cfg: EasyDict) -> gym.Wrapper: + """ Overview: - Create an env wrapper according to env_wrapper_cfg and env instance. + Create an environment wrapper according to the environment wrapper configuration and the environment instance. Arguments: - - env (:obj:`gym.Env`): An env instance to be wrapped. - - env_wrapper_cfg (:obj:`EasyDict`): Env wrapper config. - ArgumentsKeys: - - `env_wrapper_cfg`'s necessary: `type` - - `env_wrapper_cfg`'s optional: `import_names`, `kwargs` + - env (:obj:`gym.Env`): The environment instance to be wrapped. + - env_wrapper_cfg (:obj:`EasyDict`): The configuration for the environment wrapper. + Returns: + - env (:obj:`gym.Wrapper`): The wrapped environment instance. """ env_wrapper_cfg = copy.deepcopy(env_wrapper_cfg) if 'import_names' in env_wrapper_cfg: import_module(env_wrapper_cfg.pop('import_names')) env_wrapper_type = env_wrapper_cfg.pop('type') return ENV_WRAPPER_REGISTRY.build(env_wrapper_type, env, **env_wrapper_cfg.get('kwargs', {})) - - -def get_env_wrapper_cls(cfg: EasyDict) -> type: - r""" - Overview: - Get an env wrapper class according to cfg. - Arguments: - - cfg (:obj:`EasyDict`): Env wrapper config. - ArgumentsKeys: - - necessary: `type` - """ - import_module(cfg.get('import_names', [])) - return ENV_WRAPPER_REGISTRY.get(cfg.type) diff --git a/ding/envs/gym_env.py b/ding/envs/gym_env.py new file mode 100644 index 0000000000..b0f3e34dd2 --- /dev/null +++ b/ding/envs/gym_env.py @@ -0,0 +1,6 @@ +from ding.envs import BaseEnv, DingEnvWrapper + + +def env(cfg, seed_api=True, caller='collector', **kwargs) -> BaseEnv: + import gym + return DingEnvWrapper(gym.make(cfg.env_id, **kwargs), cfg=cfg, seed_api=seed_api, caller=caller) diff --git a/dizoo/board_games/go/__init__.py b/ding/example/__init__.py similarity index 100% rename from dizoo/board_games/go/__init__.py rename to ding/example/__init__.py diff --git a/ding/example/bcq.py b/ding/example/bcq.py new file mode 100755 index 0000000000..d0114d120e --- /dev/null +++ b/ding/example/bcq.py @@ -0,0 +1,42 @@ +import gym +from ditk import logging +from ding.model import BCQ +from ding.policy import BCQPolicy +from ding.envs import DingEnvWrapper, BaseEnvManagerV2 +from ding.data import create_dataset +from ding.config import compile_config +from ding.framework import task, ding_init +from ding.framework.context import OfflineRLContext +from ding.framework.middleware import interaction_evaluator, trainer, CkptSaver, offline_data_fetcher, offline_logger +from ding.utils import set_pkg_seed +from dizoo.d4rl.envs import D4RLEnv +from dizoo.d4rl.config.halfcheetah_medium_bcq_config import main_config, create_config + + +def main(): + # If you don't have offline data, you need to prepare if first and set the data_path in config + # For demonstration, we also can train a RL policy (e.g. SAC) and collect some data + logging.getLogger().setLevel(logging.INFO) + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + ding_init(cfg) + with task.start(async_mode=False, ctx=OfflineRLContext()): + evaluator_env = BaseEnvManagerV2( + env_fn=[lambda: D4RLEnv(cfg.env) for _ in range(cfg.env.evaluator_env_num)], cfg=cfg.env.manager + ) + + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + + dataset = create_dataset(cfg) + model = BCQ(**cfg.policy.model) + policy = BCQPolicy(cfg.policy, model=model) + + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) + task.use(offline_data_fetcher(cfg, dataset)) + task.use(trainer(cfg, policy.learn_mode)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=10000000)) + task.use(offline_logger()) + task.run() + + +if __name__ == "__main__": + main() diff --git a/ding/example/c51_nstep.py b/ding/example/c51_nstep.py index 0c975957e1..2b98ece213 100644 --- a/ding/example/c51_nstep.py +++ b/ding/example/c51_nstep.py @@ -40,7 +40,7 @@ def main(): task.use(nstep_reward_enhancer(cfg)) task.use(data_pusher(cfg, buffer_)) task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) - task.use(CkptSaver(cfg, policy, train_freq=100)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) task.run() diff --git a/ding/example/collect_demo_data.py b/ding/example/collect_demo_data.py index da626934cc..53e37b928c 100644 --- a/ding/example/collect_demo_data.py +++ b/ding/example/collect_demo_data.py @@ -1,7 +1,7 @@ import gym from ditk import logging import torch -from ding.model import QAC +from ding.model import ContinuousQAC from ding.policy import SACPolicy from ding.envs import DingEnvWrapper, BaseEnvManagerV2 from ding.data import offline_data_save_type @@ -22,13 +22,13 @@ def main(): set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) - model = QAC(**cfg.policy.model) + model = ContinuousQAC(**cfg.policy.model) policy = SACPolicy(cfg.policy, model=model, enable_field=['collect']) state_dict = torch.load(cfg.policy.collect.state_dict_path, map_location='cpu') policy.collect_mode.load_state_dict(state_dict) task.use(StepCollector(cfg, policy.collect_mode, collector_env)) - task.use(offline_data_saver(cfg, cfg.policy.collect.save_path, data_type='hdf5')) + task.use(offline_data_saver(cfg.policy.collect.save_path, data_type='hdf5')) task.run(max_step=1) diff --git a/ding/example/cql.py b/ding/example/cql.py index 5651121d8d..5af78dabd3 100644 --- a/ding/example/cql.py +++ b/ding/example/cql.py @@ -5,9 +5,9 @@ from ding.envs import DingEnvWrapper, BaseEnvManagerV2 from ding.data import create_dataset from ding.config import compile_config -from ding.framework import task +from ding.framework import task, ding_init from ding.framework.context import OfflineRLContext -from ding.framework.middleware import interaction_evaluator, trainer, CkptSaver, offline_data_fetcher +from ding.framework.middleware import interaction_evaluator, trainer, CkptSaver, offline_data_fetcher, offline_logger from ding.utils import set_pkg_seed from dizoo.classic_control.pendulum.envs.pendulum_env import PendulumEnv from dizoo.classic_control.pendulum.config.pendulum_cql_config import main_config, create_config @@ -18,6 +18,7 @@ def main(): # For demostration, we also can train a RL policy (e.g. SAC) and collect some data logging.getLogger().setLevel(logging.INFO) cfg = compile_config(main_config, create_cfg=create_config, auto=True) + ding_init(cfg) with task.start(async_mode=False, ctx=OfflineRLContext()): evaluator_env = BaseEnvManagerV2( env_fn=[lambda: PendulumEnv(cfg.env) for _ in range(cfg.env.evaluator_env_num)], cfg=cfg.env.manager @@ -32,7 +33,8 @@ def main(): task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) task.use(offline_data_fetcher(cfg, dataset)) task.use(trainer(cfg, policy.learn_mode)) - task.use(CkptSaver(cfg, policy, train_freq=100)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) + task.use(offline_logger()) task.run() diff --git a/ding/example/d4pg.py b/ding/example/d4pg.py index e004c0bed0..39806f166d 100644 --- a/ding/example/d4pg.py +++ b/ding/example/d4pg.py @@ -40,7 +40,7 @@ def main(): task.use(nstep_reward_enhancer(cfg)) task.use(data_pusher(cfg, buffer_)) task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) - task.use(CkptSaver(cfg, policy, train_freq=100)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) task.run() diff --git a/ding/example/ddpg.py b/ding/example/ddpg.py index 9caa9b22f2..3fa9c18db2 100644 --- a/ding/example/ddpg.py +++ b/ding/example/ddpg.py @@ -1,6 +1,6 @@ import gym from ditk import logging -from ding.model.template.qac import QAC +from ding.model.template.qac import ContinuousQAC from ding.policy import DDPGPolicy from ding.envs import DingEnvWrapper, BaseEnvManagerV2 from ding.data import DequeBuffer @@ -27,7 +27,7 @@ def main(): set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) - model = QAC(**cfg.policy.model) + model = ContinuousQAC(**cfg.policy.model) buffer_ = DequeBuffer(size=cfg.policy.other.replay_buffer.replay_buffer_size) policy = DDPGPolicy(cfg.policy, model=model) @@ -37,7 +37,7 @@ def main(): ) task.use(data_pusher(cfg, buffer_)) task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) - task.use(CkptSaver(cfg, policy, train_freq=100)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) task.use(termination_checker(max_train_iter=10000)) task.run() diff --git a/ding/example/dqn.py b/ding/example/dqn.py index 27947914d5..0959b3ab22 100644 --- a/ding/example/dqn.py +++ b/ding/example/dqn.py @@ -1,5 +1,43 @@ +""" +# Example of DQN pipeline + +Use the pipeline on a single process: + +> python3 -u ding/example/dqn.py + +Use the pipeline on multiple processes: + +We surpose there are N processes (workers) = 1 learner + 1 evaluator + (N-2) collectors + +## First Example —— Execute on one machine with multi processes. + +Execute 4 processes with 1 learner + 1 evaluator + 2 collectors +Remember to keep them connected by mesh to ensure that they can exchange information with each other. + +> ditask --package . --main ding.example.dqn.main --parallel-workers 4 --topology mesh + +## Second Example —— Execute on multiple machines. + +1. Execute 1 learner + 1 evaluator on one machine. + +> ditask --package . --main ding.example.dqn.main --parallel-workers 2 --topology mesh --node-ids 0 --ports 50515 + +2. Execute 2 collectors on another machine. (Suppose the ip of the first machine is 127.0.0.1). + Here we use `alone` topology instead of `mesh` because the collectors do not need communicate with each other. + Remember the `node_ids` cannot be duplicated with the learner, evaluator processes. + And remember to set the `ports` (should not conflict with others) and `attach_to` parameters. + The value of the `attach_to` parameter should be obtained from the log of the + process started earlier (e.g. 'NNG listen on tcp://10.0.0.4:50515'). + +> ditask --package . --main ding.example.dqn.main --parallel-workers 2 --topology alone --node-ids 2 \ + --ports 50517 --attach-to tcp://10.0.0.4:50515,tcp://127.0.0.1:50516 + +3. You can repeat step 2 to start more collectors on other machines. +""" import gym from ditk import logging +from ding.data.model_loader import FileModelLoader +from ding.data.storage_loader import FileStorageLoader from ding.model import DQN from ding.policy import DQNPolicy from ding.envs import DingEnvWrapper, BaseEnvManagerV2 @@ -8,14 +46,14 @@ from ding.framework import task, ding_init from ding.framework.context import OnlineRLContext from ding.framework.middleware import OffPolicyLearner, StepCollector, interaction_evaluator, data_pusher, \ - eps_greedy_handler, CkptSaver, online_logger + eps_greedy_handler, CkptSaver, ContextExchanger, ModelExchanger, online_logger from ding.utils import set_pkg_seed from dizoo.classic_control.cartpole.config.cartpole_dqn_config import main_config, create_config def main(): logging.getLogger().setLevel(logging.INFO) - cfg = compile_config(main_config, create_cfg=create_config, auto=True) + cfg = compile_config(main_config, create_cfg=create_config, auto=True, save_cfg=task.router.node_id == 0) ding_init(cfg) with task.start(async_mode=False, ctx=OnlineRLContext()): collector_env = BaseEnvManagerV2( @@ -33,13 +71,30 @@ def main(): buffer_ = DequeBuffer(size=cfg.policy.other.replay_buffer.replay_buffer_size) policy = DQNPolicy(cfg.policy, model=model) + # Consider the case with multiple processes + if task.router.is_active: + # You can use labels to distinguish between workers with different roles, + # here we use node_id to distinguish. + if task.router.node_id == 0: + task.add_role(task.role.LEARNER) + elif task.router.node_id == 1: + task.add_role(task.role.EVALUATOR) + else: + task.add_role(task.role.COLLECTOR) + + # Sync their context and model between each worker. + task.use(ContextExchanger(skip_n_iter=1)) + task.use(ModelExchanger(model)) + + # Here is the part of single process pipeline. task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) task.use(eps_greedy_handler(cfg)) task.use(StepCollector(cfg, policy.collect_mode, collector_env)) task.use(data_pusher(cfg, buffer_)) task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) task.use(online_logger(train_show_freq=10)) - task.use(CkptSaver(cfg, policy, train_freq=100)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) + task.run() diff --git a/ding/example/dqn_eval.py b/ding/example/dqn_eval.py new file mode 100644 index 0000000000..296d8b0b8f --- /dev/null +++ b/ding/example/dqn_eval.py @@ -0,0 +1,42 @@ +import gym +import torch +from ditk import logging +from ding.model import DQN +from ding.policy import DQNPolicy +from ding.envs import DingEnvWrapper, BaseEnvManagerV2 +from ding.config import compile_config +from ding.framework import task +from ding.framework.context import OnlineRLContext +from ding.framework.middleware import interaction_evaluator +from ding.utils import set_pkg_seed +from dizoo.classic_control.cartpole.config.cartpole_dqn_config import main_config, create_config + + +def main(): + logging.getLogger().setLevel(logging.INFO) + main_config.exp_name = 'cartpole_dqn_eval' + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + with task.start(async_mode=False, ctx=OnlineRLContext()): + evaluator_env = BaseEnvManagerV2( + env_fn=[lambda: DingEnvWrapper(gym.make("CartPole-v0")) for _ in range(cfg.env.evaluator_env_num)], + cfg=cfg.env.manager + ) + + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + model = DQN(**cfg.policy.model) + + # Load the pretrained weights. + # First, you should get a pretrained network weights. + # For example, you can run ``python3 -u ding/examples/dqn.py``. + pretrained_state_dict = torch.load('cartpole_dqn_seed0/ckpt/final.pth.tar', map_location='cpu')['model'] + model.load_state_dict(pretrained_state_dict) + + policy = DQNPolicy(cfg.policy, model=model) + + # Define the evaluator middleware. + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) + task.run(max_step=1) + + +if __name__ == "__main__": + main() diff --git a/ding/example/dqn_frozen_lake.py b/ding/example/dqn_frozen_lake.py new file mode 100644 index 0000000000..ec4b856339 --- /dev/null +++ b/ding/example/dqn_frozen_lake.py @@ -0,0 +1,45 @@ +from ditk import logging +from ding.model import DQN +from ding.policy import DQNPolicy +from ding.envs import DingEnvWrapper, BaseEnvManagerV2 +from ding.data import DequeBuffer +from ding.config import compile_config +from ding.framework import task +from ding.framework.context import OnlineRLContext +from ding.framework.middleware import OffPolicyLearner, StepCollector, interaction_evaluator, data_pusher, \ + eps_greedy_handler, CkptSaver, nstep_reward_enhancer, final_ctx_saver +from ding.utils import set_pkg_seed +from dizoo.frozen_lake.config.frozen_lake_dqn_config import main_config, create_config +from dizoo.frozen_lake.envs import FrozenLakeEnv + + +def main(): + logging.getLogger().setLevel(logging.INFO) + main_config.policy.nstep = 5 + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + with task.start(async_mode=False, ctx=OnlineRLContext()): + collector_env = BaseEnvManagerV2( + env_fn=[lambda: FrozenLakeEnv(cfg=cfg.env) for _ in range(cfg.env.collector_env_num)], cfg=cfg.env.manager + ) + evaluator_env = BaseEnvManagerV2( + env_fn=[lambda: FrozenLakeEnv(cfg=cfg.env) for _ in range(cfg.env.evaluator_env_num)], cfg=cfg.env.manager + ) + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + + model = DQN(**cfg.policy.model) + buffer_ = DequeBuffer(size=cfg.policy.other.replay_buffer.replay_buffer_size) + policy = DQNPolicy(cfg.policy, model=model) + + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) + task.use(eps_greedy_handler(cfg)) + task.use(StepCollector(cfg, policy.collect_mode, collector_env)) + task.use(nstep_reward_enhancer(cfg)) + task.use(data_pusher(cfg, buffer_)) + task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) + task.use(final_ctx_saver(cfg.exp_name)) + task.run() + + +if __name__ == "__main__": + main() diff --git a/ding/example/dqn_her.py b/ding/example/dqn_her.py index d600938af0..b88458aa33 100644 --- a/ding/example/dqn_her.py +++ b/ding/example/dqn_her.py @@ -38,7 +38,7 @@ def main(): task.use(EpisodeCollector(cfg, policy.collect_mode, collector_env)) task.use(data_pusher(cfg, buffer_)) task.use(HERLearner(cfg, policy.learn_mode, buffer_, her_reward_model)) - task.use(CkptSaver(cfg, policy, train_freq=100)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) task.run() diff --git a/ding/example/dqn_new_env.py b/ding/example/dqn_new_env.py index 97b6085aaf..e43a9a8187 100644 --- a/ding/example/dqn_new_env.py +++ b/ding/example/dqn_new_env.py @@ -40,7 +40,7 @@ def main(): task.use(StepCollector(cfg, policy.collect_mode, collector_env)) task.use(data_pusher(cfg, buffer_)) task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) - task.use(CkptSaver(cfg, policy, train_freq=100)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) task.run() diff --git a/ding/example/dqn_nstep.py b/ding/example/dqn_nstep.py index 896f2d03e5..09dc786d22 100644 --- a/ding/example/dqn_nstep.py +++ b/ding/example/dqn_nstep.py @@ -40,7 +40,7 @@ def main(): task.use(nstep_reward_enhancer(cfg)) task.use(data_pusher(cfg, buffer_)) task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) - task.use(CkptSaver(cfg, policy, train_freq=100)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) task.use(final_ctx_saver(cfg.exp_name)) task.run() diff --git a/ding/example/dqn_nstep_gymnasium.py b/ding/example/dqn_nstep_gymnasium.py new file mode 100644 index 0000000000..f712ce8cca --- /dev/null +++ b/ding/example/dqn_nstep_gymnasium.py @@ -0,0 +1,49 @@ +import gymnasium as gym +from ditk import logging +from ding.model import DQN +from ding.policy import DQNPolicy +from ding.envs import DingEnvWrapper, BaseEnvManagerV2 +from ding.data import DequeBuffer +from ding.config import compile_config +from ding.framework import task +from ding.framework.context import OnlineRLContext +from ding.framework.middleware import OffPolicyLearner, StepCollector, interaction_evaluator, data_pusher, \ + eps_greedy_handler, CkptSaver, nstep_reward_enhancer, final_ctx_saver +from ding.utils import set_pkg_seed +from dizoo.classic_control.cartpole.config.cartpole_dqn_config import main_config, create_config + + +def main(): + logging.getLogger().setLevel(logging.INFO) + main_config.exp_name = 'cartpole_dqn_nstep_gymnasium' + main_config.policy.nstep = 3 + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + with task.start(async_mode=False, ctx=OnlineRLContext()): + collector_env = BaseEnvManagerV2( + env_fn=[lambda: DingEnvWrapper(gym.make("CartPole-v0")) for _ in range(cfg.env.collector_env_num)], + cfg=cfg.env.manager + ) + evaluator_env = BaseEnvManagerV2( + env_fn=[lambda: DingEnvWrapper(gym.make("CartPole-v0")) for _ in range(cfg.env.evaluator_env_num)], + cfg=cfg.env.manager + ) + + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + + model = DQN(**cfg.policy.model) + buffer_ = DequeBuffer(size=cfg.policy.other.replay_buffer.replay_buffer_size) + policy = DQNPolicy(cfg.policy, model=model) + + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) + task.use(eps_greedy_handler(cfg)) + task.use(StepCollector(cfg, policy.collect_mode, collector_env)) + task.use(nstep_reward_enhancer(cfg)) + task.use(data_pusher(cfg, buffer_)) + task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) + task.use(final_ctx_saver(cfg.exp_name)) + task.run() + + +if __name__ == "__main__": + main() diff --git a/ding/example/dqn_per.py b/ding/example/dqn_per.py index 31196b1552..fd6d736f8b 100644 --- a/ding/example/dqn_per.py +++ b/ding/example/dqn_per.py @@ -42,7 +42,7 @@ def main(): task.use(StepCollector(cfg, policy.collect_mode, collector_env)) task.use(data_pusher(cfg, buffer_)) task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) - task.use(CkptSaver(cfg, policy, train_freq=100)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) task.run() diff --git a/ding/example/dqn_rnd.py b/ding/example/dqn_rnd.py index bbf605ed83..2d5e1b93c3 100644 --- a/ding/example/dqn_rnd.py +++ b/ding/example/dqn_rnd.py @@ -40,7 +40,7 @@ def main(): task.use(trainer(cfg, reward_model)) task.use(data_pusher(cfg, buffer_)) task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_, reward_model=reward_model)) - task.use(CkptSaver(cfg, policy, train_freq=100)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) task.run() diff --git a/ding/example/dt.py b/ding/example/dt.py new file mode 100644 index 0000000000..407ea01d6b --- /dev/null +++ b/ding/example/dt.py @@ -0,0 +1,47 @@ +import gym +from ditk import logging +from ding.model import DecisionTransformer +from ding.policy import DTPolicy +from ding.envs import DingEnvWrapper, BaseEnvManager, BaseEnvManagerV2 +from ding.envs.env_wrappers.env_wrappers import AllinObsWrapper +from ding.data import create_dataset +from ding.config import compile_config +from ding.framework import task, ding_init +from ding.framework.context import OfflineRLContext +from ding.framework.middleware import interaction_evaluator, trainer, CkptSaver, \ + offline_data_fetcher, offline_logger, termination_checker, final_ctx_saver +from ding.utils import set_pkg_seed +from dizoo.box2d.lunarlander.envs.lunarlander_env import LunarLanderEnv +from dizoo.box2d.lunarlander.config.lunarlander_dt_config import main_config, create_config + + +def main(): + # If you don't have offline data, you need to prepare if first and set the data_path in config + # For demostration, we also can train a RL policy (e.g. SAC) and collect some data + logging.getLogger().setLevel(logging.INFO) + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + ding_init(cfg) + with task.start(async_mode=False, ctx=OfflineRLContext()): + evaluator_env = BaseEnvManagerV2( + env_fn=[lambda: AllinObsWrapper(LunarLanderEnv(cfg.env)) for _ in range(cfg.env.evaluator_env_num)], + cfg=cfg.env.manager + ) + + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + + dataset = create_dataset(cfg) + cfg.policy.state_mean, cfg.policy.state_std = dataset.get_state_stats() + model = DecisionTransformer(**cfg.policy.model) + policy = DTPolicy(cfg.policy, model=model) + + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) + task.use(offline_data_fetcher(cfg, dataset)) + task.use(trainer(cfg, policy.learn_mode)) + task.use(termination_checker(max_train_iter=1e5)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) + task.use(offline_logger()) + task.run() + + +if __name__ == "__main__": + main() diff --git a/ding/example/edac.py b/ding/example/edac.py new file mode 100755 index 0000000000..40230f3008 --- /dev/null +++ b/ding/example/edac.py @@ -0,0 +1,42 @@ +import gym +from ditk import logging +from ding.model import QACEnsemble +from ding.policy import EDACPolicy +from ding.envs import DingEnvWrapper, BaseEnvManagerV2 +from ding.data import create_dataset +from ding.config import compile_config +from ding.framework import task, ding_init +from ding.framework.context import OfflineRLContext +from ding.framework.middleware import interaction_evaluator, trainer, CkptSaver, offline_data_fetcher, offline_logger +from ding.utils import set_pkg_seed +from dizoo.d4rl.envs import D4RLEnv +from dizoo.d4rl.config.halfcheetah_medium_edac_config import main_config, create_config + + +def main(): + # If you don't have offline data, you need to prepare if first and set the data_path in config + # For demostration, we also can train a RL policy (e.g. SAC) and collect some data + logging.getLogger().setLevel(logging.INFO) + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + ding_init(cfg) + with task.start(async_mode=False, ctx=OfflineRLContext()): + evaluator_env = BaseEnvManagerV2( + env_fn=[lambda: D4RLEnv(cfg.env) for _ in range(cfg.env.evaluator_env_num)], cfg=cfg.env.manager + ) + + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + + dataset = create_dataset(cfg) + model = QACEnsemble(**cfg.policy.model) + policy = EDACPolicy(cfg.policy, model=model) + + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) + task.use(offline_data_fetcher(cfg, dataset)) + task.use(trainer(cfg, policy.learn_mode)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=1e4)) + task.use(offline_logger()) + task.run() + + +if __name__ == "__main__": + main() diff --git a/ding/example/impala.py b/ding/example/impala.py new file mode 100644 index 0000000000..11602012af --- /dev/null +++ b/ding/example/impala.py @@ -0,0 +1,47 @@ +import gym +from ditk import logging +from ding.model import VAC +from ding.policy import IMPALAPolicy +from ding.envs import SubprocessEnvManagerV2 +from ding.data import DequeBuffer +from ding.config import compile_config +from ding.framework import task, ding_init +from ding.framework.context import OnlineRLContext +from ding.framework.middleware import OffPolicyLearner, StepCollector, interaction_evaluator, data_pusher, \ + CkptSaver, online_logger, termination_checker +from ding.utils import set_pkg_seed +from dizoo.box2d.lunarlander.config.lunarlander_impala_config import main_config, create_config +from dizoo.box2d.lunarlander.envs import LunarLanderEnv + + +def main(): + logging.getLogger().setLevel(logging.INFO) + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + ding_init(cfg) + with task.start(async_mode=False, ctx=OnlineRLContext()): + collector_env = SubprocessEnvManagerV2( + env_fn=[lambda: LunarLanderEnv(cfg.env) for _ in range(cfg.env.collector_env_num)], cfg=cfg.env.manager + ) + evaluator_env = SubprocessEnvManagerV2( + env_fn=[lambda: LunarLanderEnv(cfg.env) for _ in range(cfg.env.evaluator_env_num)], cfg=cfg.env.manager + ) + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + + model = VAC(**cfg.policy.model) + buffer_ = DequeBuffer( + size=cfg.policy.other.replay_buffer.replay_buffer_size, sliced=cfg.policy.other.replay_buffer.sliced + ) + policy = IMPALAPolicy(cfg.policy, model=model) + + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) + task.use(StepCollector(cfg, policy.collect_mode, collector_env, random_collect_size=1024)) + task.use(data_pusher(cfg, buffer_, group_by_env=True)) + task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) + task.use(online_logger(train_show_freq=300)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=10000)) + task.use(termination_checker(max_env_step=2e6)) + task.run() + + +if __name__ == "__main__": + main() diff --git a/ding/example/iqn_nstep.py b/ding/example/iqn_nstep.py index 49e5a99f84..eff6df85bf 100644 --- a/ding/example/iqn_nstep.py +++ b/ding/example/iqn_nstep.py @@ -40,7 +40,7 @@ def main(): task.use(nstep_reward_enhancer(cfg)) task.use(data_pusher(cfg, buffer_)) task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) - task.use(CkptSaver(cfg, policy, train_freq=100)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) task.run() diff --git a/ding/example/mappo.py b/ding/example/mappo.py new file mode 100644 index 0000000000..53ca5dff3c --- /dev/null +++ b/ding/example/mappo.py @@ -0,0 +1,45 @@ +import gym +from ditk import logging +from ding.model import MAVAC +from ding.policy import PPOPolicy +from ding.envs import DingEnvWrapper, SubprocessEnvManagerV2 +from ding.data import DequeBuffer +from ding.config import compile_config +from ding.framework import task, ding_init +from ding.framework.context import OnlineRLContext +from ding.framework.middleware import multistep_trainer, StepCollector, interaction_evaluator, CkptSaver, \ + gae_estimator, online_logger, termination_checker +from ding.utils import set_pkg_seed +from dizoo.petting_zoo.config.ptz_simple_spread_mappo_config import main_config, create_config +from dizoo.petting_zoo.envs.petting_zoo_simple_spread_env import PettingZooEnv + + +def main(): + logging.getLogger().setLevel(logging.INFO) + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + ding_init(cfg) + with task.start(async_mode=False, ctx=OnlineRLContext()): + collector_env = SubprocessEnvManagerV2( + env_fn=[lambda: PettingZooEnv(cfg.env) for _ in range(cfg.env.collector_env_num)], cfg=cfg.env.manager + ) + evaluator_env = SubprocessEnvManagerV2( + env_fn=[lambda: PettingZooEnv(cfg.env) for _ in range(cfg.env.evaluator_env_num)], cfg=cfg.env.manager + ) + + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + + model = MAVAC(**cfg.policy.model) + policy = PPOPolicy(cfg.policy, model=model) + + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) + task.use(StepCollector(cfg, policy.collect_mode, collector_env)) + task.use(gae_estimator(cfg, policy.collect_mode)) + task.use(multistep_trainer(policy.learn_mode, log_freq=100)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=1000)) + task.use(online_logger(train_show_freq=10)) + task.use(termination_checker(max_env_step=int(1e6))) + task.run() + + +if __name__ == "__main__": + main() diff --git a/ding/example/masac.py b/ding/example/masac.py new file mode 100644 index 0000000000..a268c7366b --- /dev/null +++ b/ding/example/masac.py @@ -0,0 +1,49 @@ +import gym +from ditk import logging +from ding.model import MAQAC +from ding.policy import SACDiscretePolicy +from ding.envs import DingEnvWrapper, BaseEnvManagerV2 +from ding.data import DequeBuffer +from ding.config import compile_config +from ding.framework import task, ding_init +from ding.framework.context import OnlineRLContext +from ding.framework.middleware import OffPolicyLearner, StepCollector, interaction_evaluator, CkptSaver, \ + data_pusher, online_logger, termination_checker, eps_greedy_handler +from ding.utils import set_pkg_seed +from dizoo.petting_zoo.config.ptz_simple_spread_masac_config import main_config, create_config +from dizoo.petting_zoo.envs.petting_zoo_simple_spread_env import PettingZooEnv + + +def main(): + logging.getLogger().setLevel(logging.INFO) + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + ding_init(cfg) + with task.start(async_mode=False, ctx=OnlineRLContext()): + collector_env = BaseEnvManagerV2( + env_fn=[lambda: PettingZooEnv(cfg.env) for _ in range(cfg.env.collector_env_num)], cfg=cfg.env.manager + ) + evaluator_env = BaseEnvManagerV2( + env_fn=[lambda: PettingZooEnv(cfg.env) for _ in range(cfg.env.evaluator_env_num)], cfg=cfg.env.manager + ) + + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + + model = MAQAC(**cfg.policy.model) + buffer_ = DequeBuffer(size=cfg.policy.other.replay_buffer.replay_buffer_size) + policy = SACDiscretePolicy(cfg.policy, model=model) + + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) + task.use(eps_greedy_handler(cfg)) + task.use( + StepCollector(cfg, policy.collect_mode, collector_env, random_collect_size=cfg.policy.random_collect_size) + ) + task.use(data_pusher(cfg, buffer_)) + task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_, log_freq=100)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=1000)) + task.use(online_logger(train_show_freq=10)) + task.use(termination_checker(max_env_step=int(1e6))) + task.run() + + +if __name__ == "__main__": + main() diff --git a/ding/example/pdqn.py b/ding/example/pdqn.py index d470f140b7..5bc173d83c 100644 --- a/ding/example/pdqn.py +++ b/ding/example/pdqn.py @@ -37,7 +37,7 @@ def main(): task.use(StepCollector(cfg, policy.collect_mode, collector_env)) task.use(data_pusher(cfg, buffer_)) task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) - task.use(CkptSaver(cfg, policy, train_freq=1000)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=1000)) task.run() diff --git a/ding/example/ppg_offpolicy.py b/ding/example/ppg_offpolicy.py index ab07efaa73..70bd211cd5 100644 --- a/ding/example/ppg_offpolicy.py +++ b/ding/example/ppg_offpolicy.py @@ -1,7 +1,7 @@ import gym from ditk import logging from ding.model import PPG -from ding.policy import PPGPolicy +from ding.policy import PPGOffPolicy from ding.envs import DingEnvWrapper, BaseEnvManagerV2 from ding.data import DequeBuffer from ding.data.buffer.middleware import use_time_check, sample_range_view @@ -11,7 +11,7 @@ from ding.framework.middleware import OffPolicyLearner, StepCollector, interaction_evaluator, data_pusher, \ CkptSaver, gae_estimator from ding.utils import set_pkg_seed -from dizoo.classic_control.cartpole.config.cartpole_ppg_config import main_config, create_config +from dizoo.classic_control.cartpole.config.cartpole_ppg_offpolicy_config import main_config, create_config def main(): @@ -39,13 +39,13 @@ def main(): value_buffer = buffer_.view() value_buffer.use(use_time_check(value_buffer, max_use=buffer_cfg.value.max_use)) value_buffer.use(sample_range_view(value_buffer, start=-buffer_cfg.value.replay_buffer_size)) - policy = PPGPolicy(cfg.policy, model=model) + policy = PPGOffPolicy(cfg.policy, model=model) task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) task.use(StepCollector(cfg, policy.collect_mode, collector_env)) task.use(gae_estimator(cfg, policy.collect_mode, buffer_)) task.use(OffPolicyLearner(cfg, policy.learn_mode, {'policy': policy_buffer, 'value': value_buffer})) - task.use(CkptSaver(cfg, policy, train_freq=100)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) task.run() diff --git a/ding/example/ppo.py b/ding/example/ppo.py index b2cd7aed01..b9807d1c41 100644 --- a/ding/example/ppo.py +++ b/ding/example/ppo.py @@ -1,3 +1,21 @@ +""" +# Example of PPO pipeline + +Use the pipeline on a single process: + +> python3 -u ding/example/ppo.py + +Use the pipeline on multiple processes: + +We surpose there are N processes (workers) = 1 learner + 1 evaluator + (N-2) collectors + +## First Example —— Execute on one machine with multi processes. + +Execute 4 processes with 1 learner + 1 evaluator + 2 collectors +Remember to keep them connected by mesh to ensure that they can exchange information with each other. + +> ditask --package . --main ding.example.ppo.main --parallel-workers 4 --topology mesh +""" import gym from ditk import logging from ding.model import VAC @@ -8,14 +26,14 @@ from ding.framework import task, ding_init from ding.framework.context import OnlineRLContext from ding.framework.middleware import multistep_trainer, StepCollector, interaction_evaluator, CkptSaver, \ - gae_estimator, online_logger + gae_estimator, online_logger, ContextExchanger, ModelExchanger from ding.utils import set_pkg_seed from dizoo.classic_control.cartpole.config.cartpole_ppo_config import main_config, create_config def main(): logging.getLogger().setLevel(logging.INFO) - cfg = compile_config(main_config, create_cfg=create_config, auto=True) + cfg = compile_config(main_config, create_cfg=create_config, auto=True, save_cfg=task.router.node_id == 0) ding_init(cfg) with task.start(async_mode=False, ctx=OnlineRLContext()): collector_env = BaseEnvManagerV2( @@ -32,11 +50,26 @@ def main(): model = VAC(**cfg.policy.model) policy = PPOPolicy(cfg.policy, model=model) + # Consider the case with multiple processes + if task.router.is_active: + # You can use labels to distinguish between workers with different roles, + # here we use node_id to distinguish. + if task.router.node_id == 0: + task.add_role(task.role.LEARNER) + elif task.router.node_id == 1: + task.add_role(task.role.EVALUATOR) + else: + task.add_role(task.role.COLLECTOR) + + # Sync their context and model between each worker. + task.use(ContextExchanger(skip_n_iter=1)) + task.use(ModelExchanger(model)) + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) task.use(StepCollector(cfg, policy.collect_mode, collector_env)) task.use(gae_estimator(cfg, policy.collect_mode)) - task.use(multistep_trainer(cfg, policy.learn_mode)) - task.use(CkptSaver(cfg, policy, train_freq=100)) + task.use(multistep_trainer(policy.learn_mode, log_freq=50)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) task.use(online_logger(train_show_freq=3)) task.run() diff --git a/ding/example/ppo_lunarlander.py b/ding/example/ppo_lunarlander.py new file mode 100644 index 0000000000..69cb8fe183 --- /dev/null +++ b/ding/example/ppo_lunarlander.py @@ -0,0 +1,44 @@ +import gym +from ditk import logging +from ding.model import VAC +from ding.policy import PPOPolicy +from ding.envs import DingEnvWrapper, BaseEnvManagerV2 +from ding.config import compile_config +from ding.framework import task, ding_init +from ding.framework.context import OnlineRLContext +from ding.framework.middleware import multistep_trainer, StepCollector, interaction_evaluator, CkptSaver, \ + gae_estimator, online_logger +from ding.utils import set_pkg_seed +from dizoo.box2d.lunarlander.config.lunarlander_ppo_config import main_config, create_config + + +def main(): + logging.getLogger().setLevel(logging.INFO) + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + ding_init(cfg) + with task.start(async_mode=False, ctx=OnlineRLContext()): + collector_env = BaseEnvManagerV2( + env_fn=[lambda: DingEnvWrapper(gym.make("LunarLander-v2")) for _ in range(cfg.env.collector_env_num)], + cfg=cfg.env.manager + ) + evaluator_env = BaseEnvManagerV2( + env_fn=[lambda: DingEnvWrapper(gym.make("LunarLander-v2")) for _ in range(cfg.env.evaluator_env_num)], + cfg=cfg.env.manager + ) + + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + + model = VAC(**cfg.policy.model) + policy = PPOPolicy(cfg.policy, model=model) + + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) + task.use(StepCollector(cfg, policy.collect_mode, collector_env)) + task.use(gae_estimator(cfg, policy.collect_mode)) + task.use(multistep_trainer(policy.learn_mode, log_freq=50)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) + task.use(online_logger(train_show_freq=3)) + task.run() + + +if __name__ == "__main__": + main() diff --git a/ding/example/ppo_lunarlander_continuous.py b/ding/example/ppo_lunarlander_continuous.py new file mode 100644 index 0000000000..7573c916eb --- /dev/null +++ b/ding/example/ppo_lunarlander_continuous.py @@ -0,0 +1,49 @@ +import gym +from ditk import logging +from ding.model import VAC +from ding.policy import PPOPolicy +from ding.envs import DingEnvWrapper, BaseEnvManagerV2 +from ding.config import compile_config +from ding.framework import task, ding_init +from ding.framework.context import OnlineRLContext +from ding.framework.middleware import multistep_trainer, StepCollector, interaction_evaluator, CkptSaver, \ + gae_estimator, online_logger +from ding.utils import set_pkg_seed +from dizoo.box2d.lunarlander.config.lunarlander_ppo_continuous_config import main_config, create_config + + +def main(): + logging.getLogger().setLevel(logging.INFO) + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + ding_init(cfg) + with task.start(async_mode=False, ctx=OnlineRLContext()): + collector_env = BaseEnvManagerV2( + env_fn=[ + lambda: DingEnvWrapper(gym.make("LunarLanderContinuous-v2")) for _ in range(cfg.env.collector_env_num) + ], + cfg=cfg.env.manager + ) + evaluator_env = BaseEnvManagerV2( + env_fn=[ + lambda: DingEnvWrapper(gym.make("LunarLanderContinuous-v2")) for _ in range(cfg.env.evaluator_env_num) + ], + cfg=cfg.env.manager + ) + + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + + model = VAC(**cfg.policy.model) + policy = PPOPolicy(cfg.policy, model=model) + logging.info(model) + + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) + task.use(StepCollector(cfg, policy.collect_mode, collector_env)) + task.use(gae_estimator(cfg, policy.collect_mode)) + task.use(multistep_trainer(policy.learn_mode, log_freq=50)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) + task.use(online_logger(train_show_freq=3)) + task.run() + + +if __name__ == "__main__": + main() diff --git a/ding/example/ppo_offpolicy.py b/ding/example/ppo_offpolicy.py index 5f3b11c579..738b27f230 100644 --- a/ding/example/ppo_offpolicy.py +++ b/ding/example/ppo_offpolicy.py @@ -37,7 +37,7 @@ def main(): task.use(StepCollector(cfg, policy.collect_mode, collector_env)) task.use(gae_estimator(cfg, policy.collect_mode, buffer_)) task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) - task.use(CkptSaver(cfg, policy, train_freq=100)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) task.run() diff --git a/ding/example/ppo_with_complex_obs.py b/ding/example/ppo_with_complex_obs.py new file mode 100644 index 0000000000..2563532a37 --- /dev/null +++ b/ding/example/ppo_with_complex_obs.py @@ -0,0 +1,205 @@ +from typing import Dict +import os +import torch +import torch.nn as nn +import numpy as np +import gym +from gym import spaces +from ditk import logging +from ding.envs import DingEnvWrapper, EvalEpisodeReturnWrapper, \ + BaseEnvManagerV2 +from ding.config import compile_config +from ding.policy import PPOPolicy +from ding.utils import set_pkg_seed +from ding.model import VAC +from ding.framework import task, ding_init +from ding.framework.context import OnlineRLContext +from ding.framework.middleware import multistep_trainer, StepCollector, interaction_evaluator, CkptSaver, \ + gae_estimator, online_logger +from easydict import EasyDict + +my_env_ppo_config = dict( + exp_name='my_env_ppo_seed0', + env=dict( + collector_env_num=4, + evaluator_env_num=4, + n_evaluator_episode=4, + stop_value=195, + ), + policy=dict( + cuda=True, + action_space='discrete', + model=dict( + obs_shape=dict( + key_0=dict(k1=(), k2=()), + key_1=(5, 10), + key_2=(10, 10, 3), + key_3=(2, ), + ), + action_shape=2, + action_space='discrete', + critic_head_hidden_size=138, + actor_head_hidden_size=138, + ), + learn=dict( + epoch_per_collect=2, + batch_size=64, + learning_rate=0.001, + value_weight=0.5, + entropy_weight=0.01, + clip_ratio=0.2, + learner=dict(hook=dict(save_ckpt_after_iter=100)), + ), + collect=dict( + n_sample=256, unroll_len=1, discount_factor=0.9, gae_lambda=0.95, collector=dict(transform_obs=True, ) + ), + eval=dict(evaluator=dict(eval_freq=100, ), ), + ), +) +my_env_ppo_config = EasyDict(my_env_ppo_config) +main_config = my_env_ppo_config +my_env_ppo_create_config = dict( + env_manager=dict(type='base'), + policy=dict(type='ppo'), +) +my_env_ppo_create_config = EasyDict(my_env_ppo_create_config) +create_config = my_env_ppo_create_config + + +class MyEnv(gym.Env): + + def __init__(self, seq_len=5, feature_dim=10, image_size=(10, 10, 3)): + super().__init__() + + # Define the action space + self.action_space = spaces.Discrete(2) + + # Define the observation space + self.observation_space = spaces.Dict( + ( + { + 'key_0': spaces.Dict( + { + 'k1': spaces.Box(low=0, high=np.inf, shape=(1, ), dtype=np.float32), + 'k2': spaces.Box(low=-1, high=1, shape=(1, ), dtype=np.float32), + } + ), + 'key_1': spaces.Box(low=-np.inf, high=np.inf, shape=(seq_len, feature_dim), dtype=np.float32), + 'key_2': spaces.Box(low=0, high=255, shape=image_size, dtype=np.uint8), + 'key_3': spaces.Box(low=0, high=np.array([np.inf, 3]), shape=(2, ), dtype=np.float32) + } + ) + ) + + def reset(self): + # Generate a random initial state + return self.observation_space.sample() + + def step(self, action): + # Compute the reward and done flag (which are not used in this example) + reward = np.random.uniform(low=0.0, high=1.0) + + done = False + if np.random.uniform(low=0.0, high=1.0) > 0.7: + done = True + + info = {} + + # Return the next state, reward, and done flag + return self.observation_space.sample(), reward, done, info + + +def ding_env_maker(): + return DingEnvWrapper( + MyEnv(), cfg={'env_wrapper': [ + lambda env: EvalEpisodeReturnWrapper(env), + ]} + ) + + +class Encoder(nn.Module): + + def __init__(self, feature_dim: int): + super(Encoder, self).__init__() + + # Define the networks for each input type + self.fc_net_1_k1 = nn.Sequential(nn.Linear(1, 8), nn.ReLU()) + self.fc_net_1_k2 = nn.Sequential(nn.Linear(1, 8), nn.ReLU()) + self.fc_net_1 = nn.Sequential(nn.Linear(16, 32), nn.ReLU()) + """ + Implementation of transformer_encoder refers to Vision Transformer (ViT) code: + https://arxiv.org/abs/2010.11929 + https://pytorch.org/vision/main/_modules/torchvision/models/vision_transformer.html + """ + self.class_token = nn.Parameter(torch.zeros(1, 1, feature_dim)) + self.encoder_layer = nn.TransformerEncoderLayer(d_model=feature_dim, nhead=2, batch_first=True) + self.transformer_encoder = nn.TransformerEncoder(self.encoder_layer, num_layers=1) + + self.conv_net = nn.Sequential( + nn.Conv2d(3, 16, kernel_size=3, padding=1), nn.ReLU(), nn.Conv2d(16, 32, kernel_size=3, padding=1), + nn.ReLU() + ) + self.conv_fc_net = nn.Sequential(nn.Flatten(), nn.Linear(3200, 64), nn.ReLU()) + + self.fc_net_2 = nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 32), nn.ReLU(), nn.Flatten()) + + def forward(self, inputs: Dict[str, torch.Tensor]) -> torch.Tensor: + # Unpack the input tuple + dict_input = inputs['key_0'] # dict{key:(B)} + transformer_input = inputs['key_1'] # (B, seq_len, feature_dim) + conv_input = inputs['key_2'] # (B, H, W, 3) + fc_input = inputs['key_3'] # (B, X) + + B = fc_input.shape[0] + + # Pass each input through its corresponding network + dict_output = self.fc_net_1( + torch.cat( + [self.fc_net_1_k1(dict_input['k1'].unsqueeze(-1)), + self.fc_net_1_k2(dict_input['k2'].unsqueeze(-1))], + dim=1 + ) + ) + + batch_class_token = self.class_token.expand(B, -1, -1) + transformer_output = self.transformer_encoder(torch.cat([batch_class_token, transformer_input], dim=1)) + transformer_output = transformer_output[:, 0] + + conv_output = self.conv_fc_net(self.conv_net(conv_input.permute(0, 3, 1, 2))) + fc_output = self.fc_net_2(fc_input) + + # Concatenate the outputs along the feature dimension + encoded_output = torch.cat([dict_output, transformer_output, conv_output, fc_output], dim=1) + + return encoded_output + + +def main(): + logging.getLogger().setLevel(logging.INFO) + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + ding_init(cfg) + with task.start(async_mode=False, ctx=OnlineRLContext()): + collector_env = BaseEnvManagerV2( + env_fn=[ding_env_maker for _ in range(cfg.env.collector_env_num)], cfg=cfg.env.manager + ) + evaluator_env = BaseEnvManagerV2( + env_fn=[ding_env_maker for _ in range(cfg.env.evaluator_env_num)], cfg=cfg.env.manager + ) + + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + + encoder = Encoder(feature_dim=10) + model = VAC(encoder=encoder, **cfg.policy.model) + policy = PPOPolicy(cfg.policy, model=model) + + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) + task.use(StepCollector(cfg, policy.collect_mode, collector_env)) + task.use(gae_estimator(cfg, policy.collect_mode)) + task.use(multistep_trainer(policy.learn_mode, log_freq=50)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) + task.use(online_logger(train_show_freq=3)) + task.run() + + +if __name__ == "__main__": + main() diff --git a/ding/example/qgpo.py b/ding/example/qgpo.py new file mode 100644 index 0000000000..ed844974be --- /dev/null +++ b/ding/example/qgpo.py @@ -0,0 +1,170 @@ +import torch +import gym +import d4rl +from easydict import EasyDict +from ditk import logging +from ding.model import QGPO +from ding.policy import QGPOPolicy +from ding.envs import DingEnvWrapper, BaseEnvManagerV2 +from ding.config import compile_config +from ding.framework import task, ding_init +from ding.framework.context import OfflineRLContext +from ding.framework.middleware import trainer, CkptSaver, offline_logger, wandb_offline_logger, termination_checker +from ding.framework.middleware.functional.evaluator import interaction_evaluator +from ding.framework.middleware.functional.data_processor import qgpo_support_data_generator, qgpo_offline_data_fetcher +from ding.utils import set_pkg_seed + +from dizoo.d4rl.config.halfcheetah_medium_expert_qgpo_config import main_config, create_config + + +class QGPOD4RLDataset(torch.utils.data.Dataset): + """ + Overview: + Dataset for QGPO algorithm. The training of QGPO algorithm is based on contrastive energy prediction, \ + which needs true action and fake action. The true action is sampled from the dataset, and the fake action \ + is sampled from the action support generated by the behavior policy. + Interface: + ``__init__``, ``__getitem__``, ``__len__``. + """ + + def __init__(self, cfg, device="cpu"): + """ + Overview: + Initialization method of QGPOD4RLDataset class + Arguments: + - cfg (:obj:`EasyDict`): Config dict + - device (:obj:`str`): Device name + """ + + self.cfg = cfg + data = d4rl.qlearning_dataset(gym.make(cfg.env_id)) + self.device = device + self.states = torch.from_numpy(data['observations']).float().to(self.device) + self.actions = torch.from_numpy(data['actions']).float().to(self.device) + self.next_states = torch.from_numpy(data['next_observations']).float().to(self.device) + reward = torch.from_numpy(data['rewards']).view(-1, 1).float().to(self.device) + self.is_finished = torch.from_numpy(data['terminals']).view(-1, 1).float().to(self.device) + + reward_tune = "iql_antmaze" if "antmaze" in cfg.env_id else "iql_locomotion" + if reward_tune == 'normalize': + reward = (reward - reward.mean()) / reward.std() + elif reward_tune == 'iql_antmaze': + reward = reward - 1.0 + elif reward_tune == 'iql_locomotion': + min_ret, max_ret = QGPOD4RLDataset.return_range(data, 1000) + reward /= (max_ret - min_ret) + reward *= 1000 + elif reward_tune == 'cql_antmaze': + reward = (reward - 0.5) * 4.0 + elif reward_tune == 'antmaze': + reward = (reward - 0.25) * 2.0 + self.rewards = reward + self.len = self.states.shape[0] + logging.info(f"{self.len} data loaded in QGPOD4RLDataset") + + def __getitem__(self, index): + """ + Overview: + Get data by index + Arguments: + - index (:obj:`int`): Index of data + Returns: + - data (:obj:`dict`): Data dict + + .. note:: + The data dict contains the following keys: + - s (:obj:`torch.Tensor`): State + - a (:obj:`torch.Tensor`): Action + - r (:obj:`torch.Tensor`): Reward + - s_ (:obj:`torch.Tensor`): Next state + - d (:obj:`torch.Tensor`): Is finished + - fake_a (:obj:`torch.Tensor`): Fake action for contrastive energy prediction and qgpo training \ + (fake action is sampled from the action support generated by the behavior policy) + - fake_a_ (:obj:`torch.Tensor`): Fake next action for contrastive energy prediction and qgpo training \ + (fake action is sampled from the action support generated by the behavior policy) + """ + + data = { + 's': self.states[index % self.len], + 'a': self.actions[index % self.len], + 'r': self.rewards[index % self.len], + 's_': self.next_states[index % self.len], + 'd': self.is_finished[index % self.len], + 'fake_a': self.fake_actions[index % self.len] + if hasattr(self, "fake_actions") else 0.0, # self.fake_actions + 'fake_a_': self.fake_next_actions[index % self.len] + if hasattr(self, "fake_next_actions") else 0.0, # self.fake_next_actions + } + return data + + def __len__(self): + return self.len + + def return_range(dataset, max_episode_steps): + returns, lengths = [], [] + ep_ret, ep_len = 0., 0 + for r, d in zip(dataset['rewards'], dataset['terminals']): + ep_ret += float(r) + ep_len += 1 + if d or ep_len == max_episode_steps: + returns.append(ep_ret) + lengths.append(ep_len) + ep_ret, ep_len = 0., 0 + # returns.append(ep_ret) # incomplete trajectory + lengths.append(ep_len) # but still keep track of number of steps + assert sum(lengths) == len(dataset['rewards']) + return min(returns), max(returns) + + +def main(): + # If you don't have offline data, you need to prepare if first and set the data_path in config + # For demostration, we also can train a RL policy (e.g. SAC) and collect some data + logging.getLogger().setLevel(logging.INFO) + cfg = compile_config(main_config, policy=QGPOPolicy) + ding_init(cfg) + with task.start(async_mode=False, ctx=OfflineRLContext()): + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + model = QGPO(cfg=cfg.policy.model) + policy = QGPOPolicy(cfg.policy, model=model) + dataset = QGPOD4RLDataset(cfg=cfg.dataset, device=policy._device) + if hasattr(cfg.policy, "load_path") and cfg.policy.load_path is not None: + policy_state_dict = torch.load(cfg.policy.load_path, map_location=torch.device("cpu")) + policy.learn_mode.load_state_dict(policy_state_dict) + + task.use(qgpo_support_data_generator(cfg, dataset, policy)) + task.use(qgpo_offline_data_fetcher(cfg, dataset, collate_fn=None)) + task.use(trainer(cfg, policy.learn_mode)) + for guidance_scale in cfg.policy.eval.guidance_scale: + evaluator_env = BaseEnvManagerV2( + env_fn=[ + lambda: DingEnvWrapper(env=gym.make(cfg.env.env_id), cfg=cfg.env, caller="evaluator") + for _ in range(cfg.env.evaluator_env_num) + ], + cfg=cfg.env.manager + ) + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env, guidance_scale=guidance_scale)) + task.use( + wandb_offline_logger( + cfg=EasyDict( + dict( + gradient_logger=False, + plot_logger=True, + video_logger=False, + action_logger=False, + return_logger=False, + vis_dataset=False, + ) + ), + exp_config=cfg, + metric_list=policy._monitor_vars_learn(), + project_name=cfg.exp_name + ) + ) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100000)) + task.use(offline_logger()) + task.use(termination_checker(max_train_iter=500000 + cfg.policy.learn.q_value_stop_training_iter)) + task.run() + + +if __name__ == "__main__": + main() diff --git a/ding/example/qrdqn_nstep.py b/ding/example/qrdqn_nstep.py index f6b6b5a95a..352828cf35 100644 --- a/ding/example/qrdqn_nstep.py +++ b/ding/example/qrdqn_nstep.py @@ -40,7 +40,7 @@ def main(): task.use(nstep_reward_enhancer(cfg)) task.use(data_pusher(cfg, buffer_)) task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) - task.use(CkptSaver(cfg, policy, train_freq=100)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) task.run() diff --git a/ding/example/r2d2.py b/ding/example/r2d2.py index f537ba990e..83fc617563 100644 --- a/ding/example/r2d2.py +++ b/ding/example/r2d2.py @@ -38,7 +38,7 @@ def main(): task.use(nstep_reward_enhancer(cfg)) task.use(data_pusher(cfg, buffer_, group_by_env=True)) task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) - task.use(CkptSaver(cfg, policy, train_freq=100)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) task.run() diff --git a/ding/example/sac.py b/ding/example/sac.py index 3b45129029..d83e552050 100644 --- a/ding/example/sac.py +++ b/ding/example/sac.py @@ -1,13 +1,13 @@ from ditk import logging -from ding.model import QAC +from ding.model import ContinuousQAC from ding.policy import SACPolicy from ding.envs import BaseEnvManagerV2 from ding.data import DequeBuffer from ding.config import compile_config -from ding.framework import task +from ding.framework import task, ding_init from ding.framework.context import OnlineRLContext from ding.framework.middleware import data_pusher, StepCollector, interaction_evaluator, \ - CkptSaver, OffPolicyLearner, termination_checker + CkptSaver, OffPolicyLearner, termination_checker, online_logger from ding.utils import set_pkg_seed from dizoo.classic_control.pendulum.envs.pendulum_env import PendulumEnv from dizoo.classic_control.pendulum.config.pendulum_sac_config import main_config, create_config @@ -16,6 +16,7 @@ def main(): logging.getLogger().setLevel(logging.INFO) cfg = compile_config(main_config, create_cfg=create_config, auto=True) + ding_init(cfg) with task.start(async_mode=False, ctx=OnlineRLContext()): collector_env = BaseEnvManagerV2( env_fn=[lambda: PendulumEnv(cfg.env) for _ in range(cfg.env.collector_env_num)], cfg=cfg.env.manager @@ -26,7 +27,7 @@ def main(): set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) - model = QAC(**cfg.policy.model) + model = ContinuousQAC(**cfg.policy.model) buffer_ = DequeBuffer(size=cfg.policy.other.replay_buffer.replay_buffer_size) policy = SACPolicy(cfg.policy, model=model) @@ -36,8 +37,9 @@ def main(): ) task.use(data_pusher(cfg, buffer_)) task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) - task.use(CkptSaver(cfg, policy, train_freq=100)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) task.use(termination_checker(max_train_iter=10000)) + task.use(online_logger()) task.run() diff --git a/ding/example/sqil.py b/ding/example/sqil.py index 2443d9d6c5..6df54a5724 100644 --- a/ding/example/sqil.py +++ b/ding/example/sqil.py @@ -57,7 +57,7 @@ def main(): task.use(StepCollector(cfg, expert_policy.collect_mode, expert_collector_env)) # expert data collector task.use(sqil_data_pusher(cfg, expert_buffer, expert=True)) task.use(OffPolicyLearner(cfg, policy.learn_mode, [(buffer_, 0.5), (expert_buffer, 0.5)])) - task.use(CkptSaver(cfg, policy, train_freq=100)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) task.run() diff --git a/ding/example/sqil_continuous.py b/ding/example/sqil_continuous.py index 930a5cbb27..ee3d36c9f3 100644 --- a/ding/example/sqil_continuous.py +++ b/ding/example/sqil_continuous.py @@ -1,6 +1,6 @@ from ditk import logging import torch -from ding.model import QAC +from ding.model import ContinuousQAC from ding.policy import SQILSACPolicy from ding.envs import BaseEnvManagerV2 from ding.data import DequeBuffer @@ -35,8 +35,8 @@ def main(): set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) - model = QAC(**cfg.policy.model) - expert_model = QAC(**cfg.policy.model) + model = ContinuousQAC(**cfg.policy.model) + expert_model = ContinuousQAC(**cfg.policy.model) buffer_ = DequeBuffer(size=cfg.policy.other.replay_buffer.replay_buffer_size) expert_buffer = DequeBuffer(size=cfg.policy.other.replay_buffer.replay_buffer_size) @@ -61,7 +61,7 @@ def main(): ) # expert data collector task.use(sqil_data_pusher(cfg, expert_buffer, expert=True)) task.use(OffPolicyLearner(cfg, policy.learn_mode, [(buffer_, 0.5), (expert_buffer, 0.5)])) - task.use(CkptSaver(cfg, policy, train_freq=100)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) task.use(termination_checker(max_train_iter=10000)) task.run() diff --git a/ding/example/sql.py b/ding/example/sql.py index a1c3a75680..2c2a968082 100644 --- a/ding/example/sql.py +++ b/ding/example/sql.py @@ -37,7 +37,7 @@ def main(): task.use(StepCollector(cfg, policy.collect_mode, collector_env)) task.use(data_pusher(cfg, buffer_)) task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) - task.use(CkptSaver(cfg, policy, train_freq=100)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) task.run() diff --git a/ding/example/td3.py b/ding/example/td3.py index 59fa0c59cc..9d6508dd6f 100644 --- a/ding/example/td3.py +++ b/ding/example/td3.py @@ -1,13 +1,13 @@ -import gym from ditk import logging -from ding.model.template.qac import QAC +from ding.model import ContinuousQAC from ding.policy import TD3Policy -from ding.envs import DingEnvWrapper, BaseEnvManagerV2 +from ding.envs import BaseEnvManagerV2 from ding.data import DequeBuffer from ding.config import compile_config -from ding.framework import task +from ding.framework import task, ding_init from ding.framework.context import OnlineRLContext -from ding.framework.middleware import OffPolicyLearner, StepCollector, interaction_evaluator, data_pusher, CkptSaver +from ding.framework.middleware import data_pusher, StepCollector, interaction_evaluator, \ + CkptSaver, OffPolicyLearner, termination_checker, online_logger from ding.utils import set_pkg_seed from dizoo.classic_control.pendulum.envs.pendulum_env import PendulumEnv from dizoo.classic_control.pendulum.config.pendulum_td3_config import main_config, create_config @@ -16,6 +16,7 @@ def main(): logging.getLogger().setLevel(logging.INFO) cfg = compile_config(main_config, create_cfg=create_config, auto=True) + ding_init(cfg) with task.start(async_mode=False, ctx=OnlineRLContext()): collector_env = BaseEnvManagerV2( env_fn=[lambda: PendulumEnv(cfg.env) for _ in range(cfg.env.collector_env_num)], cfg=cfg.env.manager @@ -26,9 +27,9 @@ def main(): set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) - model = QAC(**cfg.policy.model) + model = ContinuousQAC(**cfg.policy.model) buffer_ = DequeBuffer(size=cfg.policy.other.replay_buffer.replay_buffer_size) - policy = TD3Policy(cfg.policy, model) + policy = TD3Policy(cfg.policy, model=model) task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) task.use( @@ -36,7 +37,9 @@ def main(): ) task.use(data_pusher(cfg, buffer_)) task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) - task.use(CkptSaver(cfg, policy, train_freq=100)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) + task.use(termination_checker(max_train_iter=10000)) + task.use(online_logger()) task.run() diff --git a/ding/example/trex.py b/ding/example/trex.py index 7c43f7fc79..97611ba6c2 100644 --- a/ding/example/trex.py +++ b/ding/example/trex.py @@ -15,16 +15,16 @@ from ding.framework.middleware import OffPolicyLearner, StepCollector, interaction_evaluator, \ eps_greedy_handler, CkptSaver, eps_greedy_masker, sqil_data_pusher, data_pusher from ding.utils import set_pkg_seed -from dizoo.classic_control.cartpole.config.cartpole_trex_dqn_config import main_config, create_config from ding.entry import trex_collecting_data from ding.reward_model import create_reward_model +from dizoo.classic_control.cartpole.config.cartpole_trex_dqn_config import main_config, create_config def main(): logging.getLogger().setLevel(logging.INFO) demo_arg = easydict.EasyDict({'cfg': [main_config, create_config], 'seed': 0}) trex_collecting_data(demo_arg) - cfg = compile_config(main_config, create_cfg=create_config, auto=True) + cfg = compile_config(main_config, create_cfg=create_config, auto=True, renew_dir=False) with task.start(async_mode=False, ctx=OnlineRLContext()): collector_env = BaseEnvManagerV2( @@ -51,7 +51,7 @@ def main(): task.use(StepCollector(cfg, policy.collect_mode, collector_env)) task.use(data_pusher(cfg, buffer_)) task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_, reward_model)) - task.use(CkptSaver(cfg, policy, train_freq=100)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) task.run() diff --git a/ding/framework/__init__.py b/ding/framework/__init__.py index 274d6f2364..72c23d0475 100644 --- a/ding/framework/__init__.py +++ b/ding/framework/__init__.py @@ -1,5 +1,5 @@ from .context import Context, OnlineRLContext, OfflineRLContext -from .task import Task, task +from .task import Task, task, VoidMiddleware from .parallel import Parallel from .event_loop import EventLoop from .supervisor import Supervisor diff --git a/ding/framework/context.py b/ding/framework/context.py index d2e4c58357..12f144766a 100644 --- a/ding/framework/context.py +++ b/ding/framework/context.py @@ -1,5 +1,6 @@ import numpy as np import dataclasses +import treetensor.torch as ttorch from typing import Union, Dict, List @@ -10,6 +11,7 @@ class Context: Context is an object that pass contextual data between middlewares, whose life cycle is only one training iteration. It is a dict that reflect itself, so you can set any properties as you wish. + Note that the initial value of the property must be equal to False. """ _kept_keys: set = dataclasses.field(default_factory=set) total_step: int = 0 @@ -49,20 +51,31 @@ class OnlineRLContext(Context): env_episode: int = 0 train_iter: int = 0 train_data: Union[Dict, List] = None + train_output: Union[Dict, List[Dict]] = None # collect collect_kwargs: Dict = dataclasses.field(default_factory=dict) + obs: ttorch.Tensor = None + action: List = None + inference_output: Dict[int, Dict] = None trajectories: List = None episodes: List = None trajectory_end_idx: List = dataclasses.field(default_factory=list) + action: Dict = None + inference_output: Dict = None # eval eval_value: float = -np.inf last_eval_iter: int = -1 + last_eval_value: int = -np.inf + eval_output: List = dataclasses.field(default_factory=dict) + # wandb + info_for_logging: Dict = dataclasses.field(default_factory=dict) + wandb_url: str = "" def __post_init__(self): # This method is called just after __init__ method. Here, concretely speaking, # this method is called just after the object initialize its fields. # We use this method here to keep the fields needed for each iteration. - self.keep('env_step', 'env_episode', 'train_iter', 'last_eval_iter') + self.keep('env_step', 'env_episode', 'train_iter', 'last_eval_iter', 'last_eval_value', 'wandb_url') @dataclasses.dataclass @@ -70,15 +83,22 @@ class OfflineRLContext(Context): # common total_step: int = 0 + trained_env_step: int = 0 train_epoch: int = 0 train_iter: int = 0 train_data: Union[Dict, List] = None + train_output: Union[Dict, List[Dict]] = None # eval eval_value: float = -np.inf last_eval_iter: int = -1 + last_eval_value: int = -np.inf + eval_output: List = dataclasses.field(default_factory=dict) + # wandb + info_for_logging: Dict = dataclasses.field(default_factory=dict) + wandb_url: str = "" def __post_init__(self): # This method is called just after __init__ method. Here, concretely speaking, # this method is called just after the object initialize its fields. # We use this method here to keep the fields needed for each iteration. - self.keep('train_iter', 'last_eval_iter') + self.keep('trained_env_step', 'train_iter', 'last_eval_iter', 'last_eval_value', 'wandb_url') diff --git a/ding/framework/event_loop.py b/ding/framework/event_loop.py index b5f58720a1..6641d07adb 100644 --- a/ding/framework/event_loop.py +++ b/ding/framework/event_loop.py @@ -1,6 +1,7 @@ from collections import defaultdict from typing import Callable, Optional from concurrent.futures import ThreadPoolExecutor +from copy import copy import fnmatch from ditk import logging @@ -35,7 +36,10 @@ def off(self, event: str, fn: Optional[Callable] = None) -> None: """ for e in fnmatch.filter(self._listeners.keys(), event): if fn: - self._listeners[e].remove(fn) + try: + self._listeners[e].remove(fn) + except: + pass else: self._listeners[e] = [] @@ -79,7 +83,7 @@ def _trigger(self, event: str, *args, **kwargs) -> None: if event not in self._listeners: logging.debug("Event {} is not registered in the callbacks of {}!".format(event, self._name)) return - for fn in self._listeners[event]: + for fn in copy(self._listeners[event]): try: fn(*args, **kwargs) except Exception as e: diff --git a/ding/framework/message_queue/nng.py b/ding/framework/message_queue/nng.py index c0b52f7090..379601b0ed 100644 --- a/ding/framework/message_queue/nng.py +++ b/ding/framework/message_queue/nng.py @@ -30,6 +30,7 @@ def listen(self) -> None: sleep(0.1) # Wait for peers to bind for contact in self.attach_to: sock.dial(contact) + logging.info("NNG listen on {}, attach to {}".format(self.listen_to, self.attach_to)) self._running = True def publish(self, topic: str, data: bytes) -> None: diff --git a/ding/framework/middleware/__init__.py b/ding/framework/middleware/__init__.py index a2c428932c..b9e3c5005d 100644 --- a/ding/framework/middleware/__init__.py +++ b/ding/framework/middleware/__init__.py @@ -1,4 +1,7 @@ from .functional import * -from .collector import StepCollector, EpisodeCollector +from .collector import StepCollector, EpisodeCollector, PPOFStepCollector from .learner import OffPolicyLearner, HERLearner from .ckpt_handler import CkptSaver +from .distributer import ContextExchanger, ModelExchanger, PeriodicalModelExchanger +from .barrier import Barrier, BarrierRuntime +from .data_fetcher import OfflineMemoryDataFetcher diff --git a/ding/framework/middleware/barrier.py b/ding/framework/middleware/barrier.py new file mode 100644 index 0000000000..f958c079ae --- /dev/null +++ b/ding/framework/middleware/barrier.py @@ -0,0 +1,227 @@ +from time import sleep, time +from ditk import logging +from ding.framework import task +from ding.utils.lock_helper import LockContext, LockContextType +from ding.utils.design_helper import SingletonMetaclass + + +class BarrierRuntime(metaclass=SingletonMetaclass): + + def __init__(self, node_id: int, max_world_size: int = 100): + """ + Overview: + 'BarrierRuntime' is a singleton class. In addition, it must be initialized before the + class 'Parallel' starts MQ, otherwise the messages sent by other nodes may be lost after + the detection is completed. We don't have a message retransmission mechanism, and losing + a message means deadlock. + Arguments: + - node_id (int): Process ID. + - max_world_size (int, optional): The maximum total number of processes that can be + synchronized, the defalut value is 100. + """ + self.node_id = node_id + self._has_detected = False + self._range_len = len(str(max_world_size)) + 1 + + self._barrier_epoch = 0 + self._barrier_recv_peers_buff = dict() + self._barrier_recv_peers = dict() + self._barrier_ack_peers = [] + self._barrier_lock = LockContext(LockContextType.THREAD_LOCK) + + self.mq_type = task.router.mq_type + self._connected_peers = dict() + self._connected_peers_lock = LockContext(LockContextType.THREAD_LOCK) + self._keep_alive_daemon = False + + self._event_name_detect = "b_det" + self.event_name_req = "b_req" + self.event_name_ack = "b_ack" + + def _alive_msg_handler(self, peer_id): + with self._connected_peers_lock: + self._connected_peers[peer_id] = time() + + def _add_barrier_req(self, msg): + peer, epoch = self._unpickle_barrier_tag(msg) + logging.debug("Node:[{}] recv barrier request from node:{}, epoch:{}".format(self.node_id, peer, epoch)) + with self._barrier_lock: + if peer not in self._barrier_recv_peers: + self._barrier_recv_peers[peer] = [] + self._barrier_recv_peers[peer].append(epoch) + + def _add_barrier_ack(self, peer): + logging.debug("Node:[{}] recv barrier ack from node:{}".format(self.node_id, peer)) + with self._barrier_lock: + self._barrier_ack_peers.append(peer) + + def _unpickle_barrier_tag(self, msg): + return msg % self._range_len, msg // self._range_len + + def pickle_barrier_tag(self): + return int(self._barrier_epoch * self._range_len + self.node_id) + + def reset_all_peers(self): + with self._barrier_lock: + for peer, q in self._barrier_recv_peers.items(): + if len(q) != 0: + assert q.pop(0) == self._barrier_epoch + self._barrier_ack_peers = [] + self._barrier_epoch += 1 + + def get_recv_num(self): + count = 0 + with self._barrier_lock: + if len(self._barrier_recv_peers) > 0: + for _, q in self._barrier_recv_peers.items(): + if len(q) > 0 and q[0] == self._barrier_epoch: + count += 1 + return count + + def get_ack_num(self): + with self._barrier_lock: + return len(self._barrier_ack_peers) + + def detect_alive(self, expected, timeout): + # The barrier can only block other nodes within the visible range of the current node. + # If the 'attch_to' list of a node is empty, it does not know how many nodes will attach to him, + # so we cannot specify the effective range of a barrier in advance. + assert task._running + task.on(self._event_name_detect, self._alive_msg_handler) + task.on(self.event_name_req, self._add_barrier_req) + task.on(self.event_name_ack, self._add_barrier_ack) + start = time() + while True: + sleep(0.1) + task.emit(self._event_name_detect, self.node_id, only_remote=True) + # In case the other node has not had time to receive our detect message, + # we will send an additional round. + if self._has_detected: + break + with self._connected_peers_lock: + if len(self._connected_peers) == expected: + self._has_detected = True + + if time() - start > timeout: + raise TimeoutError("Node-[{}] timeout when waiting barrier! ".format(task.router.node_id)) + + task.off(self._event_name_detect) + logging.info( + "Barrier detect node done, node-[{}] has connected with {} active nodes!".format(self.node_id, expected) + ) + + +class BarrierContext: + + def __init__(self, runtime: BarrierRuntime, detect_timeout, expected_peer_num: int = 0): + self._runtime = runtime + self._expected_peer_num = expected_peer_num + self._timeout = detect_timeout + + def __enter__(self): + if not self._runtime._has_detected: + self._runtime.detect_alive(self._expected_peer_num, self._timeout) + + def __exit__(self, exc_type, exc_value, tb): + if exc_type is not None: + import traceback + traceback.print_exception(exc_type, exc_value, tb) + self._runtime.reset_all_peers() + + +class Barrier: + + def __init__(self, attch_from_nums: int, timeout: int = 60): + """ + Overview: + Barrier() is a middleware for debug or profiling. It can synchronize the task step of each + process within the scope of all visible processes. When using Barrier(), you need to pay + attention to the following points: + + 1. All processes must call the same number of Barrier(), otherwise a deadlock occurs. + + 2. 'attch_from_nums' is a very important variable, This value indicates the number of times + the current process will be attached to by other processes (the number of connections + established). + For example: + Node0: address: 127.0.0.1:12345, attach_to = [] + Node1: address: 127.0.0.1:12346, attach_to = ["tcp://127.0.0.1:12345"] + For Node0, the 'attch_from_nums' value is 1. (It will be acttched by Node1) + For Node1, the 'attch_from_nums' value is 0. (No one will attach to Node1) + Please note that this value must be given correctly, otherwise, for a node whose 'attach_to' + list is empty, it cannot perceive how many processes will establish connections with it, + resulting in any form of synchronization cannot be performed. + + 3. Barrier() is thread-safe, but it is not recommended to use barrier in multithreading. You need + to carefully calculate the number of times each thread calls Barrier() to avoid deadlock. + + 4. In normal training tasks, please do not use Barrier(), which will force the step synchronization + between each process, so it will greatly damage the training efficiency. In addition, if your + training task has dynamic processes, do not use Barrier() to prevent deadlock. + + Arguments: + - attch_from_nums (int): [description] + - timeout (int, optional): The timeout for successful detection of 'expected_peer_num' + number of nodes, the default value is 60 seconds. + """ + self.node_id = task.router.node_id + self.timeout = timeout + self._runtime: BarrierRuntime = task.router.barrier_runtime + self._barrier_peers_nums = task.get_attch_to_len() + attch_from_nums + + logging.info( + "Node:[{}], attach to num is:{}, attach from num is:{}".format( + self.node_id, task.get_attch_to_len(), attch_from_nums + ) + ) + + def __call__(self, ctx): + self._wait_barrier(ctx) + yield + self._wait_barrier(ctx) + + def _wait_barrier(self, ctx): + self_ready = False + with BarrierContext(self._runtime, self.timeout, self._barrier_peers_nums): + logging.debug("Node:[{}] enter barrier".format(self.node_id)) + # Step1: Notifies all the attached nodes that we have reached the barrier. + task.emit(self._runtime.event_name_req, self._runtime.pickle_barrier_tag(), only_remote=True) + logging.debug("Node:[{}] sended barrier request".format(self.node_id)) + + # Step2: We check the number of flags we have received. + # In the current CI design of DI-engine, there will always be a node whose 'attach_to' list is empty, + # so there will always be a node that will send ACK unconditionally, so deadlock will not occur. + if self._runtime.get_recv_num() == self._barrier_peers_nums: + self_ready = True + + # Step3: Waiting for our own to be ready. + # Even if the current process has reached the barrier, we will not send an ack immediately, + # we need to wait for the slowest directly connected or indirectly connected peer to + # reach the barrier. + start = time() + if not self_ready: + while True: + if time() - start > self.timeout: + raise TimeoutError("Node-[{}] timeout when waiting barrier! ".format(task.router.node_id)) + + if self._runtime.get_recv_num() != self._barrier_peers_nums: + sleep(0.1) + else: + break + + # Step4: Notifies all attached nodes that we are ready. + task.emit(self._runtime.event_name_ack, self.node_id, only_remote=True) + logging.debug("Node:[{}] sended barrier ack".format(self.node_id)) + + # Step5: Wait until all directly or indirectly connected nodes are ready. + start = time() + while True: + if time() - start > self.timeout: + raise TimeoutError("Node-[{}] timeout when waiting barrier! ".format(task.router.node_id)) + + if self._runtime.get_ack_num() != self._barrier_peers_nums: + sleep(0.1) + else: + break + + logging.info("Node-[{}] env_step:[{}] barrier finish".format(self.node_id, ctx.env_step)) diff --git a/ding/framework/middleware/ckpt_handler.py b/ding/framework/middleware/ckpt_handler.py index 27ffd3fa2c..ca75f16618 100644 --- a/ding/framework/middleware/ckpt_handler.py +++ b/ding/framework/middleware/ckpt_handler.py @@ -17,22 +17,32 @@ class CkptSaver: The class used to save checkpoint data. """ - def __init__(self, cfg: EasyDict, policy: Policy, train_freq: Optional[int] = None): + def __new__(cls, *args, **kwargs): + if task.router.is_active and not (task.has_role(task.role.LEARNER) or task.has_role(task.role.EVALUATOR)): + return task.void() + return super(CkptSaver, cls).__new__(cls) + + def __init__(self, policy: Policy, save_dir: str, train_freq: Optional[int] = None, save_finish: bool = True): """ Overview: Initialize the `CkptSaver`. Arguments: - - cfg (:obj:`EasyDict`): Config which should contain the following keys: `cfg.exp_name`. - policy (:obj:`Policy`): Policy used to save the checkpoint. + - save_dir (:obj:`str`): The directory path to save ckpt. - train_freq (:obj:`int`): Number of training iterations between each saving checkpoint data. + - save_finish (:obj:`bool`): Whether save final ckpt when ``task.finish = True``. """ self.policy = policy self.train_freq = train_freq - self.prefix = '{}/ckpt'.format(cfg.exp_name) + if str(os.path.basename(os.path.normpath(save_dir))) != "ckpt": + self.prefix = '{}/ckpt'.format(os.path.normpath(save_dir)) + else: + self.prefix = '{}/'.format(os.path.normpath(save_dir)) if not os.path.exists(self.prefix): os.makedirs(self.prefix) self.last_save_iter = 0 self.max_eval_value = -np.inf + self.save_finish = save_finish def __call__(self, ctx: Union["OnlineRLContext", "OfflineRLContext"]) -> None: """ @@ -40,24 +50,25 @@ def __call__(self, ctx: Union["OnlineRLContext", "OfflineRLContext"]) -> None: The method used to save checkpoint data. \ The checkpoint data will be saved in a file in following 3 cases: \ - When a multiple of `self.train_freq` iterations have elapsed since the beginning of training; \ - - When the evaluation reward is the best eval reward so far; \ + - When the evaluation episode return is the best so far; \ - When `task.finish` is True. Input of ctx: - train_iter (:obj:`int`): Number of training iteration, i.e. the number of updating policy related network. - - eval_value (:obj:`float`): The eval reward of current iteration. + - eval_value (:obj:`float`): The episode return of current iteration. """ # train enough iteration - if self.train_freq and ctx.train_iter - self.last_save_iter >= self.train_freq: - save_file( - "{}/iteration_{}.pth.tar".format(self.prefix, ctx.train_iter), self.policy.learn_mode.state_dict() - ) - self.last_save_iter = ctx.train_iter - - # best eval reward so far - if ctx.eval_value > self.max_eval_value: + if self.train_freq: + if ctx.train_iter == 0 or ctx.train_iter - self.last_save_iter >= self.train_freq: + save_file( + "{}/iteration_{}.pth.tar".format(self.prefix, ctx.train_iter), self.policy.learn_mode.state_dict() + ) + self.last_save_iter = ctx.train_iter + + # best episode return so far + if ctx.eval_value is not None and ctx.eval_value > self.max_eval_value: save_file("{}/eval.pth.tar".format(self.prefix), self.policy.learn_mode.state_dict()) self.max_eval_value = ctx.eval_value # finish - if task.finish: + if task.finish and self.save_finish: save_file("{}/final.pth.tar".format(self.prefix), self.policy.learn_mode.state_dict()) diff --git a/ding/framework/middleware/collector.py b/ding/framework/middleware/collector.py index fa70a00766..d40ba312eb 100644 --- a/ding/framework/middleware/collector.py +++ b/ding/framework/middleware/collector.py @@ -1,5 +1,6 @@ -from typing import TYPE_CHECKING, Callable, List +from typing import TYPE_CHECKING from easydict import EasyDict +import treetensor.torch as ttorch from ding.policy import get_random_policy from ding.envs import BaseEnvManager @@ -17,6 +18,11 @@ class StepCollector: process. Use the `__call__` method to execute the whole collection process. """ + def __new__(cls, *args, **kwargs): + if task.router.is_active and not task.has_role(task.role.COLLECTOR): + return task.void() + return super(StepCollector, cls).__new__(cls) + def __init__(self, cfg: EasyDict, policy, env: BaseEnvManager, random_collect_size: int = 0) -> None: """ Arguments: @@ -32,8 +38,8 @@ def __init__(self, cfg: EasyDict, policy, env: BaseEnvManager, random_collect_si self.policy = policy self.random_collect_size = random_collect_size self._transitions = TransitionList(self.env.env_num) - self._inferencer = task.wrap(inferencer(cfg, policy, env)) - self._rolloutor = task.wrap(rolloutor(cfg, policy, env, self._transitions)) + self._inferencer = task.wrap(inferencer(cfg.seed, policy, env)) + self._rolloutor = task.wrap(rolloutor(policy, env, self._transitions)) def __call__(self, ctx: "OnlineRLContext") -> None: """ @@ -47,7 +53,7 @@ def __call__(self, ctx: "OnlineRLContext") -> None: if self.random_collect_size > 0 and old < self.random_collect_size: target_size = self.random_collect_size - old random_policy = get_random_policy(self.cfg, self.policy, self.env) - current_inferencer = task.wrap(inferencer(self.cfg, random_policy, self.env)) + current_inferencer = task.wrap(inferencer(self.cfg.seed, random_policy, self.env)) else: # compatible with old config, a train sample = unroll_len step target_size = self.cfg.policy.collect.n_sample * self.cfg.policy.collect.unroll_len @@ -62,11 +68,82 @@ def __call__(self, ctx: "OnlineRLContext") -> None: break +class PPOFStepCollector: + """ + Overview: + The class of the collector running by steps, including model inference and transition \ + process. Use the `__call__` method to execute the whole collection process. + """ + + def __new__(cls, *args, **kwargs): + if task.router.is_active and not task.has_role(task.role.COLLECTOR): + return task.void() + return super(PPOFStepCollector, cls).__new__(cls) + + def __init__(self, seed: int, policy, env: BaseEnvManager, n_sample: int, unroll_len: int = 1) -> None: + """ + Arguments: + - seed (:obj:`int`): Random seed. + - policy (:obj:`Policy`): The policy to be collected. + - env (:obj:`BaseEnvManager`): The env for the collection, the BaseEnvManager object or \ + its derivatives are supported. + """ + self.env = env + self.env.seed(seed) + self.policy = policy + self.n_sample = n_sample + self.unroll_len = unroll_len + self._transitions = TransitionList(self.env.env_num) + self._env_episode_id = [_ for _ in range(env.env_num)] + self._current_id = env.env_num + + def __call__(self, ctx: "OnlineRLContext") -> None: + """ + Overview: + An encapsulation of inference and rollout middleware. Stop when completing \ + the target number of steps. + Input of ctx: + - env_step (:obj:`int`): The env steps which will increase during collection. + """ + device = self.policy._device + old = ctx.env_step + target_size = self.n_sample * self.unroll_len + + if self.env.closed: + self.env.launch() + + while True: + obs = ttorch.as_tensor(self.env.ready_obs).to(dtype=ttorch.float32) + obs = obs.to(device) + inference_output = self.policy.collect(obs, **ctx.collect_kwargs) + inference_output = inference_output.cpu() + action = inference_output.action.numpy() + timesteps = self.env.step(action) + ctx.env_step += len(timesteps) + + obs = obs.cpu() + for i, timestep in enumerate(timesteps): + transition = self.policy.process_transition(obs[i], inference_output[i], timestep) + transition.collect_train_iter = ttorch.as_tensor([ctx.train_iter]) + transition.env_data_id = ttorch.as_tensor([self._env_episode_id[timestep.env_id]]) + self._transitions.append(timestep.env_id, transition) + if timestep.done: + self.policy.reset([timestep.env_id]) + self._env_episode_id[timestep.env_id] = self._current_id + self._current_id += 1 + ctx.env_episode += 1 + + if ctx.env_step - old >= target_size: + ctx.trajectories, ctx.trajectory_end_idx = self._transitions.to_trajectories() + self._transitions.clear() + break + + class EpisodeCollector: """ Overview: The class of the collector running by episodes, including model inference and transition \ - process. Use the `__call__` method to execute the whole collection process. + process. Use the `__call__` method to execute the whole collection process. """ def __init__(self, cfg: EasyDict, policy, env: BaseEnvManager, random_collect_size: int = 0) -> None: @@ -84,14 +161,14 @@ def __init__(self, cfg: EasyDict, policy, env: BaseEnvManager, random_collect_si self.policy = policy self.random_collect_size = random_collect_size self._transitions = TransitionList(self.env.env_num) - self._inferencer = task.wrap(inferencer(cfg, policy, env)) - self._rolloutor = task.wrap(rolloutor(cfg, policy, env, self._transitions)) + self._inferencer = task.wrap(inferencer(cfg.seed, policy, env)) + self._rolloutor = task.wrap(rolloutor(policy, env, self._transitions)) def __call__(self, ctx: "OnlineRLContext") -> None: """ Overview: An encapsulation of inference and rollout middleware. Stop when completing the \ - target number of episodes. + target number of episodes. Input of ctx: - env_episode (:obj:`int`): The env env_episode which will increase during collection. """ diff --git a/ding/framework/middleware/data_fetcher.py b/ding/framework/middleware/data_fetcher.py new file mode 100644 index 0000000000..2103a8668d --- /dev/null +++ b/ding/framework/middleware/data_fetcher.py @@ -0,0 +1,99 @@ +from typing import TYPE_CHECKING +from threading import Thread, Event +from queue import Queue +import time +import numpy as np +import torch +from easydict import EasyDict +from ding.framework import task +from ding.data import Dataset, DataLoader +from ding.utils import get_rank, get_world_size + +if TYPE_CHECKING: + from ding.framework import OfflineRLContext + + +class OfflineMemoryDataFetcher: + + def __new__(cls, *args, **kwargs): + if task.router.is_active and not task.has_role(task.role.FETCHER): + return task.void() + return super(OfflineMemoryDataFetcher, cls).__new__(cls) + + def __init__(self, cfg: EasyDict, dataset: Dataset): + device = 'cuda:{}'.format(get_rank() % torch.cuda.device_count()) if cfg.policy.cuda else 'cpu' + if device != 'cpu': + stream = torch.cuda.Stream() + + def producer(queue, dataset, batch_size, device, event): + torch.set_num_threads(4) + if device != 'cpu': + nonlocal stream + sbatch_size = batch_size * get_world_size() + rank = get_rank() + idx_list = np.random.permutation(len(dataset)) + temp_idx_list = [] + for i in range(len(dataset) // sbatch_size): + temp_idx_list.extend(idx_list[i + rank * batch_size:i + (rank + 1) * batch_size]) + idx_iter = iter(temp_idx_list) + + if device != 'cpu': + with torch.cuda.stream(stream): + while True: + if queue.full(): + time.sleep(0.1) + else: + data = [] + for _ in range(batch_size): + try: + data.append(dataset.__getitem__(next(idx_iter))) + except StopIteration: + del idx_iter + idx_list = np.random.permutation(len(dataset)) + idx_iter = iter(idx_list) + data.append(dataset.__getitem__(next(idx_iter))) + data = [[i[j] for i in data] for j in range(len(data[0]))] + data = [torch.stack(x).to(device) for x in data] + queue.put(data) + if event.is_set(): + break + else: + while True: + if queue.full(): + time.sleep(0.1) + else: + data = [] + for _ in range(batch_size): + try: + data.append(dataset.__getitem__(next(idx_iter))) + except StopIteration: + del idx_iter + idx_list = np.random.permutation(len(dataset)) + idx_iter = iter(idx_list) + data.append(dataset.__getitem__(next(idx_iter))) + data = [[i[j] for i in data] for j in range(len(data[0]))] + data = [torch.stack(x) for x in data] + queue.put(data) + if event.is_set(): + break + + self.queue = Queue(maxsize=50) + self.event = Event() + self.producer_thread = Thread( + target=producer, + args=(self.queue, dataset, cfg.policy.batch_size, device, self.event), + name='cuda_fetcher_producer' + ) + + def __call__(self, ctx: "OfflineRLContext"): + if not self.producer_thread.is_alive(): + time.sleep(5) + self.producer_thread.start() + while self.queue.empty(): + time.sleep(0.001) + ctx.train_data = self.queue.get() + + def __del__(self): + if self.producer_thread.is_alive(): + self.event.set() + del self.queue diff --git a/ding/framework/middleware/distributer.py b/ding/framework/middleware/distributer.py new file mode 100644 index 0000000000..d2f5e36402 --- /dev/null +++ b/ding/framework/middleware/distributer.py @@ -0,0 +1,415 @@ +import numpy as np +from time import sleep, time +from dataclasses import fields +from typing import TYPE_CHECKING, List, Dict, Any, Optional, Union +from ditk import logging +from ding.framework import task +from ding.data import StorageLoader, Storage, ModelLoader +if TYPE_CHECKING: + from ding.framework.context import Context + from torch.nn import Module + + +class ContextExchanger: + + def __init__(self, skip_n_iter: int = 1, storage_loader: Optional[StorageLoader] = None) -> None: + """ + Overview: + Exchange context between processes, + support properties: trajectories, episodes, env_step, env_episode, train_iter + Arguments: + - skip_n_iter (:obj:`int`): For collectors, it may be necessary to skip waiting \ + for the first n iterations to collect data for the learner to learn. This parameter \ + will not work on learner. + - storage_loader (:obj:`Optional[StorageLoader]`): Turn data into storage class to reduce \ + the network overhead. + """ + if not task.router.is_active: + raise RuntimeError("ContextHandler should be used in parallel mode!") + self._state = {} + self._local_state = {} # just save local state, not send to remote node + if task.has_role(task.role.COLLECTOR): + self._local_state['env_step'] = 0 + self._local_state['env_episode'] = 0 + self._event_name = "context_exchanger_{role}" + self._skip_n_iter = skip_n_iter + self._storage_loader = storage_loader + for role in task.role: # Only subscribe to other roles + if not task.has_role(role): + task.on(self._event_name.format(role=role), self.put) + if storage_loader: + task.once("finish", lambda _: storage_loader.shutdown()) + + def __new__(cls, *args, **kwargs): + if not task.router.is_active: + return task.void() + + if len(task.roles) == 0: + logging.warning("The task does not have any roles defined, the ContextExchanger will not work.") + return task.void() + + if len(task.roles) > 1: + logging.warning( + "Use multiple roles in one exchanger may lead to unexpected result, please check your code." + ) + + return super(ContextExchanger, cls).__new__(cls) + + def __call__(self, ctx: "Context"): + self.merge(ctx) + yield + payload = self.fetch(ctx) + if payload: + if self._storage_loader and task.has_role(task.role.COLLECTOR): + payload = self._storage_loader.save(payload) + for role in task.roles: + task.emit(self._event_name.format(role=role), payload, only_remote=True) + + def __del__(self): + if self._storage_loader: + self._storage_loader.shutdown() + + def put(self, payload: Union[Dict, Storage]): + """ + Overview: + Get attributes from ctx on the callback of event. + Each attribute should have a standalone put handler, which named `_put_{key}` + """ + + def callback(payload: Dict): + for key, item in payload.items(): + fn_name = "_put_{}".format(key) + if hasattr(self, fn_name): + getattr(self, fn_name)(item) + else: + logging.warning("Receive unexpected key ({}) in context exchanger".format(key)) + + if isinstance(payload, Storage): + assert self._storage_loader is not None, "Storage loader is not defined when data is a storage object." + self._storage_loader.load(payload, callback) + else: + callback(payload) + + def fetch(self, ctx: "Context") -> Dict[str, Any]: + """ + Overview: + Fetch attributes from ctx before emit them to the event bus. + Each attribute should have a standalone fetch handler, which named `_fetch_{key}` + """ + payload = {} + for field in fields(ctx): + key, item = field.name, getattr(ctx, field.name) + fn_name = "_fetch_{}".format(key) + if hasattr(self, fn_name): + value = getattr(self, fn_name)(item) + if value is not None: + payload[key] = value + return payload + + def merge(self, ctx: "Context"): + if task.has_role(task.role.LEARNER): + # Learner should always wait for trajs. + # TODO: Automaticlly wait based on properties, not roles. + while len(self._state) == 0: + sleep(0.01) + elif ctx.total_step >= self._skip_n_iter: + start = time() + while len(self._state) == 0: + if time() - start > 60: + logging.warning("Timeout when waiting for new context! Node id: {}".format(task.router.node_id)) + break + sleep(0.01) + + for k, v in self._state.items(): + if not task.has_role(task.role.COLLECTOR) and k.startswith('increment_'): + pure_k = k.split('increment_')[-1] + setattr(ctx, pure_k, getattr(ctx, pure_k) + v) + else: + setattr(ctx, k, v) + self._state = {} + + # Handle each attibute of context + def _put_trajectories(self, traj: List[Any]): + if not task.has_role(task.role.LEARNER): + return + if "trajectories" not in self._state: + self._state["trajectories"] = [] + self._state["trajectories"].extend(traj) + + def _fetch_trajectories(self, traj: List[Any]): + if task.has_role(task.role.COLLECTOR): + return traj + + def _put_episodes(self, episodes: List[Any]): + if not task.has_role(task.role.LEARNER): + return + if "episodes" not in self._state: + self._state["episodes"] = [] + self._state["episodes"].extend(episodes) + + def _fetch_episodes(self, episodes: List[Any]): + if task.has_role(task.role.COLLECTOR): + return episodes + + def _put_trajectory_end_idx(self, trajectory_end_idx: List[str]): + if not task.has_role(task.role.LEARNER): + return + if "trajectory_end_idx" not in self._state: + self._state["trajectory_end_idx"] = [] + self._state["trajectory_end_idx"].extend(trajectory_end_idx) + + def _fetch_trajectory_end_idx(self, trajectory_end_idx: List[str]): + if task.has_role(task.role.COLLECTOR): + return trajectory_end_idx + + def _put_env_step(self, increment_env_step: int): + if not task.has_role(task.role.COLLECTOR): + if 'increment_env_step' not in self._state: + self._state['increment_env_step'] = 0 + self._state["increment_env_step"] += increment_env_step + + def _fetch_env_step(self, env_step: int): + if task.has_role(task.role.COLLECTOR): + increment_env_step = env_step - self._local_state['env_step'] + self._local_state['env_step'] = env_step + return increment_env_step + + def _put_env_episode(self, increment_env_episode: int): + if not task.has_role(task.role.COLLECTOR): + if 'increment_env_episode' not in self._state: + self._state['increment_env_episode'] = 0 + self._state["increment_env_episode"] += increment_env_episode + + def _fetch_env_episode(self, env_episode: int): + if task.has_role(task.role.COLLECTOR): + increment_env_episode = env_episode - self._local_state['env_episode'] + self._local_state['env_episode'] = env_episode + return increment_env_episode + + def _put_train_iter(self, train_iter: int): + if not task.has_role(task.role.LEARNER): + self._state["train_iter"] = train_iter + + def _fetch_train_iter(self, train_iter: int): + if task.has_role(task.role.LEARNER): + return train_iter + + +class ModelExchanger: + + def __init__(self, model: "Module", model_loader: Optional[ModelLoader] = None) -> None: + """ + Overview: + Exchange model between processes, only the learner will send the model, + otherwise the model will only be received. + If you are using a shared model on a single host, there is no need to use this middleware. + Arguments: + - model (:obj:`torch.nn.Module`): Pytorch module. + - model_loader (:obj:`ModelLoader`): Encode model in subprocess. + """ + self._model = model + self._model_loader = model_loader + self._event_name = "model_exchanger" + self._state_dict_cache: Optional[Union[object, Storage]] = None + self._is_learner = task.has_role(task.role.LEARNER) + if not self._is_learner: + task.on(self._event_name, self._cache_state_dict) + if model_loader: + task.once("finish", lambda _: model_loader.shutdown()) + + def _cache_state_dict(self, state_dict: Union[object, Storage]): + self._state_dict_cache = state_dict + + def __new__(cls, *args, **kwargs): + if not task.router.is_active: + return task.void() + + if len(task.roles) == 0: + logging.warning("The task does not have any roles defined, the ModelExchanger will not work.") + return task.void() + + if len(task.roles) > 1: + logging.warning( + "Use multiple roles in one exchanger may lead to unexpected result, please check your code." + ) + + return super(ModelExchanger, cls).__new__(cls) + + def __call__(self, ctx: "Context") -> Any: + if self._model_loader: + self._model_loader.start() + + if not self._is_learner: + if ctx.total_step != 0: # Skip first iteration + self._update_model() + else: + yield + self._send_model() + + def _update_model(self): + start = time() + while True: + if task.finish: + return + if time() - start > 60: + logging.warning("Timeout when waiting for new model! Node id: {}".format(task.router.node_id)) + break + if self._state_dict_cache is None: + sleep(0.01) + else: + if isinstance(self._state_dict_cache, Storage) and self._model_loader is not None: + try: + self._model.load_state_dict(self._model_loader.load(self._state_dict_cache)) + self._state_dict_cache = None + break + except FileNotFoundError as e: + logging.warning( + "Model file has been deleted on node {}, maybe you can increase the ttl.".format( + task.router.node_id + ) + ) + self._state_dict_cache = None + continue + else: + self._model.load_state_dict(self._state_dict_cache) + self._state_dict_cache = None + break + + def _send_model(self): + if self._model_loader: + self._model_loader.save(self._send_callback) + else: + task.emit(self._event_name, self._model.state_dict(), only_remote=True) + + def _send_callback(self, storage: Storage): + if task.running: + task.emit(self._event_name, storage, only_remote=True) + + def __del__(self): + if self._model_loader: + self._model_loader.shutdown() + + +class PeriodicalModelExchanger: + + def __init__( + self, + model: "Module", + mode: str, + period: int = 1, + delay_toleration: float = np.inf, + stale_toleration: int = 1, + event_name: str = "model_exchanger", + model_loader: Optional[ModelLoader] = None + ) -> None: + """ + Overview: + Exchange model between processes, set the mode to "send" or "receive" to specify the role of the process. + If you are using a shared model on a single host, there is no need to use this middleware. + Arguments: + - model (:obj:`torch.nn.Module`): Pytorch module. + - mode (:obj:`str`): "send" or "receive". + - period (:obj:`int`): The period of model exchange. + - delay_toleration (:obj:`float`): The permitted time interval for receiving model after being sent. + - stale_toleration (:obj:`int`): The permitted number of iterations for receiving model after being sent. + - event_name (:obj:`str`): The event name for model exchange. + - model_loader (:obj:`ModelLoader`): ModelLoader for this PeriodicalModelExchanger to use. + """ + self._model = model + self._model_loader = model_loader + self._event_name = event_name + self._period = period + self._mode = mode + if self._mode == "receive": + self._id_counter = -1 + self._model_id = -1 + else: + self._id_counter = 0 + self._stale_toleration = stale_toleration + self._model_stale = stale_toleration + self._delay_toleration = delay_toleration + self._state_dict_cache: Optional[Union[object, Storage]] = None + + if self._mode == "receive": + task.on(self._event_name, self._cache_state_dict) + if model_loader: + task.once("finish", lambda _: model_loader.shutdown()) + + def _cache_state_dict(self, msg: Dict[str, Any]): + if msg['id'] % self._period == 0: + self._state_dict_cache = msg['model'] + self._id_counter = msg['id'] + self._time = msg['time'] + + def __new__(cls, *args, **kwargs): + return super(PeriodicalModelExchanger, cls).__new__(cls) + + def __call__(self, ctx: "Context") -> Any: + if self._model_loader: + self._model_loader.start() + + if self._mode == "receive": + if ctx.total_step != 0: # Skip first iteration + self._update_model() + elif self._mode == "send": + yield + if self._id_counter % self._period == 0: + self._send_model(id=self._id_counter) + self._id_counter += 1 + else: + raise NotImplementedError + + def _update_model(self): + start = time() + while True: + if task.finish: + return + if time() - start > 60: + logging.warning("Timeout when waiting for new model! Node id: {}".format(task.router.node_id)) + self._model_stale += 1 + break + if self._state_dict_cache is None: + if self._model_stale < self._stale_toleration and time() - self._time < self._delay_toleration: + self._model_stale += 1 + break + else: + sleep(0.01) + else: + if self._id_counter > self._model_id and time() - self._time < self._delay_toleration: + if isinstance(self._state_dict_cache, Storage) and self._model_loader is not None: + try: + self._model.load_state_dict(self._model_loader.load(self._state_dict_cache)) + self._state_dict_cache = None + self._model_id = self._id_counter + self._model_stale = 1 + break + except FileNotFoundError as e: + logging.warning( + "Model file has been deleted on node {}, maybe you can increase the ttl.".format( + task.router.node_id + ) + ) + self._state_dict_cache = None + continue + else: + self._model.load_state_dict(self._state_dict_cache) + self._state_dict_cache = None + self._model_id = self._id_counter + self._model_stale = 1 + break + else: + self._model_stale += 1 + + def _send_model(self, id: int): + if self._model_loader: + self._model_loader.save(self._send_callback) + else: + task.emit(self._event_name, {'id': id, 'model': self._model.state_dict(), 'time': time()}, only_remote=True) + + def _send_callback(self, storage: Storage): + if task.running: + task.emit(self._event_name, storage, only_remote=True) + + def __del__(self): + if self._model_loader: + self._model_loader.shutdown() diff --git a/ding/framework/middleware/functional/__init__.py b/ding/framework/middleware/functional/__init__.py index f352790be6..8474f2626e 100644 --- a/ding/framework/middleware/functional/__init__.py +++ b/ding/framework/middleware/functional/__init__.py @@ -1,13 +1,15 @@ from .trainer import trainer, multistep_trainer from .data_processor import offpolicy_data_fetcher, data_pusher, offline_data_fetcher, offline_data_saver, \ - sqil_data_pusher + offline_data_fetcher_from_mem, sqil_data_pusher, buffer_saver from .collector import inferencer, rolloutor, TransitionList -from .evaluator import interaction_evaluator -from .termination_checker import termination_checker -from .logger import online_logger, offline_logger +from .evaluator import interaction_evaluator, interaction_evaluator_ttorch +from .termination_checker import termination_checker, ddp_termination_checker +from .logger import online_logger, offline_logger, wandb_online_logger, wandb_offline_logger from .ctx_helper import final_ctx_saver # algorithm from .explorer import eps_greedy_handler, eps_greedy_masker -from .advantage_estimator import gae_estimator +from .advantage_estimator import gae_estimator, ppof_adv_estimator, montecarlo_return_estimator from .enhancer import reward_estimator, her_data_enhancer, nstep_reward_enhancer +from .priority import priority_calculator +from .timer import epoch_timer diff --git a/ding/framework/middleware/functional/advantage_estimator.py b/ding/framework/middleware/functional/advantage_estimator.py index 56d9a9b01d..6daf8d4528 100644 --- a/ding/framework/middleware/functional/advantage_estimator.py +++ b/ding/framework/middleware/functional/advantage_estimator.py @@ -5,9 +5,10 @@ import treetensor.torch as ttorch from ding.policy import Policy from ding.data import Buffer -from ding.rl_utils import gae, gae_data +from ding.rl_utils import gae, gae_data, get_train_sample from ding.framework import task from ding.utils.data import ttorch_collate +from ding.utils.dict_helper import convert_easy_dict_to_dict from ding.torch_utils import to_device if TYPE_CHECKING: @@ -26,8 +27,16 @@ def gae_estimator(cfg: EasyDict, policy: Policy, buffer_: Optional[Buffer] = Non - policy (:obj:`Policy`): Policy in `policy.collect_mode`, used to get model to calculate value. - buffer\_ (:obj:`Optional[Buffer]`): The `buffer_` to push the processed data in if `buffer_` is not None. """ + if task.router.is_active and not task.has_role(task.role.LEARNER): + return task.void() model = policy.get_attribute('model') + if buffer_ is not None: + # Unify the shape of obs and action + obs_shape = cfg['policy']['model']['obs_shape'] + obs_shape = torch.Size(torch.tensor(obs_shape)) if isinstance(obs_shape, list) \ + else ttorch.size.Size(convert_easy_dict_to_dict(obs_shape)) if isinstance(obs_shape, dict) \ + else torch.Size(torch.tensor(obs_shape).unsqueeze(0)) def _gae(ctx: "OnlineRLContext"): """ @@ -40,26 +49,108 @@ def _gae(ctx: "OnlineRLContext"): Output of ctx: - train_data (:obj:`List[treetensor.torch.Tensor]`): The processed data if `buffer_` is None. """ - data = ctx.trajectories # list - data = ttorch_collate(data) + cuda = cfg.policy.cuda and torch.cuda.is_available() + + # action shape (B,) for discete action, (B, D,) for continuous action + # reward shape (B,) done shape (B,) value shape (B,) + data = ttorch_collate(ctx.trajectories, cat_1dim=True) + if data['action'].dtype in [torch.float16, torch.float32, torch.double] \ + and data['action'].dim() == 1: + # action shape + data['action'] = data['action'].unsqueeze(-1) + with torch.no_grad(): - value = model.forward(data.obs, mode='compute_critic')['value'] - next_value = model.forward(data.next_obs, mode='compute_critic')['value'] + if cuda: + data = data.cuda() + value = model.forward(data.obs.to(dtype=ttorch.float32), mode='compute_critic')['value'] + next_value = model.forward(data.next_obs.to(dtype=ttorch.float32), mode='compute_critic')['value'] data.value = value - traj_flag = data.done.clone() - traj_flag[ctx.trajectory_end_idx] = True - data.traj_flag = traj_flag + traj_flag = data.done.clone() + traj_flag[ctx.trajectory_end_idx] = True + data.traj_flag = traj_flag - # done is bool type when acquired from env.step - data_ = gae_data(data.value, next_value, data.reward, data.done.float(), traj_flag.float()) - data.adv = gae(data_, cfg.policy.collect.discount_factor, cfg.policy.collect.gae_lambda) + # done is bool type when acquired from env.step + data_ = gae_data(data.value, next_value, data.reward, data.done.float(), traj_flag.float()) + data.adv = gae(data_, cfg.policy.collect.discount_factor, cfg.policy.collect.gae_lambda) if buffer_ is None: ctx.train_data = data else: + data = data.cpu() data = ttorch.split(data, 1) + # To ensure the shape of obs is same as config + if data[0]['obs'].shape == obs_shape: + pass + elif data[0]['obs'].shape[0] == 1 and data[0]['obs'].shape[1:] == obs_shape: + for d in data: + d['obs'] = d['obs'].squeeze(0) + d['next_obs'] = d['next_obs'].squeeze(0) + if 'logit' in data[0]: + for d in data: + d['logit'] = d['logit'].squeeze(0) + if 'log_prob' in data[0]: + for d in data: + d['log_prob'] = d['log_prob'].squeeze(0) + else: + raise RuntimeError("The shape of obs is {}, which is not same as config.".format(data[0]['obs'].shape)) + + if data[0]['action'].dtype in [torch.float16, torch.float32, torch.double] \ + and data[0]['action'].dim() == 2: + for d in data: + d['action'] = d['action'].squeeze(0) for d in data: buffer_.push(d) ctx.trajectories = None return _gae + + +def ppof_adv_estimator(policy: Policy) -> Callable: + if task.router.is_active and not task.has_role(task.role.LEARNER): + return task.void() + + def _estimator(ctx: "OnlineRLContext"): + data = ttorch_collate(ctx.trajectories, cat_1dim=True) + if data['action'].dtype == torch.float32 and data['action'].dim() == 1: + data['action'] = data['action'].unsqueeze(-1) + traj_flag = data.done.clone() + traj_flag[ctx.trajectory_end_idx] = True + data.traj_flag = traj_flag + ctx.train_data = data + + return _estimator + + +def montecarlo_return_estimator(policy: Policy) -> Callable: + if task.router.is_active and not task.has_role(task.role.LEARNER): + return task.void() + + def pg_policy_get_train_sample(data): + assert data[-1]['done'], "PG needs a complete epsiode" + + if policy._cfg.learn.ignore_done: + raise NotImplementedError + + R = 0. + if isinstance(data, ttorch.Tensor): + data_size = data['done'].shape[0] + data['return'] = ttorch.Tensor([0.0 for i in range(data_size)]) + for i in reversed(range(data_size)): + R = policy._gamma * R + data['reward'][i] + data['return'][i] = R + return get_train_sample(data, policy._unroll_len) + else: + raise ValueError + + def _estimator(ctx: "OnlineRLContext"): + train_data = [] + for episode in ctx.episodes: + data = ttorch_collate(episode, cat_1dim=True) + if data['action'].dtype in [torch.float16, torch.float32, torch.double] \ + and data['action'].dim() == 1: + data['action'] = data['action'].unsqueeze(-1) + data = pg_policy_get_train_sample(data) + train_data.append(data) + ctx.train_data = ttorch.cat(train_data, dim=0) + + return _estimator diff --git a/ding/framework/middleware/functional/collector.py b/ding/framework/middleware/functional/collector.py index fa47e7cf83..d87d1e0997 100644 --- a/ding/framework/middleware/functional/collector.py +++ b/ding/framework/middleware/functional/collector.py @@ -1,10 +1,12 @@ from typing import TYPE_CHECKING, Callable, List, Tuple, Any -from easydict import EasyDict from functools import reduce import treetensor.torch as ttorch +import numpy as np +from ditk import logging +from ding.utils import EasyTimer from ding.envs import BaseEnvManager from ding.policy import Policy -from ding.torch_utils import to_ndarray +from ding.torch_utils import to_ndarray, get_shape0 if TYPE_CHECKING: from ding.framework import OnlineRLContext @@ -45,36 +47,38 @@ def clear(self): item.clear() -def inferencer(cfg: EasyDict, policy: Policy, env: BaseEnvManager) -> Callable: +def inferencer(seed: int, policy: Policy, env: BaseEnvManager) -> Callable: """ Overview: The middleware that executes the inference process. Arguments: - - cfg (:obj:`EasyDict`): Config. + - seed (:obj:`int`): Random seed. - policy (:obj:`Policy`): The policy to be inferred. - env (:obj:`BaseEnvManager`): The env where the inference process is performed. \ The env.ready_obs (:obj:`tnp.array`) will be used as model input. """ - env.seed(cfg.seed) + env.seed(seed) def _inference(ctx: "OnlineRLContext"): """ Output of ctx: - - obs (:obj:`Dict[Tensor]`): The input states fed into the model. + - obs (:obj:`Union[torch.Tensor, Dict[torch.Tensor]]`): The input observations collected \ + from all collector environments. - action: (:obj:`List[np.ndarray]`): The inferred actions listed by env_id. - - inference_output (:obj:`Dict[int, Dict]`): The dict that contains env_id (int) \ - and inference result (Dict). + - inference_output (:obj:`Dict[int, Dict]`): The dict of which the key is env_id (int), \ + and the value is inference result (Dict). """ if env.closed: env.launch() - obs = ttorch.as_tensor(env.ready_obs).to(dtype=ttorch.float32) + obs = ttorch.as_tensor(env.ready_obs) ctx.obs = obs + obs = obs.to(dtype=ttorch.float32) # TODO mask necessary rollout - obs = {i: obs[i] for i in range(obs.shape[0])} # TBD + obs = {i: obs[i] for i in range(get_shape0(obs))} # TBD inference_output = policy.forward(obs, **ctx.collect_kwargs) ctx.action = [to_ndarray(v['action']) for v in inference_output.values()] # TBD ctx.inference_output = inference_output @@ -82,12 +86,16 @@ def _inference(ctx: "OnlineRLContext"): return _inference -def rolloutor(cfg: EasyDict, policy: Policy, env: BaseEnvManager, transitions: TransitionList) -> Callable: +def rolloutor( + policy: Policy, + env: BaseEnvManager, + transitions: TransitionList, + collect_print_freq=100, +) -> Callable: """ Overview: The middleware that executes the transition process in the env. Arguments: - - cfg (:obj:`EasyDict`): Config. - policy (:obj:`Policy`): The policy to be used during transition. - env (:obj:`BaseEnvManager`): The env for the collection, the BaseEnvManager object or \ its derivatives are supported. @@ -98,6 +106,13 @@ def rolloutor(cfg: EasyDict, policy: Policy, env: BaseEnvManager, transitions: T env_episode_id = [_ for _ in range(env.env_num)] current_id = env.env_num + timer = EasyTimer() + last_train_iter = 0 + total_envstep_count = 0 + total_episode_count = 0 + total_train_sample_count = 0 + env_info = {env_id: {'time': 0., 'step': 0, 'train_sample': 0} for env_id in range(env.env_num)} + episode_info = [] def _rollout(ctx: "OnlineRLContext"): """ @@ -113,22 +128,88 @@ def _rollout(ctx: "OnlineRLContext"): trajectory stops. """ - nonlocal current_id + nonlocal current_id, env_info, episode_info, timer, \ + total_episode_count, total_envstep_count, total_train_sample_count, last_train_iter timesteps = env.step(ctx.action) ctx.env_step += len(timesteps) timesteps = [t.tensor() for t in timesteps] - # TODO abnormal env step + + collected_sample = 0 + collected_step = 0 + collected_episode = 0 + interaction_duration = timer.value / len(timesteps) for i, timestep in enumerate(timesteps): - transition = policy.process_transition(ctx.obs[i], ctx.inference_output[i], timestep) - transition = ttorch.as_tensor(transition) # TBD - transition.collect_train_iter = ttorch.as_tensor([ctx.train_iter]) - transition.env_data_id = ttorch.as_tensor([env_episode_id[timestep.env_id]]) - transitions.append(timestep.env_id, transition) + with timer: + transition = policy.process_transition(ctx.obs[i], ctx.inference_output[i], timestep) + transition = ttorch.as_tensor(transition) + transition.collect_train_iter = ttorch.as_tensor([ctx.train_iter]) + transition.env_data_id = ttorch.as_tensor([env_episode_id[timestep.env_id]]) + transitions.append(timestep.env_id, transition) + + collected_step += 1 + collected_sample += len(transition.obs) + env_info[timestep.env_id.item()]['step'] += 1 + env_info[timestep.env_id.item()]['train_sample'] += len(transition.obs) + + env_info[timestep.env_id.item()]['time'] += timer.value + interaction_duration if timestep.done: - policy.reset([timestep.env_id]) - env_episode_id[timestep.env_id] = current_id + info = { + 'reward': timestep.info['eval_episode_return'], + 'time': env_info[timestep.env_id.item()]['time'], + 'step': env_info[timestep.env_id.item()]['step'], + 'train_sample': env_info[timestep.env_id.item()]['train_sample'], + } + # reset corresponding env info + env_info[timestep.env_id.item()] = {'time': 0., 'step': 0, 'train_sample': 0} + + episode_info.append(info) + policy.reset([timestep.env_id.item()]) + env_episode_id[timestep.env_id.item()] = current_id + collected_episode += 1 current_id += 1 ctx.env_episode += 1 - # TODO log + + total_envstep_count += collected_step + total_episode_count += collected_episode + total_train_sample_count += collected_sample + + if (ctx.train_iter - last_train_iter) >= collect_print_freq and len(episode_info) > 0: + output_log(episode_info, total_episode_count, total_envstep_count, total_train_sample_count) + last_train_iter = ctx.train_iter return _rollout + + +def output_log(episode_info, total_episode_count, total_envstep_count, total_train_sample_count) -> None: + """ + Overview: + Print the output log information. You can refer to the docs of `Best Practice` to understand \ + the training generated logs and tensorboards. + Arguments: + - train_iter (:obj:`int`): the number of training iteration. + """ + episode_count = len(episode_info) + envstep_count = sum([d['step'] for d in episode_info]) + train_sample_count = sum([d['train_sample'] for d in episode_info]) + duration = sum([d['time'] for d in episode_info]) + episode_return = [d['reward'].item() for d in episode_info] + info = { + 'episode_count': episode_count, + 'envstep_count': envstep_count, + 'train_sample_count': train_sample_count, + 'avg_envstep_per_episode': envstep_count / episode_count, + 'avg_sample_per_episode': train_sample_count / episode_count, + 'avg_envstep_per_sec': envstep_count / duration, + 'avg_train_sample_per_sec': train_sample_count / duration, + 'avg_episode_per_sec': episode_count / duration, + 'reward_mean': np.mean(episode_return), + 'reward_std': np.std(episode_return), + 'reward_max': np.max(episode_return), + 'reward_min': np.min(episode_return), + 'total_envstep_count': total_envstep_count, + 'total_train_sample_count': total_train_sample_count, + 'total_episode_count': total_episode_count, + # 'each_reward': episode_return, + } + episode_info.clear() + logging.info("collect end:\n{}".format('\n'.join(['{}: {}'.format(k, v) for k, v in info.items()]))) diff --git a/ding/framework/middleware/functional/ctx_helper.py b/ding/framework/middleware/functional/ctx_helper.py index 2994889eb5..8c3254079b 100644 --- a/ding/framework/middleware/functional/ctx_helper.py +++ b/ding/framework/middleware/functional/ctx_helper.py @@ -11,15 +11,19 @@ def final_ctx_saver(name: str) -> Callable: def _save(ctx: "Context"): if task.finish: + # make sure the items to be recorded are all kept in the context with open(os.path.join(name, 'result.pkl'), 'wb') as f: final_data = { 'total_step': ctx.total_step, 'train_iter': ctx.train_iter, - 'eval_value': ctx.eval_value, + 'last_eval_iter': ctx.last_eval_iter, + 'eval_value': ctx.last_eval_value, } if ctx.has_attr('env_step'): final_data['env_step'] = ctx.env_step final_data['env_episode'] = ctx.env_episode + if ctx.has_attr('trained_env_step'): + final_data['trained_env_step'] = ctx.trained_env_step pickle.dump(final_data, f) return _save diff --git a/ding/framework/middleware/functional/data_processor.py b/ding/framework/middleware/functional/data_processor.py index 0ec60de3b0..5882716ef3 100644 --- a/ding/framework/middleware/functional/data_processor.py +++ b/ding/framework/middleware/functional/data_processor.py @@ -1,9 +1,14 @@ +import os from typing import TYPE_CHECKING, Callable, List, Union, Tuple, Dict, Optional from easydict import EasyDict from ditk import logging +import numpy as np import torch +import tqdm from ding.data import Buffer, Dataset, DataLoader, offline_data_save_type from ding.data.buffer.middleware import PriorityExperienceReplay +from ding.framework import task +from ding.utils import get_rank if TYPE_CHECKING: from ding.framework import OnlineRLContext, OfflineRLContext @@ -17,6 +22,8 @@ def data_pusher(cfg: EasyDict, buffer_: Buffer, group_by_env: Optional[bool] = N - cfg (:obj:`EasyDict`): Config. - buffer (:obj:`Buffer`): Buffer to push the data in. """ + if task.router.is_active and not task.has_role(task.role.LEARNER): + return task.void() def _push(ctx: "OnlineRLContext"): """ @@ -45,6 +52,42 @@ def _push(ctx: "OnlineRLContext"): return _push +def buffer_saver(cfg: EasyDict, buffer_: Buffer, every_envstep: int = 1000, replace: bool = False): + """ + Overview: + Save current buffer data. + Arguments: + - cfg (:obj:`EasyDict`): Config. + - buffer (:obj:`Buffer`): Buffer to push the data in. + - every_envstep (:obj:`int`): save at every env step. + - replace (:obj:`bool`): Whether replace the last file. + """ + + buffer_saver_env_counter = -every_envstep + + def _save(ctx: "OnlineRLContext"): + """ + Overview: + In ctx, `ctx.env_step` should not be None. + Input of ctx: + - env_step (:obj:`int`): env step. + """ + nonlocal buffer_saver_env_counter + if ctx.env_step is not None: + if ctx.env_step >= every_envstep + buffer_saver_env_counter: + buffer_saver_env_counter = ctx.env_step + if replace: + buffer_.save_data(os.path.join(cfg.exp_name, "replaybuffer", "data_latest.hkl")) + else: + buffer_.save_data( + os.path.join(cfg.exp_name, "replaybuffer", "data_envstep_{}.hkl".format(ctx.env_step)) + ) + else: + raise RuntimeError("buffer_saver only supports collecting data by step rather than episode.") + + return _save + + def offpolicy_data_fetcher( cfg: EasyDict, buffer_: Union[Buffer, List[Tuple[Buffer, float]], Dict[str, Buffer]], @@ -140,7 +183,55 @@ def _fetch(ctx: "OnlineRLContext"): return _fetch -def offline_data_fetcher(cfg: EasyDict, dataset: Dataset) -> Callable: +def offline_data_fetcher_from_mem(cfg: EasyDict, dataset: Dataset) -> Callable: + + from threading import Thread + from queue import Queue + import time + stream = torch.cuda.Stream() + + def producer(queue, dataset, batch_size, device): + torch.set_num_threads(4) + nonlocal stream + idx_iter = iter(range(len(dataset) - batch_size)) + + if len(dataset) < batch_size: + logging.warning('batch_size is too large!!!!') + with torch.cuda.stream(stream): + while True: + if queue.full(): + time.sleep(0.1) + else: + try: + start_idx = next(idx_iter) + except StopIteration: + del idx_iter + idx_iter = iter(range(len(dataset) - batch_size)) + start_idx = next(idx_iter) + data = [dataset.__getitem__(idx) for idx in range(start_idx, start_idx + batch_size)] + data = [[i[j] for i in data] for j in range(len(data[0]))] + data = [torch.stack(x).to(device) for x in data] + queue.put(data) + + queue = Queue(maxsize=50) + device = 'cuda:{}'.format(get_rank() % torch.cuda.device_count()) if cfg.policy.cuda else 'cpu' + producer_thread = Thread( + target=producer, args=(queue, dataset, cfg.policy.learn.batch_size, device), name='cuda_fetcher_producer' + ) + + def _fetch(ctx: "OfflineRLContext"): + nonlocal queue, producer_thread + if not producer_thread.is_alive(): + time.sleep(5) + producer_thread.start() + while queue.empty(): + time.sleep(0.001) + ctx.train_data = queue.get() + + return _fetch + + +def offline_data_fetcher(cfg: EasyDict, dataset: Dataset, collate_fn=lambda x: x) -> Callable: """ Overview: The outer function transforms a Pytorch `Dataset` to `DataLoader`. \ @@ -152,7 +243,8 @@ def offline_data_fetcher(cfg: EasyDict, dataset: Dataset) -> Callable: - dataset (:obj:`Dataset`): The dataset of type `torch.utils.data.Dataset` which stores the data. """ # collate_fn is executed in policy now - dataloader = DataLoader(dataset, batch_size=cfg.policy.learn.batch_size, shuffle=True, collate_fn=lambda x: x) + dataloader = DataLoader(dataset, batch_size=cfg.policy.learn.batch_size, shuffle=True, collate_fn=collate_fn) + dataloader = iter(dataloader) def _fetch(ctx: "OfflineRLContext"): """ @@ -164,22 +256,28 @@ def _fetch(ctx: "OfflineRLContext"): Output of ctx: - train_data (:obj:`List[Tensor]`): The fetched data batch. """ - while True: - for i, data in enumerate(dataloader): - ctx.train_data = data - yield + nonlocal dataloader + try: + ctx.train_data = next(dataloader) # noqa + except StopIteration: ctx.train_epoch += 1 + del dataloader + dataloader = DataLoader( + dataset, batch_size=cfg.policy.learn.batch_size, shuffle=True, collate_fn=collate_fn + ) + dataloader = iter(dataloader) + ctx.train_data = next(dataloader) # TODO apply data update (e.g. priority) in offline setting when necessary + ctx.trained_env_step += len(ctx.train_data) return _fetch -def offline_data_saver(cfg: EasyDict, data_path: str, data_type: str = 'hdf5') -> Callable: +def offline_data_saver(data_path: str, data_type: str = 'hdf5') -> Callable: """ Overview: Save the expert data of offline RL in a directory. Arguments: - - cfg (:obj:`EasyDict`): Config. - data_path (:obj:`str`): File path where the expert data will be written into, which is usually ./expert.pkl'. - data_type (:obj:`str`): Define the type of the saved data. \ The type of saved data is pkl if `data_type == 'naive'`. \ @@ -223,3 +321,125 @@ def _pusher(ctx: "OnlineRLContext"): ctx.trajectories = None return _pusher + + +def qgpo_support_data_generator(cfg, dataset, policy) -> Callable: + + behavior_policy_stop_training_iter = cfg.policy.learn.behavior_policy_stop_training_iter if hasattr( + cfg.policy.learn, 'behavior_policy_stop_training_iter' + ) else np.inf + energy_guided_policy_begin_training_iter = cfg.policy.learn.energy_guided_policy_begin_training_iter if hasattr( + cfg.policy.learn, 'energy_guided_policy_begin_training_iter' + ) else 0 + actions_generated = False + + def generate_fake_actions(): + allstates = dataset.states[:].cpu().numpy() + actions_sampled = [] + for states in tqdm.tqdm(np.array_split(allstates, allstates.shape[0] // 4096 + 1)): + actions_sampled.append( + policy._model.sample( + states, + sample_per_state=cfg.policy.learn.M, + diffusion_steps=cfg.policy.learn.diffusion_steps, + guidance_scale=0.0, + ) + ) + actions = np.concatenate(actions_sampled) + + allnextstates = dataset.next_states[:].cpu().numpy() + actions_next_states_sampled = [] + for next_states in tqdm.tqdm(np.array_split(allnextstates, allnextstates.shape[0] // 4096 + 1)): + actions_next_states_sampled.append( + policy._model.sample( + next_states, + sample_per_state=cfg.policy.learn.M, + diffusion_steps=cfg.policy.learn.diffusion_steps, + guidance_scale=0.0, + ) + ) + actions_next_states = np.concatenate(actions_next_states_sampled) + return actions, actions_next_states + + def _data_generator(ctx: "OfflineRLContext"): + nonlocal actions_generated + + if ctx.train_iter >= energy_guided_policy_begin_training_iter: + if ctx.train_iter > behavior_policy_stop_training_iter: + # no need to generate fake actions if fake actions are already generated + if actions_generated: + pass + else: + actions, actions_next_states = generate_fake_actions() + dataset.fake_actions = torch.Tensor(actions.astype(np.float32)).to(cfg.policy.model.device) + dataset.fake_next_actions = torch.Tensor(actions_next_states.astype(np.float32) + ).to(cfg.policy.model.device) + actions_generated = True + else: + # generate fake actions + actions, actions_next_states = generate_fake_actions() + dataset.fake_actions = torch.Tensor(actions.astype(np.float32)).to(cfg.policy.model.device) + dataset.fake_next_actions = torch.Tensor(actions_next_states.astype(np.float32) + ).to(cfg.policy.model.device) + actions_generated = True + else: + # no need to generate fake actions + pass + + return _data_generator + + +def qgpo_offline_data_fetcher(cfg: EasyDict, dataset: Dataset, collate_fn=lambda x: x) -> Callable: + """ + Overview: + The outer function transforms a Pytorch `Dataset` to `DataLoader`. \ + The return function is a generator which each time fetches a batch of data from the previous `DataLoader`.\ + Please refer to the link https://pytorch.org/tutorials/beginner/basics/data_tutorial.html \ + and https://pytorch.org/docs/stable/data.html for more details. + Arguments: + - cfg (:obj:`EasyDict`): Config which should contain the following keys: `cfg.policy.learn.batch_size`. + - dataset (:obj:`Dataset`): The dataset of type `torch.utils.data.Dataset` which stores the data. + """ + # collate_fn is executed in policy now + dataloader = DataLoader(dataset, batch_size=cfg.policy.learn.batch_size, shuffle=True, collate_fn=collate_fn) + dataloader_q = DataLoader(dataset, batch_size=cfg.policy.learn.batch_size_q, shuffle=True, collate_fn=collate_fn) + + behavior_policy_stop_training_iter = cfg.policy.learn.behavior_policy_stop_training_iter if hasattr( + cfg.policy.learn, 'behavior_policy_stop_training_iter' + ) else np.inf + energy_guided_policy_begin_training_iter = cfg.policy.learn.energy_guided_policy_begin_training_iter if hasattr( + cfg.policy.learn, 'energy_guided_policy_begin_training_iter' + ) else 0 + + def get_behavior_policy_training_data(): + while True: + yield from dataloader + + data = get_behavior_policy_training_data() + + def get_q_training_data(): + while True: + yield from dataloader_q + + data_q = get_q_training_data() + + def _fetch(ctx: "OfflineRLContext"): + """ + Overview: + Every time this generator is iterated, the fetched data will be assigned to ctx.train_data. \ + After the dataloader is empty, the attribute `ctx.train_epoch` will be incremented by 1. + Input of ctx: + - train_epoch (:obj:`int`): Number of `train_epoch`. + Output of ctx: + - train_data (:obj:`List[Tensor]`): The fetched data batch. + """ + + if ctx.train_iter >= energy_guided_policy_begin_training_iter: + ctx.train_data = next(data_q) + else: + ctx.train_data = next(data) + + # TODO apply data update (e.g. priority) in offline setting when necessary + ctx.trained_env_step += len(ctx.train_data) + + return _fetch diff --git a/ding/framework/middleware/functional/enhancer.py b/ding/framework/middleware/functional/enhancer.py index 4bd9ad45e4..597a086850 100644 --- a/ding/framework/middleware/functional/enhancer.py +++ b/ding/framework/middleware/functional/enhancer.py @@ -2,7 +2,7 @@ from easydict import EasyDict from ditk import logging import torch -from ding.policy import Policy +from ding.framework import task if TYPE_CHECKING: from ding.framework import OnlineRLContext from ding.reward_model import BaseRewardModel, HerRewardModel @@ -17,6 +17,8 @@ def reward_estimator(cfg: EasyDict, reward_model: "BaseRewardModel") -> Callable - cfg (:obj:`EasyDict`): Config. - reward_model (:obj:`BaseRewardModel`): Reward model. """ + if task.router.is_active and not task.has_role(task.role.LEARNER): + return task.void() def _enhance(ctx: "OnlineRLContext"): """ @@ -40,6 +42,8 @@ def her_data_enhancer(cfg: EasyDict, buffer_: "Buffer", her_reward_model: "HerRe - her_reward_model (:obj:`HerRewardModel`): Hindsight Experience Replay (HER) model \ which is used to process episodes. """ + if task.router.is_active and not task.has_role(task.role.LEARNER): + return task.void() def _fetch_and_enhance(ctx: "OnlineRLContext"): """ @@ -69,6 +73,9 @@ def _fetch_and_enhance(ctx: "OnlineRLContext"): def nstep_reward_enhancer(cfg: EasyDict) -> Callable: + if task.router.is_active and (not task.has_role(task.role.LEARNER) and not task.has_role(task.role.COLLECTOR)): + return task.void() + def _enhance(ctx: "OnlineRLContext"): nstep = cfg.policy.nstep gamma = cfg.policy.discount_factor diff --git a/ding/framework/middleware/functional/evaluator.py b/ding/framework/middleware/functional/evaluator.py index 00281dc0a4..3fa407b6de 100644 --- a/ding/framework/middleware/functional/evaluator.py +++ b/ding/framework/middleware/functional/evaluator.py @@ -1,22 +1,20 @@ -from typing import TYPE_CHECKING, Callable, Any, List, Union +from typing import Callable, Any, List, Union, Optional from abc import ABC, abstractmethod from collections import deque from ditk import logging import numpy as np import torch +import treetensor.numpy as tnp import treetensor.torch as ttorch from easydict import EasyDict from ding.envs import BaseEnvManager -from ding.framework.context import OfflineRLContext +from ding.framework.context import Context, OfflineRLContext, OnlineRLContext from ding.policy import Policy from ding.data import Dataset, DataLoader from ding.framework import task -from ding.torch_utils import tensor_to_list, to_ndarray +from ding.torch_utils import to_ndarray, get_shape0 from ding.utils import lists_to_dicts -if TYPE_CHECKING: - from ding.framework import Context, OnlineRLContext - class IMetric(ABC): @@ -48,8 +46,8 @@ class VectorEvalMonitor(object): any thing, it is likely that we will get more short episodes than long episodes. As a result, \ our average reward will have a bias and may not be accurate. we use VectorEvalMonitor to solve the problem. Interfaces: - __init__, is_finished, update_info, update_reward, get_episode_reward, get_latest_reward, get_current_episode,\ - get_episode_info + __init__, is_finished, update_info, update_reward, get_episode_return, get_latest_reward, get_current_episode,\ + get_episode_info, update_video, get_episode_video """ def __init__(self, env_num: int, n_episode: int) -> None: @@ -70,6 +68,14 @@ def __init__(self, env_num: int, n_episode: int) -> None: each_env_episode[i] += 1 self._reward = {env_id: deque(maxlen=maxlen) for env_id, maxlen in enumerate(each_env_episode)} self._info = {env_id: deque(maxlen=maxlen) for env_id, maxlen in enumerate(each_env_episode)} + self._video = { + env_id: deque([[] for _ in range(maxlen)], maxlen=maxlen) + for env_id, maxlen in enumerate(each_env_episode) + } + self._output = { + env_id: deque([[] for _ in range(maxlen)], maxlen=maxlen) + for env_id, maxlen in enumerate(each_env_episode) + } def is_finished(self) -> bool: """ @@ -88,7 +94,6 @@ def update_info(self, env_id: int, info: Any) -> None: - env_id: (:obj:`int`): the id of the environment we need to update information - info: (:obj:`Any`): the information we need to update """ - info = tensor_to_list(info) self._info[env_id].append(info) def update_reward(self, env_id: Union[int, np.ndarray], reward: Any) -> None: @@ -105,10 +110,10 @@ def update_reward(self, env_id: Union[int, np.ndarray], reward: Any) -> None: env_id = env_id.item() self._reward[env_id].append(reward) - def get_episode_reward(self) -> list: + def get_episode_return(self) -> list: """ Overview: - Get the total reward of one episode. + Sum up all reward and get the total return of one episode. """ return sum([list(v) for v in self._reward.values()], []) # sum(iterable, start) @@ -131,22 +136,83 @@ def get_current_episode(self) -> int: def get_episode_info(self) -> dict: """ Overview: - Get all episode information, such as total reward of one episode. + Get all episode information, such as total return of one episode. """ if len(self._info[0]) == 0: return None else: + # sum among all envs total_info = sum([list(v) for v in self._info.values()], []) + if isinstance(total_info[0], tnp.ndarray): + total_info = [t.json() for t in total_info] total_info = lists_to_dicts(total_info) new_dict = {} for k in total_info.keys(): - if np.isscalar(total_info[k][0]): - new_dict[k + '_mean'] = np.mean(total_info[k]) - total_info.update(new_dict) - return total_info - - -def interaction_evaluator(cfg: EasyDict, policy: Policy, env: BaseEnvManager) -> Callable: + try: + if np.isscalar(total_info[k][0].item()): + new_dict[k + '_mean'] = np.mean(total_info[k]) + except: # noqa + pass + return new_dict + + def _select_idx(self): + reward = [t.item() for t in self.get_episode_return()] + sortarg = np.argsort(reward) + # worst, median(s), best + if len(sortarg) == 1: + idxs = [sortarg[0]] + elif len(sortarg) == 2: + idxs = [sortarg[0], sortarg[-1]] + elif len(sortarg) == 3: + idxs = [sortarg[0], sortarg[len(sortarg) // 2], sortarg[-1]] + else: + # TensorboardX pad the number of videos to even numbers with black frames, + # therefore providing even number of videos prevents black frames being rendered. + idxs = [sortarg[0], sortarg[len(sortarg) // 2 - 1], sortarg[len(sortarg) // 2], sortarg[-1]] + return idxs + + def update_video(self, imgs): + for env_id, img in imgs.items(): + if len(self._reward[env_id]) == self._reward[env_id].maxlen: + continue + self._video[env_id][len(self._reward[env_id])].append(img) + + def get_episode_video(self): + """ + Overview: + Convert list of videos into [N, T, C, H, W] tensor, containing + worst, median, best evaluation trajectories for video logging. + """ + videos = sum([list(v) for v in self._video.values()], []) + videos = [np.transpose(np.stack(video, 0), [0, 3, 1, 2]) for video in videos] + idxs = self._select_idx() + videos = [videos[idx] for idx in idxs] + # pad videos to the same length with last frames + max_length = max(video.shape[0] for video in videos) + for i in range(len(videos)): + if videos[i].shape[0] < max_length: + padding = np.tile([videos[i][-1]], (max_length - videos[i].shape[0], 1, 1, 1)) + videos[i] = np.concatenate([videos[i], padding], 0) + videos = np.stack(videos, 0) + assert len(videos.shape) == 5, 'Need [N, T, C, H, W] input tensor for video logging!' + return videos + + def update_output(self, output): + for env_id, o in output.items(): + if len(self._reward[env_id]) == self._reward[env_id].maxlen: + continue + self._output[env_id][len(self._reward[env_id])].append(to_ndarray(o)) + + def get_episode_output(self): + output = sum([list(v) for v in self._output.values()], []) + idxs = self._select_idx() + output = [output[idx] for idx in idxs] + return output + + +def interaction_evaluator( + cfg: EasyDict, policy: Policy, env: BaseEnvManager, render: bool = False, **kwargs +) -> Callable: """ Overview: The middleware that executes the evaluation. @@ -154,11 +220,15 @@ def interaction_evaluator(cfg: EasyDict, policy: Policy, env: BaseEnvManager) -> - cfg (:obj:`EasyDict`): Config. - policy (:obj:`Policy`): The policy to be evaluated. - env (:obj:`BaseEnvManager`): The env for the evaluation. + - render (:obj:`bool`): Whether to render env images and policy logits. + - kwargs: (:obj:`Any`): Other arguments for specific evaluation. """ + if task.router.is_active and not task.has_role(task.role.EVALUATOR): + return task.void() env.seed(cfg.seed, dynamic_seed=False) - def _evaluate(ctx: "OnlineRLContext"): + def _evaluate(ctx: Union["OnlineRLContext", "OfflineRLContext"]): """ Overview: - The evaluation will be executed if the task begins and enough train_iter passed \ @@ -170,9 +240,15 @@ def _evaluate(ctx: "OnlineRLContext"): - eval_value (:obj:`float`): The average reward in the current evaluation. """ + # evaluation will be executed if the task begins or enough train_iter after last evaluation if ctx.last_eval_iter != -1 and \ - (ctx.train_iter - ctx.last_eval_iter < cfg.policy.eval.evaluator.eval_freq): - return + (ctx.train_iter - ctx.last_eval_iter < cfg.policy.eval.evaluator.eval_freq): + if ctx.train_iter != ctx.last_eval_iter: + return + if len(kwargs) > 0: + kwargs_str = '/'.join([f'{k}({v})' for k, v in kwargs.items()]) + else: + kwargs_str = '' if env.closed: env.launch() @@ -183,29 +259,163 @@ def _evaluate(ctx: "OnlineRLContext"): while not eval_monitor.is_finished(): obs = ttorch.as_tensor(env.ready_obs).to(dtype=ttorch.float32) - obs = {i: obs[i] for i in range(obs.shape[0])} # TBD - inference_output = policy.forward(obs) - action = [to_ndarray(v['action']) for v in inference_output.values()] # TBD + obs = {i: obs[i] for i in range(get_shape0(obs))} # TBD + if len(kwargs) > 0: + inference_output = policy.forward(obs, **kwargs) + else: + inference_output = policy.forward(obs) + if render: + eval_monitor.update_video(env.ready_imgs) + eval_monitor.update_output(inference_output) + output = [v for v in inference_output.values()] + action = [to_ndarray(v['action']) for v in output] # TBD timesteps = env.step(action) for timestep in timesteps: env_id = timestep.env_id.item() if timestep.done: policy.reset([env_id]) - reward = timestep.info.final_eval_reward + reward = timestep.info.eval_episode_return eval_monitor.update_reward(env_id, reward) - episode_reward = eval_monitor.get_episode_reward() - eval_reward = np.mean(episode_reward) - stop_flag = eval_reward >= cfg.env.stop_value and ctx.train_iter > 0 - if isinstance(ctx, OfflineRLContext): - logging.info('Evaluation: Train Iter({})\tEval Reward({:.3f})'.format(ctx.train_iter, eval_reward)) - else: + if 'episode_info' in timestep.info: + eval_monitor.update_info(env_id, timestep.info.episode_info) + episode_return = eval_monitor.get_episode_return() + episode_return_min = np.min(episode_return) + episode_return_max = np.max(episode_return) + episode_return_std = np.std(episode_return) + episode_return = np.mean(episode_return) + stop_flag = episode_return >= cfg.env.stop_value and ctx.train_iter > 0 + if isinstance(ctx, OnlineRLContext): logging.info( - 'Evaluation: Train Iter({})\tEnv Step({})\tEval Reward({:.3f})'.format( - ctx.train_iter, ctx.env_step, eval_reward + 'Evaluation: Train Iter({}) Env Step({}) Episode Return({:.3f}) {}'.format( + ctx.train_iter, ctx.env_step, episode_return, kwargs_str ) ) + elif isinstance(ctx, OfflineRLContext): + logging.info( + 'Evaluation: Train Iter({}) Eval Return({:.3f}) {}'.format(ctx.train_iter, episode_return, kwargs_str) + ) + else: + raise TypeError("not supported ctx type: {}".format(type(ctx))) ctx.last_eval_iter = ctx.train_iter - ctx.eval_value = eval_reward + ctx.eval_value = episode_return + ctx.eval_value_min = episode_return_min + ctx.eval_value_max = episode_return_max + ctx.eval_value_std = episode_return_std + ctx.last_eval_value = ctx.eval_value + ctx.eval_output = {'episode_return': episode_return} + episode_info = eval_monitor.get_episode_info() + if episode_info is not None: + ctx.eval_output['episode_info'] = episode_info + if render: + ctx.eval_output['replay_video'] = eval_monitor.get_episode_video() + ctx.eval_output['output'] = eval_monitor.get_episode_output() + else: + ctx.eval_output['output'] = output # for compatibility + + if len(kwargs) > 0: + ctx.info_for_logging.update( + { + f'{kwargs_str}/eval_episode_return': episode_return, + f'{kwargs_str}/eval_episode_return_min': episode_return_min, + f'{kwargs_str}/eval_episode_return_max': episode_return_max, + f'{kwargs_str}/eval_episode_return_std': episode_return_std, + } + ) + + if stop_flag: + task.finish = True + + return _evaluate + + +def interaction_evaluator_ttorch( + seed: int, + policy: Policy, + env: BaseEnvManager, + n_evaluator_episode: Optional[int] = None, + stop_value: float = np.inf, + eval_freq: int = 1000, + render: bool = False, +) -> Callable: + """ + Overview: + The middleware that executes the evaluation with ttorch data. + Arguments: + - policy (:obj:`Policy`): The policy to be evaluated. + - env (:obj:`BaseEnvManager`): The env for the evaluation. + - render (:obj:`bool`): Whether to render env images and policy logits. + """ + if task.router.is_active and not task.has_role(task.role.EVALUATOR): + return task.void() + + env.seed(seed, dynamic_seed=False) + if n_evaluator_episode is None: + n_evaluator_episode = env.env_num + + def _evaluate(ctx: "OnlineRLContext"): + """ + Overview: + - The evaluation will be executed if the task begins and enough train_iter passed \ + since last evaluation. + Input of ctx: + - last_eval_iter (:obj:`int`): Last evaluation iteration. + - train_iter (:obj:`int`): Current train iteration. + Output of ctx: + - eval_value (:obj:`float`): The average reward in the current evaluation. + """ + + # evaluation will be executed if the task begins or enough train_iter after last evaluation + if ctx.last_eval_iter != -1 and (ctx.train_iter - ctx.last_eval_iter < eval_freq): + return + + if env.closed: + env.launch() + else: + env.reset() + policy.reset() + device = policy._device + eval_monitor = VectorEvalMonitor(env.env_num, n_evaluator_episode) + + while not eval_monitor.is_finished(): + obs = ttorch.as_tensor(env.ready_obs).to(dtype=ttorch.float32) + obs = obs.to(device) + inference_output = policy.eval(obs) + inference_output = inference_output.cpu() + if render: + eval_monitor.update_video(env.ready_imgs) + # eval_monitor.update_output(inference_output) + action = inference_output.action.numpy() + timesteps = env.step(action) + for timestep in timesteps: + env_id = timestep.env_id.item() + if timestep.done: + policy.reset([env_id]) + reward = timestep.info.eval_episode_return + eval_monitor.update_reward(env_id, reward) + if 'episode_info' in timestep.info: + eval_monitor.update_info(env_id, timestep.info.episode_info) + episode_return = eval_monitor.get_episode_return() + episode_return_std = np.std(episode_return) + episode_return_mean = np.mean(episode_return) + stop_flag = episode_return_mean >= stop_value and ctx.train_iter > 0 + logging.info( + 'Evaluation: Train Iter({})\tEnv Step({})\tMean Episode Return({:.3f})'.format( + ctx.train_iter, ctx.env_step, episode_return_mean + ) + ) + ctx.last_eval_iter = ctx.train_iter + ctx.eval_value = episode_return_mean + ctx.eval_value_std = episode_return_std + ctx.last_eval_value = ctx.eval_value + ctx.eval_output = {'episode_return': episode_return} + episode_info = eval_monitor.get_episode_info() + if episode_info is not None: + ctx.eval_output['episode_info'] = episode_info + if render: + ctx.eval_output['replay_video'] = eval_monitor.get_episode_video() + ctx.eval_output['output'] = eval_monitor.get_episode_output() + else: + ctx.eval_output['output'] = inference_output.numpy() # for compatibility if stop_flag: task.finish = True @@ -233,7 +443,7 @@ def _evaluate(ctx: "Context"): avg_eval_output = metric.reduce_mean(eval_output) stop_flag = metric.gt(avg_eval_output, cfg.env.stop_value) and ctx.train_iter > 0 logging.info( - 'Evaluation: Train Iter({})\tEnv Step({})\tEval Reward({:.3f})'.format( + 'Evaluation: Train Iter({})\tEnv Step({})\tEpisode Return({:.3f})'.format( ctx.train_iter, ctx.env_step, avg_eval_output ) ) diff --git a/ding/framework/middleware/functional/explorer.py b/ding/framework/middleware/functional/explorer.py index 4c7364004d..45aa9bd24a 100644 --- a/ding/framework/middleware/functional/explorer.py +++ b/ding/framework/middleware/functional/explorer.py @@ -1,6 +1,7 @@ -from typing import TYPE_CHECKING, Callable, List +from typing import TYPE_CHECKING, Callable from easydict import EasyDict from ding.rl_utils import get_epsilon_greedy_fn +from ding.framework import task if TYPE_CHECKING: from ding.framework import OnlineRLContext @@ -13,6 +14,8 @@ def eps_greedy_handler(cfg: EasyDict) -> Callable: Arguments: - cfg (:obj:`EasyDict`): Config. """ + if task.router.is_active and not task.has_role(task.role.COLLECTOR): + return task.void() eps_cfg = cfg.policy.other.eps handle = get_epsilon_greedy_fn(eps_cfg.start, eps_cfg.end, eps_cfg.decay, eps_cfg.type) diff --git a/ding/framework/middleware/functional/logger.py b/ding/framework/middleware/functional/logger.py index b700a19071..5f968a5fd3 100644 --- a/ding/framework/middleware/functional/logger.py +++ b/ding/framework/middleware/functional/logger.py @@ -1,31 +1,66 @@ -from typing import TYPE_CHECKING, Callable, Dict, List +from typing import TYPE_CHECKING, Optional, Callable, Dict, List, Union +from ditk import logging +from easydict import EasyDict +from matplotlib import pyplot as plt +from matplotlib import animation +import os import numpy as np +import torch +import wandb +import pickle +import treetensor.numpy as tnp +from ding.framework import task +from ding.envs import BaseEnvManagerV2 from ding.utils import DistributedWriter +from ding.torch_utils import to_ndarray +from ding.utils.default_helper import one_time_warning if TYPE_CHECKING: from ding.framework import OnlineRLContext, OfflineRLContext def online_logger(record_train_iter: bool = False, train_show_freq: int = 100) -> Callable: + """ + Overview: + Create an online RL tensorboard logger for recording training and evaluation metrics. + Arguments: + - record_train_iter (:obj:`bool`): Whether to record training iteration. Default is False. + - train_show_freq (:obj:`int`): Frequency of showing training logs. Default is 100. + Returns: + - _logger (:obj:`Callable`): A logger function that takes an OnlineRLContext object as input. + Raises: + - RuntimeError: If writer is None. + - NotImplementedError: If the key of train_output is not supported, such as "scalars". + + Examples: + >>> task.use(online_logger(record_train_iter=False, train_show_freq=1000)) + """ + if task.router.is_active and not task.has_role(task.role.LEARNER): + return task.void() writer = DistributedWriter.get_instance() + if writer is None: + raise RuntimeError("logger writer is None, you should call `ding_init(cfg)` at the beginning of training.") last_train_show_iter = -1 def _logger(ctx: "OnlineRLContext"): + if task.finish: + writer.close() nonlocal last_train_show_iter + if not np.isinf(ctx.eval_value): if record_train_iter: - writer.add_scalar('basic/eval_episode_reward_mean-env_step', ctx.eval_value, ctx.env_step) - writer.add_scalar('basic/eval_episode_reward_mean-train_iter', ctx.eval_value, ctx.train_iter) + writer.add_scalar('basic/eval_episode_return_mean-env_step', ctx.eval_value, ctx.env_step) + writer.add_scalar('basic/eval_episode_return_mean-train_iter', ctx.eval_value, ctx.train_iter) else: - writer.add_scalar('basic/eval_episode_reward_mean', ctx.eval_value, ctx.env_step) + writer.add_scalar('basic/eval_episode_return_mean', ctx.eval_value, ctx.env_step) if ctx.train_output is not None and ctx.train_iter - last_train_show_iter >= train_show_freq: last_train_show_iter = ctx.train_iter if isinstance(ctx.train_output, List): - output = ctx.train_output.pop() # only use latest output + output = ctx.train_output.pop() # only use latest output for some algorithms, like PPO else: output = ctx.train_output for k, v in output.items(): - if k in ['priority']: + if k in ['priority', 'td_error_priority']: continue if "[scalars]" in k: new_k = k.split(']')[-1] @@ -45,13 +80,36 @@ def _logger(ctx: "OnlineRLContext"): return _logger -def offline_logger() -> Callable: +def offline_logger(train_show_freq: int = 100) -> Callable: + """ + Overview: + Create an offline RL tensorboard logger for recording training and evaluation metrics. + Arguments: + - train_show_freq (:obj:`int`): Frequency of showing training logs. Defaults to 100. + Returns: + - _logger (:obj:`Callable`): A logger function that takes an OfflineRLContext object as input. + Raises: + - RuntimeError: If writer is None. + - NotImplementedError: If the key of train_output is not supported, such as "scalars". + + Examples: + >>> task.use(offline_logger(train_show_freq=1000)) + """ + if task.router.is_active and not task.has_role(task.role.LEARNER): + return task.void() writer = DistributedWriter.get_instance() + if writer is None: + raise RuntimeError("logger writer is None, you should call `ding_init(cfg)` at the beginning of training.") + last_train_show_iter = -1 def _logger(ctx: "OfflineRLContext"): + nonlocal last_train_show_iter + if task.finish: + writer.close() if not np.isinf(ctx.eval_value): - writer.add_scalar('basic/eval_episode_reward_mean-train_iter', ctx.eval_value, ctx.train_iter) - if ctx.train_output is not None: + writer.add_scalar('basic/eval_episode_return_mean-train_iter', ctx.eval_value, ctx.train_iter) + if ctx.train_output is not None and ctx.train_iter - last_train_show_iter >= train_show_freq: + last_train_show_iter = ctx.train_iter output = ctx.train_output for k, v in output.items(): if k in ['priority']: @@ -66,3 +124,593 @@ def _logger(ctx: "OfflineRLContext"): writer.add_scalar('basic/train_{}-train_iter'.format(k), v, ctx.train_iter) return _logger + + +# four utility functions for wandb logger +def softmax(logit: np.ndarray) -> np.ndarray: + v = np.exp(logit) + return v / v.sum(axis=-1, keepdims=True) + + +def action_prob(num, action_prob, ln): + ax = plt.gca() + ax.set_ylim([0, 1]) + for rect, x in zip(ln, action_prob[num]): + rect.set_height(x) + return ln + + +def return_prob(num, return_prob, ln): + return ln + + +def return_distribution(episode_return): + num = len(episode_return) + max_return = max(episode_return) + min_return = min(episode_return) + hist, bins = np.histogram(episode_return, bins=np.linspace(min_return - 50, max_return + 50, 6)) + gap = (max_return - min_return + 100) / 5 + x_dim = ['{:.1f}'.format(min_return - 50 + gap * x) for x in range(5)] + return hist / num, x_dim + + +def wandb_online_logger( + record_path: str = None, + cfg: Union[dict, EasyDict] = None, + exp_config: Union[dict, EasyDict] = None, + metric_list: Optional[List[str]] = None, + env: Optional[BaseEnvManagerV2] = None, + model: Optional[torch.nn.Module] = None, + anonymous: bool = False, + project_name: str = 'default-project', + run_name: str = None, + wandb_sweep: bool = False, +) -> Callable: + """ + Overview: + Wandb visualizer to track the experiment. + Arguments: + - record_path (:obj:`str`): The path to save the replay of simulation. + - cfg (:obj:`Union[dict, EasyDict]`): Config, a dict of following settings: + - gradient_logger: boolean. Whether to track the gradient. + - plot_logger: boolean. Whether to track the metrics like reward and loss. + - video_logger: boolean. Whether to upload the rendering video replay. + - action_logger: boolean. `q_value` or `action probability`. + - return_logger: boolean. Whether to track the return value. + - metric_list (:obj:`Optional[List[str]]`): Logged metric list, specialized by different policies. + - env (:obj:`BaseEnvManagerV2`): Evaluator environment. + - model (:obj:`nn.Module`): Policy neural network model. + - anonymous (:obj:`bool`): Open the anonymous mode of wandb or not. The anonymous mode allows visualization \ + of data without wandb count. + - project_name (:obj:`str`): The name of wandb project. + - run_name (:obj:`str`): The name of wandb run. + - wandb_sweep (:obj:`bool`): Whether to use wandb sweep. + ''' + Returns: + - _plot (:obj:`Callable`): A logger function that takes an OnlineRLContext object as input. + """ + if task.router.is_active and not task.has_role(task.role.LEARNER): + return task.void() + color_list = ["orange", "red", "blue", "purple", "green", "darkcyan"] + if metric_list is None: + metric_list = ["q_value", "target q_value", "loss", "lr", "entropy", "target_q_value", "td_error"] + # Initialize wandb with default settings + # Settings can be covered by calling wandb.init() at the top of the script + if exp_config: + if not wandb_sweep: + if run_name is not None: + if anonymous: + wandb.init(project=project_name, config=exp_config, reinit=True, name=run_name, anonymous="must") + else: + wandb.init(project=project_name, config=exp_config, reinit=True, name=run_name) + else: + if anonymous: + wandb.init(project=project_name, config=exp_config, reinit=True, anonymous="must") + else: + wandb.init(project=project_name, config=exp_config, reinit=True) + else: + if run_name is not None: + if anonymous: + wandb.init(project=project_name, config=exp_config, name=run_name, anonymous="must") + else: + wandb.init(project=project_name, config=exp_config, name=run_name) + else: + if anonymous: + wandb.init(project=project_name, config=exp_config, anonymous="must") + else: + wandb.init(project=project_name, config=exp_config) + else: + if not wandb_sweep: + if run_name is not None: + if anonymous: + wandb.init(project=project_name, reinit=True, name=run_name, anonymous="must") + else: + wandb.init(project=project_name, reinit=True, name=run_name) + else: + if anonymous: + wandb.init(project=project_name, reinit=True, anonymous="must") + else: + wandb.init(project=project_name, reinit=True) + else: + if run_name is not None: + if anonymous: + wandb.init(project=project_name, name=run_name, anonymous="must") + else: + wandb.init(project=project_name, name=run_name) + else: + if anonymous: + wandb.init(project=project_name, anonymous="must") + else: + wandb.init(project=project_name) + plt.switch_backend('agg') + if cfg is None: + cfg = EasyDict( + dict( + gradient_logger=False, + plot_logger=True, + video_logger=False, + action_logger=False, + return_logger=False, + ) + ) + else: + if not isinstance(cfg, EasyDict): + cfg = EasyDict(cfg) + for key in ["gradient_logger", "plot_logger", "video_logger", "action_logger", "return_logger", "vis_dataset"]: + if key not in cfg.keys(): + cfg[key] = False + + # The visualizer is called to save the replay of the simulation + # which will be uploaded to wandb later + if env is not None and cfg.video_logger is True and record_path is not None: + env.enable_save_replay(replay_path=record_path) + if cfg.gradient_logger: + wandb.watch(model, log="all", log_freq=100, log_graph=True) + else: + one_time_warning( + "If you want to use wandb to visualize the gradient, please set gradient_logger = True in the config." + ) + + first_plot = True + + def _plot(ctx: "OnlineRLContext"): + nonlocal first_plot + if first_plot: + first_plot = False + ctx.wandb_url = wandb.run.get_project_url() + + info_for_logging = {} + + if cfg.plot_logger: + for metric in metric_list: + if isinstance(ctx.train_output, Dict) and metric in ctx.train_output: + if isinstance(ctx.train_output[metric], torch.Tensor): + info_for_logging.update({metric: ctx.train_output[metric].cpu().detach().numpy()}) + else: + info_for_logging.update({metric: ctx.train_output[metric]}) + elif isinstance(ctx.train_output, List) and len(ctx.train_output) > 0 and metric in ctx.train_output[0]: + metric_value_list = [] + for item in ctx.train_output: + if isinstance(item[metric], torch.Tensor): + metric_value_list.append(item[metric].cpu().detach().numpy()) + else: + metric_value_list.append(item[metric]) + metric_value = np.mean(metric_value_list) + info_for_logging.update({metric: metric_value}) + else: + one_time_warning( + "If you want to use wandb to visualize the result, please set plot_logger = True in the config." + ) + + if ctx.eval_value != -np.inf: + if hasattr(ctx, "eval_value_min"): + info_for_logging.update({ + "episode return min": ctx.eval_value_min, + }) + if hasattr(ctx, "eval_value_max"): + info_for_logging.update({ + "episode return max": ctx.eval_value_max, + }) + if hasattr(ctx, "eval_value_std"): + info_for_logging.update({ + "episode return std": ctx.eval_value_std, + }) + if hasattr(ctx, "eval_value"): + info_for_logging.update({ + "episode return mean": ctx.eval_value, + }) + if hasattr(ctx, "train_iter"): + info_for_logging.update({ + "train iter": ctx.train_iter, + }) + if hasattr(ctx, "env_step"): + info_for_logging.update({ + "env step": ctx.env_step, + }) + + eval_output = ctx.eval_output['output'] + episode_return = ctx.eval_output['episode_return'] + episode_return = np.array(episode_return) + if len(episode_return.shape) == 2: + episode_return = episode_return.squeeze(1) + + if cfg.video_logger: + if 'replay_video' in ctx.eval_output: + # save numpy array "images" of shape (N,1212,3,224,320) to N video files in mp4 format + # The numpy tensor must be either 4 dimensional or 5 dimensional. + # Channels should be (time, channel, height, width) or (batch, time, channel, height width) + video_images = ctx.eval_output['replay_video'] + video_images = video_images.astype(np.uint8) + info_for_logging.update({"replay_video": wandb.Video(video_images, fps=60)}) + elif record_path is not None: + file_list = [] + for p in os.listdir(record_path): + if os.path.splitext(p)[-1] == ".mp4": + file_list.append(p) + file_list.sort(key=lambda fn: os.path.getmtime(os.path.join(record_path, fn))) + video_path = os.path.join(record_path, file_list[-2]) + info_for_logging.update({"video": wandb.Video(video_path, format="mp4")}) + + if cfg.action_logger: + action_path = os.path.join(record_path, (str(ctx.env_step) + "_action.gif")) + if all(['logit' in v for v in eval_output]) or hasattr(eval_output, "logit"): + if isinstance(eval_output, tnp.ndarray): + action_prob = softmax(eval_output.logit) + else: + action_prob = [softmax(to_ndarray(v['logit'])) for v in eval_output] + fig, ax = plt.subplots() + plt.ylim([-1, 1]) + action_dim = len(action_prob[1]) + x_range = [str(x + 1) for x in range(action_dim)] + ln = ax.bar(x_range, [0 for x in range(action_dim)], color=color_list[:action_dim]) + ani = animation.FuncAnimation( + fig, action_prob, fargs=(action_prob, ln), blit=True, save_count=len(action_prob) + ) + ani.save(action_path, writer='pillow') + info_for_logging.update({"action": wandb.Video(action_path, format="gif")}) + + elif all(['action' in v for v in eval_output[0]]): + for i, action_trajectory in enumerate(eval_output): + fig, ax = plt.subplots() + fig_data = np.array([[i + 1, *v['action']] for i, v in enumerate(action_trajectory)]) + steps = fig_data[:, 0] + actions = fig_data[:, 1:] + plt.ylim([-1, 1]) + for j in range(actions.shape[1]): + ax.scatter(steps, actions[:, j]) + info_for_logging.update({"actions_of_trajectory_{}".format(i): fig}) + + if cfg.return_logger: + return_path = os.path.join(record_path, (str(ctx.env_step) + "_return.gif")) + fig, ax = plt.subplots() + ax = plt.gca() + ax.set_ylim([0, 1]) + hist, x_dim = return_distribution(episode_return) + assert len(hist) == len(x_dim) + ln_return = ax.bar(x_dim, hist, width=1, color='r', linewidth=0.7) + ani = animation.FuncAnimation(fig, return_prob, fargs=(hist, ln_return), blit=True, save_count=1) + ani.save(return_path, writer='pillow') + info_for_logging.update({"return distribution": wandb.Video(return_path, format="gif")}) + + if bool(info_for_logging): + wandb.log(data=info_for_logging, step=ctx.env_step) + plt.clf() + + return _plot + + +def wandb_offline_logger( + record_path: str = None, + cfg: Union[dict, EasyDict] = None, + exp_config: Union[dict, EasyDict] = None, + metric_list: Optional[List[str]] = None, + env: Optional[BaseEnvManagerV2] = None, + model: Optional[torch.nn.Module] = None, + anonymous: bool = False, + project_name: str = 'default-project', + run_name: str = None, + wandb_sweep: bool = False, +) -> Callable: + """ + Overview: + Wandb visualizer to track the experiment. + Arguments: + - record_path (:obj:`str`): The path to save the replay of simulation. + - cfg (:obj:`Union[dict, EasyDict]`): Config, a dict of following settings: + - gradient_logger: boolean. Whether to track the gradient. + - plot_logger: boolean. Whether to track the metrics like reward and loss. + - video_logger: boolean. Whether to upload the rendering video replay. + - action_logger: boolean. `q_value` or `action probability`. + - return_logger: boolean. Whether to track the return value. + - vis_dataset: boolean. Whether to visualize the dataset. + - metric_list (:obj:`Optional[List[str]]`): Logged metric list, specialized by different policies. + - env (:obj:`BaseEnvManagerV2`): Evaluator environment. + - model (:obj:`nn.Module`): Policy neural network model. + - anonymous (:obj:`bool`): Open the anonymous mode of wandb or not. The anonymous mode allows visualization \ + of data without wandb count. + - project_name (:obj:`str`): The name of wandb project. + - run_name (:obj:`str`): The name of wandb run. + - wandb_sweep (:obj:`bool`): Whether to use wandb sweep. + ''' + Returns: + - _plot (:obj:`Callable`): A logger function that takes an OfflineRLContext object as input. + """ + if task.router.is_active and not task.has_role(task.role.LEARNER): + return task.void() + color_list = ["orange", "red", "blue", "purple", "green", "darkcyan"] + if metric_list is None: + metric_list = ["q_value", "target q_value", "loss", "lr", "entropy", "target_q_value", "td_error"] + # Initialize wandb with default settings + # Settings can be covered by calling wandb.init() at the top of the script + if exp_config: + if not wandb_sweep: + if run_name is not None: + if anonymous: + wandb.init(project=project_name, config=exp_config, reinit=True, name=run_name, anonymous="must") + else: + wandb.init(project=project_name, config=exp_config, reinit=True, name=run_name) + else: + if anonymous: + wandb.init(project=project_name, config=exp_config, reinit=True, anonymous="must") + else: + wandb.init(project=project_name, config=exp_config, reinit=True) + else: + if run_name is not None: + if anonymous: + wandb.init(project=project_name, config=exp_config, name=run_name, anonymous="must") + else: + wandb.init(project=project_name, config=exp_config, name=run_name) + else: + if anonymous: + wandb.init(project=project_name, config=exp_config, anonymous="must") + else: + wandb.init(project=project_name, config=exp_config) + else: + if not wandb_sweep: + if run_name is not None: + if anonymous: + wandb.init(project=project_name, reinit=True, name=run_name, anonymous="must") + else: + wandb.init(project=project_name, reinit=True, name=run_name) + else: + if anonymous: + wandb.init(project=project_name, reinit=True, anonymous="must") + else: + wandb.init(project=project_name, reinit=True) + else: + if run_name is not None: + if anonymous: + wandb.init(project=project_name, name=run_name, anonymous="must") + else: + wandb.init(project=project_name, name=run_name) + else: + if anonymous: + wandb.init(project=project_name, anonymous="must") + else: + wandb.init(project=project_name) + plt.switch_backend('agg') + plt.switch_backend('agg') + if cfg is None: + cfg = EasyDict( + dict( + gradient_logger=False, + plot_logger=True, + video_logger=False, + action_logger=False, + return_logger=False, + vis_dataset=True, + ) + ) + else: + if not isinstance(cfg, EasyDict): + cfg = EasyDict(cfg) + for key in ["gradient_logger", "plot_logger", "video_logger", "action_logger", "return_logger", "vis_dataset"]: + if key not in cfg.keys(): + cfg[key] = False + + # The visualizer is called to save the replay of the simulation + # which will be uploaded to wandb later + if env is not None and cfg.video_logger is True and record_path is not None: + env.enable_save_replay(replay_path=record_path) + if cfg.gradient_logger: + wandb.watch(model, log="all", log_freq=100, log_graph=True) + else: + one_time_warning( + "If you want to use wandb to visualize the gradient, please set gradient_logger = True in the config." + ) + + first_plot = True + + def _vis_dataset(datasetpath: str): + try: + from sklearn.manifold import TSNE + except ImportError: + import sys + logging.warning("Please install sklearn first, such as `pip3 install scikit-learn`.") + sys.exit(1) + try: + import h5py + except ImportError: + import sys + logging.warning("Please install h5py first, such as `pip3 install h5py`.") + sys.exit(1) + assert os.path.splitext(datasetpath)[-1] in ['.pkl', '.h5', '.hdf5'] + if os.path.splitext(datasetpath)[-1] == '.pkl': + with open(datasetpath, 'rb') as f: + data = pickle.load(f) + obs = [] + action = [] + reward = [] + for i in range(len(data)): + obs.extend(data[i]['observations']) + action.extend(data[i]['actions']) + reward.extend(data[i]['rewards']) + elif os.path.splitext(datasetpath)[-1] in ['.h5', '.hdf5']: + with h5py.File(datasetpath, 'r') as f: + obs = f['obs'][()] + action = f['action'][()] + reward = f['reward'][()] + + cmap = plt.cm.hsv + obs = np.array(obs) + reward = np.array(reward) + obs_action = np.hstack((obs, np.array(action))) + reward = reward / (max(reward) - min(reward)) + + embedded_obs = TSNE(n_components=2).fit_transform(obs) + embedded_obs_action = TSNE(n_components=2).fit_transform(obs_action) + x_min, x_max = np.min(embedded_obs, 0), np.max(embedded_obs, 0) + embedded_obs = embedded_obs / (x_max - x_min) + + x_min, x_max = np.min(embedded_obs_action, 0), np.max(embedded_obs_action, 0) + embedded_obs_action = embedded_obs_action / (x_max - x_min) + + fig = plt.figure() + f, axes = plt.subplots(nrows=1, ncols=3) + + axes[0].scatter(embedded_obs[:, 0], embedded_obs[:, 1], c=cmap(reward)) + axes[1].scatter(embedded_obs[:, 0], embedded_obs[:, 1], c=cmap(action)) + axes[2].scatter(embedded_obs_action[:, 0], embedded_obs_action[:, 1], c=cmap(reward)) + axes[0].set_title('state-reward') + axes[1].set_title('state-action') + axes[2].set_title('stateAction-reward') + plt.savefig('dataset.png') + + wandb.log({"dataset": wandb.Image("dataset.png")}) + + if cfg.vis_dataset is True: + _vis_dataset(exp_config.dataset_path) + + def _plot(ctx: "OfflineRLContext"): + nonlocal first_plot + if first_plot: + first_plot = False + ctx.wandb_url = wandb.run.get_project_url() + + info_for_logging = {} + + if cfg.plot_logger: + for metric in metric_list: + if isinstance(ctx.train_output, Dict) and metric in ctx.train_output: + if isinstance(ctx.train_output[metric], torch.Tensor): + info_for_logging.update({metric: ctx.train_output[metric].cpu().detach().numpy()}) + else: + info_for_logging.update({metric: ctx.train_output[metric]}) + elif isinstance(ctx.train_output, List) and len(ctx.train_output) > 0 and metric in ctx.train_output[0]: + metric_value_list = [] + for item in ctx.train_output: + if isinstance(item[metric], torch.Tensor): + metric_value_list.append(item[metric].cpu().detach().numpy()) + else: + metric_value_list.append(item[metric]) + metric_value = np.mean(metric_value_list) + info_for_logging.update({metric: metric_value}) + else: + one_time_warning( + "If you want to use wandb to visualize the result, please set plot_logger = True in the config." + ) + + if ctx.eval_value != -np.inf: + if hasattr(ctx, "info_for_logging"): + """ + .. note:: + The info_for_logging is a dict that contains the information to be logged. + Users can add their own information to the dict. + All the information in the dict will be logged to wandb. + """ + info_for_logging.update(ctx.info_for_logging) + + if hasattr(ctx, "eval_value_min"): + info_for_logging.update({ + "episode return min": ctx.eval_value_min, + }) + if hasattr(ctx, "eval_value_max"): + info_for_logging.update({ + "episode return max": ctx.eval_value_max, + }) + if hasattr(ctx, "eval_value_std"): + info_for_logging.update({ + "episode return std": ctx.eval_value_std, + }) + if hasattr(ctx, "eval_value"): + info_for_logging.update({ + "episode return mean": ctx.eval_value, + }) + if hasattr(ctx, "train_iter"): + info_for_logging.update({ + "train iter": ctx.train_iter, + }) + if hasattr(ctx, "train_epoch"): + info_for_logging.update({ + "train_epoch": ctx.train_epoch, + }) + + eval_output = ctx.eval_output['output'] + episode_return = ctx.eval_output['episode_return'] + episode_return = np.array(episode_return) + if len(episode_return.shape) == 2: + episode_return = episode_return.squeeze(1) + + if cfg.video_logger: + if 'replay_video' in ctx.eval_output: + # save numpy array "images" of shape (N,1212,3,224,320) to N video files in mp4 format + # The numpy tensor must be either 4 dimensional or 5 dimensional. + # Channels should be (time, channel, height, width) or (batch, time, channel, height width) + video_images = ctx.eval_output['replay_video'] + video_images = video_images.astype(np.uint8) + info_for_logging.update({"replay_video": wandb.Video(video_images, fps=60)}) + elif record_path is not None: + file_list = [] + for p in os.listdir(record_path): + if os.path.splitext(p)[-1] == ".mp4": + file_list.append(p) + file_list.sort(key=lambda fn: os.path.getmtime(os.path.join(record_path, fn))) + video_path = os.path.join(record_path, file_list[-2]) + info_for_logging.update({"video": wandb.Video(video_path, format="mp4")}) + + if cfg.action_logger: + action_path = os.path.join(record_path, (str(ctx.trained_env_step) + "_action.gif")) + if all(['logit' in v for v in eval_output]) or hasattr(eval_output, "logit"): + if isinstance(eval_output, tnp.ndarray): + action_prob = softmax(eval_output.logit) + else: + action_prob = [softmax(to_ndarray(v['logit'])) for v in eval_output] + fig, ax = plt.subplots() + plt.ylim([-1, 1]) + action_dim = len(action_prob[1]) + x_range = [str(x + 1) for x in range(action_dim)] + ln = ax.bar(x_range, [0 for x in range(action_dim)], color=color_list[:action_dim]) + ani = animation.FuncAnimation( + fig, action_prob, fargs=(action_prob, ln), blit=True, save_count=len(action_prob) + ) + ani.save(action_path, writer='pillow') + info_for_logging.update({"action": wandb.Video(action_path, format="gif")}) + + elif all(['action' in v for v in eval_output[0]]): + for i, action_trajectory in enumerate(eval_output): + fig, ax = plt.subplots() + fig_data = np.array([[i + 1, *v['action']] for i, v in enumerate(action_trajectory)]) + steps = fig_data[:, 0] + actions = fig_data[:, 1:] + plt.ylim([-1, 1]) + for j in range(actions.shape[1]): + ax.scatter(steps, actions[:, j]) + info_for_logging.update({"actions_of_trajectory_{}".format(i): fig}) + + if cfg.return_logger: + return_path = os.path.join(record_path, (str(ctx.trained_env_step) + "_return.gif")) + fig, ax = plt.subplots() + ax = plt.gca() + ax.set_ylim([0, 1]) + hist, x_dim = return_distribution(episode_return) + assert len(hist) == len(x_dim) + ln_return = ax.bar(x_dim, hist, width=1, color='r', linewidth=0.7) + ani = animation.FuncAnimation(fig, return_prob, fargs=(hist, ln_return), blit=True, save_count=1) + ani.save(return_path, writer='pillow') + info_for_logging.update({"return distribution": wandb.Video(return_path, format="gif")}) + + if bool(info_for_logging): + wandb.log(data=info_for_logging, step=ctx.trained_env_step) + plt.clf() + + return _plot diff --git a/ding/framework/middleware/functional/priority.py b/ding/framework/middleware/functional/priority.py new file mode 100644 index 0000000000..e62afbb5c9 --- /dev/null +++ b/ding/framework/middleware/functional/priority.py @@ -0,0 +1,24 @@ +from typing import TYPE_CHECKING, Callable +from ding.framework import task +if TYPE_CHECKING: + from ding.framework import OnlineRLContext + + +def priority_calculator(priority_calculation_fn: Callable) -> Callable: + """ + Overview: + The middleware that calculates the priority of the collected data. + Arguments: + - priority_calculation_fn (:obj:`Callable`): The function that calculates the priority of the collected data. + """ + + if task.router.is_active and not task.has_role(task.role.COLLECTOR): + return task.void() + + def _priority_calculator(ctx: "OnlineRLContext") -> None: + + priority = priority_calculation_fn(ctx.trajectories) + for i in range(len(priority)): + ctx.trajectories[i]['priority'] = priority[i] + + return _priority_calculator diff --git a/ding/framework/middleware/functional/termination_checker.py b/ding/framework/middleware/functional/termination_checker.py index 58c371d57b..3e0ed51887 100644 --- a/ding/framework/middleware/functional/termination_checker.py +++ b/ding/framework/middleware/functional/termination_checker.py @@ -1,5 +1,8 @@ from typing import TYPE_CHECKING, Union, Callable, Optional +from ditk import logging import numpy as np +import torch +from ding.utils import broadcast from ding.framework import task if TYPE_CHECKING: @@ -14,9 +17,38 @@ def termination_checker(max_env_step: Optional[int] = None, max_train_iter: Opti def _check(ctx: Union["OnlineRLContext", "OfflineRLContext"]): # ">" is better than ">=" when taking logger result into consideration - if ctx.env_step > max_env_step: + assert hasattr(ctx, "env_step") or hasattr(ctx, "train_iter"), "Context must have env_step or train_iter" + if hasattr(ctx, "env_step") and ctx.env_step > max_env_step: task.finish = True - if ctx.train_iter > max_train_iter: + logging.info('Exceeded maximum number of env_step({}), program is terminated'.format(ctx.env_step)) + elif hasattr(ctx, "train_iter") and ctx.train_iter > max_train_iter: task.finish = True + logging.info('Exceeded maximum number of train_iter({}), program is terminated'.format(ctx.train_iter)) + + return _check + + +def ddp_termination_checker(max_env_step=None, max_train_iter=None, rank=0): + if rank == 0: + if max_env_step is None: + max_env_step = np.inf + if max_train_iter is None: + max_train_iter = np.inf + + def _check(ctx): + if rank == 0: + if ctx.env_step > max_env_step: + finish = torch.ones(1).long().cuda() + logging.info('Exceeded maximum number of env_step({}), program is terminated'.format(ctx.env_step)) + elif ctx.train_iter > max_train_iter: + finish = torch.ones(1).long().cuda() + logging.info('Exceeded maximum number of train_iter({}), program is terminated'.format(ctx.train_iter)) + else: + finish = torch.LongTensor([task.finish]).cuda() + else: + finish = torch.zeros(1).long().cuda() + # broadcast finish result to other DDP workers + broadcast(finish, 0) + task.finish = finish.cpu().bool().item() return _check diff --git a/ding/framework/middleware/functional/timer.py b/ding/framework/middleware/functional/timer.py new file mode 100644 index 0000000000..db8a2c0056 --- /dev/null +++ b/ding/framework/middleware/functional/timer.py @@ -0,0 +1,35 @@ +import numpy as np +from collections import deque +from ditk import logging +from time import time + +from ding.framework import task +from typing import TYPE_CHECKING +if TYPE_CHECKING: + from ding.framework.context import Context + + +def epoch_timer(print_per: int = 1, smooth_window: int = 10): + """ + Overview: + Print time cost of each epoch. + Arguments: + - print_per (:obj:`int`): Print each N epoch. + - smooth_window (:obj:`int`): The window size to smooth the mean. + """ + records = deque(maxlen=print_per * smooth_window) + + def _epoch_timer(ctx: "Context"): + start = time() + yield + time_cost = time() - start + records.append(time_cost) + if ctx.total_step % print_per == 0: + logging.info( + "[Epoch Timer][Node:{:>2}]: Cost: {:.2f}ms, Mean: {:.2f}ms".format( + task.router.node_id or 0, time_cost * 1000, + np.mean(records) * 1000 + ) + ) + + return _epoch_timer diff --git a/ding/framework/middleware/functional/trainer.py b/ding/framework/middleware/functional/trainer.py index 29d58d664a..11c281c1a1 100644 --- a/ding/framework/middleware/functional/trainer.py +++ b/ding/framework/middleware/functional/trainer.py @@ -1,22 +1,23 @@ from typing import TYPE_CHECKING, Callable, Union from easydict import EasyDict +import treetensor.torch as ttorch from ditk import logging import numpy as np from ding.policy import Policy -from ding.framework import task, OfflineRLContext +from ding.framework import task, OfflineRLContext, OnlineRLContext -if TYPE_CHECKING: - from ding.framework import OnlineRLContext - -def trainer(cfg: EasyDict, policy: Policy) -> Callable: +def trainer(cfg: EasyDict, policy: Policy, log_freq: int = 100) -> Callable: """ Overview: The middleware that executes a single training process. Arguments: - cfg (:obj:`EasyDict`): Config. - policy (:obj:`Policy`): The policy to be trained in step-by-step mode. + - log_freq (:obj:`int`): The frequency (iteration) of showing log. """ + if task.router.is_active and not task.has_role(task.role.LEARNER): + return task.void() def _train(ctx: Union["OnlineRLContext", "OfflineRLContext"]): """ @@ -32,32 +33,38 @@ def _train(ctx: Union["OnlineRLContext", "OfflineRLContext"]): if ctx.train_data is None: return train_output = policy.forward(ctx.train_data) - if ctx.train_iter % cfg.policy.learn.learner.hook.log_show_after_iter == 0: - if isinstance(ctx, OfflineRLContext): - logging.info( - 'Training: Train Iter({})\tLoss({:.3f})'.format(ctx.train_iter, train_output['total_loss']) - ) + if ctx.train_iter % log_freq == 0: + if isinstance(train_output, list): + train_output_loss = np.mean([item['total_loss'] for item in train_output]) else: + train_output_loss = train_output['total_loss'] + if isinstance(ctx, OnlineRLContext): logging.info( 'Training: Train Iter({})\tEnv Step({})\tLoss({:.3f})'.format( - ctx.train_iter, ctx.env_step, train_output['total_loss'] + ctx.train_iter, ctx.env_step, train_output_loss ) ) - + elif isinstance(ctx, OfflineRLContext): + logging.info('Training: Train Iter({})\tLoss({:.3f})'.format(ctx.train_iter, train_output_loss)) + else: + raise TypeError("not supported ctx type: {}".format(type(ctx))) ctx.train_iter += 1 ctx.train_output = train_output return _train -def multistep_trainer(cfg: EasyDict, policy: Policy) -> Callable: +def multistep_trainer(policy: Policy, log_freq: int = 100) -> Callable: """ Overview: The middleware that executes training for a target num of steps. Arguments: - - cfg (:obj:`EasyDict`): Config. - policy (:obj:`Policy`): The policy specialized for multi-step training. + - log_freq (:obj:`int`): The frequency (iteration) of showing log. """ + if task.router.is_active and not task.has_role(task.role.LEARNER): + return task.void() + last_log_iter = -1 def _train(ctx: Union["OnlineRLContext", "OfflineRLContext"]): """ @@ -72,8 +79,15 @@ def _train(ctx: Union["OnlineRLContext", "OfflineRLContext"]): if ctx.train_data is None: # no enough data from data fetcher return - train_output = policy.forward(ctx.train_data) - if ctx.train_iter % cfg.policy.learn.learner.hook.log_show_after_iter == 0: + if hasattr(policy, "_device"): # For ppof policy + data = ctx.train_data.to(policy._device) + elif hasattr(policy, "get_attribute"): # For other policy + data = ctx.train_data.to(policy.get_attribute("device")) + else: + assert AttributeError("Policy should have attribution '_device'.") + train_output = policy.forward(data) + nonlocal last_log_iter + if ctx.train_iter - last_log_iter >= log_freq: loss = np.mean([o['total_loss'] for o in train_output]) if isinstance(ctx, OfflineRLContext): logging.info('Training: Train Iter({})\tLoss({:.3f})'.format(ctx.train_iter, loss)) @@ -81,6 +95,7 @@ def _train(ctx: Union["OnlineRLContext", "OfflineRLContext"]): logging.info( 'Training: Train Iter({})\tEnv Step({})\tLoss({:.3f})'.format(ctx.train_iter, ctx.env_step, loss) ) + last_log_iter = ctx.train_iter ctx.train_iter += len(train_output) ctx.train_output = train_output diff --git a/ding/framework/middleware/learner.py b/ding/framework/middleware/learner.py index 4d60f117bd..9abf88e9b3 100644 --- a/ding/framework/middleware/learner.py +++ b/ding/framework/middleware/learner.py @@ -8,6 +8,8 @@ if TYPE_CHECKING: from ding.framework import Context, OnlineRLContext + from ding.policy import Policy + from ding.reward_model import BaseRewardModel class OffPolicyLearner: @@ -17,24 +19,31 @@ class OffPolicyLearner: the `__call__` method to execute the whole learning process. """ + def __new__(cls, *args, **kwargs): + if task.router.is_active and not task.has_role(task.role.LEARNER): + return task.void() + return super(OffPolicyLearner, cls).__new__(cls) + def __init__( self, cfg: EasyDict, - policy, + policy: 'Policy', buffer_: Union[Buffer, List[Tuple[Buffer, float]], Dict[str, Buffer]], - reward_model=None + reward_model: Optional['BaseRewardModel'] = None, + log_freq: int = 100, ) -> None: """ Arguments: - cfg (:obj:`EasyDict`): Config. - policy (:obj:`Policy`): The policy to be trained. - - buffer\_ (:obj:`Buffer`): The replay buffer to store the data for training. - - reward_model (:obj:`nn.Module`): Additional reward estimator likes RND, ICM, etc. \ + - buffer (:obj:`Buffer`): The replay buffer to store the data for training. + - reward_model (:obj:`BaseRewardModel`): Additional reward estimator likes RND, ICM, etc. \ default to None. + - log_freq (:obj:`int`): The frequency (iteration) of showing log. """ self.cfg = cfg self._fetcher = task.wrap(offpolicy_data_fetcher(cfg, buffer_)) - self._trainer = task.wrap(trainer(cfg, policy)) + self._trainer = task.wrap(trainer(cfg, policy, log_freq=log_freq)) if reward_model is not None: self._reward_estimator = task.wrap(reward_estimator(cfg, reward_model)) else: diff --git a/ding/framework/middleware/tests/mock_for_test.py b/ding/framework/middleware/tests/mock_for_test.py index 7fa6cde7aa..9bffc1875e 100644 --- a/ding/framework/middleware/tests/mock_for_test.py +++ b/ding/framework/middleware/tests/mock_for_test.py @@ -1,5 +1,6 @@ from typing import Union, Any, List, Callable, Dict, Optional from collections import namedtuple +import random import torch import treetensor.numpy as tnp from easydict import EasyDict @@ -62,7 +63,7 @@ def process_transition(self, obs: Any, model_output: dict, timestep: namedtuple) 'logit': 1.0, 'value': 2.0, 'reward': 0.1, - 'done': True, + 'done': timestep.done, } return transition @@ -74,7 +75,7 @@ def __init__(self) -> None: self.env_num = env_num self.obs_dim = obs_dim self.closed = False - self._reward_grow_indicator = 1 + self._steps = [0 for _ in range(self.env_num)] @property def ready_obs(self) -> tnp.array: @@ -90,20 +91,29 @@ def launch(self, reset_param: Optional[Dict] = None) -> None: return def reset(self, reset_param: Optional[Dict] = None) -> None: - return + self._steps = [0 for _ in range(self.env_num)] def step(self, actions: tnp.ndarray) -> List[tnp.ndarray]: timesteps = [] for i in range(self.env_num): + if self._steps[i] < 5: + done = False + elif self._steps[i] < 10: + done = random.random() > 0.5 + else: + done = True + if done: + self._steps[i] = 0 + else: + self._steps[i] += 1 timestep = dict( obs=torch.rand(self.obs_dim), reward=1.0, - done=True, - info={'final_eval_reward': self._reward_grow_indicator * 1.0}, + done=done, + info={'eval_episode_return': 10.0} if done else {}, env_id=i, ) timesteps.append(tnp.array(timestep)) - self._reward_grow_indicator += 1 # final_eval_reward will increase as step method is called return timesteps diff --git a/ding/framework/middleware/tests/test_advantage_estimator.py b/ding/framework/middleware/tests/test_advantage_estimator.py index cc898b2546..dd2a12c19b 100644 --- a/ding/framework/middleware/tests/test_advantage_estimator.py +++ b/ding/framework/middleware/tests/test_advantage_estimator.py @@ -9,6 +9,7 @@ import copy from ding.framework.middleware.functional.advantage_estimator import gae_estimator +from ding.framework.middleware.functional.advantage_estimator import montecarlo_return_estimator from ding.utils.data import ttorch_collate from typing import Any, List, Dict, Optional @@ -33,7 +34,21 @@ def get_attribute(self, name: str) -> Any: def call_gae_estimator(batch_size: int = 32, trajectory_end_idx_size: int = 5, buffer: Optional[Buffer] = None): - cfg = EasyDict({'policy': {'collect': {'discount_factor': 0.9, 'gae_lambda': 0.95}}}) + cfg = EasyDict( + { + 'policy': { + 'model': { + 'obs_shape': 4, + 'action_shape': 2, + }, + 'collect': { + 'discount_factor': 0.9, + 'gae_lambda': 0.95 + }, + 'cuda': False + } + } + ) ctx = OnlineRLContext() assert trajectory_end_idx_size <= batch_size @@ -53,7 +68,7 @@ def call_gae_estimator(batch_size: int = 32, trajectory_end_idx_size: int = 5, b } ) for _ in range(batch_size) ] - ctx.trajectories_copy = ttorch_collate(copy.deepcopy(ctx.trajectories)) + ctx.trajectories_copy = ttorch_collate(copy.deepcopy(ctx.trajectories), cat_1dim=True) traj_flag = ctx.trajectories_copy.done.clone() traj_flag[ctx.trajectory_end_idx] = True ctx.trajectories_copy.traj_flag = traj_flag @@ -62,12 +77,12 @@ def call_gae_estimator(batch_size: int = 32, trajectory_end_idx_size: int = 5, b gae_estimator(cfg, MockPolicy(TheModelClass()), buffer)(ctx) if buffer is not None: - train_data = [d.data for d in buffer.export_data()] + train_data = [d.data for d in list(buffer.storage)] for d in train_data: - d.logit = d.logit[0] - d.next_obs = d.next_obs[0] - d.obs = d.obs[0] - ctx.train_data = ttorch_collate(train_data) + d.logit = d.logit + d.next_obs = d.next_obs + d.obs = d.obs + ctx.train_data = ttorch_collate(train_data, cat_1dim=True) assert ctx.trajectories is None assert torch.equal(ctx.trajectories_copy.action, ctx.train_data.action) @@ -85,3 +100,61 @@ def test_gae_estimator(): trajectory_end_idx_size = 5 call_gae_estimator(batch_size, trajectory_end_idx_size) call_gae_estimator(batch_size, trajectory_end_idx_size, DequeBuffer(size=batch_size)) + + +class MockPGPolicy(Mock): + + def __init__(self, cfg) -> None: + super(MockPGPolicy, self).__init__() + self._cfg = EasyDict(cfg) + self._gamma = self._cfg.collect.discount_factor + self._unroll_len = self._cfg.collect.unroll_len + + def get_attribute(self, name: str) -> Any: + return self._model + + +def call_montecarlo_return_estimator(batch_size: int = 32): + + cfg = dict( + learn=dict(ignore_done=False, ), + collect=dict( + unroll_len=1, + discount_factor=0.9, + ), + ) + ctx = OnlineRLContext() + ctx.episodes = [ + [ + treetensor.torch.Tensor( + { + 'action': treetensor.torch.randint(low=0, high=2, size=(1, )), + 'collect_train_iter': [0], + 'done': False if i != batch_size - 1 else True, + 'logit': treetensor.torch.randn(2), + 'next_obs': treetensor.torch.randn(4), + 'obs': treetensor.torch.randn(4), + 'reward': [1.0], + 'value': torch.distributions.uniform.Uniform(0, 4).sample([1]) + } + ) for i in range(batch_size) + ] + ] + ctx.episodes_copy = treetensor.torch.concat( + [ttorch_collate(copy.deepcopy(episode), cat_1dim=True) for episode in ctx.episodes], dim=0 + ) + with patch("ding.policy.Policy", MockPGPolicy): + montecarlo_return_estimator(MockPGPolicy(cfg))(ctx) + + assert torch.equal(ctx.episodes_copy.action, ctx.train_data.action) + assert torch.equal(ctx.episodes_copy.collect_train_iter, ctx.train_data.collect_train_iter) + assert torch.equal(ctx.episodes_copy.logit, ctx.train_data.logit) + assert torch.equal(ctx.episodes_copy.next_obs, ctx.train_data.next_obs) + assert torch.equal(ctx.episodes_copy.obs, ctx.train_data.obs) + assert torch.equal(ctx.episodes_copy.reward, ctx.train_data.reward) + + +@pytest.mark.unittest +def test_montecarlo_return_estimator(): + batch_size = 32 + call_montecarlo_return_estimator(batch_size) diff --git a/ding/framework/middleware/tests/test_barrier.py b/ding/framework/middleware/tests/test_barrier.py new file mode 100644 index 0000000000..de176eda2d --- /dev/null +++ b/ding/framework/middleware/tests/test_barrier.py @@ -0,0 +1,144 @@ +import random +import time +import socket +import pytest +import multiprocessing as mp +from ditk import logging +from ding.framework import task +from ding.framework.parallel import Parallel +from ding.framework.context import OnlineRLContext +from ding.framework.middleware.barrier import Barrier + +PORTS_LIST = ["1235", "1236", "1237"] + + +class EnvStepMiddleware: + + def __call__(self, ctx): + yield + ctx.env_step += 1 + + +class SleepMiddleware: + + def __init__(self, node_id): + self.node_id = node_id + + def random_sleep(self, diection, step): + random.seed(self.node_id + step) + sleep_second = random.randint(1, 5) + logging.info("Node:[{}] env_step:[{}]-{} will sleep:{}s".format(self.node_id, step, diection, sleep_second)) + for i in range(sleep_second): + time.sleep(1) + print("Node:[{}] sleepping...".format(self.node_id)) + logging.info("Node:[{}] env_step:[{}]-{} wake up!".format(self.node_id, step, diection)) + + def __call__(self, ctx): + self.random_sleep("forward", ctx.env_step) + yield + self.random_sleep("backward", ctx.env_step) + + +def star_barrier(): + with task.start(ctx=OnlineRLContext()): + node_id = task.router.node_id + if node_id == 0: + attch_from_nums = 3 + else: + attch_from_nums = 0 + barrier = Barrier(attch_from_nums) + task.use(barrier, lock=False) + task.use(SleepMiddleware(node_id), lock=False) + task.use(barrier, lock=False) + task.use(EnvStepMiddleware(), lock=False) + try: + task.run(2) + except Exception as e: + logging.error(e) + assert False + + +def mesh_barrier(): + with task.start(ctx=OnlineRLContext()): + node_id = task.router.node_id + attch_from_nums = 3 - task.router.node_id + barrier = Barrier(attch_from_nums) + task.use(barrier, lock=False) + task.use(SleepMiddleware(node_id), lock=False) + task.use(barrier, lock=False) + task.use(EnvStepMiddleware(), lock=False) + try: + task.run(2) + except Exception as e: + logging.error(e) + assert False + + +def unmatch_barrier(): + with task.start(ctx=OnlineRLContext()): + node_id = task.router.node_id + attch_from_nums = 3 - task.router.node_id + task.use(Barrier(attch_from_nums, 5), lock=False) + if node_id != 2: + task.use(Barrier(attch_from_nums, 5), lock=False) + try: + task.run(2) + except TimeoutError as e: + assert node_id != 2 + logging.info("Node:[{}] timeout with barrier".format(node_id)) + else: + time.sleep(5) + assert node_id == 2 + logging.info("Node:[{}] finish barrier".format(node_id)) + + +def launch_barrier(args): + i, topo, fn, test_id = args + address = socket.gethostbyname(socket.gethostname()) + topology = "alone" + attach_to = [] + port_base = PORTS_LIST[test_id] + port = port_base + str(i) + if topo == 'star': + if i != 0: + attach_to = ['tcp://{}:{}{}'.format(address, port_base, 0)] + elif topo == 'mesh': + for j in range(i): + attach_to.append('tcp://{}:{}{}'.format(address, port_base, j)) + + Parallel.runner( + node_ids=i, + ports=int(port), + attach_to=attach_to, + topology=topology, + protocol="tcp", + n_parallel_workers=1, + startup_interval=0 + )(fn) + + +@pytest.mark.unittest +def test_star_topology_barrier(): + ctx = mp.get_context("spawn") + with ctx.Pool(processes=4) as pool: + pool.map(launch_barrier, [[i, 'star', star_barrier, 0] for i in range(4)]) + pool.close() + pool.join() + + +@pytest.mark.unittest +def test_mesh_topology_barrier(): + ctx = mp.get_context("spawn") + with ctx.Pool(processes=4) as pool: + pool.map(launch_barrier, [[i, 'mesh', mesh_barrier, 1] for i in range(4)]) + pool.close() + pool.join() + + +@pytest.mark.unittest +def test_unmatch_barrier(): + ctx = mp.get_context("spawn") + with ctx.Pool(processes=4) as pool: + pool.map(launch_barrier, [[i, 'mesh', unmatch_barrier, 2] for i in range(4)]) + pool.close() + pool.join() diff --git a/ding/framework/middleware/tests/test_ckpt_handler.py b/ding/framework/middleware/tests/test_ckpt_handler.py index 9a22ffece0..f0d81f0a33 100644 --- a/ding/framework/middleware/tests/test_ckpt_handler.py +++ b/ding/framework/middleware/tests/test_ckpt_handler.py @@ -11,7 +11,7 @@ from unittest.mock import Mock, patch from ding.framework import task -from ding.utils import save_file +from ding.policy.base_policy import Policy class TheModelClass(nn.Module): @@ -22,24 +22,28 @@ def state_dict(self): class MockPolicy(Mock): - def __init__(self, model) -> None: - super(MockPolicy, self).__init__() + def __init__(self, model, **kwargs) -> None: + super(MockPolicy, self).__init__(model) self.learn_mode = model + @property + def eval_mode(self): + return EasyDict({"state_dict": lambda: {}}) + @pytest.mark.unittest def test_ckpt_saver(): - cfg = EasyDict({'exp_name': 'test_ckpt_saver_exp'}) + exp_name = 'test_ckpt_saver_exp' ctx = OnlineRLContext() train_freq = 100 model = TheModelClass() - if not os.path.exists(cfg.exp_name): - os.makedirs(cfg.exp_name) + if not os.path.exists(exp_name): + os.makedirs(exp_name) - prefix = '{}/ckpt'.format(cfg.exp_name) + prefix = '{}/ckpt'.format(exp_name) with patch("ding.policy.Policy", MockPolicy), task.start(): policy = MockPolicy(model) @@ -48,9 +52,9 @@ def mock_save_file(path, data, fs_type=None, use_lock=False): assert path == "{}/eval.pth.tar".format(prefix) with patch("ding.framework.middleware.ckpt_handler.save_file", mock_save_file): - ctx.train_iter = 0 + ctx.train_iter = 1 ctx.eval_value = 9.4 - ckpt_saver = CkptSaver(cfg, policy, train_freq) + ckpt_saver = CkptSaver(policy, exp_name, train_freq) ckpt_saver(ctx) def mock_save_file(path, data, fs_type=None, use_lock=False): @@ -68,4 +72,4 @@ def mock_save_file(path, data, fs_type=None, use_lock=False): task.finish = True ckpt_saver(ctx) - shutil.rmtree(cfg.exp_name) + shutil.rmtree(exp_name) diff --git a/ding/framework/middleware/tests/test_collector.py b/ding/framework/middleware/tests/test_collector.py index 16c577ae2b..40dc6dae7f 100644 --- a/ding/framework/middleware/tests/test_collector.py +++ b/ding/framework/middleware/tests/test_collector.py @@ -10,12 +10,11 @@ @pytest.mark.unittest def test_inferencer(): - cfg = copy.deepcopy(CONFIG) ctx = OnlineRLContext() with patch("ding.policy.Policy", MockPolicy), patch("ding.envs.BaseEnvManagerV2", MockEnv): policy = MockPolicy() env = MockEnv() - inferencer(cfg, policy, env)(ctx) + inferencer(0, policy, env)(ctx) assert isinstance(ctx.inference_output, dict) assert ctx.inference_output[0] == {'action': torch.Tensor([0.])} # sum of zeros([2, 2]) assert ctx.inference_output[1] == {'action': torch.Tensor([4.])} # sum of ones([2, 2]) @@ -23,17 +22,19 @@ def test_inferencer(): @pytest.mark.unittest def test_rolloutor(): - cfg = copy.deepcopy(CONFIG) + N = 20 ctx = OnlineRLContext() transitions = TransitionList(2) with patch("ding.policy.Policy", MockPolicy), patch("ding.envs.BaseEnvManagerV2", MockEnv): policy = MockPolicy() env = MockEnv() - for _ in range(10): - inferencer(cfg, policy, env)(ctx) - rolloutor(cfg, policy, env, transitions)(ctx) - assert ctx.env_episode == 20 # 10 * env_num - assert ctx.env_step == 20 # 10 * env_num + i = inferencer(0, policy, env) + r = rolloutor(policy, env, transitions) + for _ in range(N): + i(ctx) + r(ctx) + assert ctx.env_step == N * 2 # N * env_num + assert ctx.env_episode >= N // 10 * 2 # N * env_num @pytest.mark.unittest diff --git a/ding/framework/middleware/tests/test_data_processor.py b/ding/framework/middleware/tests/test_data_processor.py index 4b67f6cc07..d63d392943 100644 --- a/ding/framework/middleware/tests/test_data_processor.py +++ b/ding/framework/middleware/tests/test_data_processor.py @@ -1,10 +1,11 @@ +import tempfile import pytest from ding.data.buffer import DequeBuffer from ding.framework import Context, OnlineRLContext, OfflineRLContext from ding.framework.middleware.functional.data_processor import \ - data_pusher, offpolicy_data_fetcher, offline_data_fetcher, offline_data_saver, sqil_data_pusher + data_pusher, offpolicy_data_fetcher, offline_data_fetcher, offline_data_saver, sqil_data_pusher, buffer_saver from ding.data.buffer.middleware import PriorityExperienceReplay @@ -61,14 +62,14 @@ def offpolicy_data_fetcher_type_buffer_helper(priority=0.5, use_list=True): assert [d['reward'] for d in ctx.train_data] == [1 for i in range(20)] assert [d['info'] for d in ctx.train_data] == ['xxx' for i in range(20)] assert [d['priority_IS'] for d in ctx.train_data] == [torch.tensor([1]) for i in range(20)] - assert buffer.export_data()[0].meta['priority'] == 1.0 + assert list(buffer.storage)[0].meta['priority'] == 1.0 # assert sorted(ctx.train_data) == [i for i in range(20)] try: next(func_generator) except StopIteration: pass - assert buffer.export_data()[0].meta['priority'] == priority + assert list(buffer.storage)[0].meta['priority'] == priority def call_offpolicy_data_fetcher_type_buffer(): @@ -152,7 +153,9 @@ def __len__(self): ctx.train_epoch = 0 data_tmp = [] - for i, _ in enumerate(offline_data_fetcher(cfg, MyDataset())(ctx)): + fetch = offline_data_fetcher(cfg, MyDataset()) + for i in range(num_batch): + fetch(ctx) assert i // num_batch == ctx.train_epoch data_tmp.extend(ctx.train_data) @@ -186,7 +189,7 @@ def mock_offline_data_save_type(exp_data, expert_data_path, data_type): with patch("ding.framework.middleware.functional.data_processor.offline_data_save_type", mock_offline_data_save_type): - offline_data_saver(cfg=None, data_path=data_path_, data_type='naive')(ctx) + offline_data_saver(data_path=data_path_, data_type='naive')(ctx) assert ctx.trajectories is None @@ -200,7 +203,7 @@ def mock_offline_data_save_type(exp_data, expert_data_path, data_type): with patch("ding.framework.middleware.functional.data_processor.offline_data_save_type", mock_offline_data_save_type): - offline_data_saver(cfg=None, data_path=data_path_, data_type='hdf5')(ctx) + offline_data_saver(data_path=data_path_, data_type='hdf5')(ctx) assert ctx.trajectories is None @@ -224,7 +227,7 @@ def test_sqil_data_pusher(): buffer = DequeBuffer(size=10) sqil_data_pusher(cfg=None, buffer_=buffer, expert=True)(ctx) assert buffer.count() == 5 - assert all(t.data.reward == 1 for t in buffer.export_data()) + assert all(t.data.reward == 1 for t in list(buffer.storage)) # expert = False ctx = OnlineRLContext() @@ -232,4 +235,26 @@ def test_sqil_data_pusher(): buffer = DequeBuffer(size=10) sqil_data_pusher(cfg=None, buffer_=buffer, expert=False)(ctx) assert buffer.count() == 5 - assert all(t.data.reward == 0 for t in buffer.export_data()) + assert all(t.data.reward == 0 for t in list(buffer.storage)) + + +@pytest.mark.unittest +def test_buffer_saver(): + with tempfile.TemporaryDirectory() as tmpdirname: + test_folder = os.path.join(tmpdirname, "test_buffer_saver") + cfg = EasyDict({"exp_name": test_folder}) + os.makedirs(test_folder) + buffer_ = DequeBuffer(size=10) + ctx = OnlineRLContext() + ctx.trajectories = [i for i in range(5)] + ctx.env_step = 0 + data_pusher(cfg=cfg, buffer_=buffer_)(ctx) + assert buffer_.count() == 5 + buffer_saver(cfg=cfg, buffer_=buffer_, replace=False)(ctx) + buffer_saver(cfg=cfg, buffer_=buffer_, replace=True)(ctx) + buffer_1 = DequeBuffer(size=10) + buffer_1.load_data(os.path.join(test_folder, "replaybuffer", "data_latest.hkl")) + assert buffer_1.count() == 5 + buffer_2 = DequeBuffer(size=10) + buffer_2.load_data(os.path.join(test_folder, "replaybuffer", "data_envstep_0.hkl")) + assert buffer_2.count() == 5 diff --git a/ding/framework/middleware/tests/test_distributer.py b/ding/framework/middleware/tests/test_distributer.py new file mode 100644 index 0000000000..7651e66ec7 --- /dev/null +++ b/ding/framework/middleware/tests/test_distributer.py @@ -0,0 +1,269 @@ +import shutil +from time import sleep +import pytest +import numpy as np +import tempfile + +import torch +from ding.data.model_loader import FileModelLoader +from ding.data.storage_loader import FileStorageLoader +from ding.framework import task +from ding.framework.context import OnlineRLContext +from ding.framework.middleware.distributer import ContextExchanger, ModelExchanger, PeriodicalModelExchanger +from ding.framework.parallel import Parallel +from ding.utils.default_helper import set_pkg_seed +from os import path + + +def context_exchanger_main(): + with task.start(ctx=OnlineRLContext()): + if task.router.node_id == 0: + task.add_role(task.role.LEARNER) + elif task.router.node_id == 1: + task.add_role(task.role.COLLECTOR) + + task.use(ContextExchanger(skip_n_iter=1)) + + if task.has_role(task.role.LEARNER): + + def learner_context(ctx: OnlineRLContext): + assert len(ctx.trajectories) == 2 + assert len(ctx.trajectory_end_idx) == 4 + assert len(ctx.episodes) == 8 + assert ctx.env_step > 0 + assert ctx.env_episode > 0 + yield + ctx.train_iter += 1 + + task.use(learner_context) + elif task.has_role(task.role.COLLECTOR): + + def collector_context(ctx: OnlineRLContext): + if ctx.total_step > 0: + assert ctx.train_iter > 0 + yield + ctx.trajectories = [np.random.rand(10, 10) for _ in range(2)] + ctx.trajectory_end_idx = [1 for _ in range(4)] + ctx.episodes = [np.random.rand(10, 10) for _ in range(8)] + ctx.env_step += 1 + ctx.env_episode += 1 + + task.use(collector_context) + + task.run(max_step=3) + + +@pytest.mark.tmp +def test_context_exchanger(): + Parallel.runner(n_parallel_workers=2)(context_exchanger_main) + + +def context_exchanger_with_storage_loader_main(): + with task.start(ctx=OnlineRLContext()): + if task.router.node_id == 0: + task.add_role(task.role.LEARNER) + elif task.router.node_id == 1: + task.add_role(task.role.COLLECTOR) + + tempdir = path.join(tempfile.gettempdir(), "test_storage_loader") + storage_loader = FileStorageLoader(dirname=tempdir) + try: + task.use(ContextExchanger(skip_n_iter=1, storage_loader=storage_loader)) + + if task.has_role(task.role.LEARNER): + + def learner_context(ctx: OnlineRLContext): + assert len(ctx.trajectories) == 2 + assert len(ctx.trajectory_end_idx) == 4 + assert len(ctx.episodes) == 8 + assert ctx.env_step > 0 + assert ctx.env_episode > 0 + yield + ctx.train_iter += 1 + + task.use(learner_context) + elif task.has_role(task.role.COLLECTOR): + + def collector_context(ctx: OnlineRLContext): + if ctx.total_step > 0: + assert ctx.train_iter > 0 + yield + ctx.trajectories = [np.random.rand(10, 10) for _ in range(2)] + ctx.trajectory_end_idx = [1 for _ in range(4)] + ctx.episodes = [np.random.rand(10, 10) for _ in range(8)] + ctx.env_step += 1 + ctx.env_episode += 1 + + task.use(collector_context) + + task.run(max_step=3) + finally: + storage_loader.shutdown() + sleep(1) + if path.exists(tempdir): + shutil.rmtree(tempdir) + + +@pytest.mark.tmp +def test_context_exchanger_with_storage_loader(): + Parallel.runner(n_parallel_workers=2)(context_exchanger_with_storage_loader_main) + + +class MockPolicy: + + def __init__(self) -> None: + self._model = self._get_model(10, 10) + + def _get_model(self, X_shape, y_shape) -> torch.nn.Module: + return torch.nn.Sequential( + torch.nn.Linear(X_shape, 24), torch.nn.ReLU(), torch.nn.Linear(24, 24), torch.nn.ReLU(), + torch.nn.Linear(24, y_shape) + ) + + def train(self, X, y): + loss_fn = torch.nn.MSELoss(reduction="mean") + optimizer = torch.optim.Adam(self._model.parameters(), lr=0.01) + y_pred = self._model(X) + loss = loss_fn(y_pred, y) + optimizer.zero_grad() + loss.backward() + optimizer.step() + + def predict(self, X): + with torch.no_grad(): + return self._model(X) + + +def model_exchanger_main(): + with task.start(ctx=OnlineRLContext()): + set_pkg_seed(0, use_cuda=False) + policy = MockPolicy() + X = torch.rand(10) + y = torch.rand(10) + + if task.router.node_id == 0: + task.add_role(task.role.LEARNER) + else: + task.add_role(task.role.COLLECTOR) + + task.use(ModelExchanger(policy._model)) + + if task.has_role(task.role.LEARNER): + + def train(ctx): + policy.train(X, y) + sleep(0.3) + + task.use(train) + else: + y_pred1 = policy.predict(X) + + def pred(ctx): + if ctx.total_step > 0: + y_pred2 = policy.predict(X) + # Ensure model is upgraded + assert any(y_pred1 != y_pred2) + sleep(0.3) + + task.use(pred) + + task.run(2) + + +@pytest.mark.tmp +def test_model_exchanger(): + Parallel.runner(n_parallel_workers=2, startup_interval=0)(model_exchanger_main) + + +def model_exchanger_main_with_model_loader(): + with task.start(ctx=OnlineRLContext()): + set_pkg_seed(0, use_cuda=False) + policy = MockPolicy() + X = torch.rand(10) + y = torch.rand(10) + + if task.router.node_id == 0: + task.add_role(task.role.LEARNER) + else: + task.add_role(task.role.COLLECTOR) + + tempdir = path.join(tempfile.gettempdir(), "test_model_loader") + model_loader = FileModelLoader(policy._model, dirname=tempdir) + task.use(ModelExchanger(policy._model, model_loader=model_loader)) + + try: + if task.has_role(task.role.LEARNER): + + def train(ctx): + policy.train(X, y) + sleep(0.3) + + task.use(train) + else: + y_pred1 = policy.predict(X) + + def pred(ctx): + if ctx.total_step > 0: + y_pred2 = policy.predict(X) + # Ensure model is upgraded + assert any(y_pred1 != y_pred2) + sleep(0.3) + + task.use(pred) + task.run(2) + finally: + model_loader.shutdown() + sleep(0.3) + if path.exists(tempdir): + shutil.rmtree(tempdir) + + +@pytest.mark.tmp +def test_model_exchanger_with_model_loader(): + Parallel.runner(n_parallel_workers=2, startup_interval=0)(model_exchanger_main_with_model_loader) + + +def periodical_model_exchanger_main(): + with task.start(ctx=OnlineRLContext()): + set_pkg_seed(0, use_cuda=False) + policy = MockPolicy() + X = torch.rand(10) + y = torch.rand(10) + + if task.router.node_id == 0: + task.add_role(task.role.LEARNER) + task.use(PeriodicalModelExchanger(policy._model, mode="send", period=3)) + else: + task.add_role(task.role.COLLECTOR) + task.use(PeriodicalModelExchanger(policy._model, mode="receive", period=1, stale_toleration=3)) + + if task.has_role(task.role.LEARNER): + + def train(ctx): + policy.train(X, y) + sleep(0.3) + + task.use(train) + else: + y_pred1 = policy.predict(X) + print("y_pred1: ", y_pred1) + stale = 1 + + def pred(ctx): + nonlocal stale + y_pred2 = policy.predict(X) + print("y_pred2: ", y_pred2) + stale += 1 + assert stale <= 3 or all(y_pred1 == y_pred2) + if any(y_pred1 != y_pred2): + stale = 1 + + sleep(0.3) + + task.use(pred) + task.run(8) + + +@pytest.mark.tmp +def test_periodical_model_exchanger(): + Parallel.runner(n_parallel_workers=2, startup_interval=0)(periodical_model_exchanger_main) diff --git a/ding/framework/middleware/tests/test_enhancer.py b/ding/framework/middleware/tests/test_enhancer.py index a1765c031d..10d34b264f 100644 --- a/ding/framework/middleware/tests/test_enhancer.py +++ b/ding/framework/middleware/tests/test_enhancer.py @@ -2,8 +2,7 @@ import torch from ding.framework import OnlineRLContext from ding.data.buffer import DequeBuffer -from easydict import EasyDict -from typing import Any, List, Dict, Optional +from typing import Any import numpy as np import copy from ding.framework.middleware.functional.enhancer import reward_estimator, her_data_enhancer diff --git a/ding/framework/middleware/tests/test_evaluator.py b/ding/framework/middleware/tests/test_evaluator.py index 4e78b150bc..6f0f88586a 100644 --- a/ding/framework/middleware/tests/test_evaluator.py +++ b/ding/framework/middleware/tests/test_evaluator.py @@ -24,4 +24,4 @@ def test_interaction_evaluator(): # there are 2 env_num and 5 episodes in the test. # so when interaction_evaluator runs the first time, reward is [[1, 2, 3], [2, 3]] and the avg = 2.2 # the second time, reward is [[4, 5, 6], [5, 6]] . . . - assert ctx.eval_value == 2.2 + i // 10 * 3.0 + assert ctx.eval_value == 10.0 diff --git a/ding/framework/middleware/tests/test_logger.py b/ding/framework/middleware/tests/test_logger.py index 1c0a772f1d..e5ba9a21fd 100644 --- a/ding/framework/middleware/tests/test_logger.py +++ b/ding/framework/middleware/tests/test_logger.py @@ -1,14 +1,19 @@ -import pytest -from ding.framework import OnlineRLContext, OfflineRLContext, ding_init -from ding.framework.middleware.functional import online_logger, offline_logger -from easydict import EasyDict -import os from os import path -import shutil +import os +from functools import partial +from easydict import EasyDict from collections import deque +import pytest +import wandb +import torch.nn as nn +from unittest.mock import MagicMock from unittest.mock import Mock, patch + from ding.utils import DistributedWriter -import copy +from ding.framework.middleware.tests import MockPolicy, CONFIG +from ding.framework import OnlineRLContext, OfflineRLContext +from ding.framework.middleware.functional import online_logger, offline_logger, wandb_online_logger, \ + wandb_offline_logger test_folder = "test_exp" test_path = path.join(os.getcwd(), test_folder) @@ -57,10 +62,10 @@ def __init__(self): self.ctx = get_online_ctx() def add_scalar(self, tag, scalar_value, global_step): - if tag in ['basic/eval_episode_reward_mean-env_step', 'basic/eval_episode_reward_mean']: + if tag in ['basic/eval_episode_return_mean-env_step', 'basic/eval_episode_return_mean']: assert scalar_value == self.ctx.eval_value assert global_step == self.ctx.env_step - elif tag == 'basic/eval_episode_reward_mean-train_iter': + elif tag == 'basic/eval_episode_return_mean-train_iter': assert scalar_value == self.ctx.eval_value assert global_step == self.ctx.train_iter elif tag in ['basic/train_td_error-env_step', 'basic/train_td_error']: @@ -77,6 +82,9 @@ def add_histogram(self, tag, values, global_step): assert values == [1, 2, 3, 4, 5, 6] assert global_step in [self.ctx.train_iter, self.ctx.env_step] + def close(self): + pass + def mock_get_online_instance(): return MockOnlineWriter() @@ -126,7 +134,7 @@ def __init__(self): def add_scalar(self, tag, scalar_value, global_step): assert global_step == self.ctx.train_iter - if tag == 'basic/eval_episode_reward_mean-train_iter': + if tag == 'basic/eval_episode_return_mean-train_iter': assert scalar_value == self.ctx.eval_value elif tag == 'basic/train_td_error-train_iter': assert scalar_value == self.ctx.train_output['td_error'] @@ -138,12 +146,14 @@ def add_histogram(self, tag, values, global_step): assert values == [1, 2, 3, 4, 5, 6] assert global_step == self.ctx.train_iter + def close(self): + pass + def mock_get_offline_instance(): return MockOfflineWriter() -@pytest.mark.unittest class TestOfflineLogger: def test_offline_logger_no_scalars(self, offline_ctx_output_dict): @@ -154,3 +164,131 @@ def test_offline_logger_scalars(self, offline_scalar_ctx): with patch.object(DistributedWriter, 'get_instance', new=mock_get_offline_instance): with pytest.raises(NotImplementedError) as exc_info: offline_logger()(offline_scalar_ctx) + + +class TheModelClass(nn.Module): + + def state_dict(self): + return 'fake_state_dict' + + +class TheEnvClass(Mock): + + def enable_save_replay(self, replay_path): + return + + +class TheObsDataClass(Mock): + + def __getitem__(self, index): + return [[1, 1, 1]] * 50 + + +class The1DDataClass(Mock): + + def __getitem__(self, index): + return [[1]] * 50 + + +@pytest.mark.unittest +def test_wandb_online_logger(): + record_path = './video_qbert_dqn' + cfg = EasyDict( + dict( + gradient_logger=True, + plot_logger=True, + action_logger=True, + return_logger=True, + video_logger=True, + ) + ) + env = TheEnvClass() + ctx = OnlineRLContext() + ctx.train_output = [{'reward': 1, 'q_value': [1.0]}] + wandb.init(config=cfg, anonymous="must") + + def mock_metric_logger(data, step): + metric_list = [ + "q_value", + "target q_value", + "loss", + "lr", + "entropy", + "reward", + "q value", + "video", + "q value distribution", + "train iter", + "episode return mean", + "env step", + "action", + "actions_of_trajectory_0", + "actions_of_trajectory_1", + "actions_of_trajectory_2", + "actions_of_trajectory_3", + "return distribution", + ] + assert set(data.keys()) <= set(metric_list) + + def mock_gradient_logger(input_model, model, log, log_freq, log_graph): + assert input_model == model + + def test_wandb_online_logger_metric(): + model = TheModelClass() + with patch.object(wandb, 'log', new=mock_metric_logger): + wandb_online_logger(record_path, cfg, env=env, model=model, anonymous=True)(ctx) + + def test_wandb_online_logger_gradient(): + model = TheModelClass() + with patch.object(wandb, 'watch', new=partial(mock_gradient_logger, model=model)): + wandb_online_logger(record_path, cfg, env=env, model=model, anonymous=True)(ctx) + + test_wandb_online_logger_metric() + test_wandb_online_logger_gradient() + + +@pytest.mark.tmp +def test_wandb_offline_logger(): + record_path = './video_pendulum_cql' + cfg = EasyDict(dict(gradient_logger=True, plot_logger=True, action_logger=True, vis_dataset=True)) + env = TheEnvClass() + ctx = OfflineRLContext() + ctx.train_output = [{'reward': 1, 'q_value': [1.0]}] + model = TheModelClass() + wandb.init(config=cfg, anonymous="must") + exp_config = EasyDict(dict(dataset_path='dataset.h5')) + + def mock_metric_logger(data, step=None): + metric_list = [ + "q_value", "target q_value", "loss", "lr", "entropy", "reward", "q value", "video", "q value distribution", + "train iter", 'dataset' + ] + assert set(data.keys()) < set(metric_list) + + def mock_gradient_logger(input_model, log, log_freq, log_graph): + assert input_model == model + + def mock_image_logger(imagepath): + assert os.path.splitext(imagepath)[-1] == '.png' + + def test_wandb_offline_logger_gradient(): + cfg.vis_dataset = False + print(cfg) + with patch.object(wandb, 'watch', new=mock_gradient_logger): + wandb_offline_logger( + record_path=record_path, cfg=cfg, exp_config=exp_config, env=env, model=model, anonymous=True + )(ctx) + + def test_wandb_offline_logger_dataset(): + cfg.vis_dataset = True + m = MagicMock() + m.__enter__.return_value = {'obs': TheObsDataClass(), 'action': The1DDataClass(), 'reward': The1DDataClass()} + with patch.object(wandb, 'log', new=mock_metric_logger): + with patch.object(wandb, 'Image', new=mock_image_logger): + with patch('h5py.File', return_value=m): + wandb_offline_logger( + record_path=record_path, cfg=cfg, exp_config=exp_config, env=env, model=model, anonymous=True + )(ctx) + + test_wandb_offline_logger_gradient() + test_wandb_offline_logger_dataset() diff --git a/ding/framework/middleware/tests/test_priority.py b/ding/framework/middleware/tests/test_priority.py new file mode 100644 index 0000000000..19261213d6 --- /dev/null +++ b/ding/framework/middleware/tests/test_priority.py @@ -0,0 +1,33 @@ +#unittest for priority_calculator + +import unittest +import pytest +import numpy as np +from unittest.mock import Mock, patch +from ding.framework import OnlineRLContext, OfflineRLContext +from ding.framework import task, Parallel +from ding.framework.middleware.functional import priority_calculator + + +class MockPolicy(Mock): + + def priority_fun(self, data): + return np.random.rand(len(data)) + + +@pytest.mark.unittest +def test_priority_calculator(): + policy = MockPolicy() + ctx = OnlineRLContext() + ctx.trajectories = [ + { + 'obs': np.random.rand(2, 2), + 'next_obs': np.random.rand(2, 2), + 'reward': np.random.rand(1), + 'info': {} + } for _ in range(10) + ] + priority_calculator_middleware = priority_calculator(priority_calculation_fn=policy.priority_fun) + priority_calculator_middleware(ctx) + assert len(ctx.trajectories) == 10 + assert all([isinstance(traj['priority'], float) for traj in ctx.trajectories]) diff --git a/ding/framework/middleware/tests/test_trainer.py b/ding/framework/middleware/tests/test_trainer.py index 97d4d187b3..b9dbf9f55c 100644 --- a/ding/framework/middleware/tests/test_trainer.py +++ b/ding/framework/middleware/tests/test_trainer.py @@ -1,7 +1,8 @@ import pytest import random -import torch import copy +import torch +import treetensor.torch as ttorch from unittest.mock import Mock, patch from ding.data.buffer import DequeBuffer from ding.framework import OnlineRLContext, task @@ -10,6 +11,8 @@ class MockPolicy(Mock): + _device = 'cpu' + # MockPolicy class for train mode def forward(self, train_data, **kwargs): res = { @@ -19,6 +22,8 @@ def forward(self, train_data, **kwargs): class MultiStepMockPolicy(Mock): + _device = 'cpu' + # MockPolicy class for multi-step train mode def forward(self, train_data, **kwargs): res = [ @@ -34,7 +39,7 @@ def forward(self, train_data, **kwargs): def get_mock_train_input(): data = {'obs': torch.rand(2, 2), 'next_obs': torch.rand(2, 2), 'reward': random.random(), 'info': {}} - return data + return ttorch.as_tensor(data) @pytest.mark.unittest @@ -74,7 +79,7 @@ def test_multistep_trainer(): with patch("ding.policy.Policy", MultiStepMockPolicy): policy = MultiStepMockPolicy() for _ in range(30): - multistep_trainer(cfg, policy)(ctx) + multistep_trainer(policy, 10)(ctx) assert ctx.train_iter == 60 assert ctx.train_output[0]["total_loss"] == 0.1 assert ctx.train_output[1]["total_loss"] == 1.0 diff --git a/ding/framework/parallel.py b/ding/framework/parallel.py index 70134f6584..df7b430a8f 100644 --- a/ding/framework/parallel.py +++ b/ding/framework/parallel.py @@ -3,8 +3,8 @@ import random import time import traceback -from mpire.pool import WorkerPool import pickle +from mpire.pool import WorkerPool from ditk import logging import tempfile import socket @@ -27,6 +27,7 @@ def __init__(self) -> None: self._listener = None self.is_active = False self.node_id = None + self.local_id = None self.labels = set() self._event_loop = EventLoop("parallel_{}".format(id(self))) self._retries = 0 # Retries in auto recovery @@ -34,20 +35,30 @@ def __init__(self) -> None: def _run( self, node_id: int, + local_id: int, + n_parallel_workers: int, labels: Optional[Set[str]] = None, auto_recover: bool = False, max_retries: int = float("inf"), mq_type: str = "nng", + startup_interval: int = 1, **kwargs ) -> None: self.node_id = node_id + self.local_id = local_id + self.startup_interval = startup_interval + self.n_parallel_workers = n_parallel_workers self.labels = labels or set() self.auto_recover = auto_recover self.max_retries = max_retries self._mq = MQ_REGISTRY.get(mq_type)(**kwargs) + time.sleep(self.local_id * self.startup_interval) self._listener = Thread(target=self.listen, name="mq_listener", daemon=True) self._listener.start() + self.mq_type = mq_type + self.barrier_runtime = Parallel.get_barrier_runtime()(self.node_id) + @classmethod def runner( cls, @@ -63,7 +74,8 @@ def runner( auto_recover: bool = False, max_retries: int = float("inf"), redis_host: Optional[str] = None, - redis_port: Optional[int] = None + redis_port: Optional[int] = None, + startup_interval: int = 1 ) -> Callable: """ Overview: @@ -85,6 +97,7 @@ def runner( - max_retries (:obj:`int`): Max retries for auto recover. - redis_host (:obj:`str`): Redis server host. - redis_port (:obj:`int`): Redis server port. + - startup_interval (:obj:`int`): Start up interval between each task. Returns: - _runner (:obj:`Callable`): The wrapper function for main. """ @@ -102,7 +115,10 @@ def _runner(main_process: Callable, *args, **kwargs) -> None: - main_process (:obj:`Callable`): The main function, your program start from here. """ runner_params = args_parsers[mq_type](**all_args) - params_group = [[runner_kwargs, (main_process, args, kwargs)] for runner_kwargs in runner_params] + params_group = [] + for i, runner_kwargs in enumerate(runner_params): + runner_kwargs["local_id"] = i + params_group.append([runner_kwargs, (main_process, args, kwargs)]) if n_parallel_workers == 1: cls._subprocess_runner(*params_group[0]) @@ -128,7 +144,6 @@ def _nng_args_parser( ) -> Dict[str, dict]: attach_to = attach_to or [] nodes = cls.get_node_addrs(n_parallel_workers, protocol=protocol, address=address, ports=ports) - logging.info("Bind subprocesses on these addresses: {}".format(nodes)) def cleanup_nodes(): for node in nodes: @@ -156,6 +171,7 @@ def topology_network(i: int) -> List[str]: "node_id": candidate_node_ids[i], "listen_to": nodes[i], "attach_to": topology_network(i), + "n_parallel_workers": n_parallel_workers, } runner_params.append(runner_kwargs) @@ -166,7 +182,7 @@ def _redis_args_parser(cls, n_parallel_workers: int, node_ids: Optional[Union[Li runner_params = [] candidate_node_ids = cls.padding_param(node_ids, n_parallel_workers, 0) for i in range(n_parallel_workers): - runner_kwargs = {**kwargs, "node_id": candidate_node_ids[i]} + runner_kwargs = {**kwargs, "n_parallel_workers": n_parallel_workers, "node_id": candidate_node_ids[i]} runner_params.append(runner_kwargs) return runner_params @@ -179,6 +195,7 @@ def _subprocess_runner(cls, runner_kwargs: dict, main_params: Tuple[Union[List, - runner_params (:obj:`Tuple[Union[List, Dict]]`): Args and kwargs for runner. - main_params (:obj:`Tuple[Union[List, Dict]]`): Args and kwargs for main function. """ + logging.getLogger().setLevel(logging.INFO) main_process, args, kwargs = main_params with Parallel() as router: @@ -320,7 +337,7 @@ def emit(self, event: str, *args, **kwargs) -> None: if self.is_active: payload = {"a": args, "k": kwargs} try: - data = pickle.dumps(payload, protocol=-1) + data = pickle.dumps(payload, protocol=pickle.HIGHEST_PROTOCOL) except AttributeError as e: logging.error("Arguments are not pickable! Event: {}, Args: {}".format(event, args)) raise e @@ -351,12 +368,24 @@ def get_ip(cls): try: # doesn't even have to be reachable s.connect(('10.255.255.255', 1)) - IP = s.getsockname()[0] + ip = s.getsockname()[0] except Exception: - IP = '127.0.0.1' + ip = '127.0.0.1' finally: s.close() - return IP + return ip + + def get_attch_to_len(self) -> int: + """ + Overview: + Get the length of the 'attach_to' list of message queue. + Returns: + int: the length of the self._mq.attach_to. Returns 0 if self._mq is not initialized + """ + if self._mq: + if hasattr(self._mq, 'attach_to'): + return len(self._mq.attach_to) + return 0 def __enter__(self) -> "Parallel": return self @@ -375,3 +404,9 @@ def stop(self): self._listener.join(timeout=1) self._listener = None self._event_loop.stop() + + @classmethod + def get_barrier_runtime(cls): + # We get the BarrierRuntime object in the closure to avoid circular import. + from ding.framework.middleware.barrier import BarrierRuntime + return BarrierRuntime diff --git a/ding/framework/supervisor.py b/ding/framework/supervisor.py index 22f67177f0..7d385c12c6 100644 --- a/ding/framework/supervisor.py +++ b/ding/framework/supervisor.py @@ -1,5 +1,7 @@ from abc import ABC, abstractmethod -import multiprocessing as mp +import functools +import torch.multiprocessing as mp +from multiprocessing.context import BaseContext import threading import queue import platform @@ -12,6 +14,13 @@ from enum import Enum +@functools.lru_cache(maxsize=1) +def get_mp_ctx() -> BaseContext: + context = 'spawn' if platform.system().lower() == 'windows' else 'fork' + mp_ctx = mp.get_context(context) + return mp_ctx + + @dataclass class SendPayload: proc_id: int @@ -29,6 +38,7 @@ class RecvPayload: method: str = None data: Any = None err: Exception = None + extra: Any = None class ReserveMethod(Enum): @@ -41,27 +51,16 @@ class ChildType(Enum): THREAD = "thread" -@dataclass -class SharedObject: - buf: Any - callback: Callable - - class Child(ABC): """ Abstract class of child process/thread. """ - def __init__( - self, proc_id: int, init: Callable, *args, shared_object: Optional[SharedObject] = None, **kwargs - ) -> None: + def __init__(self, proc_id: int, init: Union[Callable, object], **kwargs) -> None: self._proc_id = proc_id self._init = init - self._args = args - self._kwargs = kwargs self._recv_queue = None self._send_queue = None - self._shared_object = shared_object @abstractmethod def start(self, recv_queue: Union[mp.Queue, queue.Queue]): @@ -82,15 +81,17 @@ def send(self, payload: SendPayload): def _target( self, proc_id: int, - init: Callable, - args: List, - kwargs: Dict[str, Any], + init: Union[Callable, object], send_queue: Union[mp.Queue, queue.Queue], recv_queue: Union[mp.Queue, queue.Queue], - shared_object: Optional[SharedObject] = None + shm_buffer: Optional[Any] = None, + shm_callback: Optional[Callable] = None ): send_payload = SendPayload(proc_id=proc_id) - child_ins = init(*args, **kwargs) + if isinstance(init, Callable): + child_ins = init() + else: + child_ins = init while True: try: send_payload: SendPayload = send_queue.get() @@ -103,8 +104,8 @@ def _target( recv_payload = RecvPayload( proc_id=proc_id, req_id=send_payload.req_id, method=send_payload.method, data=data ) - if shared_object: - shared_object.callback(recv_payload, shared_object.buf) + if shm_callback is not None and shm_buffer is not None: + shm_callback(recv_payload, shm_buffer) recv_queue.put(recv_payload) except Exception as e: logging.warning(traceback.format_exc()) @@ -121,27 +122,35 @@ def __del__(self): class ChildProcess(Child): def __init__( - self, proc_id: int, init: Callable, *args, shared_object: Optional[SharedObject] = None, **kwargs + self, + proc_id: int, + init: Union[Callable, object], + shm_buffer: Optional[Any] = None, + shm_callback: Optional[Callable] = None, + mp_ctx: Optional[BaseContext] = None, + **kwargs ) -> None: - super().__init__(proc_id, init, *args, shared_object=shared_object, **kwargs) + super().__init__(proc_id, init, **kwargs) self._proc = None + self._mp_ctx = mp_ctx + self._shm_buffer = shm_buffer + self._shm_callback = shm_callback def start(self, recv_queue: mp.Queue): - self._recv_queue = recv_queue - context = 'spawn' if platform.system().lower() == 'windows' else 'fork' - ctx = mp.get_context(context) - self._send_queue = ctx.Queue() - proc = ctx.Process( - target=self._target, - args=( - self._proc_id, self._init, self._args, self._kwargs, self._send_queue, self._recv_queue, - self._shared_object - ), - name="supervisor_child_{}_{}".format(self._proc_id, time.time()), - daemon=True - ) - proc.start() - self._proc = proc + if self._proc is None: + self._recv_queue = recv_queue + ctx = self._mp_ctx or get_mp_ctx() + self._send_queue = ctx.Queue() + proc = ctx.Process( + target=self._target, + args=( + self._proc_id, self._init, self._send_queue, self._recv_queue, self._shm_buffer, self._shm_callback + ), + name="supervisor_child_{}_{}".format(self._proc_id, time.time()), + daemon=True + ) + proc.start() + self._proc = proc def shutdown(self, timeout: Optional[float] = None): if self._proc: @@ -156,28 +165,30 @@ def shutdown(self, timeout: Optional[float] = None): self._send_queue = None def send(self, payload: SendPayload): + if self._send_queue is None: + logging.warning("Child worker has been terminated or not started.") + return self._send_queue.put(payload) class ChildThread(Child): - def __init__( - self, proc_id: int, init: Callable, *args, shared_object: Optional[SharedObject] = None, **kwargs - ) -> None: - super().__init__(proc_id, init, *args, shared_object=shared_object, **kwargs) + def __init__(self, proc_id: int, init: Union[Callable, object], *args, **kwargs) -> None: + super().__init__(proc_id, init, *args, **kwargs) self._thread = None def start(self, recv_queue: queue.Queue): - self._recv_queue = recv_queue - self._send_queue = queue.Queue() - thread = threading.Thread( - target=self._target, - args=(self._proc_id, self._init, self._args, self._kwargs, self._send_queue, self._recv_queue), - name="supervisor_child_{}_{}".format(self._proc_id, time.time()), - daemon=True - ) - thread.start() - self._thread = thread + if self._thread is None: + self._recv_queue = recv_queue + self._send_queue = queue.Queue() + thread = threading.Thread( + target=self._target, + args=(self._proc_id, self._init, self._send_queue, self._recv_queue), + name="supervisor_child_{}_{}".format(self._proc_id, time.time()), + daemon=True + ) + thread.start() + self._thread = thread def shutdown(self, timeout: Optional[float] = None): if self._thread: @@ -187,6 +198,9 @@ def shutdown(self, timeout: Optional[float] = None): self._send_queue = None def send(self, payload: SendPayload): + if self._send_queue is None: + logging.warning("Child worker has been terminated or not started.") + return self._send_queue.put(payload) @@ -194,26 +208,32 @@ class Supervisor: TYPE_MAPPING = {ChildType.PROCESS: ChildProcess, ChildType.THREAD: ChildThread} - QUEUE_MAPPING = { - ChildType.PROCESS: mp.get_context('spawn' if platform.system().lower() == 'windows' else 'fork').Queue, - ChildType.THREAD: queue.Queue - } - - def __init__(self, type_: ChildType) -> None: + def __init__(self, type_: ChildType, mp_ctx: Optional[BaseContext] = None) -> None: self._children: List[Child] = [] self._type = type_ self._child_class = self.TYPE_MAPPING[self._type] self._running = False self.__queue = None + self._mp_ctx = mp_ctx or get_mp_ctx() - def register(self, init: Callable, *args, shared_object: Optional[SharedObject] = None, **kwargs) -> None: + def register( + self, + init: Union[Callable, object], + shm_buffer: Optional[Any] = None, + shm_callback: Optional[Callable] = None + ) -> None: proc_id = len(self._children) - self._children.append(self._child_class(proc_id, init, *args, shared_object=shared_object, **kwargs)) + self._children.append( + self._child_class(proc_id, init, shm_buffer=shm_buffer, shm_callback=shm_callback, mp_ctx=self._mp_ctx) + ) @property def _recv_queue(self) -> Union[queue.Queue, mp.Queue]: if not self.__queue: - self.__queue = self.QUEUE_MAPPING[self._type]() + if self._type is ChildType.PROCESS: + self.__queue = self._mp_ctx.Queue() + elif self._type is ChildType.THREAD: + self.__queue = queue.Queue() return self.__queue @_recv_queue.setter @@ -233,6 +253,9 @@ def send(self, payload: SendPayload) -> None: Arguments: - payload (:obj:`SendPayload`): Send payload. """ + if not self._running: + logging.warning("Please call start_link before sending any payload to child process.") + return self._children[payload.proc_id].send(payload) def recv(self, ignore_err: bool = False, timeout: float = None) -> RecvPayload: diff --git a/ding/framework/task.py b/ding/framework/task.py index 53e95716b0..e5f8e7d9f7 100644 --- a/ding/framework/task.py +++ b/ding/framework/task.py @@ -7,8 +7,11 @@ import concurrent.futures import fnmatch import math +import enum from types import GeneratorType from typing import Any, Awaitable, Callable, Dict, Generator, Iterable, List, Optional, Set, Union +import inspect + from ding.framework.context import Context from ding.framework.parallel import Parallel from ding.framework.event_loop import EventLoop @@ -50,11 +53,29 @@ def runtime_handler(task: "Task", *args, async_mode: Optional[bool] = None, **kw return runtime_handler +class Role(str, enum.Enum): + LEARNER = "learner" + COLLECTOR = "collector" + EVALUATOR = "evaluator" + FETCHER = 'fetcher' + + +class VoidMiddleware: + + def __call__(self, _): + return + + class Task: """ - Tash will manage the execution order of the entire pipeline, register new middleware, + Task will manage the execution order of the entire pipeline, register new middleware, and generate new context objects. """ + role = Role + + def __init__(self) -> None: + self.router = Parallel() + self._finish = False def start( self, @@ -71,6 +92,7 @@ def start( self._wrappers = [] self.ctx = ctx or Context() self._backward_stack = OrderedDict() + self._roles = set() # Bind event loop functions self._event_loop = EventLoop("task_{}".format(id(self))) @@ -85,7 +107,6 @@ def start( self.labels = labels or set() # Parallel segment - self.router = Parallel() if async_mode or self.router.is_active: self._activate_async() @@ -99,6 +120,21 @@ def sync_finish(value): self.init_labels() return self + def add_role(self, role: Role): + self._roles.add(role) + + def has_role(self, role: Role) -> bool: + if len(self._roles) == 0: + return True + return role in self._roles + + @property + def roles(self) -> Set[Role]: + return self._roles + + def void(self): + return VoidMiddleware() + def init_labels(self): if self.async_mode: self.labels.add("async") @@ -120,6 +156,9 @@ def use(self, fn: Callable, lock: Union[bool, Lock] = False) -> 'Task': Returns: - task (:obj:`Task`): The task. """ + assert isinstance(fn, Callable), "Middleware function should be a callable object, current fn {}".format(fn) + if isinstance(fn, VoidMiddleware): # Skip void function + return self for wrapper in self._wrappers: fn = wrapper(fn) self._middleware.append(self.wrap(fn, lock=lock)) @@ -192,7 +231,6 @@ def wrap(self, fn: Callable, lock: Union[bool, Lock] = False) -> Callable: if lock is True: lock = self._thread_lock - @wraps(fn) def forward(ctx: Context): if lock: with lock: @@ -212,6 +250,11 @@ def backward(): return backward + if hasattr(fn, "__name__"): + forward = wraps(fn)(forward) + else: + forward = wraps(fn.__class__)(forward) + return forward @enable_async @@ -258,6 +301,10 @@ def backward(self, backward_stack: Optional[Dict[str, Generator]] = None) -> Non except StopIteration: continue + @property + def running(self): + return self._running + def serial(self, *fns: List[Callable]) -> Callable: """ Overview: @@ -330,6 +377,8 @@ def stop(self) -> None: Overview: Stop and cleanup every thing in the runtime of task. """ + if self.router.is_active: + self.emit("finish", True) if self._thread_pool: self._thread_pool.shutdown() self._event_loop.stop() @@ -472,8 +521,6 @@ def finish(self): @finish.setter def finish(self, value: bool): self._finish = value - if self.router.is_active and value is True: - self.emit("finish", value) def _wrap_event_name(self, event: str) -> str: """ @@ -490,5 +537,17 @@ def _activate_async(self): if not self._async_loop: self._async_loop = asyncio.new_event_loop() + def get_attch_to_len(self) -> int: + """ + Overview: + Get the length of the 'attach_to' list in Parallel._mq. + Returns: + int: the length of the Parallel._mq. + """ + if self.router.is_active: + return self.router.get_attch_to_len() + else: + raise RuntimeError("The router is inactive, failed to be obtained the length of 'attch_to' list.") + task = Task() diff --git a/ding/framework/tests/context_fake_data.py b/ding/framework/tests/context_fake_data.py index 0f7f55db9f..ee65048e6c 100644 --- a/ding/framework/tests/context_fake_data.py +++ b/ding/framework/tests/context_fake_data.py @@ -8,22 +8,25 @@ n_sample = 8 action_dim = 1 obs_dim = 4 +logit_dim = 2 n_episodes = 2 n_episode_length = 16 +update_per_collect = 4 +collector_env_num = 8 # the range here is meaningless and just for test def fake_train_data(): train_data = ttorch.as_tensor( { - 'action': torch.randint(0, 1, size=(action_dim, )), + 'action': torch.randint(0, 2, size=(action_dim, )), 'collect_train_iter': torch.randint(0, 100, size=(1, )), 'done': torch.tensor(False), 'env_data_id': torch.tensor([2]), 'next_obs': torch.randn(obs_dim), 'obs': torch.randn(obs_dim), - 'reward': torch.randint(0, 1, size=(1, )), + 'reward': torch.randint(0, 2, size=(1, )), } ) return train_data @@ -35,6 +38,19 @@ def fake_online_rl_context(): env_episode=random.randint(0, 100), train_iter=random.randint(0, 100), train_data=[fake_train_data() for _ in range(batch_size)], + train_output=[{ + 'cur_lr': 0.001, + 'total_loss': random.uniform(0, 2) + } for _ in range(update_per_collect)], + obs=torch.randn(collector_env_num, obs_dim), + action=[np.random.randint(low=0, high=1, size=(action_dim), dtype=np.int64) for _ in range(collector_env_num)], + inference_output={ + env_id: { + 'logit': torch.randn(logit_dim), + 'action': torch.randint(0, 2, size=(action_dim, )) + } + for env_id in range(collector_env_num) + }, collect_kwargs={'eps': random.uniform(0, 1)}, trajectories=[fake_train_data() for _ in range(n_sample)], episodes=[[fake_train_data() for _ in range(n_episode_length)] for _ in range(n_episodes)], @@ -50,6 +66,10 @@ def fake_offline_rl_context(): train_epoch=random.randint(0, 100), train_iter=random.randint(0, 100), train_data=[fake_train_data() for _ in range(batch_size)], + train_output=[{ + 'cur_lr': 0.001, + 'total_loss': random.uniform(0, 2) + } for _ in range(update_per_collect)], eval_value=random.uniform(-1.0, 1.0), last_eval_iter=random.randint(0, 100), ) diff --git a/ding/framework/tests/test_event_loop.py b/ding/framework/tests/test_event_loop.py index 1dddae6164..2f3545f3f5 100644 --- a/ding/framework/tests/test_event_loop.py +++ b/ding/framework/tests/test_event_loop.py @@ -31,6 +31,8 @@ def callback(n, lock): assert counter == 10 # Test once + counter = 0 + loop.once("count", callback) loop.once("count", callback) loop.emit("count", 10, lock) sleep(0.1) diff --git a/ding/framework/tests/test_parallel.py b/ding/framework/tests/test_parallel.py index b042cb3a57..7bdf6ea343 100644 --- a/ding/framework/tests/test_parallel.py +++ b/ding/framework/tests/test_parallel.py @@ -1,9 +1,7 @@ from collections import defaultdict import pytest import time -import os from ding.framework import Parallel -from ding.utils.design_helper import SingletonMetaclass def parallel_main(): @@ -26,10 +24,10 @@ def test_callback(key): time.sleep(0.7) -@pytest.mark.unittest +@pytest.mark.tmp def test_parallel_run(): - Parallel.runner(n_parallel_workers=2)(parallel_main) - Parallel.runner(n_parallel_workers=2, protocol="tcp")(parallel_main) + Parallel.runner(n_parallel_workers=2, startup_interval=0.1)(parallel_main) + Parallel.runner(n_parallel_workers=2, protocol="tcp", startup_interval=0.1)(parallel_main) def uncaught_exception_main(): @@ -41,11 +39,11 @@ def uncaught_exception_main(): time.sleep(0.2) -@pytest.mark.unittest +@pytest.mark.tmp def test_uncaught_exception(): # Make one process crash, then the parent process will also crash and output the stack of the wrong process. with pytest.raises(Exception) as exc_info: - Parallel.runner(n_parallel_workers=2, topology="mesh")(uncaught_exception_main) + Parallel.runner(n_parallel_workers=2, topology="mesh", startup_interval=0.1)(uncaught_exception_main) e = exc_info._excinfo[1] assert "uncaught exception" in str(e) @@ -54,6 +52,7 @@ def disconnected_main(): router = Parallel() if router.node_id == 0: + time.sleep(0.1) # Receive two messages then exit greets = [] router.on("greeting", lambda: greets.append(".")) @@ -71,11 +70,11 @@ def disconnected_main(): assert i == 9 -@pytest.mark.unittest +@pytest.mark.tmp def test_disconnected(): # Make one process exit normally and the rest will still run, even if the network request # is not received by other processes. - Parallel.runner(n_parallel_workers=2, topology="mesh")(disconnected_main) + Parallel.runner(n_parallel_workers=2, topology="mesh", startup_interval=0.1)(disconnected_main) class AutoRecover: @@ -142,12 +141,16 @@ def main(cls): raise Exception("Invalid node id") -@pytest.mark.unittest +@pytest.mark.tmp def test_auto_recover(): # With max_retries=1 - Parallel.runner(n_parallel_workers=3, topology="mesh", auto_recover=True, max_retries=1)(AutoRecover.main) + Parallel.runner( + n_parallel_workers=3, topology="mesh", auto_recover=True, max_retries=1, startup_interval=0.1 + )(AutoRecover.main) # With max_retries=0 with pytest.raises(Exception) as exc_info: - Parallel.runner(n_parallel_workers=3, topology="mesh", auto_recover=True, max_retries=0)(AutoRecover.main) + Parallel.runner( + n_parallel_workers=3, topology="mesh", auto_recover=True, max_retries=0, startup_interval=0.1 + )(AutoRecover.main) e = exc_info._excinfo[1] assert "P1 Error" in str(e) diff --git a/ding/framework/tests/test_supervisor.py b/ding/framework/tests/test_supervisor.py index 57a0f0d49f..d6f4c646fa 100644 --- a/ding/framework/tests/test_supervisor.py +++ b/ding/framework/tests/test_supervisor.py @@ -1,9 +1,9 @@ import multiprocessing as mp import ctypes -from time import sleep +from time import sleep, time from typing import Any, Dict, List import pytest -from ding.framework.supervisor import RecvPayload, SendPayload, Supervisor, ChildType, SharedObject +from ding.framework.supervisor import RecvPayload, SendPayload, Supervisor, ChildType class MockEnv(): @@ -25,13 +25,16 @@ def block(self): def block_reset(self): sleep(10) + def sleep1(self): + sleep(1) -@pytest.mark.unittest + +@pytest.mark.tmp @pytest.mark.parametrize("type_", [ChildType.PROCESS, ChildType.THREAD]) def test_supervisor(type_): sv = Supervisor(type_=type_) for _ in range(3): - sv.register(MockEnv, "AnyArgs") + sv.register(lambda: MockEnv("AnyArgs")) sv.start_link() for env_id in range(len(sv._children)): @@ -71,6 +74,25 @@ def test_supervisor(type_): sv.shutdown() +@pytest.mark.tmp +def test_supervisor_spawn(): + sv = Supervisor(type_=ChildType.PROCESS, mp_ctx=mp.get_context("spawn")) + for _ in range(3): + sv.register(MockEnv("AnyArgs")) + sv.start_link() + + for env_id in range(len(sv._children)): + sv.send(SendPayload(proc_id=env_id, method="step", args=["any action"])) + + recv_states: List[RecvPayload] = [] + for _ in range(3): + recv_states.append(sv.recv()) + + assert sum([payload.proc_id for payload in recv_states]) == 3 + assert all([payload.data == 1 for payload in recv_states]) + sv.shutdown() + + class MockCrashEnv(MockEnv): def step(self, _): @@ -81,13 +103,13 @@ def step(self, _): return self._counter -# @pytest.mark.unittest +@pytest.mark.tmp @pytest.mark.parametrize("type_", [ChildType.PROCESS, ChildType.THREAD]) def test_crash_supervisor(type_): sv = Supervisor(type_=type_) for _ in range(2): - sv.register(MockEnv, "AnyArgs") - sv.register(MockCrashEnv, "AnyArgs") + sv.register(lambda: MockEnv("AnyArgs")) + sv.register(lambda: MockCrashEnv("AnyArgs")) sv.start_link() # Send 6 messages, will cause the third subprocess crash @@ -121,12 +143,12 @@ def test_crash_supervisor(type_): sv.shutdown() -@pytest.mark.unittest +@pytest.mark.tmp @pytest.mark.parametrize("type_", [ChildType.PROCESS, ChildType.THREAD]) def test_recv_all(type_): sv = Supervisor(type_=type_) for _ in range(3): - sv.register(MockEnv, "AnyArgs") + sv.register(lambda: MockEnv("AnyArgs")) sv.start_link() # Test recv_all @@ -162,7 +184,7 @@ def recv_callback(recv_payload: RecvPayload, remain_payloads: Dict[str, SendPayl def test_timeout(type_): sv = Supervisor(type_=type_) for _ in range(3): - sv.register(MockEnv, "AnyArgs") + sv.register(lambda: MockEnv("AnyArgs")) sv.start_link() send_payloads = [] @@ -202,7 +224,7 @@ def test_timeout(type_): def test_timeout_with_callback(type_): sv = Supervisor(type_=type_) for _ in range(3): - sv.register(MockEnv, "AnyArgs") + sv.register(lambda: MockEnv("AnyArgs")) sv.start_link() send_payloads = [] @@ -239,25 +261,50 @@ def recv_callback(recv_payload: RecvPayload, remain_payloads: Dict[str, SendPayl sv.shutdown(timeout=1) -@pytest.mark.unittest +@pytest.mark.tmp # gitlab ci and local test pass, github always fail def test_shared_memory(): sv = Supervisor(type_=ChildType.PROCESS) def shm_callback(payload: RecvPayload, shm: Any): - shm[payload.proc_id] = payload.data + shm[payload.proc_id] = payload.req_id payload.data = 0 shm = mp.Array(ctypes.c_uint8, 3) for i in range(3): - sv.register(MockEnv, "AnyArgs", shared_object=SharedObject(buf=shm, callback=shm_callback)) + sv.register(lambda: MockEnv("AnyArgs"), shm_buffer=shm, shm_callback=shm_callback) sv.start_link() + # Send init request for env_id in range(len(sv._children)): - sv.send(SendPayload(proc_id=env_id, method="step", args=["any action"])) + sv.send(SendPayload(proc_id=env_id, req_id=env_id, method="sleep1", args=[])) - for i in range(3): + start = time() + for i in range(6): payload = sv.recv() assert payload.data == 0 - assert shm[payload.proc_id] == 1 + assert shm[payload.proc_id] == payload.req_id + sv.send(SendPayload(proc_id=payload.proc_id, req_id=i, method="sleep1", args=[])) + + # Non blocking + assert time() - start < 3 sv.shutdown() + + +@pytest.mark.benchmark +@pytest.mark.parametrize("type_", [ChildType.PROCESS, ChildType.THREAD]) +def test_supervisor_benchmark(type_): + sv = Supervisor(type_=type_) + for _ in range(3): + sv.register(lambda: MockEnv("AnyArgs")) + sv.start_link() + + for env_id in range(len(sv._children)): + sv.send(SendPayload(proc_id=env_id, method="step", args=[""])) + + start = time() + for _ in range(1000): + payload = sv.recv() + sv.send(SendPayload(proc_id=payload.proc_id, method="step", args=[""])) + + assert time() - start < 1 diff --git a/ding/framework/tests/test_task.py b/ding/framework/tests/test_task.py index 7ad13977bc..67f3dc34c7 100644 --- a/ding/framework/tests/test_task.py +++ b/ding/framework/tests/test_task.py @@ -1,6 +1,7 @@ +import multiprocessing as mp import pytest from threading import Lock -from time import sleep +from time import sleep, time import random import dataclasses from ding.framework import task, Context, Parallel @@ -123,12 +124,12 @@ def _counter(ctx): assert sync_count > 0 -@pytest.mark.unittest +@pytest.mark.tmp def test_parallel_pipeline(): - Parallel.runner(n_parallel_workers=2)(parallel_main) + Parallel.runner(n_parallel_workers=2, startup_interval=0.1)(parallel_main) -@pytest.mark.unittest +@pytest.mark.tmp def test_emit(): with task.start(): greets = [] @@ -160,12 +161,12 @@ def emit_remote_main(): assert len(greets) == 0 -@pytest.mark.unittest +@pytest.mark.tmp def test_emit_remote(): - Parallel.runner(n_parallel_workers=2)(emit_remote_main) + Parallel.runner(n_parallel_workers=2, startup_interval=0.1)(emit_remote_main) -@pytest.mark.unittest +@pytest.mark.tmp def test_wait_for(): # Wait for will only work in async or parallel mode with task.start(async_mode=True, n_async_workers=2): @@ -197,7 +198,7 @@ def step1(_): task.run(max_step=1) -@pytest.mark.unittest +@pytest.mark.tmp def test_async_exception(): with task.start(async_mode=True, n_async_workers=2): @@ -226,12 +227,12 @@ def early_stop_main(): assert task.ctx.total_step < 7 -@pytest.mark.unittest +@pytest.mark.tmp def test_early_stop(): - Parallel.runner(n_parallel_workers=2)(early_stop_main) + Parallel.runner(n_parallel_workers=2, startup_interval=0.1)(early_stop_main) -@pytest.mark.unittest +@pytest.mark.tmp def test_parallel_in_sequencial(): result = [] @@ -249,7 +250,7 @@ def slow(_): assert result == ["begin", "fast", "slow"] -@pytest.mark.unittest +@pytest.mark.tmp def test_serial_in_parallel(): result = [] @@ -333,3 +334,49 @@ def slowest(ctx): task.use(fast, lock=lock) task.run(1) assert task.ctx.result == "slowest" + + +def broadcast_finish_main(): + with task.start(): + + def tick(ctx: Context): + if task.router.node_id == 1 and ctx.total_step == 1: + task.finish = True + sleep(1) + + task.use(tick) + task.run(20) + + +def broadcast_main_target(): + Parallel.runner( + n_parallel_workers=1, protocol="tcp", address="127.0.0.1", topology="mesh", ports=50555, startup_interval=0.1 + )(broadcast_finish_main) + + +def broadcast_secondary_target(): + "Start two standalone processes and connect to the main process." + Parallel.runner( + n_parallel_workers=2, + protocol="tcp", + address="127.0.0.1", + topology="alone", + ports=50556, + attach_to=["tcp://127.0.0.1:50555"], + node_ids=[1, 2], + startup_interval=0.1 + )(broadcast_finish_main) + + +@pytest.mark.tmp # gitlab ci and local test pass, github always fail +@pytest.mark.timeout(10) +def test_broadcast_finish(): + start = time() + ctx = mp.get_context("spawn") + main_process = ctx.Process(target=broadcast_main_target) + secondary_process = ctx.Process(target=broadcast_secondary_target) + main_process.start() + secondary_process.start() + main_process.join() + secondary_process.join() + assert (time() - start) < 10 diff --git a/ding/framework/wrapper/step_timer.py b/ding/framework/wrapper/step_timer.py index f7d123bc62..dfabdd1476 100644 --- a/ding/framework/wrapper/step_timer.py +++ b/ding/framework/wrapper/step_timer.py @@ -5,11 +5,20 @@ import numpy as np import time from ditk import logging +from ding.framework import task class StepTimer: def __init__(self, print_per_step: int = 1, smooth_window: int = 10) -> None: + """ + Overview: + Print time cost of each step (execute one middleware). + Arguments: + - print_per_step (:obj:`int`): Print each N step. + - smooth_window (:obj:`int`): The window size to smooth the mean. + """ + self.print_per_step = print_per_step self.records = defaultdict(lambda: deque(maxlen=print_per_step * smooth_window)) @@ -36,11 +45,12 @@ def executor(ctx): time_cost += time.time() - start_time else: time_cost = time.time() - start_time - self.records[step_name].append(time_cost * 1000) + self.records[step_name].append(time_cost) if ctx.total_step % self.print_per_step == 0: logging.info( - "[Step Timer] {}: Cost: {:.2f}ms, Mean: {:.2f}ms".format( - step_name, time_cost * 1000, np.mean(self.records[step_name]) + "[Step Timer][Node:{:>2}] {}: Cost: {:.2f}ms, Mean: {:.2f}ms".format( + task.router.node_id or 0, step_name, time_cost * 1000, + np.mean(self.records[step_name]) * 1000 ) ) diff --git a/ding/hpc_rl/tests/test_vtrace.py b/ding/hpc_rl/tests/test_vtrace.py index bbfeba0c1c..c26ab4f407 100644 --- a/ding/hpc_rl/tests/test_vtrace.py +++ b/ding/hpc_rl/tests/test_vtrace.py @@ -1,7 +1,7 @@ import time import torch import torch.nn.functional as F -from hpc_rll.origin.vtrace import vtrace_error, vtrace_data +from hpc_rll.origin.vtrace import vtrace_error_discrete_action, vtrace_data from hpc_rll.rl_utils.vtrace import VTrace from testbase import mean_relative_error, times @@ -48,7 +48,7 @@ def vtrace_val(): ori_target_output.requires_grad_(True) ori_value.requires_grad_(True) - ori_loss = vtrace_error( + ori_loss = vtrace_error_discrete_action( vtrace_data(ori_target_output, ori_behaviour_output, ori_action, ori_value, ori_reward, None) ) ori_loss = sum(ori_loss) @@ -114,7 +114,7 @@ def vtrace_perf(): ori_value.requires_grad_(True) for i in range(times): t = time.time() - ori_loss = vtrace_error( + ori_loss = vtrace_error_discrete_action( vtrace_data(ori_target_output, ori_behaviour_output, ori_action, ori_value, ori_reward, None) ) ori_loss = sum(ori_loss) diff --git a/ding/hpc_rl/wrapper.py b/ding/hpc_rl/wrapper.py index 07bd958d06..2e4ac0bf95 100644 --- a/ding/hpc_rl/wrapper.py +++ b/ding/hpc_rl/wrapper.py @@ -69,7 +69,7 @@ def register_runtime_fn(fn_name, runtime_name, shape): 'ScatterConnection': ['hpc_rll.torch_utils.network.scatter_connection', 'ScatterConnection'], 'td_lambda_error': ['hpc_rll.rl_utils.td', 'TDLambda'], 'upgo_loss': ['hpc_rll.rl_utils.upgo', 'UPGO'], - 'vtrace_error': ['hpc_rll.rl_utils.vtrace', 'VTrace'], + 'vtrace_error_discrete_action': ['hpc_rll.rl_utils.vtrace', 'VTrace'], } fn_str = fn_name_mapping[fn_name] cls = getattr(importlib.import_module(fn_str[0]), fn_str[1]) diff --git a/ding/league/base_league.py b/ding/league/base_league.py index 8a5ed333e5..d59e7f084b 100644 --- a/ding/league/base_league.py +++ b/ding/league/base_league.py @@ -86,7 +86,7 @@ def __init__(self, cfg: EasyDict) -> None: self.payoff = create_payoff(self.cfg.payoff) metric_cfg = self.cfg.metric self.metric_env = LeagueMetricEnv(metric_cfg.mu, metric_cfg.sigma, metric_cfg.tau, metric_cfg.draw_probability) - self._active_players_lock = LockContext(type_=LockContextType.THREAD_LOCK) + self._active_players_lock = LockContext(lock_type=LockContextType.THREAD_LOCK) self._init_players() def _init_players(self) -> None: diff --git a/ding/league/metric.py b/ding/league/metric.py index 41c587cd02..be675898f4 100644 --- a/ding/league/metric.py +++ b/ding/league/metric.py @@ -5,6 +5,17 @@ class EloCalculator(object): + """ + Overview: + A class that calculates Elo ratings for players based on game results. + + Attributes: + - score (:obj:`dict`): A dictionary that maps game results to scores. + + Interfaces: + ``__init__``, ``get_new_rating``, ``get_new_rating_array``. + """ + score = { 1: 1.0, # win 0: 0.5, # draw @@ -18,6 +29,20 @@ def get_new_rating(cls, result: int, k_factor: int = 32, beta: int = 200) -> Tuple[int, int]: + """ + Overview: + Calculates the new ratings for two players based on their current ratings and game result. + + Arguments: + - rating_a (:obj:`int`): The current rating of player A. + - rating_b (:obj:`int`): The current rating of player B. + - result (:obj:`int`): The result of the game: 1 for player A win, 0 for draw, -1 for player B win. + - k_factor (:obj:`int`): The K-factor used in the Elo rating system. Defaults to 32. + - beta (:obj:`int`): The beta value used in the Elo rating system. Defaults to 200. + + Returns: + -ret (:obj:`Tuple[int, int]`): The new ratings for player A and player B, respectively. + """ assert result in [1, 0, -1] expect_a = 1. / (1. + math.pow(10, (rating_b - rating_a) / (2. * beta))) expect_b = 1. / (1. + math.pow(10, (rating_a - rating_b) / (2. * beta))) @@ -35,10 +60,25 @@ def get_new_rating_array( beta: int = 200 ) -> np.ndarray: """ + Overview: + Calculates the new ratings for multiple players based on their current ratings, game results, \ + and game counts. + + Arguments: + - rating (obj:`np.ndarray`): An array of current ratings for each player. + - result (obj:`np.ndarray`): An array of game results, where 1 represents a win, 0 represents a draw, \ + and -1 represents a loss. + - game_count (obj:`np.ndarray`): An array of game counts for each player. + - k_factor (obj:`int`): The K-factor used in the Elo rating system. Defaults to 32. + - beta (obj:`int`): The beta value used in the Elo rating system. Defaults to 200. + + Returns: + -ret(obj:`np.ndarray`): An array of new ratings for each player. + Shapes: - rating: :math:`(N, )`, N is the number of player - result: :math:`(N, N)` - game_count: :math:`(N, N)` + - rating (obj:`np.ndarray`): :math:`(N, )`, N is the number of player + - result (obj:`np.ndarray`): :math:`(N, N)` + - game_count (obj:`np.ndarray`): :math:`(N, N)` """ rating_diff = np.expand_dims(rating, 0) - np.expand_dims(rating, 1) expect = 1. / (1. + np.power(10, rating_diff / (2. * beta))) * game_count @@ -48,6 +88,13 @@ def get_new_rating_array( class PlayerRating(Rating): + """ + Overview: + Represents the rating of a player. + + Interfaces: + ``__init__``, ``__repr__``. + """ def __init__(self, mu: float = None, sigma: float = None, elo_init: int = None) -> None: super(PlayerRating, self).__init__(mu, sigma) @@ -62,7 +109,11 @@ def __repr__(self) -> str: class LeagueMetricEnv(TrueSkill): """ Overview: - TrueSkill rating system among game players, for more details pleas refer to ``https://trueskill.org/`` + A class that represents a TrueSkill rating system for game players. Inherits from the TrueSkill class. \ + For more details, please refer to https://trueskill.org/. + + Interfaces: + ``__init__``, ``create_rating``, ``rate_1vs1``, ``rate_1vsC``. """ def __init__(self, *args, elo_init: int = 1200, **kwargs) -> None: @@ -70,6 +121,21 @@ def __init__(self, *args, elo_init: int = 1200, **kwargs) -> None: self.elo_init = elo_init def create_rating(self, mu: float = None, sigma: float = None, elo_init: int = None) -> PlayerRating: + """ + Overview: + Creates a new player rating object with the specified mean, standard deviation, and Elo rating. + + Arguments: + - mu (:obj:`float`): The mean value of the player's skill rating. If not provided, the default \ + TrueSkill mean is used. + - sigma (:obj:`float`): The standard deviation of the player's skill rating. If not provided, \ + the default TrueSkill sigma is used. + - elo_init (:obj:int`): The initial Elo rating value for the player. If not provided, the default \ + elo_init value of the LeagueMetricEnv class is used. + + Returns: + - PlayerRating: A player rating object with the specified mean, standard deviation, and Elo rating. + """ if mu is None: mu = self.mu if sigma is None: @@ -91,11 +157,23 @@ def _rate_1vs1(t1, t2, **kwargs): t2 = PlayerRating(t2.mu, t2.sigma, t2_elo) return t1, t2 - def rate_1vs1(self, - team1: PlayerRating, - team2: PlayerRating, - result: List[str] = None, - **kwargs) -> Tuple[PlayerRating, PlayerRating]: + def rate_1vs1(self, team1: PlayerRating, team2: PlayerRating, result: List[str] = None, **kwargs) \ + -> Tuple[PlayerRating, PlayerRating]: + """ + Overview: + Rates two teams of players against each other in a 1 vs 1 match and returns the updated ratings \ + for both teams. + + Arguments: + - team1 (:obj:`PlayerRating`): The rating object representing the first team of players. + - team2 (:obj:`PlayerRating`): The rating object representing the second team of players. + - result (:obj:`List[str]`): The result of the match. Can be 'wins', 'draws', or 'losses'. If \ + not provided, the default behavior is to rate the match as a win for team1. + + Returns: + - ret (:obj:`Tuple[PlayerRating, PlayerRating]`): A tuple containing the updated ratings for team1 \ + and team2. + """ if result is None: return self._rate_1vs1(team1, team2, **kwargs) else: @@ -111,6 +189,19 @@ def rate_1vs1(self, return team1, team2 def rate_1vsC(self, team1: PlayerRating, team2: PlayerRating, result: List[str]) -> PlayerRating: + """ + Overview: + Rates a team of players against a single player in a 1 vs C match and returns the updated rating \ + for the team. + + Arguments: + - team1 (:obj:`PlayerRating`): The rating object representing the team of players. + - team2 (:obj:`PlayerRating`): The rating object representing the single player. + - result (:obj:`List[str]`): The result of the match. Can be 'wins', 'draws', or 'losses'. + + Returns: + - PlayerRating: The updated rating for the team of players. + """ for r in result: if r == 'wins': team1, _ = self._rate_1vs1(team1, team2) diff --git a/ding/league/shared_payoff.py b/ding/league/shared_payoff.py index 7576d441c0..82381468ad 100644 --- a/ding/league/shared_payoff.py +++ b/ding/league/shared_payoff.py @@ -75,7 +75,7 @@ def __init__(self, cfg: EasyDict): # ``_min_win_rate_games``` is used in ``self._win_rate`` method for calculating win rate between two players. self._min_win_rate_games = cfg.get('min_win_rate_games', 8) # Thread lock. - self._lock = LockContext(type_=LockContextType.THREAD_LOCK) + self._lock = LockContext(lock_type=LockContextType.THREAD_LOCK) def __repr__(self) -> str: headers = ["Home Player", "Away Player", "Wins", "Draws", "Losses", "Naive Win Rate"] diff --git a/ding/model/common/__init__.py b/ding/model/common/__init__.py old mode 100644 new mode 100755 index 8a1a0eb1b9..b63aedfcd9 --- a/ding/model/common/__init__.py +++ b/ding/model/common/__init__.py @@ -1,4 +1,5 @@ -from .head import DiscreteHead, DuelingHead, DistributionHead, RainbowHead, QRDQNHead, \ - QuantileHead, FQFHead, RegressionHead, ReparameterizationHead, MultiHead, head_cls_map -from .encoder import ConvEncoder, FCEncoder, IMPALAConvEncoder +from .head import DiscreteHead, DuelingHead, DistributionHead, RainbowHead, QRDQNHead, StochasticDuelingHead, \ + QuantileHead, FQFHead, RegressionHead, ReparameterizationHead, MultiHead, BranchingHead, head_cls_map, \ + independent_normal_dist, AttentionPolicyHead, PopArtVHead, EnsembleHead +from .encoder import ConvEncoder, FCEncoder, IMPALAConvEncoder, GaussianFourierProjectionTimeEncoder from .utils import create_model diff --git a/ding/model/common/encoder.py b/ding/model/common/encoder.py index 35cde749d4..8936f6af4d 100644 --- a/ding/model/common/encoder.py +++ b/ding/model/common/encoder.py @@ -1,20 +1,22 @@ -from typing import Optional +from typing import Optional, Dict, Union, List from functools import reduce import operator import math +import numpy as np import torch import torch.nn as nn from torch.nn import functional as F from ding.torch_utils import ResFCBlock, ResBlock, Flatten, normed_linear, normed_conv2d +from ding.torch_utils.network.dreamer import Conv2dSame, DreamerLayerNorm from ding.utils import SequenceType def prod(iterable): """ - Product of all elements.(To be deprecated soon.) - This function denifition is for supporting python version that under 3.8. - In python3.8 and larger, 'math.prod()' is recommended. + Overview: + Product of all elements.(To be deprecated soon.) This function denifition is for supporting python version \ + that under 3.8. In Python3.8 and larger, 'math.prod()' is recommended. """ return reduce(operator.mul, iterable, 1) @@ -22,9 +24,9 @@ def prod(iterable): class ConvEncoder(nn.Module): """ Overview: - The ``Convolution Encoder`` used to encode raw 2-dim observations. + The Convolution Encoder is used to encode 2-dim image observations. Interfaces: - ``__init__``, ``_get_flatten_size``, ``forward``. + ``__init__``, ``forward``. """ def __init__( @@ -35,11 +37,12 @@ def __init__( kernel_size: SequenceType = [8, 4, 3], stride: SequenceType = [4, 2, 1], padding: Optional[SequenceType] = None, + layer_norm: Optional[bool] = False, norm_type: Optional[str] = None ) -> None: """ Overview: - Init the ``Convolution Encoder`` according to the provided arguments. + Initialize the ``Convolution Encoder`` according to the provided arguments. Arguments: - obs_shape (:obj:`SequenceType`): Sequence of ``in_channel``, plus one or more ``input size``. - hidden_size_list (:obj:`SequenceType`): Sequence of ``hidden_size`` of subsequent conv layers \ @@ -50,6 +53,8 @@ def __init__( - stride (:obj:`SequenceType`): Sequence of ``stride`` of subsequent conv layers. - padding (:obj:`SequenceType`): Padding added to all four sides of the input for each conv layer. \ See ``nn.Conv2d`` for more details. Default is ``None``. + - layer_norm (:obj:`bool`): Whether to use ``DreamerLayerNorm``, which is kind of special trick \ + proposed in DreamerV3. - norm_type (:obj:`str`): Type of normalization to use. See ``ding.torch_utils.network.ResBlock`` \ for more details. Default is ``None``. """ @@ -63,16 +68,35 @@ def __init__( layers = [] input_size = obs_shape[0] # in_channel for i in range(len(kernel_size)): - layers.append(nn.Conv2d(input_size, hidden_size_list[i], kernel_size[i], stride[i], padding[i])) - layers.append(self.act) + if layer_norm: + layers.append( + Conv2dSame( + in_channels=input_size, + out_channels=hidden_size_list[i], + kernel_size=(kernel_size[i], kernel_size[i]), + stride=(2, 2), + bias=False, + ) + ) + layers.append(DreamerLayerNorm(hidden_size_list[i])) + layers.append(self.act) + else: + layers.append(nn.Conv2d(input_size, hidden_size_list[i], kernel_size[i], stride[i], padding[i])) + layers.append(self.act) input_size = hidden_size_list[i] - assert len(set(hidden_size_list[3:-1])) <= 1, "Please indicate the same hidden size for res block parts" - for i in range(3, len(self.hidden_size_list) - 1): - layers.append(ResBlock(self.hidden_size_list[i], activation=self.act, norm_type=norm_type)) + if len(self.hidden_size_list) >= len(kernel_size) + 2: + assert self.hidden_size_list[len(kernel_size) - 1] == self.hidden_size_list[ + len(kernel_size)], "Please indicate the same hidden size between conv and res block" + assert len( + set(hidden_size_list[len(kernel_size):-1]) + ) <= 1, "Please indicate the same hidden size for res block parts" + for i in range(len(kernel_size), len(self.hidden_size_list) - 1): + layers.append(ResBlock(self.hidden_size_list[i - 1], activation=self.act, norm_type=norm_type)) layers.append(Flatten()) self.main = nn.Sequential(*layers) flatten_size = self._get_flatten_size() + self.output_size = hidden_size_list[-1] # outside to use self.mid = nn.Linear(flatten_size, hidden_size_list[-1]) def _get_flatten_size(self) -> int: @@ -83,6 +107,18 @@ def _get_flatten_size(self) -> int: - outputs (:obj:`torch.Tensor`): Size ``int`` Tensor representing the number of ``in-features``. Shapes: - outputs: :math:`(1,)`. + Examples: + >>> conv = ConvEncoder( + >>> obs_shape=(4, 84, 84), + >>> hidden_size_list=[32, 64, 64, 128], + >>> activation=nn.ReLU(), + >>> kernel_size=[8, 4, 3], + >>> stride=[4, 2, 1], + >>> padding=None, + >>> layer_norm=False, + >>> norm_type=None + >>> ) + >>> flatten_size = conv._get_flatten_size() """ test_data = torch.randn(1, *self.obs_shape) with torch.no_grad(): @@ -92,13 +128,27 @@ def _get_flatten_size(self) -> int: def forward(self, x: torch.Tensor) -> torch.Tensor: """ Overview: - Return output embedding tensor of the env observation. + Return output 1D embedding tensor of the env's 2D image observation. Arguments: - - x (:obj:`torch.Tensor`): Env raw observation. + - x (:obj:`torch.Tensor`): Raw 2D observation of the environment. Returns: - outputs (:obj:`torch.Tensor`): Output embedding tensor. Shapes: - - outputs: :math:`(B, N)`, where ``N = hidden_size_list[-1]``. + - x : :math:`(B, C, H, W)`, where ``B`` is batch size, ``C`` is channel, ``H`` is height, ``W`` is width. + - outputs: :math:`(B, N)`, where ``N = hidden_size_list[-1]`` . + Examples: + >>> conv = ConvEncoder( + >>> obs_shape=(4, 84, 84), + >>> hidden_size_list=[32, 64, 64, 128], + >>> activation=nn.ReLU(), + >>> kernel_size=[8, 4, 3], + >>> stride=[4, 2, 1], + >>> padding=None, + >>> layer_norm=False, + >>> norm_type=None + >>> ) + >>> x = torch.randn(1, 4, 84, 84) + >>> output = conv(x) """ x = self.main(x) x = self.mid(x) @@ -108,7 +158,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class FCEncoder(nn.Module): """ Overview: - The ``FCEncoder`` used in models to encode raw 1-dim observations. + The full connected encoder is used to encode 1-dim input variable. Interfaces: ``__init__``, ``forward``. """ @@ -119,11 +169,12 @@ def __init__( hidden_size_list: SequenceType, res_block: bool = False, activation: Optional[nn.Module] = nn.ReLU(), - norm_type: Optional[str] = None + norm_type: Optional[str] = None, + dropout: Optional[float] = None ) -> None: """ Overview: - Init the FC Encoder according to arguments. + Initialize the FC Encoder according to arguments. Arguments: - obs_shape (:obj:`int`): Observation shape. - hidden_size_list (:obj:`SequenceType`): Sequence of ``hidden_size`` of subsequent FC layers. @@ -131,6 +182,7 @@ def __init__( - activation (:obj:`nn.Module`): Type of activation to use in ``ResFCBlock``. Default is ``nn.ReLU()``. - norm_type (:obj:`str`): Type of normalization to use. See ``ding.torch_utils.network.ResFCBlock`` \ for more details. Default is ``None``. + - dropout (:obj:`float`): Dropout rate of the dropout layer. If ``None`` then default no dropout layer. """ super(FCEncoder, self).__init__() self.obs_shape = obs_shape @@ -140,17 +192,21 @@ def __init__( if res_block: assert len(set(hidden_size_list)) == 1, "Please indicate the same hidden size for res block parts" if len(hidden_size_list) == 1: - self.main = ResFCBlock(hidden_size_list[0], activation=self.act, norm_type=norm_type) + self.main = ResFCBlock(hidden_size_list[0], activation=self.act, norm_type=norm_type, dropout=dropout) else: layers = [] for i in range(len(hidden_size_list)): - layers.append(ResFCBlock(hidden_size_list[0], activation=self.act, norm_type=norm_type)) + layers.append( + ResFCBlock(hidden_size_list[0], activation=self.act, norm_type=norm_type, dropout=dropout) + ) self.main = nn.Sequential(*layers) else: layers = [] for i in range(len(hidden_size_list) - 1): layers.append(nn.Linear(hidden_size_list[i], hidden_size_list[i + 1])) layers.append(self.act) + if dropout is not None: + layers.append(nn.Dropout(dropout)) self.main = nn.Sequential(*layers) def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -162,7 +218,18 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: Returns: - outputs (:obj:`torch.Tensor`): Output embedding tensor. Shapes: + - x : :math:`(B, M)`, where ``M = obs_shape``. - outputs: :math:`(B, N)`, where ``N = hidden_size_list[-1]``. + Examples: + >>> fc = FCEncoder( + >>> obs_shape=4, + >>> hidden_size_list=[32, 64, 64, 128], + >>> activation=nn.ReLU(), + >>> norm_type=None, + >>> dropout=None + >>> ) + >>> x = torch.randn(1, 4) + >>> output = fc(x) """ x = self.act(self.init(x)) x = self.main(x) @@ -170,23 +237,32 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class StructEncoder(nn.Module): - # TODO(nyz) - pass + + def __init__(self, obs_shape: Dict[str, Union[int, List[int]]]) -> None: + super(StructEncoder, self).__init__() + # TODO concrete implementation + raise NotImplementedError class IMPALACnnResidualBlock(nn.Module): """ - Residual basic block (without batchnorm) in IMPALA cnn encoder. - Preserves channel number and shape + Overview: + This CNN encoder residual block is residual basic block used in IMPALA algorithm, + which preserves the channel number and shape. + IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures + https://arxiv.org/pdf/1802.01561.pdf + Interfaces: + ``__init__``, ``forward``. """ - def __init__(self, in_channnel, scale=1, batch_norm=False): + def __init__(self, in_channnel: int, scale: float = 1, batch_norm: bool = False): """ Overview: - Init every impala cnn residual block. + Initialize the IMPALA CNN residual block according to arguments. Arguments: - in_channnel (:obj:`int`): Channel number of input features. - - scale (:obj:`float`): Scale of module. + - scale (:obj:`float`): Scale of module, defaults to 1. + - batch_norm (:obj:`bool`): Whether use batch normalization, defaults to False. """ super().__init__() self.in_channnel = in_channnel @@ -198,10 +274,17 @@ def __init__(self, in_channnel, scale=1, batch_norm=False): self.bn0 = nn.BatchNorm2d(self.in_channnel) self.bn1 = nn.BatchNorm2d(self.in_channnel) - def residual(self, x): - # inplace should be False for the first relu, so that it does not change the input, - # which will be used for skip connection. - # getattr is for backwards compatibility with loaded models + def residual(self, x: torch.Tensor) -> torch.Tensor: + """ + Overview: + Return output tensor of the residual block, keep the shape and channel number unchanged. + The inplace of activation function should be False for the first relu, + so that it does not change the origin input tensor of the residual block. + Arguments: + - x (:obj:`torch.Tensor`): Input tensor. + Returns: + - output (:obj:`torch.Tensor`): Output tensor. + """ if self.batch_norm: x = self.bn0(x) x = F.relu(x, inplace=False) @@ -212,19 +295,38 @@ def residual(self, x): x = self.conv1(x) return x - def forward(self, x): + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Overview: + Return output tensor of the residual block, keep the shape and channel number unchanged. + Arguments: + - x (:obj:`torch.Tensor`): Input tensor. + Returns: + - output (:obj:`torch.Tensor`): Output tensor. + Examples: + >>> block = IMPALACnnResidualBlock(16) + >>> x = torch.randn(1, 16, 84, 84) + >>> output = block(x) + """ return x + self.residual(x) class IMPALACnnDownStack(nn.Module): """ - Downsampling stack from Impala CNN + Overview: + Downsampling stack of CNN encoder used in IMPALA algorithmn. + Every IMPALACnnDownStack consists n IMPALACnnResidualBlock, + which reduces the spatial size by 2 with maxpooling. + IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures + https://arxiv.org/pdf/1802.01561.pdf + Interfaces: + ``__init__``, ``forward``. """ def __init__(self, in_channnel, nblock, out_channel, scale=1, pool=True, **kwargs): """ Overview: - Init every impala cnn block of the Impala Cnn Encoder. + Initialize every impala cnn block of the Impala Cnn Encoder. Arguments: - in_channnel (:obj:`int`): Channel number of input features. - nblock (:obj:`int`): Residual Block number in each block. @@ -240,7 +342,20 @@ def __init__(self, in_channnel, nblock, out_channel, scale=1, pool=True, **kwarg s = scale / math.sqrt(nblock) self.blocks = nn.ModuleList([IMPALACnnResidualBlock(out_channel, scale=s, **kwargs) for _ in range(nblock)]) - def forward(self, x): + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Overview: + Return output tensor of the downsampling stack. The output shape is different from input shape. And you \ + can refer to the ``output_shape`` method to get the output shape. + Arguments: + - x (:obj:`torch.Tensor`): Input tensor. + Returns: + - output (:obj:`torch.Tensor`): Output tensor. + Examples: + >>> stack = IMPALACnnDownStack(16, 2, 32) + >>> x = torch.randn(1, 16, 84, 84) + >>> output = stack(x) + """ x = self.firstconv(x) if self.pool: x = F.max_pool2d(x, kernel_size=3, stride=2, padding=1) @@ -248,7 +363,22 @@ def forward(self, x): x = block(x) return x - def output_shape(self, inshape): + def output_shape(self, inshape: tuple) -> tuple: + """ + Overview: + Calculate the output shape of the downsampling stack according to input shape and related arguments. + Arguments: + - inshape (:obj:`tuple`): Input shape. + Returns: + - output_shape (:obj:`tuple`): Output shape. + Shapes: + - inshape (:obj:`tuple`): :math:`(C, H, W)`, where C is channel number, H is height and W is width. + - output_shape (:obj:`tuple`): :math:`(C, H, W)`, where C is channel number, H is height and W is width. + Examples: + >>> stack = IMPALACnnDownStack(16, 2, 32) + >>> inshape = (16, 84, 84) + >>> output_shape = stack.output_shape(inshape) + """ c, h, w = inshape assert c == self.in_channnel if self.pool: @@ -258,24 +388,40 @@ def output_shape(self, inshape): class IMPALAConvEncoder(nn.Module): + """ + Overview: + IMPALA CNN encoder, which is used in IMPALA algorithm. + IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures, \ + https://arxiv.org/pdf/1802.01561.pdf, + Interface: + ``__init__``, ``forward``, ``output_shape``. + """ name = "IMPALAConvEncoder" # put it here to preserve pickle compat def __init__( - self, obs_shape, channels=(16, 32, 32), outsize=256, scale_ob=255.0, nblock=2, final_relu=True, **kwargs - ): + self, + obs_shape: SequenceType, + channels: SequenceType = (16, 32, 32), + outsize: int = 256, + scale_ob: float = 255.0, + nblock: int = 2, + final_relu: bool = True, + **kwargs + ) -> None: """ Overview: - Init the Encoder described in paper, \ - IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures, \ - https://arxiv.org/pdf/1802.01561.pdf, + Initialize the IMPALA CNN encoder according to arguments. Arguments: - - obs_shape (:obj:`int`): Observation shape. - - channels (:obj:`SequenceType`): Channel number of each impala cnn block. \ - The size of it is the number of impala cnn blocks in the encoder - - outsize (:obj:`int`): Out feature of the encoder. - - scale_ob (:obj:`float`): Scale of each pixel. - - nblock (:obj:`int`): Residual Block number in each block. - - final_relu (:obj:`bool`): Whether to use Relu in the end of encoder. + - obs_shape (:obj:`SequenceType`): 2D image observation shape. + - channels (:obj:`SequenceType`): The channel number of a series of impala cnn blocks. \ + Each element of the sequence is the output channel number of a impala cnn block. + - outsize (:obj:`int`): The output size the final linear layer, which means the dimension of the \ + 1D embedding vector. + - scale_ob (:obj:`float`): The scale of the input observation, which is used to normalize the input \ + observation, such as dividing 255.0 for the raw image observation. + - nblock (:obj:`int`): The number of Residual Block in each block. + - final_relu (:obj:`bool`): Whether to use ReLU activation in the final output of encoder. + - kwargs (:obj:`Dict[str, Any]`): Other arguments for ``IMPALACnnDownStack``. """ super().__init__() self.scale_ob = scale_ob @@ -291,7 +437,30 @@ def __init__( self.outsize = outsize self.final_relu = final_relu - def forward(self, x): + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Overview: + Return the 1D embedding vector of the input 2D observation. + Arguments: + - x (:obj:`torch.Tensor`): Input 2D observation tensor. + Returns: + - output (:obj:`torch.Tensor`): Output 1D embedding vector. + Shapes: + - x (:obj:`torch.Tensor`): :math:`(B, C, H, W)`, where B is batch size, C is channel number, H is height \ + and W is width. + - output (:obj:`torch.Tensor`): :math:`(B, outsize)`, where B is batch size. + Examples: + >>> encoder = IMPALAConvEncoder( + >>> obs_shape=(4, 84, 84), + >>> channels=(16, 32, 32), + >>> outsize=256, + >>> scale_ob=255.0, + >>> nblock=2, + >>> final_relu=True, + >>> ) + >>> x = torch.randn(1, 4, 84, 84) + >>> output = encoder(x) + """ x = x / self.scale_ob for (i, layer) in enumerate(self.stacks): x = layer(x) @@ -302,3 +471,46 @@ def forward(self, x): if self.final_relu: x = torch.relu(x) return x + + +class GaussianFourierProjectionTimeEncoder(nn.Module): + """ + Overview: + Gaussian random features for encoding time steps. + This module is used as the encoder of time in generative models such as diffusion model. + Interfaces: + ``__init__``, ``forward``. + """ + + def __init__(self, embed_dim, scale=30.): + """ + Overview: + Initialize the Gaussian Fourier Projection Time Encoder according to arguments. + Arguments: + - embed_dim (:obj:`int`): The dimension of the output embedding vector. + - scale (:obj:`float`): The scale of the Gaussian random features. + """ + super().__init__() + # Randomly sample weights during initialization. These weights are fixed + # during optimization and are not trainable. + self.W = nn.Parameter(torch.randn(embed_dim // 2) * scale * 2 * np.pi, requires_grad=False) + + def forward(self, x): + """ + Overview: + Return the output embedding vector of the input time step. + Arguments: + - x (:obj:`torch.Tensor`): Input time step tensor. + Returns: + - output (:obj:`torch.Tensor`): Output embedding vector. + Shapes: + - x (:obj:`torch.Tensor`): :math:`(B,)`, where B is batch size. + - output (:obj:`torch.Tensor`): :math:`(B, embed_dim)`, where B is batch size, embed_dim is the \ + dimension of the output embedding vector. + Examples: + >>> encoder = GaussianFourierProjectionTimeEncoder(128) + >>> x = torch.randn(100) + >>> output = encoder(x) + """ + x_proj = x[..., None] * self.W[None, :] + return torch.cat([torch.sin(x_proj), torch.cos(x_proj)], dim=-1) diff --git a/ding/model/common/head.py b/ding/model/common/head.py old mode 100644 new mode 100755 index ddfb542c9a..1131e8a2e8 --- a/ding/model/common/head.py +++ b/ding/model/common/head.py @@ -1,4 +1,4 @@ -from typing import Optional, Dict +from typing import Optional, Dict, Union, List import math import torch @@ -6,19 +6,18 @@ import torch.nn.functional as F from torch.distributions import Normal, Independent -from ding.torch_utils import fc_block, noise_block, NoiseLinearLayer, MLP +from ding.torch_utils import fc_block, noise_block, NoiseLinearLayer, MLP, PopArt, conv1d_block from ding.rl_utils import beta_function_map from ding.utils import lists_to_dicts, SequenceType class DiscreteHead(nn.Module): """ - Overview: - The ``DiscreteHead`` used to output discrete actions logit. \ - Input is a (:obj:`torch.Tensor`) of shape ``(B, N)`` and returns a (:obj:`Dict`) containing \ - output ``logit``. - Interfaces: - ``__init__``, ``forward``. + Overview: + The ``DiscreteHead`` is used to generate discrete actions logit or Q-value logit, \ + which is often used in q-learning algorithms or actor-critic algorithms for discrete action space. + Interfaces: + ``__init__``, ``forward``. """ def __init__( @@ -28,6 +27,7 @@ def __init__( layer_num: int = 1, activation: Optional[nn.Module] = nn.ReLU(), norm_type: Optional[str] = None, + dropout: Optional[float] = None, noise: Optional[bool] = False, ) -> None: """ @@ -41,6 +41,7 @@ def __init__( If ``None``, then default set activation to ``nn.ReLU()``. Default ``None``. - norm_type (:obj:`str`): The type of normalization to use. See ``ding.torch_utils.network.fc_block`` \ for more details. Default ``None``. + - dropout (:obj:`float`): The dropout rate, default set to None. - noise (:obj:`bool`): Whether use ``NoiseLinearLayer`` as ``layer_fn`` in Q networks' MLP. \ Default ``False``. """ @@ -55,6 +56,8 @@ def __init__( layer_num, layer_fn=layer, activation=activation, + use_dropout=dropout is not None, + dropout_probability=dropout, norm_type=norm_type ), block(hidden_size, output_size) ) @@ -70,7 +73,6 @@ def forward(self, x: torch.Tensor) -> Dict: Shapes: - x: :math:`(B, N)`, where ``B = batch_size`` and ``N = hidden_size``. - logit: :math:`(B, M)`, where ``M = output_size``. - Examples: >>> head = DiscreteHead(64, 64) >>> inputs = torch.randn(4, 64) @@ -83,13 +85,11 @@ def forward(self, x: torch.Tensor) -> Dict: class DistributionHead(nn.Module): """ - Overview: - The ``DistributionHead`` used to output Q-value distribution. \ - Input is a (:obj:`torch.Tensor`) of shape ``(B, N)`` and returns a (:obj:`Dict`) containing \ - outputs ``logit`` and ``distribution``. - - Interfaces: - ``__init__``, ``forward``. + Overview: + The ``DistributionHead`` is used to generate distribution for Q-value. + This module is used in C51 algorithm. + Interfaces: + ``__init__``, ``forward``. """ def __init__( @@ -156,7 +156,6 @@ def forward(self, x: torch.Tensor) -> Dict: - x: :math:`(B, N)`, where ``B = batch_size`` and ``N = hidden_size``. - logit: :math:`(B, M)`, where ``M = output_size``. - distribution: :math:`(B, M, n_atom)`. - Examples: >>> head = DistributionHead(64, 64) >>> inputs = torch.randn(4, 64) @@ -174,14 +173,126 @@ def forward(self, x: torch.Tensor) -> Dict: return {'logit': q, 'distribution': dist} -class RainbowHead(nn.Module): +class BranchingHead(nn.Module): """ + Overview: + The ``BranchingHead`` is used to generate Q-value with different branches. + This module is used in Branch DQN. + Interfaces: + ``__init__``, ``forward``. + """ + + def __init__( + self, + hidden_size: int, + num_branches: int = 0, + action_bins_per_branch: int = 2, + layer_num: int = 1, + a_layer_num: Optional[int] = None, + v_layer_num: Optional[int] = None, + norm_type: Optional[str] = None, + activation: Optional[nn.Module] = nn.ReLU(), + noise: Optional[bool] = False, + ) -> None: + """ + Overview: + Init the ``BranchingHead`` layers according to the provided arguments. \ + This head achieves a linear increase of the number of network outputs \ + with the number of degrees of freedom by allowing a level of independence for each individual action. + Therefore, this head is suitable for high dimensional action Spaces. + Arguments: + - hidden_size (:obj:`int`): The ``hidden_size`` of the MLP connected to ``BranchingHead``. + - num_branches (:obj:`int`): The number of branches, which is equivalent to the action dimension. + - action_bins_per_branch (:obj:int): The number of action bins in each dimension. + - layer_num (:obj:`int`): The number of layers used in the network to compute Advantage and Value output. + - a_layer_num (:obj:`int`): The number of layers used in the network to compute Advantage output. + - v_layer_num (:obj:`int`): The number of layers used in the network to compute Value output. + - output_size (:obj:`int`): The number of outputs. + - norm_type (:obj:`str`): The type of normalization to use. See ``ding.torch_utils.network.fc_block`` \ + for more details. Default ``None``. + - activation (:obj:`nn.Module`): The type of activation function to use in MLP. \ + If ``None``, then default set activation to ``nn.ReLU()``. Default ``None``. + - noise (:obj:`bool`): Whether use ``NoiseLinearLayer`` as ``layer_fn`` in Q networks' MLP. \ + Default ``False``. + """ + super(BranchingHead, self).__init__() + if a_layer_num is None: + a_layer_num = layer_num + if v_layer_num is None: + v_layer_num = layer_num + self.num_branches = num_branches + self.action_bins_per_branch = action_bins_per_branch + + layer = NoiseLinearLayer if noise else nn.Linear + block = noise_block if noise else fc_block + # value network + + self.V = nn.Sequential( + MLP( + hidden_size, + hidden_size, + hidden_size, + v_layer_num, + layer_fn=layer, + activation=activation, + norm_type=norm_type + ), block(hidden_size, 1) + ) + # action branching network + action_output_dim = action_bins_per_branch + self.branches = nn.ModuleList( + [ + nn.Sequential( + MLP( + hidden_size, + hidden_size, + hidden_size, + a_layer_num, + layer_fn=layer, + activation=activation, + norm_type=norm_type + ), block(hidden_size, action_output_dim) + ) for _ in range(self.num_branches) + ] + ) + + def forward(self, x: torch.Tensor) -> Dict: + """ Overview: - The ``RainbowHead`` used to output Q-value distribution. \ - Input is a (:obj:`torch.Tensor`) of shape ``(B, N)`` and returns a (:obj:`Dict`) containing \ - outputs ``logit`` and ``distribution``. - Interfaces: - ``__init__``, ``forward``. + Use encoded embedding tensor to run MLP with ``BranchingHead`` and return the prediction dictionary. + Arguments: + - x (:obj:`torch.Tensor`): Tensor containing input embedding. + Returns: + - outputs (:obj:`Dict`): Dict containing keyword ``logit`` (:obj:`torch.Tensor`). + Shapes: + - x: :math:`(B, N)`, where ``B = batch_size`` and ``N = hidden_size``. + - logit: :math:`(B, M)`, where ``M = output_size``. + Examples: + >>> head = BranchingHead(64, 5, 2) + >>> inputs = torch.randn(4, 64) + >>> outputs = head(inputs) + >>> assert isinstance(outputs, dict) and outputs['logit'].shape == torch.Size([4, 5, 2]) + """ + value_out = self.V(x) + value_out = torch.unsqueeze(value_out, 1) + action_out = [] + for b in self.branches: + action_out.append(b(x)) + action_scores = torch.stack(action_out, 1) + # From the paper, this implementation performs better than both the naive alternative (Q = V + A) \ + # and the local maximum reduction method (Q = V + max(A)). + action_scores = action_scores - torch.mean(action_scores, 2, keepdim=True) + logits = value_out + action_scores + return {'logit': logits} + + +class RainbowHead(nn.Module): + """ + Overview: + The ``RainbowHead`` is used to generate distribution of Q-value. + This module is used in Rainbow DQN. + Interfaces: + ``__init__``, ``forward``. """ def __init__( @@ -259,7 +370,6 @@ def forward(self, x: torch.Tensor) -> Dict: - x: :math:`(B, N)`, where ``B = batch_size`` and ``N = hidden_size``. - logit: :math:`(B, M)`, where ``M = output_size``. - distribution: :math:`(B, M, n_atom)`. - Examples: >>> head = RainbowHead(64, 64) >>> inputs = torch.randn(4, 64) @@ -282,12 +392,10 @@ def forward(self, x: torch.Tensor) -> Dict: class QRDQNHead(nn.Module): """ - Overview: - The ``QRDQNHead`` (Quantile Regression DQN) used to output action quantiles. \ - Input is a (:obj:`torch.Tensor`) of shape ``(B, N)`` and returns a (:obj:`Dict`) containing \ - output ``logit``, ``q``, and ``tau``. - Interfaces: - ``__init__``, ``forward``. + Overview: + The ``QRDQNHead`` (Quantile Regression DQN) is used to output action quantiles. + Interfaces: + ``__init__``, ``forward``. """ def __init__( @@ -346,7 +454,6 @@ def forward(self, x: torch.Tensor) -> Dict: - logit: :math:`(B, M)`, where ``M = output_size``. - q: :math:`(B, M, num_quantiles)`. - tau: :math:`(B, M, 1)`. - Examples: >>> head = QRDQNHead(64, 64) >>> inputs = torch.randn(4, 64) @@ -368,12 +475,16 @@ def forward(self, x: torch.Tensor) -> Dict: class QuantileHead(nn.Module): """ - Overview: - The ``QuantileHead`` used to output action quantiles. \ - Input is a (:obj:`torch.Tensor`) of shape ``(B, N)`` and returns a (:obj:`Dict`) containing \ - output ``logit``, ``q``, and ``quantiles``. - Interfaces: - ``__init__``, ``forward``, ``quantile_net``. + Overview: + The ``QuantileHead`` is used to output action quantiles. + This module is used in IQN. + Interfaces: + ``__init__``, ``forward``, ``quantile_net``. + + .. note:: + The difference between ``QuantileHead`` and ``QRDQNHead`` is that ``QuantileHead`` models the \ + state-action quantile function as a mapping from state-actions and samples from some base distribution \ + while ``QRDQNHead`` approximates random returns by a uniform mixture of Diracs functions. """ def __init__( @@ -467,7 +578,6 @@ def forward(self, x: torch.Tensor, num_quantiles: Optional[int] = None) -> Dict: - logit: :math:`(B, M)`, where ``M = output_size``. - q: :math:`(num_quantiles, B, M)`. - quantiles: :math:`(quantile_embedding_size, 1)`. - Examples: >>> head = QuantileHead(64, 64) >>> inputs = torch.randn(4, 64) @@ -501,12 +611,19 @@ def forward(self, x: torch.Tensor, num_quantiles: Optional[int] = None) -> Dict: class FQFHead(nn.Module): """ - Overview: - The ``FQFHead`` used to output action quantiles. \ - Input is a (:obj:`torch.Tensor`) of shape ``(B, N)`` and returns a (:obj:`Dict`) containing \ - output ``logit``, ``q``, ``quantiles``, ``quantiles_hats``, ``q_tau_i`` and ``entropies``. - Interfaces: - ``__init__``, ``forward``, ``quantile_net``. + Overview: + The ``FQFHead`` is used to output action quantiles. + This module is used in FQF. + Interfaces: + ``__init__``, ``forward``, ``quantile_net``. + + .. note:: + The implementation of FQFHead is based on the paper https://arxiv.org/abs/1911.02140. + The difference between FQFHead and QuantileHead is that, in FQF, \ + N adjustable quantile values for N adjustable quantile fractions are estimated to approximate \ + the quantile function. The distribution of the return is approximated by a weighted mixture of N \ + Diracs functions. While in IQN, the state-action quantile function is modeled as a mapping from \ + state-actions and samples from some base distribution. """ def __init__( @@ -673,12 +790,11 @@ def forward(self, x: torch.Tensor, num_quantiles: Optional[int] = None) -> Dict: class DuelingHead(nn.Module): """ - Overview: - The ``DuelingHead`` used to output discrete actions logit. \ - Input is a (:obj:`torch.Tensor`) of shape ``(B, N)`` and returns a (:obj:`Dict`) containing \ - output ``logit``. - Interfaces: - ``__init__``, ``forward``. + Overview: + The ``DuelingHead`` is used to output discrete actions logit. + This module is used in Dueling DQN. + Interfaces: + ``__init__``, ``forward``. """ def __init__( @@ -690,6 +806,7 @@ def __init__( v_layer_num: Optional[int] = None, activation: Optional[nn.Module] = nn.ReLU(), norm_type: Optional[str] = None, + dropout: Optional[float] = None, noise: Optional[bool] = False, ) -> None: """ @@ -704,6 +821,7 @@ def __init__( If ``None``, then default set activation to ``nn.ReLU()``. Default ``None``. - norm_type (:obj:`str`): The type of normalization to use. See ``ding.torch_utils.network.fc_block`` \ for more details. Default ``None``. + - dropout (:obj:`float`): The dropout rate of dropout layer. Default ``None``. - noise (:obj:`bool`): Whether use ``NoiseLinearLayer`` as ``layer_fn`` in Q networks' MLP. \ Default ``False``. """ @@ -722,6 +840,8 @@ def __init__( a_layer_num, layer_fn=layer, activation=activation, + use_dropout=dropout is not None, + dropout_probability=dropout, norm_type=norm_type ), block(hidden_size, output_size) ) @@ -733,6 +853,8 @@ def __init__( v_layer_num, layer_fn=layer, activation=activation, + use_dropout=dropout is not None, + dropout_probability=dropout, norm_type=norm_type ), block(hidden_size, 1) ) @@ -748,7 +870,6 @@ def forward(self, x: torch.Tensor) -> Dict: Shapes: - x: :math:`(B, N)`, where ``B = batch_size`` and ``N = hidden_size``. - logit: :math:`(B, M)`, where ``M = output_size``. - Examples: >>> head = DuelingHead(64, 64) >>> inputs = torch.randn(4, 64) @@ -764,13 +885,11 @@ def forward(self, x: torch.Tensor) -> Dict: class StochasticDuelingHead(nn.Module): """ - Overview: - The ``Stochastic Dueling Network`` proposed in paper ACER (arxiv 1611.01224). \ - Dueling network architecture in continuous action space. \ - Input is a (:obj:`torch.Tensor`) of shape ``(B, N)`` and returns a (:obj:`Dict`) containing \ - outputs ``q_value`` and ``v_value``. - Interfaces: - ``__init__``, ``forward``. + Overview: + The ``Stochastic Dueling Network`` is proposed in paper ACER (arxiv 1611.01224). \ + That is to say, dueling network architecture in continuous action space. + Interfaces: + ``__init__``, ``forward``. """ def __init__( @@ -868,6 +987,16 @@ def forward( - sigma: :math:`(B, A)`. - q_value: :math:`(B, 1)`. - v_value: :math:`(B, 1)`. + Examples: + >>> head = StochasticDuelingHead(64, 64) + >>> inputs = torch.randn(4, 64) + >>> a = torch.randn(4, 64) + >>> mu = torch.randn(4, 64) + >>> sigma = torch.ones(4, 64) + >>> outputs = head(inputs, a, mu, sigma) + >>> assert isinstance(outputs, dict) + >>> assert outputs['q_value'].shape == torch.Size([4, 1]) + >>> assert outputs['v_value'].shape == torch.Size([4, 1]) """ batch_size = s.shape[0] # batch_size or batch_size * T @@ -897,22 +1026,23 @@ def forward( class RegressionHead(nn.Module): """ - Overview: - The ``RegressionHead`` used to output actions Q-value. - Input is a (:obj:`torch.Tensor`) of shape ``(B, N)`` and returns a (:obj:`Dict`) containing \ - output ``pred``. - Interfaces: - ``__init__``, ``forward``. + Overview: + The ``RegressionHead`` is used to regress continuous variables. + This module is used for generating Q-value (DDPG critic) of continuous actions, \ + or state value (A2C/PPO), or directly predicting continuous action (DDPG actor). + Interfaces: + ``__init__``, ``forward``. """ def __init__( - self, - hidden_size: int, - output_size: int, - layer_num: int = 2, - final_tanh: Optional[bool] = False, - activation: Optional[nn.Module] = nn.ReLU(), - norm_type: Optional[str] = None + self, + input_size: int, + output_size: int, + layer_num: int = 2, + final_tanh: Optional[bool] = False, + activation: Optional[nn.Module] = nn.ReLU(), + norm_type: Optional[str] = None, + hidden_size: int = None, ) -> None: """ Overview: @@ -928,7 +1058,9 @@ def __init__( for more details. Default ``None``. """ super(RegressionHead, self).__init__() - self.main = MLP(hidden_size, hidden_size, hidden_size, layer_num, activation=activation, norm_type=norm_type) + if hidden_size is None: + hidden_size = input_size + self.main = MLP(input_size, hidden_size, hidden_size, layer_num, activation=activation, norm_type=norm_type) self.last = nn.Linear(hidden_size, output_size) # for convenience of special initialization self.final_tanh = final_tanh if self.final_tanh: @@ -945,7 +1077,6 @@ def forward(self, x: torch.Tensor) -> Dict: Shapes: - x: :math:`(B, N)`, where ``B = batch_size`` and ``N = hidden_size``. - pred: :math:`(B, M)`, where ``M = output_size``. - Examples: >>> head = RegressionHead(64, 64) >>> inputs = torch.randn(4, 64) @@ -964,27 +1095,29 @@ def forward(self, x: torch.Tensor) -> Dict: class ReparameterizationHead(nn.Module): """ - Overview: - The ``ReparameterizationHead`` used to output action ``mu`` and ``sigma``. - Input is a (:obj:`torch.Tensor`) of shape ``(B, N)`` and returns a (:obj:`Dict`) containing \ - outputs ``mu`` and ``sigma``. - Interfaces: - ``__init__``, ``forward``. + Overview: + The ``ReparameterizationHead`` is used to generate Gaussian distribution of continuous variable, \ + which is parameterized by ``mu`` and ``sigma``. + This module is often used in stochastic policies, such as PPO and SAC. + Interfaces: + ``__init__``, ``forward``. """ - - default_sigma_type = ['fixed', 'independent', 'conditioned'] + # The "happo" type here is to align with the sigma initialization method of the network in the original HAPPO \ + # paper. The code here needs to be optimized later. + default_sigma_type = ['fixed', 'independent', 'conditioned', 'happo'] default_bound_type = ['tanh', None] def __init__( - self, - hidden_size: int, - output_size: int, - layer_num: int = 2, - sigma_type: Optional[str] = None, - fixed_sigma_value: Optional[float] = 1.0, - activation: Optional[nn.Module] = nn.ReLU(), - norm_type: Optional[str] = None, - bound_type: Optional[str] = None, + self, + input_size: int, + output_size: int, + layer_num: int = 2, + sigma_type: Optional[str] = None, + fixed_sigma_value: Optional[float] = 1.0, + activation: Optional[nn.Module] = nn.ReLU(), + norm_type: Optional[str] = None, + bound_type: Optional[str] = None, + hidden_size: int = None ) -> None: """ Overview: @@ -1005,6 +1138,8 @@ def __init__( Default is ``None``. """ super(ReparameterizationHead, self).__init__() + if hidden_size is None: + hidden_size = input_size self.sigma_type = sigma_type assert sigma_type in self.default_sigma_type, "Please indicate sigma_type as one of {}".format( self.default_sigma_type @@ -1013,7 +1148,7 @@ def __init__( assert bound_type in self.default_bound_type, "Please indicate bound_type as one of {}".format( self.default_bound_type ) - self.main = MLP(hidden_size, hidden_size, hidden_size, layer_num, activation=activation, norm_type=norm_type) + self.main = MLP(input_size, hidden_size, hidden_size, layer_num, activation=activation, norm_type=norm_type) self.mu = nn.Linear(hidden_size, output_size) if self.sigma_type == 'fixed': self.sigma = torch.full((1, output_size), fixed_sigma_value) @@ -1021,6 +1156,11 @@ def __init__( self.log_sigma_param = nn.Parameter(torch.zeros(1, output_size)) elif self.sigma_type == 'conditioned': self.log_sigma_layer = nn.Linear(hidden_size, output_size) + elif self.sigma_type == 'happo': + self.sigma_x_coef = 1. + self.sigma_y_coef = 0.5 + # This parameter (x_coef, y_coef) refers to the HAPPO paper http://arxiv.org/abs/2109.11251. + self.log_sigma_param = nn.Parameter(torch.ones(1, output_size) * self.sigma_x_coef) def forward(self, x: torch.Tensor) -> Dict: """ @@ -1036,7 +1176,6 @@ def forward(self, x: torch.Tensor) -> Dict: - x: :math:`(B, N)`, where ``B = batch_size`` and ``N = hidden_size``. - mu: :math:`(B, M)`, where ``M = output_size``. - sigma: :math:`(B, M)`. - Examples: >>> head = ReparameterizationHead(64, 64, sigma_type='fixed') >>> inputs = torch.randn(4, 64) @@ -1057,17 +1196,129 @@ def forward(self, x: torch.Tensor) -> Dict: elif self.sigma_type == 'conditioned': log_sigma = self.log_sigma_layer(x) sigma = torch.exp(torch.clamp(log_sigma, -20, 2)) + elif self.sigma_type == 'happo': + log_sigma = self.log_sigma_param + torch.zeros_like(mu) + sigma = torch.sigmoid(log_sigma / self.sigma_x_coef) * self.sigma_y_coef return {'mu': mu, 'sigma': sigma} -class MultiHead(nn.Module): +class PopArtVHead(nn.Module): + """ + Overview: + The ``PopArtVHead`` is used to generate adaptive normalized state value. More information can be found in \ + paper Multi-task Deep Reinforcement Learning with PopArt. \ + https://arxiv.org/abs/1809.04474 \ + This module is used in PPO or IMPALA. + Interfaces: + ``__init__``, ``forward``. """ + + def __init__( + self, + hidden_size: int, + output_size: int, + layer_num: int = 1, + activation: Optional[nn.Module] = nn.ReLU(), + norm_type: Optional[str] = None, + ) -> None: + """ + Overview: + Init the ``PopArtVHead`` layers according to the provided arguments. + Arguments: + - hidden_size (:obj:`int`): The ``hidden_size`` of the MLP connected to ``PopArtVHead``. + - output_size (:obj:`int`): The number of outputs. + - layer_num (:obj:`int`): The number of layers used in the network to compute Q value output. + - activation (:obj:`nn.Module`): The type of activation function to use in MLP. \ + If ``None``, then default set activation to ``nn.ReLU()``. Default ``None``. + - norm_type (:obj:`str`): The type of normalization to use. See ``ding.torch_utils.network.fc_block`` \ + for more details. Default ``None``. + """ + super(PopArtVHead, self).__init__() + self.popart = PopArt(hidden_size, output_size) + self.Q = nn.Sequential( + MLP( + hidden_size, + hidden_size, + hidden_size, + layer_num, + layer_fn=nn.Linear, + activation=activation, + norm_type=norm_type + ), self.popart + ) + + def forward(self, x: torch.Tensor) -> Dict: + """ Overview: - The ``MultiHead`` used to output actions logit. \ - Input is a (:obj:`torch.Tensor`) of shape ``(B, N)`` and returns a (:obj:`Dict`) containing \ - output ``logit``. - Interfaces: - ``__init__``, ``forward``. + Use encoded embedding tensor to run MLP with ``PopArtVHead`` and return the normalized prediction and \ + the unnormalized prediction dictionary. + Arguments: + - x (:obj:`torch.Tensor`): Tensor containing input embedding. + Returns: + - outputs (:obj:`Dict`): Dict containing keyword ``pred`` (:obj:`torch.Tensor`) \ + and ``unnormalized_pred`` (:obj:`torch.Tensor`). + Shapes: + - x: :math:`(B, N)`, where ``B = batch_size`` and ``N = hidden_size``. + - logit: :math:`(B, M)`, where ``M = output_size``. + Examples: + >>> head = PopArtVHead(64, 64) + >>> inputs = torch.randn(4, 64) + >>> outputs = head(inputs) + >>> assert isinstance(outputs, dict) and outputs['pred'].shape == torch.Size([4, 64]) and \ + outputs['unnormalized_pred'].shape == torch.Size([4, 64]) + """ + x = self.Q(x) + return x + + +class AttentionPolicyHead(nn.Module): + """ + Overview: + Cross-attention-type discrete action policy head, which is often used in variable discrete action space. + Interfaces: + ``__init__``, ``forward``. + """ + + def __init__(self) -> None: + super(AttentionPolicyHead, self).__init__() + + def forward(self, key: torch.Tensor, query: torch.Tensor) -> torch.Tensor: + """ + Overview: + Use attention-like mechanism to combine key and query tensor to output discrete action logit. + Arguments: + - key (:obj:`torch.Tensor`): Tensor containing key embedding. + - query (:obj:`torch.Tensor`): Tensor containing query embedding. + Returns: + - logit (:obj:`torch.Tensor`): Tensor containing output discrete action logit. + Shapes: + - key: :math:`(B, N, K)`, where ``B = batch_size``, ``N = possible discrete action choices`` and \ + ``K = hidden_size``. + - query: :math:`(B, K)`. + - logit: :math:`(B, N)`. + Examples: + >>> head = AttentionPolicyHead() + >>> key = torch.randn(4, 5, 64) + >>> query = torch.randn(4, 64) + >>> logit = head(key, query) + >>> assert logit.shape == torch.Size([4, 5]) + + .. note:: + In this head, we assume that the ``key`` and ``query`` tensor are both normalized. + """ + if len(query.shape) == 2 and len(key.shape) == 3: + query = query.unsqueeze(1) + logit = (key * query).sum(-1) + return logit + + +class MultiHead(nn.Module): + """ + Overview: + The ``MultiHead`` is used to generate multiple similar results. + For example, we can combine ``Distribution`` and ``MultiHead`` to generate multi-discrete action space logit. + Interfaces: + ``__init__``, ``forward``. """ def __init__(self, head_cls: type, hidden_size: int, output_size_list: SequenceType, **head_kwargs) -> None: @@ -1098,7 +1349,6 @@ def forward(self, x: torch.Tensor) -> Dict: Shapes: - x: :math:`(B, N)`, where ``B = batch_size`` and ``N = hidden_size``. - logit: :math:`(B, Mi)`, where ``Mi = output_size`` corresponding to output ``i``. - Examples: >>> head = MultiHead(DuelingHead, 64, [2, 3, 5], v_layer_num=2) >>> inputs = torch.randn(4, 64) @@ -1116,6 +1366,103 @@ def forward(self, x: torch.Tensor) -> Dict: return lists_to_dicts([m(x) for m in self.pred]) +class EnsembleHead(nn.Module): + """ + Overview: + The ``EnsembleHead`` is used to generate Q-value for Q-ensemble in model-based RL algorithms. + Interfaces: + ``__init__``, ``forward``. + """ + + def __init__( + self, + input_size: int, + output_size: int, + hidden_size: int, + layer_num: int, + ensemble_num: int, + activation: Optional[nn.Module] = nn.ReLU(), + norm_type: Optional[str] = None + ) -> None: + super(EnsembleHead, self).__init__() + d = input_size + layers = [] + for _ in range(layer_num): + layers.append( + conv1d_block( + d * ensemble_num, + hidden_size * ensemble_num, + kernel_size=1, + stride=1, + groups=ensemble_num, + activation=activation, + norm_type=norm_type + ) + ) + d = hidden_size + + # Adding activation for last layer will lead to train fail + layers.append( + conv1d_block( + hidden_size * ensemble_num, + output_size * ensemble_num, + kernel_size=1, + stride=1, + groups=ensemble_num, + activation=None, + norm_type=None + ) + ) + self.pred = nn.Sequential(*layers) + + def forward(self, x: torch.Tensor) -> Dict: + """ + Overview: + Use encoded embedding tensor to run MLP with ``EnsembleHead`` and return the prediction dictionary. + Arguments: + - x (:obj:`torch.Tensor`): Tensor containing input embedding. + Returns: + - outputs (:obj:`Dict`): Dict containing keyword ``pred`` (:obj:`torch.Tensor`). + Shapes: + - x: :math:`(B, N * ensemble_num, 1)`, where ``B = batch_size`` and ``N = hidden_size``. + - pred: :math:`(B, M * ensemble_num, 1)`, where ``M = output_size``. + Examples: + >>> head = EnsembleHead(64 * 10, 64 * 10) + >>> inputs = torch.randn(4, 64 * 10, 1) ` + >>> outputs = head(inputs) + >>> assert isinstance(outputs, dict) + >>> assert outputs['pred'].shape == torch.Size([10, 64 * 10]) + """ + x = self.pred(x).squeeze(-1) + return {'pred': x} + + +def independent_normal_dist(logits: Union[List, Dict]) -> torch.distributions.Distribution: + """ + Overview: + Convert different types logit to independent normal distribution. + Arguments: + - logits (:obj:`Union[List, Dict]`): The logits to be converted. + Returns: + - dist (:obj:`torch.distributions.Distribution`): The converted normal distribution. + Examples: + >>> logits = [torch.randn(4, 5), torch.ones(4, 5)] + >>> dist = independent_normal_dist(logits) + >>> assert isinstance(dist, torch.distributions.Independent) + >>> assert isinstance(dist.base_dist, torch.distributions.Normal) + >>> assert dist.base_dist.loc.shape == torch.Size([4, 5]) + >>> assert dist.base_dist.scale.shape == torch.Size([4, 5]) + Raises: + - TypeError: If the type of logits is not ``list`` or ``dict``. + """ + if isinstance(logits, (list, tuple)): + return Independent(Normal(*logits), 1) + elif isinstance(logits, dict): + return Independent(Normal(logits['mu'], logits['sigma']), 1) + else: + raise TypeError("invalid logits type: {}".format(type(logits))) + + head_cls_map = { # discrete 'discrete': DiscreteHead, @@ -1125,9 +1472,15 @@ def forward(self, x: torch.Tensor) -> Dict: 'rainbow': RainbowHead, 'qrdqn': QRDQNHead, 'quantile': QuantileHead, + 'fqf': FQFHead, + 'branch': BranchingHead, + 'attention_policy': AttentionPolicyHead, # continuous 'regression': RegressionHead, 'reparameterization': ReparameterizationHead, + 'popart': PopArtVHead, + 'sdn': StochasticDuelingHead, # multi 'multi': MultiHead, + 'ensemble': EnsembleHead, } diff --git a/ding/model/common/tests/test_encoder.py b/ding/model/common/tests/test_encoder.py index 26fae7b22c..d73f10561a 100644 --- a/ding/model/common/tests/test_encoder.py +++ b/ding/model/common/tests/test_encoder.py @@ -2,7 +2,7 @@ import numpy as np import pytest -from ding.model import ConvEncoder, FCEncoder, IMPALAConvEncoder +from ding.model import ConvEncoder, FCEncoder, IMPALAConvEncoder, GaussianFourierProjectionTimeEncoder from ding.torch_utils import is_differentiable B = 4 @@ -24,6 +24,20 @@ def test_conv_encoder(self): self.output_check(model, outputs) assert outputs.shape == (B, 128) + def test_dreamer_conv_encoder(self): + inputs = torch.randn(B, C, H, W) + model = ConvEncoder( + (C, H, W), + hidden_size_list=[32, 64, 128, 256, 128], + activation=torch.nn.SiLU(), + kernel_size=[4, 4, 4, 4], + layer_norm=True + ) + print(model) + outputs = model(inputs) + self.output_check(model, outputs) + assert outputs.shape == (B, 128) + def test_fc_encoder(self): inputs = torch.randn(B, 32) hidden_size_list = [128 for _ in range(3)] @@ -47,3 +61,10 @@ def test_impalaconv_encoder(self): outputs = model(inputs) self.output_check(model, outputs) assert outputs.shape == (B, 256) + + def test_GaussianFourierProjectionTimeEncoder(self): + inputs = torch.randn(B) + model = GaussianFourierProjectionTimeEncoder(128) + print(model) + outputs = model(inputs) + assert outputs.shape == (B, 128) diff --git a/ding/model/common/tests/test_head.py b/ding/model/common/tests/test_head.py index 3c36640047..044ff75cbd 100644 --- a/ding/model/common/tests/test_head.py +++ b/ding/model/common/tests/test_head.py @@ -2,7 +2,7 @@ import numpy as np import pytest -from ding.model.common.head import DuelingHead, ReparameterizationHead, MultiHead, StochasticDuelingHead +from ding.model.common.head import DuelingHead, ReparameterizationHead, MultiHead, StochasticDuelingHead, EnsembleHead from ding.torch_utils import is_differentiable B = 4 @@ -84,3 +84,10 @@ def test_stochastic_dueling(self): assert isinstance(sigma.grad, torch.Tensor) assert outputs['q_value'].shape == (B, 1) assert outputs['v_value'].shape == (B, 1) + + def test_ensemble(self): + inputs = torch.randn(B, embedding_dim * 3, 1) + model = EnsembleHead(embedding_dim, action_shape, 3, 3, 3) + outputs = model(inputs)['pred'] + self.output_check(model, outputs) + assert outputs.shape == (B, action_shape * 3) diff --git a/ding/model/common/utils.py b/ding/model/common/utils.py index 494481fa8b..f74a179962 100644 --- a/ding/model/common/utils.py +++ b/ding/model/common/utils.py @@ -1,3 +1,4 @@ +import copy import torch from easydict import EasyDict from ding.utils import import_module, MODEL_REGISTRY @@ -6,13 +7,25 @@ def create_model(cfg: EasyDict) -> torch.nn.Module: """ Overview: - Create a model given a config dictionary. + Create a neural network model according to the given EasyDict-type ``cfg``. Arguments: - - cfg: (:obj:`dict`): Training configuration. The key ``import_name`` is \ - used to import modules, and they key ``type`` is used to build the model. + - cfg: (:obj:`EasyDict`): User's model config. The key ``import_name`` is \ + used to import modules, and they key ``type`` is used to indicate the model. Returns: - - (:obj:`torch.nn.Module`) Training configuration corresponding model. + - (:obj:`torch.nn.Module`): The created neural network model. + Examples: + >>> cfg = EasyDict({ + >>> 'import_names': ['ding.model.template.q_learning'], + >>> 'type': 'dqn', + >>> 'obs_shape': 4, + >>> 'action_shape': 2, + >>> }) + >>> model = create_model(cfg) + + .. tip:: + This method will not modify the ``cfg`` , it will deepcopy the ``cfg`` and then modify it. """ + cfg = copy.deepcopy(cfg) import_module(cfg.pop('import_names', [])) - # must use pop + # here we must use the pop opeartion to ensure compatibility return MODEL_REGISTRY.build(cfg.pop("type"), **cfg) diff --git a/ding/model/template/__init__.py b/ding/model/template/__init__.py old mode 100644 new mode 100755 index 00a5a9099a..d9b05b59dd --- a/ding/model/template/__init__.py +++ b/ding/model/template/__init__.py @@ -1,10 +1,12 @@ # general -from .q_learning import DQN, RainbowDQN, QRDQN, IQN, FQF, DRQN, C51DQN -from .qac import QAC, DiscreteQAC +from .q_learning import DQN, RainbowDQN, QRDQN, IQN, FQF, DRQN, C51DQN, BDQ, GTrXLDQN +from .qac import DiscreteQAC, ContinuousQAC from .pdqn import PDQN -from .vac import VAC +from .vac import VAC, DREAMERVAC from .bc import DiscreteBC, ContinuousBC +from .language_transformer import LanguageTransformer # algorithm-specific +from .pg import PG from .ppg import PPG from .qmix import Mixer, QMix from .collaq import CollaQ @@ -17,6 +19,14 @@ from .mavac import MAVAC from .ngu import NGU from .qac_dist import QACDIST -from .maqac import MAQAC, ContinuousMAQAC +from .maqac import DiscreteMAQAC, ContinuousMAQAC +from .madqn import MADQN from .vae import VanillaVAE from .decision_transformer import DecisionTransformer +from .procedure_cloning import ProcedureCloningMCTS, ProcedureCloningBFS +from .bcq import BCQ +from .edac import EDAC +from .hpt import HPT +from .qgpo import QGPO +from .ebm import EBM, AutoregressiveEBM +from .havac import HAVAC diff --git a/ding/model/template/acer.py b/ding/model/template/acer.py index 2e28ef0b2c..44bb386cba 100644 --- a/ding/model/template/acer.py +++ b/ding/model/template/acer.py @@ -9,9 +9,11 @@ @MODEL_REGISTRY.register('acer') class ACER(nn.Module): - r""" + """ Overview: - The ACER model. + The model of algorithmn ACER(Actor Critic with Experience Replay) + Sample Efficient Actor-Critic with Experience Replay. + https://arxiv.org/abs/1611.01224 Interfaces: ``__init__``, ``forward``, ``compute_actor``, ``compute_critic`` """ @@ -29,7 +31,7 @@ def __init__( activation: Optional[nn.Module] = nn.ReLU(), norm_type: Optional[str] = None, ) -> None: - r""" + """ Overview: Init the ACER Model according to arguments. Arguments: @@ -78,64 +80,29 @@ def __init__( self.critic = nn.ModuleList(self.critic) def forward(self, inputs: Union[torch.Tensor, Dict], mode: str) -> Dict: - r""" + """ Overview: - Use observation to predict output. - Parameter updates with ACER's MLPs forward setup. + Use observation to predict output. + Parameter updates with ACER's MLPs forward setup. Arguments: - Forward with ``'compute_actor'``: - - inputs (:obj:`torch.Tensor`): - The encoded embedding tensor, determined with given ``hidden_size``, i.e. ``(B, N=hidden_size)``. - Whether ``actor_head_hidden_size`` or ``critic_head_hidden_size`` depend on ``mode``. - - Forward with ``'compute_critic'``, inputs:`torch.Tensor` Necessary Keys: - - ``obs`` encoded tensors. - - mode (:obj:`str`): Name of the forward mode. Returns: - outputs (:obj:`Dict`): Outputs of network forward. - - Forward with ``'compute_actor'``, Necessary Keys (either): - - logit (:obj:`torch.Tensor`): - - logit (:obj:`torch.Tensor`): Logit encoding tensor. - - Forward with ``'compute_critic'``, Necessary Keys: - - q_value (:obj:`torch.Tensor`): Q value tensor. - - Actor Shapes: + Shapes (Actor): - obs (:obj:`torch.Tensor`): :math:`(B, N1)`, where B is batch size and N1 is ``obs_shape`` - logit (:obj:`torch.FloatTensor`): :math:`(B, N2)`, where B is batch size and N2 is ``action_shape`` - - Critic Shapes: + Shapes (Critic): - inputs (:obj:`torch.Tensor`): :math:`(B, N1)`, B is batch size and N1 corresponds to ``obs_shape`` - q_value (:obj:`torch.FloatTensor`): :math:`(B, N2)`, where B is batch size and N2 is ``action_shape`` - Actor Examples: - >>> # Regression mode - >>> model = ACER(64, 64) - >>> inputs = torch.randn(4, 64) - >>> actor_outputs = model(inputs,'compute_actor') - >>> assert actor_outputs['logit'].shape == torch.Size([4, 64]) - - Critic Examples: - >>> inputs = torch.randn(4,N) - >>> model = ACER(obs_shape=(N, ),action_shape=5) - >>> model(inputs, mode='compute_critic')['q_value'] # q value - tensor([[-0.0681, -0.0431, -0.0530, 0.1454, -0.1093], - [-0.0647, -0.0281, -0.0527, 0.1409, -0.1162], - [-0.0596, -0.0321, -0.0676, 0.1386, -0.1113], - [-0.0874, -0.0406, -0.0487, 0.1346, -0.1135]], - grad_fn=) - - """ assert mode in self.mode, "not support forward mode: {}/{}".format(mode, self.mode) return getattr(self, mode)(inputs) def compute_actor(self, inputs: torch.Tensor) -> Dict: - r""" + """ Overview: Use encoded embedding tensor to predict output. - Execute parameter updates with ``'compute_actor'`` mode + Execute parameter updates with ``compute_actor`` mode Use encoded embedding tensor to predict output. Arguments: - inputs (:obj:`torch.Tensor`): @@ -144,7 +111,6 @@ def compute_actor(self, inputs: torch.Tensor) -> Dict: - mode (:obj:`str`): Name of the forward mode. Returns: - outputs (:obj:`Dict`): Outputs of forward pass encoder and head. - ReturnsKeys (either): - logit (:obj:`torch.FloatTensor`): :math:`(B, N1)`, where B is batch size and N1 is ``action_shape`` Shapes: @@ -163,31 +129,24 @@ def compute_actor(self, inputs: torch.Tensor) -> Dict: return x def compute_critic(self, inputs: torch.Tensor) -> Dict: - r""" + """ Overview: - Execute parameter updates with ``'compute_critic'`` mode + Execute parameter updates with ``compute_critic`` mode Use encoded embedding tensor to predict output. Arguments: - ``obs``, ``action`` encoded tensors. - mode (:obj:`str`): Name of the forward mode. Returns: - outputs (:obj:`Dict`): Q-value output. - ReturnKeys: - q_value (:obj:`torch.Tensor`): Q value tensor with same size as batch size. Shapes: - obs (:obj:`torch.Tensor`): :math:`(B, N1)`, where B is batch size and N1 is ``obs_shape`` - q_value (:obj:`torch.FloatTensor`): :math:`(B, N2)`, where B is batch size and N2 is ``action_shape``. - Examples: >>> inputs =torch.randn(4, N) >>> model = ACER(obs_shape=(N, ),action_shape=5) - >>> model(inputs, mode='compute_critic')['q_value'] # q value - tensor([[-0.0681, -0.0431, -0.0530, 0.1454, -0.1093], - [-0.0647, -0.0281, -0.0527, 0.1409, -0.1162], - [-0.0596, -0.0321, -0.0676, 0.1386, -0.1113], - [-0.0874, -0.0406, -0.0487, 0.1346, -0.1135]], - grad_fn=) + >>> model(inputs, mode='compute_critic')['q_value'] """ obs = inputs diff --git a/ding/model/template/atoc.py b/ding/model/template/atoc.py index f0863481c2..a06f536aef 100644 --- a/ding/model/template/atoc.py +++ b/ding/model/template/atoc.py @@ -5,28 +5,24 @@ from ding.utils import squeeze, MODEL_REGISTRY, SequenceType from ding.torch_utils import MLP -from ..common import RegressionHead +from ding.model.common import RegressionHead class ATOCAttentionUnit(nn.Module): - r""" + """ Overview: - the attention unit of the atoc network. We now implement it as two-layer MLP, same as the original paper - + The attention unit of the ATOC network. We now implement it as two-layer MLP, same as the original paper. Interface: - __init__, forward + ``__init__``, ``forward`` .. note:: - "ATOC paper: We use two-layer MLP to implement the attention unit but it is also can be realized by RNN." - """ def __init__(self, thought_size: int, embedding_size: int) -> None: - r""" + """ Overview: - init the attention unit according to the size of input args - + Initialize the attention unit according to the size of input arguments. Arguments: - thought_size (:obj:`int`): the size of input thought - embedding_size (:obj:`int`): the size of hidden layers @@ -42,15 +38,19 @@ def __init__(self, thought_size: int, embedding_size: int) -> None: self._act2 = nn.Sigmoid() def forward(self, data: Union[Dict, torch.Tensor]) -> torch.Tensor: - r""" + """ Overview: - forward method take the thought of agents as input and output the prob of these agent\ - being initiator - + Take the thought of agents as input and generate the probability of these agent being initiator Arguments: - x (:obj:`Union[Dict, torch.Tensor`): the input tensor or dict contain the thoughts tensor - - ret (:obj:`torch.Tensor`): the output initiator prob - + - ret (:obj:`torch.Tensor`): the output initiator probability + Shapes: + - data['thought']: :math:`(M, B, N)`, M is the num of thoughts to integrate,\ + B is batch_size and N is thought size + Examples: + >>> attention_unit = ATOCAttentionUnit(64, 64) + >>> thought = torch.randn(2, 3, 64) + >>> attention_unit(thought) """ x = data if isinstance(data, Dict): @@ -61,24 +61,21 @@ def forward(self, data: Union[Dict, torch.Tensor]) -> torch.Tensor: x = self._act1(x) x = self._fc3(x) x = self._act2(x) - # return {'initiator': x} return x.squeeze(-1) class ATOCCommunicationNet(nn.Module): - r""" + """ Overview: - atoc commnication net is a bi-direction LSTM, so it can integrate all the thoughts in the group - + This ATOC commnication net is a bi-direction LSTM, so it can integrate all the thoughts in the group. Interface: - __init__, forward + ``__init__``, ``forward`` """ def __init__(self, thought_size: int) -> None: - r""" + """ Overview: - init method of the communication network - + Initialize the communication network according to the size of input arguments. Arguments: - thought_size (:obj:`int`): the size of input thought @@ -93,32 +90,34 @@ def __init__(self, thought_size: int) -> None: self._bi_lstm = nn.LSTM(self._thought_size, self._comm_hidden_size, bidirectional=True) def forward(self, data: Union[Dict, torch.Tensor]): - r""" + """ Overview: - the forward method that integrate thoughts + The forward of ATOCCommunicationNet integrates thoughts in the group. Arguments: - x (:obj:`Union[Dict, torch.Tensor`): the input tensor or dict contain the thoughts tensor - out (:obj:`torch.Tensor`): the integrated thoughts Shapes: - data['thoughts']: :math:`(M, B, N)`, M is the num of thoughts to integrate,\ B is batch_size and N is thought size + Examples: + >>> comm_net = ATOCCommunicationNet(64) + >>> thoughts = torch.randn(2, 3, 64) + >>> comm_net(thoughts) """ self._bi_lstm.flatten_parameters() x = data if isinstance(data, Dict): x = data['thoughts'] out, _ = self._bi_lstm(x) - # return {'thoughts': out} return out class ATOCActorNet(nn.Module): - r""" + """ Overview: - the overall ATOC actor network - + The actor network of ATOC. Interface: - __init__, forward + ``__init__``, ``forward`` .. note:: "ATOC paper: The neural networks use ReLU and batch normalization for some hidden layers." @@ -139,10 +138,9 @@ def __init__( activation: Optional[nn.Module] = nn.ReLU(), norm_type: Optional[str] = None, ): - r""" + """ Overview: - the init method of atoc actor network - + Initialize the actor network of ATOC Arguments: - obs_shape(:obj:`Union[Tuple, int]`): the observation size - thought_size (:obj:`int`): the size of thoughts @@ -194,11 +192,9 @@ def __init__( self.comm_net = ATOCCommunicationNet(self._thought_size) def forward(self, obs: torch.Tensor) -> Dict: - r""" + """ Overview: - the forward method of actor network, take the input obs, and calculate the corresponding action, group, \ - initiator_prob, thoughts, etc... - + Take the input obs, and calculate the corresponding action, group, initiator_prob, thoughts, etc... Arguments: - obs (:obj:`Dict`): the input obs containing the observation Returns: @@ -207,6 +203,18 @@ def forward(self, obs: torch.Tensor) -> Dict: ReturnsKeys: - necessary: ``action`` - optional: ``group``, ``initiator_prob``, ``is_initiator``, ``new_thoughts``, ``old_thoughts`` + Shapes: + - obs (:obj:`torch.Tensor`): :math:`(B, A, N)`, where B is batch size, A is agent num, N is obs size + - action (:obj:`torch.Tensor`): :math:`(B, A, M)`, where M is action size + - group (:obj:`torch.Tensor`): :math:`(B, A, A)` + - initiator_prob (:obj:`torch.Tensor`): :math:`(B, A)` + - is_initiator (:obj:`torch.Tensor`): :math:`(B, A)` + - new_thoughts (:obj:`torch.Tensor`): :math:`(B, A, M)` + - old_thoughts (:obj:`torch.Tensor`): :math:`(B, A, M)` + Examples: + >>> actor_net = ATOCActorNet(64, 64, 64, 3) + >>> obs = torch.randn(2, 3, 64) + >>> actor_net(obs) """ assert len(obs.shape) == 3 self._cur_batch_size = obs.shape[0] @@ -238,6 +246,25 @@ def forward(self, obs: torch.Tensor) -> Dict: return {'action': action} def _get_initiate_group(self, current_thoughts): + """ + Overview: + Calculate the initiator probability, group and is_initiator + Arguments: + - current_thoughts (:obj:`torch.Tensor`): tensor of current thoughts + Returns: + - init_prob (:obj:`torch.Tensor`): tesnor of initiator probability + - is_initiator (:obj:`torch.Tensor`): tensor of is initiator + - group (:obj:`torch.Tensor`): tensor of group + Shapes: + - current_thoughts (:obj:`torch.Tensor`): :math:`(B, A, M)`, where M is thought size + - init_prob (:obj:`torch.Tensor`): :math:`(B, A)` + - is_initiator (:obj:`torch.Tensor`): :math:`(B, A)` + - group (:obj:`torch.Tensor`): :math:`(B, A, A)` + Examples: + >>> actor_net = ATOCActorNet(64, 64, 64, 3) + >>> current_thoughts = torch.randn(2, 3, 64) + >>> actor_net._get_initiate_group(current_thoughts) + """ if not self._communication: raise NotImplementedError init_prob = self.attention(current_thoughts) # B, A @@ -267,10 +294,25 @@ def _get_initiate_group(self, current_thoughts): def _get_new_thoughts(self, current_thoughts, group, is_initiator): """ + Overview: + Calculate the new thoughts according to current thoughts, group and is_initiator + Arguments: + - current_thoughts (:obj:`torch.Tensor`): tensor of current thoughts + - group (:obj:`torch.Tensor`): tensor of group + - is_initiator (:obj:`torch.Tensor`): tensor of is initiator + Returns: + - new_thoughts (:obj:`torch.Tensor`): tensor of new thoughts Shapes: - current_thoughts (:obj:`torch.Tensor`): :math:`(B, A, M)`, where M is thought size - group: (:obj:`torch.Tensor`): :math:`(B, A, A)` - is_initiator (:obj:`torch.Tensor`): :math:`(B, A)` + - new_thoughts (:obj:`torch.Tensor`): :math:`(B, A, M)` + Examples: + >>> actor_net = ATOCActorNet(64, 64, 64, 3) + >>> current_thoughts = torch.randn(2, 3, 64) + >>> group = torch.randn(2, 3, 3) + >>> is_initiator = torch.randn(2, 3) + >>> actor_net._get_new_thoughts(current_thoughts, group, is_initiator) """ if not self._communication: raise NotImplementedError @@ -306,12 +348,13 @@ def _get_new_thoughts(self, current_thoughts, group, is_initiator): @MODEL_REGISTRY.register('atoc') class ATOC(nn.Module): - r""" + """ Overview: The QAC network of ATOC, a kind of extension of DDPG for MARL. - + Learning Attentional Communication for Multi-Agent Cooperation + https://arxiv.org/abs/1805.07733 Interface: - __init__, forward, compute_critic, compute_actor, optimize_actor_attention + ``__init__``, ``forward``, ``compute_critic``, ``compute_actor``, ``optimize_actor_attention`` """ mode = ['compute_actor', 'compute_critic', 'optimize_actor_attention'] @@ -330,10 +373,9 @@ def __init__( activation: Optional[nn.Module] = nn.ReLU(), norm_type: Optional[str] = None, ) -> None: - r""" + """ Overview: - init the atoc QAC network - + Initialize the ATOC QAC network Arguments: - obs_shape(:obj:`Union[Tuple, int]`): the observation space shape - thought_size (:obj:`int`): the size of thoughts @@ -367,16 +409,33 @@ def __init__( ) def _compute_delta_q(self, obs: torch.Tensor, actor_outputs: Dict) -> torch.Tensor: - r""" + """ Overview: calculate the delta_q according to obs and actor_outputs - Arguments: - obs (:obj:`torch.Tensor`): the observations - actor_outputs (:obj:`dict`): the output of actors - delta_q (:obj:`Dict`): the calculated delta_q + Returns: + - delta_q (:obj:`Dict`): the calculated delta_q ArgumentsKeys: - necessary: ``new_thoughts``, ``old_thoughts``, ``group``, ``is_initiator`` + Shapes: + - obs (:obj:`torch.Tensor`): :math:`(B, A, N)`, where B is batch size, A is agent num, N is obs size + - actor_outputs (:obj:`Dict`): the output of actor network, including ``action``, ``new_thoughts``, \ + ``old_thoughts``, ``group``, ``initiator_prob``, ``is_initiator`` + - action (:obj:`torch.Tensor`): :math:`(B, A, M)` where M is action size + - new_thoughts (:obj:`torch.Tensor`): :math:`(B, A, M)` where M is thought size + - old_thoughts (:obj:`torch.Tensor`): :math:`(B, A, M)` where M is thought size + - group (:obj:`torch.Tensor`): :math:`(B, A, A)` + - initiator_prob (:obj:`torch.Tensor`): :math:`(B, A)` + - is_initiator (:obj:`torch.Tensor`): :math:`(B, A)` + - delta_q (:obj:`torch.Tensor`): :math:`(B, A)` + Examples: + >>> net = ATOC(64, 64, 64, 3) + >>> obs = torch.randn(2, 3, 64) + >>> actor_outputs = net.compute_actor(obs) + >>> net._compute_delta_q(obs, actor_outputs) """ if not self._communication: raise NotImplementedError @@ -413,10 +472,9 @@ def _compute_delta_q(self, obs: torch.Tensor, actor_outputs: Dict) -> torch.Tens return curr_delta_q def compute_actor(self, obs: torch.Tensor, get_delta_q: bool = False) -> Dict[str, torch.Tensor]: - r''' + ''' Overview: compute the action according to inputs, call the _compute_delta_q function to compute delta_q - Arguments: - obs (:obj:`torch.Tensor`): observation - get_delta_q (:obj:`bool`) : whether need to get delta_q @@ -425,7 +483,19 @@ def compute_actor(self, obs: torch.Tensor, get_delta_q: bool = False) -> Dict[st ReturnsKeys: - necessary: ``action`` - optional: ``group``, ``initiator_prob``, ``is_initiator``, ``new_thoughts``, ``old_thoughts``, ``delta_q`` - + Shapes: + - obs (:obj:`torch.Tensor`): :math:`(B, A, N)`, where B is batch size, A is agent num, N is obs size + - action (:obj:`torch.Tensor`): :math:`(B, A, M)`, where M is action size + - group (:obj:`torch.Tensor`): :math:`(B, A, A)` + - initiator_prob (:obj:`torch.Tensor`): :math:`(B, A)` + - is_initiator (:obj:`torch.Tensor`): :math:`(B, A)` + - new_thoughts (:obj:`torch.Tensor`): :math:`(B, A, M)` + - old_thoughts (:obj:`torch.Tensor`): :math:`(B, A, M)` + - delta_q (:obj:`torch.Tensor`): :math:`(B, A)` + Examples: + >>> net = ATOC(64, 64, 64, 3) + >>> obs = torch.randn(2, 3, 64) + >>> net.compute_actor(obs) ''' outputs = self.actor(obs) if get_delta_q and self._communication: @@ -435,10 +505,25 @@ def compute_actor(self, obs: torch.Tensor, get_delta_q: bool = False) -> Dict[st def compute_critic(self, inputs: Dict) -> Dict: """ + Overview: + compute the q_value according to inputs + Arguments: + - inputs (:obj:`Dict`): the inputs contain the obs and action + Returns: + - outputs (:obj:`Dict`): the output of critic network ArgumentsKeys: - necessary: ``obs``, ``action`` ReturnsKeys: - necessary: ``q_value`` + Shapes: + - obs (:obj:`torch.Tensor`): :math:`(B, A, N)`, where B is batch size, A is agent num, N is obs size + - action (:obj:`torch.Tensor`): :math:`(B, A, M)`, where M is action size + - q_value (:obj:`torch.Tensor`): :math:`(B, A)` + Examples: + >>> net = ATOC(64, 64, 64, 3) + >>> obs = torch.randn(2, 3, 64) + >>> action = torch.randn(2, 3, 64) + >>> net.compute_critic({'obs': obs, 'action': action}) """ obs, action = inputs['obs'], inputs['action'] if len(action.shape) == 2: # (B, A) -> (B, A, 1) @@ -448,14 +533,31 @@ def compute_critic(self, inputs: Dict) -> Dict: return {'q_value': x} def optimize_actor_attention(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: - r""" + """ Overview: return the actor attention loss - Arguments: - inputs (:obj:`Dict`): the inputs contain the delta_q, initiator_prob, and is_initiator Returns - loss (:obj:`Dict`): the loss of actor attention unit + ArgumentsKeys: + - necessary: ``delta_q``, ``initiator_prob``, ``is_initiator`` + ReturnsKeys: + - necessary: ``loss`` + Shapes: + - delta_q (:obj:`torch.Tensor`): :math:`(B, A)` + - initiator_prob (:obj:`torch.Tensor`): :math:`(B, A)` + - is_initiator (:obj:`torch.Tensor`): :math:`(B, A)` + - loss (:obj:`torch.Tensor`): :math:`(1)` + Examples: + >>> net = ATOC(64, 64, 64, 3) + >>> delta_q = torch.randn(2, 3) + >>> initiator_prob = torch.randn(2, 3) + >>> is_initiator = torch.randn(2, 3) + >>> net.optimize_actor_attention( + >>> {'delta_q': delta_q, + >>> 'initiator_prob': initiator_prob, + >>> 'is_initiator': is_initiator}) """ if not self._communication: raise NotImplementedError diff --git a/ding/model/template/bc.py b/ding/model/template/bc.py index 0da1f3e171..5348c750a6 100644 --- a/ding/model/template/bc.py +++ b/ding/model/template/bc.py @@ -1,9 +1,8 @@ -from typing import Union, Optional, Dict, Callable, List +from typing import Union, Optional, Dict import torch import torch.nn as nn from easydict import EasyDict -from ding.torch_utils import get_lstm from ding.utils import MODEL_REGISTRY, SequenceType, squeeze from ..common import FCEncoder, ConvEncoder, DiscreteHead, DuelingHead, \ MultiHead, RegressionHead, ReparameterizationHead @@ -11,17 +10,24 @@ @MODEL_REGISTRY.register('discrete_bc') class DiscreteBC(nn.Module): + """ + Overview: + The DiscreteBC network. + Interfaces: + ``__init__``, ``forward`` + """ def __init__( - self, - obs_shape: Union[int, SequenceType], - action_shape: Union[int, SequenceType], - encoder_hidden_size_list: SequenceType = [128, 128, 64], - dueling: bool = True, - head_hidden_size: Optional[int] = None, - head_layer_num: int = 1, - activation: Optional[nn.Module] = nn.ReLU(), - norm_type: Optional[str] = None + self, + obs_shape: Union[int, SequenceType], + action_shape: Union[int, SequenceType], + encoder_hidden_size_list: SequenceType = [128, 128, 64], + dueling: bool = True, + head_hidden_size: Optional[int] = None, + head_layer_num: int = 1, + activation: Optional[nn.Module] = nn.ReLU(), + norm_type: Optional[str] = None, + strides: Optional[list] = None, ) -> None: """ Overview: @@ -35,9 +41,11 @@ def __init__( - head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` of head network. - head_layer_num (:obj:`int`): The number of layers used in the head network to compute Q value output - activation (:obj:`Optional[nn.Module]`): The type of activation function in networks \ - if ``None`` then default set it to ``nn.ReLU()`` + if ``None`` then default set it to ``nn.ReLU()``. - norm_type (:obj:`Optional[str]`): The type of normalization in networks, see \ ``ding.torch_utils.fc_block`` for more details. + - strides (:obj:`Optional[list]`): The strides for each convolution layers, such as [2, 2, 2]. The length \ + of this argument should be the same as ``encoder_hidden_size_list``. """ super(DiscreteBC, self).__init__() # For compatibility: 1, (1, ), [4, 32, 32] @@ -49,7 +57,14 @@ def __init__( self.encoder = FCEncoder(obs_shape, encoder_hidden_size_list, activation=activation, norm_type=norm_type) # Conv Encoder elif len(obs_shape) == 3: - self.encoder = ConvEncoder(obs_shape, encoder_hidden_size_list, activation=activation, norm_type=norm_type) + if not strides: + self.encoder = ConvEncoder( + obs_shape, encoder_hidden_size_list, activation=activation, norm_type=norm_type + ) + else: + self.encoder = ConvEncoder( + obs_shape, encoder_hidden_size_list, activation=activation, norm_type=norm_type, stride=strides + ) else: raise RuntimeError( "not support obs_shape for pre-defined encoder: {}, please customize your own BC".format(obs_shape) @@ -75,7 +90,7 @@ def __init__( ) def forward(self, x: torch.Tensor) -> Dict: - r""" + """ Overview: DiscreteBC forward computation graph, input observation tensor to predict q_value. Arguments: @@ -100,7 +115,7 @@ def forward(self, x: torch.Tensor) -> Dict: @MODEL_REGISTRY.register('continuous_bc') class ContinuousBC(nn.Module): - r""" + """ Overview: The ContinuousBC network. Interfaces: @@ -119,7 +134,7 @@ def __init__( ) -> None: """ Overview: - Initailize the ContinuousBC Model according to input arguments. + Initialize the ContinuousBC Model according to input arguments. Arguments: - obs_shape (:obj:`Union[int, SequenceType]`): Observation's shape, such as 128, (156, ). - action_shape (:obj:`Union[int, SequenceType, EasyDict]`): Action's shape, such as 4, (3, ), \ @@ -165,14 +180,34 @@ def __init__( ) ) - def forward(self, inputs: Union[torch.Tensor, Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]: + def forward(self, inputs: Union[torch.Tensor, Dict[str, torch.Tensor]]) -> Dict: """ Overview: - The unique execution (forward) method of ContinuousBC method. - Arguments: - - inputs (:obj:`torch.Tensor`): Observation data, defaults to tensor. - Returns: - - output (:obj:`Dict`): Output dict data, including differnet key-values among distinct action_space. + The unique execution (forward) method of ContinuousBC. + Arguments: + - inputs (:obj:`torch.Tensor`): Observation data, defaults to tensor. + Returns: + - output (:obj:`Dict`): Output dict data, including different key-values among distinct action_space. + ReturnsKeys: + - action (:obj:`torch.Tensor`): action output of actor network, \ + with shape :math:`(B, action_shape)`. + - logit (:obj:`List[torch.Tensor]`): reparameterized action output of actor network, \ + with shape :math:`(B, action_shape)`. + Shapes: + - inputs (:obj:`torch.Tensor`): :math:`(B, N)`, where B is batch size and N is ``obs_shape`` + - action (:obj:`torch.FloatTensor`): :math:`(B, M)`, where B is batch size and M is ``action_shape`` + - logit (:obj:`List[torch.FloatTensor]`): :math:`(B, M)`, where B is batch size and M is ``action_shape`` + Examples (Regression): + >>> model = ContinuousBC(32, 6, action_space='regression') + >>> inputs = torch.randn(4, 32) + >>> outputs = model(inputs) + >>> assert isinstance(outputs, dict) and outputs['action'].shape == torch.Size([4, 6]) + Examples (Reparameterization): + >>> model = ContinuousBC(32, 6, action_space='reparameterization') + >>> inputs = torch.randn(4, 32) + >>> outputs = model(inputs) + >>> assert isinstance(outputs, dict) and outputs['logit'][0].shape == torch.Size([4, 6]) + >>> assert outputs['logit'][1].shape == torch.Size([4, 6]) """ if self.action_space == 'regression': x = self.actor(inputs) diff --git a/ding/model/template/bcq.py b/ding/model/template/bcq.py new file mode 100755 index 0000000000..0e72927a76 --- /dev/null +++ b/ding/model/template/bcq.py @@ -0,0 +1,210 @@ +from typing import Union, Dict, Optional, List +from easydict import EasyDict +import numpy as np +import torch +import torch.nn as nn + +from ding.utils import SequenceType, squeeze, MODEL_REGISTRY +from ..common import RegressionHead, ReparameterizationHead +from .vae import VanillaVAE + + +@MODEL_REGISTRY.register('bcq') +class BCQ(nn.Module): + """ + Overview: + Model of BCQ (Batch-Constrained deep Q-learning). + Off-Policy Deep Reinforcement Learning without Exploration. + https://arxiv.org/abs/1812.02900 + Interface: + ``forward``, ``compute_actor``, ``compute_critic``, ``compute_vae``, ``compute_eval`` + Property: + ``mode`` + """ + + mode = ['compute_actor', 'compute_critic', 'compute_vae', 'compute_eval'] + + def __init__( + self, + obs_shape: Union[int, SequenceType], + action_shape: Union[int, SequenceType, EasyDict], + actor_head_hidden_size: List = [400, 300], + critic_head_hidden_size: List = [400, 300], + activation: Optional[nn.Module] = nn.ReLU(), + vae_hidden_dims: List = [750, 750], + phi: float = 0.05 + ) -> None: + """ + Overview: + Initialize neural network, i.e. agent Q network and actor. + Arguments: + - obs_shape (:obj:`int`): the dimension of observation state + - action_shape (:obj:`int`): the dimension of action shape + - actor_hidden_size (:obj:`list`): the list of hidden size of actor + - critic_hidden_size (:obj:'list'): the list of hidden size of critic + - activation (:obj:`nn.Module`): Activation function in network, defaults to nn.ReLU(). + - vae_hidden_dims (:obj:`list`): the list of hidden size of vae + """ + super(BCQ, self).__init__() + obs_shape: int = squeeze(obs_shape) + action_shape = squeeze(action_shape) + self.action_shape = action_shape + self.input_size = obs_shape + self.phi = phi + + critic_input_size = self.input_size + action_shape + self.critic = nn.ModuleList() + for _ in range(2): + net = [] + d = critic_input_size + for dim in critic_head_hidden_size: + net.append(nn.Linear(d, dim)) + net.append(activation) + d = dim + net.append(nn.Linear(d, 1)) + self.critic.append(nn.Sequential(*net)) + + net = [] + d = critic_input_size + for dim in actor_head_hidden_size: + net.append(nn.Linear(d, dim)) + net.append(activation) + d = dim + net.append(nn.Linear(d, 1)) + self.actor = nn.Sequential(*net) + + self.vae = VanillaVAE(action_shape, obs_shape, action_shape * 2, vae_hidden_dims) + + def forward(self, inputs: Dict[str, torch.Tensor], mode: str) -> Dict[str, torch.Tensor]: + """ + Overview: + The unique execution (forward) method of BCQ method, and one can indicate different modes to implement \ + different computation graph, including ``compute_actor`` and ``compute_critic`` in BCQ. + Mode compute_actor: + Arguments: + - inputs (:obj:`Dict`): Input dict data, including obs and action tensor. + Returns: + - output (:obj:`Dict`): Output dict data, including action tensor. + Mode compute_critic: + Arguments: + - inputs (:obj:`Dict`): Input dict data, including obs and action tensor. + Returns: + - output (:obj:`Dict`): Output dict data, including q_value tensor. + Mode compute_vae: + Arguments: + - inputs (:obj:`Dict`): Input dict data, including obs and action tensor. + Returns: + - outputs (:obj:`Dict`): Dict containing keywords ``recons_action`` \ + (:obj:`torch.Tensor`), ``prediction_residual`` (:obj:`torch.Tensor`), \ + ``input`` (:obj:`torch.Tensor`), ``mu`` (:obj:`torch.Tensor`), \ + ``log_var`` (:obj:`torch.Tensor`) and ``z`` (:obj:`torch.Tensor`). + Mode compute_eval: + Arguments: + - inputs (:obj:`Dict`): Input dict data, including obs and action tensor. + Returns: + - output (:obj:`Dict`): Output dict data, including action tensor. + Examples: + >>> inputs = {'obs': torch.randn(4, 32), 'action': torch.randn(4, 6)} + >>> model = BCQ(32, 6) + >>> outputs = model(inputs, mode='compute_actor') + >>> outputs = model(inputs, mode='compute_critic') + >>> outputs = model(inputs, mode='compute_vae') + >>> outputs = model(inputs, mode='compute_eval') + + .. note:: + For specific examples, one can refer to API doc of ``compute_actor`` and ``compute_critic`` respectively. + """ + assert mode in self.mode, "not support forward mode: {}/{}".format(mode, self.mode) + return getattr(self, mode)(inputs) + + def compute_critic(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: + """ + Overview: + Use critic network to compute q value. + Arguments: + - inputs (:obj:`Dict`): Input dict data, including obs and action tensor. + Returns: + - outputs (:obj:`Dict`): Dict containing keywords ``q_value`` (:obj:`torch.Tensor`). + Shapes: + - inputs (:obj:`Dict`): :math:`(B, N, D)`, where B is batch size, N is sample number, D is input dimension. + - outputs (:obj:`Dict`): :math:`(B, N)`. + Examples: + >>> inputs = {'obs': torch.randn(4, 32), 'action': torch.randn(4, 6)} + >>> model = BCQ(32, 6) + >>> outputs = model.compute_critic(inputs) + """ + obs, action = inputs['obs'], inputs['action'] + if len(action.shape) == 1: # (B, ) -> (B, 1) + action = action.unsqueeze(1) + x = torch.cat([obs, action], dim=-1) + x = [m(x).squeeze() for m in self.critic] + return {'q_value': x} + + def compute_actor(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, Union[torch.Tensor, Dict[str, torch.Tensor]]]: + """ + Overview: + Use actor network to compute action. + Arguments: + - inputs (:obj:`Dict`): Input dict data, including obs and action tensor. + Returns: + - outputs (:obj:`Dict`): Dict containing keywords ``action`` (:obj:`torch.Tensor`). + Shapes: + - inputs (:obj:`Dict`): :math:`(B, N, D)`, where B is batch size, N is sample number, D is input dimension. + - outputs (:obj:`Dict`): :math:`(B, N)`. + Examples: + >>> inputs = {'obs': torch.randn(4, 32), 'action': torch.randn(4, 6)} + >>> model = BCQ(32, 6) + >>> outputs = model.compute_actor(inputs) + """ + input = torch.cat([inputs['obs'], inputs['action']], -1) + x = self.actor(input) + action = self.phi * 1 * torch.tanh(x) + action = (action + inputs['action']).clamp(-1, 1) + return {'action': action} + + def compute_vae(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: + """ + Overview: + Use vae network to compute action. + Arguments: + - inputs (:obj:`Dict`): Input dict data, including obs and action tensor. + Returns: + - outputs (:obj:`Dict`): Dict containing keywords ``recons_action`` (:obj:`torch.Tensor`), \ + ``prediction_residual`` (:obj:`torch.Tensor`), ``input`` (:obj:`torch.Tensor`), \ + ``mu`` (:obj:`torch.Tensor`), ``log_var`` (:obj:`torch.Tensor`) and ``z`` (:obj:`torch.Tensor`). + Shapes: + - inputs (:obj:`Dict`): :math:`(B, N, D)`, where B is batch size, N is sample number, D is input dimension. + - outputs (:obj:`Dict`): :math:`(B, N)`. + Examples: + >>> inputs = {'obs': torch.randn(4, 32), 'action': torch.randn(4, 6)} + >>> model = BCQ(32, 6) + >>> outputs = model.compute_vae(inputs) + """ + return self.vae.forward(inputs) + + def compute_eval(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: + """ + Overview: + Use actor network to compute action. + Arguments: + - inputs (:obj:`Dict`): Input dict data, including obs and action tensor. + Returns: + - outputs (:obj:`Dict`): Dict containing keywords ``action`` (:obj:`torch.Tensor`). + Shapes: + - inputs (:obj:`Dict`): :math:`(B, N, D)`, where B is batch size, N is sample number, D is input dimension. + - outputs (:obj:`Dict`): :math:`(B, N)`. + Examples: + >>> inputs = {'obs': torch.randn(4, 32), 'action': torch.randn(4, 6)} + >>> model = BCQ(32, 6) + >>> outputs = model.compute_eval(inputs) + """ + obs = inputs['obs'] + obs_rep = obs.clone().unsqueeze(0).repeat_interleave(100, dim=0) + z = torch.randn((obs_rep.shape[0], obs_rep.shape[1], self.action_shape * 2)).to(obs.device).clamp(-0.5, 0.5) + sample_action = self.vae.decode_with_obs(z, obs_rep)['reconstruction_action'] + action = self.compute_actor({'obs': obs_rep, 'action': sample_action})['action'] + q = self.compute_critic({'obs': obs_rep, 'action': action})['q_value'][0] + idx = q.argmax(dim=0).unsqueeze(0).unsqueeze(-1) + idx = idx.repeat_interleave(action.shape[-1], dim=-1) + action = action.gather(0, idx).squeeze() + return {'action': action} diff --git a/ding/model/template/collaq.py b/ding/model/template/collaq.py index dd8db59986..3a1e4fec08 100644 --- a/ding/model/template/collaq.py +++ b/ding/model/template/collaq.py @@ -14,7 +14,7 @@ class CollaQMultiHeadAttention(nn.Module): Overview: The head of collaq attention module. Interface: - __init__, forward + ``__init__``, ``forward`` """ def __init__( @@ -69,8 +69,26 @@ def forward(self, q, k, v, mask=None): - q (:obj:`torch.nn.Sequential`): the transformer information q - k (:obj:`torch.nn.Sequential`): the transformer information k - v (:obj:`torch.nn.Sequential`): the transformer information v - Output: + Returns: - q (:obj:`torch.nn.Sequential`): the transformer output q + - residual (:obj:`torch.nn.Sequential`): the transformer output residual + Shapes: + - q (:obj:`torch.nn.Sequential`): :math:`(B, L, N)` where B is batch_size, L is sequence length, \ + N is the size of input q + - k (:obj:`torch.nn.Sequential`): :math:`(B, L, N)` where B is batch_size, L is sequence length, \ + N is the size of input k + - v (:obj:`torch.nn.Sequential`): :math:`(B, L, N)` where B is batch_size, L is sequence length, \ + N is the size of input v + - q (:obj:`torch.nn.Sequential`): :math:`(B, L, N)` where B is batch_size, L is sequence length, \ + N is the size of output q + - residual (:obj:`torch.nn.Sequential`): :math:`(B, L, N)` where B is batch_size, L is sequence length, \ + N is the size of output residual + Examples: + >>> net = CollaQMultiHeadAttention(1, 2, 3, 4, 5, 6) + >>> q = torch.randn(1, 2, 2) + >>> k = torch.randn(1, 3, 3) + >>> v = torch.randn(1, 3, 3) + >>> q, residual = net(q, k, v) """ d_k, d_v, n_head = self.d_k, self.d_v, self.n_head batch_size, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1) @@ -104,7 +122,7 @@ class CollaQSMACAttentionModule(nn.Module): Collaq attention module. Used to get agent's attention observation. It includes agent's observation\ and agent's part of the observation information of the agent's concerned allies Interface: - __init__, _cut_obs, forward + ``__init__``, ``_cut_obs``, ``forward`` """ def __init__( @@ -140,9 +158,16 @@ def _cut_obs(self, obs: torch.Tensor): cut the observed information into self's observation and allay's observation Arguments: - obs (:obj:`torch.Tensor`): input each agent's observation - Return: + Returns: - self_features (:obj:`torch.Tensor`): output self agent's attention observation - ally_features (:obj:`torch.Tensor`): output ally agent's attention observation + Shapes: + - obs (:obj:`torch.Tensor`): :math:`(T, B, A, N)` where T is timestep, B is batch_size, \ + A is agent_num, N is obs_shape + - self_features (:obj:`torch.Tensor`): :math:`(T, B, A, N)` where T is timestep, B is batch_size, \ + A is agent_num, N is self_feature_range[1] - self_feature_range[0] + - ally_features (:obj:`torch.Tensor`): :math:`(T, B, A, N)` where T is timestep, B is batch_size, \ + A is agent_num, N is ally_feature_range[1] - ally_feature_range[0] """ # obs shape = (T, B, A, obs_shape) self_features = obs[:, :, :, self.self_feature_range[0]:self.self_feature_range[1]] @@ -155,8 +180,11 @@ def forward(self, inputs: torch.Tensor): forward computation to get agent's attention observation information Arguments: - obs (:obj:`torch.Tensor`): input each agent's observation - Return: + Returns: - obs (:obj:`torch.Tensor`): output agent's attention observation + Shapes: + - obs (:obj:`torch.Tensor`): :math:`(T, B, A, N)` where T is timestep, B is batch_size, \ + A is agent_num, N is obs_shape """ # obs shape = (T, B ,A, obs_shape) obs = inputs @@ -182,9 +210,16 @@ def forward(self, inputs: torch.Tensor): class CollaQ(nn.Module): """ Overview: - CollaQ network + The network of CollaQ (Collaborative Q-learning) algorithm. + It includes two parts: q_network and q_alone_network. + The q_network is used to get the q_value of the agent's observation and \ + the agent's part of the observation information of the agent's concerned allies. + The q_alone_network is used to get the q_value of the agent's observation and \ + the agent's observation information without the agent's concerned allies. + Multi-Agent Collaboration via Reward Attribution Decomposition + https://arxiv.org/abs/2010.08531 Interface: - __init__, forward, _setup_global_encoder + ``__init__``, ``forward``, ``_setup_global_encoder`` """ def __init__( @@ -275,7 +310,8 @@ def __init__( def forward(self, data: dict, single_step: bool = True) -> dict: """ Overview: - forward computation graph of collaQ network + The forward method calculates the q_value of each agent and the total q_value of all agents. + The q_value of each agent is calculated by the q_network, and the total q_value is calculated by the mixer. Arguments: - data (:obj:`dict`): input data dict with keys ['obs', 'prev_state', 'action'] - agent_state (:obj:`torch.Tensor`): each agent local state(obs) @@ -302,6 +338,32 @@ def forward(self, data: dict, single_step: bool = True) -> dict: - total_q (:obj:`torch.Tensor`): :math:`(T, B)` - agent_q (:obj:`torch.Tensor`): :math:`(T, B, A, P)`, where P is action_shape - next_state (:obj:`list`): math:`(B, A)`, a list of length B, and each element is a list of length A + Examples: + >>> collaQ_model = CollaQ( + >>> agent_num=4, + >>> obs_shape=32, + >>> alone_obs_shape=24, + >>> global_obs_shape=32 * 4, + >>> action_shape=9, + >>> hidden_size_list=[128, 64], + >>> self_feature_range=[8, 10], + >>> ally_feature_range=[10, 16], + >>> attention_size=64, + >>> mixer=True, + >>> activation=torch.nn.Tanh() + >>> ) + >>> data={ + >>> 'obs': { + >>> 'agent_state': torch.randn(8, 4, 4, 32), + >>> 'agent_alone_state': torch.randn(8, 4, 4, 24), + >>> 'agent_alone_padding_state': torch.randn(8, 4, 4, 32), + >>> 'global_state': torch.randn(8, 4, 32 * 4), + >>> 'action_mask': torch.randint(0, 2, size=(8, 4, 4, 9)) + >>> }, + >>> 'prev_state': [[[None for _ in range(4)] for _ in range(3)] for _ in range(4)], + >>> 'action': torch.randint(0, 9, size=(8, 4, 4)) + >>> } + >>> output = collaQ_model(data, single_step=False) """ agent_state, agent_alone_state = data['obs']['agent_state'], data['obs']['agent_alone_state'] agent_alone_padding_state = data['obs']['agent_alone_padding_state'] @@ -349,27 +411,20 @@ def forward(self, data: dict, single_step: bool = True) -> dict: agent_alone_state = agent_alone_state.reshape(T, -1, *agent_alone_state.shape[3:]) agent_alone_padding_state = agent_alone_padding_state.reshape(T, -1, *agent_alone_padding_state.shape[3:]) - colla_output = self._q_network( - { - 'obs': agent_state, - 'prev_state': colla_prev_state, - 'enable_fast_timestep': True - } - ) + colla_output = self._q_network({ + 'obs': agent_state, + 'prev_state': colla_prev_state, + }) colla_alone_output = self._q_network( { 'obs': agent_alone_padding_state, 'prev_state': colla_alone_prev_state, - 'enable_fast_timestep': True - } - ) - alone_output = self._q_alone_network( - { - 'obs': agent_alone_state, - 'prev_state': alone_prev_state, - 'enable_fast_timestep': True } ) + alone_output = self._q_alone_network({ + 'obs': agent_alone_state, + 'prev_state': alone_prev_state, + }) agent_alone_q, alone_next_state = alone_output['logit'], alone_output['next_state'] agent_colla_alone_q, colla_alone_next_state = colla_alone_output['logit'], colla_alone_output['next_state'] @@ -426,7 +481,7 @@ def _setup_global_encoder(self, global_obs_shape: int, embedding_size: int) -> t Arguments: - global_obs_shape (:obj:`int`): the dimension of global observation state - embedding_size (:obj:`int`): the dimension of state emdedding - Return: + Returns: - outputs (:obj:`torch.nn.Module`): Global observation encoding network """ return MLP(global_obs_shape, embedding_size, embedding_size, 2, activation=self._act) diff --git a/ding/model/template/coma.py b/ding/model/template/coma.py index c120dfd78d..5ed7973635 100644 --- a/ding/model/template/coma.py +++ b/ding/model/template/coma.py @@ -11,9 +11,9 @@ class COMAActorNetwork(nn.Module): """ Overview: - Decentralized actor network in COMA + Decentralized actor network in COMA algorithm. Interface: - __init__, forward + ``__init__``, ``forward`` """ def __init__( @@ -24,7 +24,7 @@ def __init__( ): """ Overview: - initialize COMA actor network + Initialize COMA actor network Arguments: - obs_shape (:obj:`int`): the dimension of each agent's observation state - action_shape (:obj:`int`): the dimension of action shape @@ -35,10 +35,30 @@ def __init__( def forward(self, inputs: Dict) -> Dict: """ + Overview: + The forward computation graph of COMA actor network + Arguments: + - inputs (:obj:`dict`): input data dict with keys ['obs', 'prev_state'] + - agent_state (:obj:`torch.Tensor`): each agent local state(obs) + - action_mask (:obj:`torch.Tensor`): the masked action + - prev_state (:obj:`torch.Tensor`): the previous hidden state + Returns: + - output (:obj:`dict`): output data dict with keys ['logit', 'next_state', 'action_mask'] ArgumentsKeys: - necessary: ``obs`` { ``agent_state``, ``action_mask`` }, ``prev_state`` ReturnsKeys: - necessary: ``logit``, ``next_state``, ``action_mask`` + Examples: + >>> T, B, A, N = 4, 8, 3, 32 + >>> embedding_dim = 64 + >>> action_dim = 6 + >>> data = torch.randn(T, B, A, N) + >>> model = COMAActorNetwork((N, ), action_dim, [128, embedding_dim]) + >>> prev_state = [[None for _ in range(A)] for _ in range(B)] + >>> for t in range(T): + >>> inputs = {'obs': {'agent_state': data[t], 'action_mask': None}, 'prev_state': prev_state} + >>> outputs = model(inputs) + >>> logit, prev_state = outputs['logit'], outputs['next_state'] """ agent_state = inputs['obs']['agent_state'] prev_state = inputs['prev_state'] @@ -50,7 +70,7 @@ def forward(self, inputs: Dict) -> Dict: T, B, A = agent_state.shape[:3] agent_state = agent_state.reshape(T, -1, *agent_state.shape[3:]) prev_state = reduce(lambda x, y: x + y, prev_state) - output = self.main({'obs': agent_state, 'prev_state': prev_state, 'enable_fast_timestep': True}) + output = self.main({'obs': agent_state, 'prev_state': prev_state}) logit, next_state = output['logit'], output['next_state'] next_state, _ = list_split(next_state, step=A) logit = logit.reshape(T, B, A, -1) @@ -62,9 +82,9 @@ def forward(self, inputs: Dict) -> Dict: class COMACriticNetwork(nn.Module): """ Overview: - Centralized critic network in COMA + Centralized critic network in COMA algorithm. Interface: - __init__, forward + ``__init__``, ``forward`` """ def __init__( @@ -80,6 +100,14 @@ def __init__( - input_size (:obj:`int`): the size of input global observation - action_shape (:obj:`int`): the dimension of action shape - hidden_size_list (:obj:`list`): the list of hidden size, default to 128 + Returns: + - output (:obj:`dict`): output data dict with keys ['q_value'] + Shapes: + - obs (:obj:`dict`): ``agent_state``: :math:`(T, B, A, N, D)`, ``action_mask``: :math:`(T, B, A, N, A)` + - prev_state (:obj:`list`): :math:`[[[h, c] for _ in range(A)] for _ in range(B)]` + - logit (:obj:`torch.Tensor`): :math:`(T, B, A, N, A)` + - next_state (:obj:`list`): :math:`[[[h, c] for _ in range(A)] for _ in range(B)]` + - action_mask (:obj:`torch.Tensor`): :math:`(T, B, A, N, A)` """ super(COMACriticNetwork, self).__init__() self.action_shape = action_shape @@ -101,6 +129,19 @@ def forward(self, data: Dict) -> Dict: - necessary: ``obs`` { ``agent_state``, ``global_state`` }, ``action``, ``prev_state`` ReturnsKeys: - necessary: ``q_value`` + Examples: + >>> agent_num, bs, T = 4, 3, 8 + >>> obs_dim, global_obs_dim, action_dim = 32, 32 * 4, 9 + >>> coma_model = COMACriticNetwork( + >>> obs_dim - action_dim + global_obs_dim + 2 * action_dim * agent_num, action_dim) + >>> data = { + >>> 'obs': { + >>> 'agent_state': torch.randn(T, bs, agent_num, obs_dim), + >>> 'global_state': torch.randn(T, bs, global_obs_dim), + >>> }, + >>> 'action': torch.randint(0, action_dim, size=(T, bs, agent_num)), + >>> } + >>> output = coma_model(data) """ x = self._preprocess_data(data) q = self.mlp(x) @@ -145,8 +186,13 @@ def _preprocess_data(self, data: Dict) -> torch.Tensor: class COMA(nn.Module): """ Overview: - COMA network is QAC-type actor-critic. + The network of COMA algorithm, which is QAC-type actor-critic. + Interface: + ``__init__``, ``forward`` + Properties: + - mode (:obj:`list`): The list of forward mode, including ``compute_actor`` and ``compute_critic`` """ + mode = ['compute_actor', 'compute_critic'] def __init__( @@ -174,12 +220,53 @@ def __init__( def forward(self, inputs: Dict, mode: str) -> Dict: """ + Overview: + forward computation graph of COMA network + Arguments: + - inputs (:obj:`dict`): input data dict with keys ['obs', 'prev_state', 'action'] + - agent_state (:obj:`torch.Tensor`): each agent local state(obs) + - global_state (:obj:`torch.Tensor`): global state(obs) + - action (:obj:`torch.Tensor`): the masked action ArgumentsKeys: - necessary: ``obs`` { ``agent_state``, ``global_state``, ``action_mask`` }, ``action``, ``prev_state`` ReturnsKeys: - necessary: - compute_critic: ``q_value`` - compute_actor: ``logit``, ``next_state``, ``action_mask`` + Shapes: + - obs (:obj:`dict`): ``agent_state``: :math:`(T, B, A, N, D)`, ``action_mask``: :math:`(T, B, A, N, A)` + - prev_state (:obj:`list`): :math:`[[[h, c] for _ in range(A)] for _ in range(B)]` + - logit (:obj:`torch.Tensor`): :math:`(T, B, A, N, A)` + - next_state (:obj:`list`): :math:`[[[h, c] for _ in range(A)] for _ in range(B)]` + - action_mask (:obj:`torch.Tensor`): :math:`(T, B, A, N, A)` + - q_value (:obj:`torch.Tensor`): :math:`(T, B, A, N, A)` + Examples: + >>> agent_num, bs, T = 4, 3, 8 + >>> agent_num, bs, T = 4, 3, 8 + >>> obs_dim, global_obs_dim, action_dim = 32, 32 * 4, 9 + >>> coma_model = COMA( + >>> agent_num=agent_num, + >>> obs_shape=dict(agent_state=(obs_dim, ), global_state=(global_obs_dim, )), + >>> action_shape=action_dim, + >>> actor_hidden_size_list=[128, 64], + >>> ) + >>> prev_state = [[None for _ in range(agent_num)] for _ in range(bs)] + >>> data = { + >>> 'obs': { + >>> 'agent_state': torch.randn(T, bs, agent_num, obs_dim), + >>> 'action_mask': None, + >>> }, + >>> 'prev_state': prev_state, + >>> } + >>> output = coma_model(data, mode='compute_actor') + >>> data= { + >>> 'obs': { + >>> 'agent_state': torch.randn(T, bs, agent_num, obs_dim), + >>> 'global_state': torch.randn(T, bs, global_obs_dim), + >>> }, + >>> 'action': torch.randint(0, action_dim, size=(T, bs, agent_num)), + >>> } + >>> output = coma_model(data, mode='compute_critic') """ assert mode in self.mode, "not support forward mode: {}/{}".format(mode, self.mode) if mode == 'compute_actor': diff --git a/ding/model/template/decision_transformer.py b/ding/model/template/decision_transformer.py index ef86657f86..3d35497383 100644 --- a/ding/model/template/decision_transformer.py +++ b/ding/model/template/decision_transformer.py @@ -1,20 +1,128 @@ """ -The code is transplanted from https://github.com/nikhilbarhate99/min-decision-transformer +this extremely minimal Decision Transformer model is based on +the following causal transformer (GPT) implementation: + +Misha Laskin's tweet: +https://twitter.com/MishaLaskin/status/1481767788775628801?cxt=HHwWgoCzmYD9pZApAAAA + +and its corresponding notebook: +https://colab.research.google.com/drive/1NUBqyboDcGte5qAJKOl8gaJC28V_73Iv?usp=sharing + +** the above colab notebook has a bug while applying masked_fill +which is fixed in the following code """ -from ding.utils import MODEL_REGISTRY -from typing import Tuple -from ding.torch_utils.network.transformer import Attention +import math +from typing import Union, Optional, Tuple + import torch import torch.nn as nn +import torch.nn.functional as F +from ding.utils import SequenceType -class Block(nn.Module): +class MaskedCausalAttention(nn.Module): + """ + Overview: + The implementation of masked causal attention in decision transformer. The input of this module is a sequence \ + of several tokens. For the calculated hidden embedding for the i-th token, it is only related the 0 to i-1 \ + input tokens by applying a mask to the attention map. Thus, this module is called masked-causal attention. + Interfaces: + ``__init__``, ``forward`` + """ def __init__(self, h_dim: int, max_T: int, n_heads: int, drop_p: float) -> None: + """ + Overview: + Initialize the MaskedCausalAttention Model according to input arguments. + Arguments: + - h_dim (:obj:`int`): The dimension of the hidden layers, such as 128. + - max_T (:obj:`int`): The max context length of the attention, such as 6. + - n_heads (:obj:`int`): The number of heads in calculating attention, such as 8. + - drop_p (:obj:`float`): The drop rate of the drop-out layer, such as 0.1. + """ super().__init__() - self.attention = Attention(h_dim, h_dim, h_dim, n_heads, nn.Dropout(drop_p)) + + self.n_heads = n_heads + self.max_T = max_T + + self.q_net = nn.Linear(h_dim, h_dim) + self.k_net = nn.Linear(h_dim, h_dim) + self.v_net = nn.Linear(h_dim, h_dim) + + self.proj_net = nn.Linear(h_dim, h_dim) + self.att_drop = nn.Dropout(drop_p) + self.proj_drop = nn.Dropout(drop_p) + + ones = torch.ones((max_T, max_T)) + mask = torch.tril(ones).view(1, 1, max_T, max_T) + + # register buffer makes sure mask does not get updated + # during backpropagation + self.register_buffer('mask', mask) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Overview: + MaskedCausalAttention forward computation graph, input a sequence tensor \ + and return a tensor with the same shape. + Arguments: + - x (:obj:`torch.Tensor`): The input tensor. + Returns: + - out (:obj:`torch.Tensor`): Output tensor, the shape is the same as the input. + Examples: + >>> inputs = torch.randn(2, 4, 64) + >>> model = MaskedCausalAttention(64, 5, 4, 0.1) + >>> outputs = model(inputs) + >>> assert outputs.shape == torch.Size([2, 4, 64]) + """ + B, T, C = x.shape # batch size, seq length, h_dim * n_heads + + N, D = self.n_heads, C // self.n_heads # N = num heads, D = attention dim + + # rearrange q, k, v as (B, N, T, D) + q = self.q_net(x).view(B, T, N, D).transpose(1, 2) + k = self.k_net(x).view(B, T, N, D).transpose(1, 2) + v = self.v_net(x).view(B, T, N, D).transpose(1, 2) + + # weights (B, N, T, T) + weights = q @ k.transpose(2, 3) / math.sqrt(D) + # causal mask applied to weights + weights = weights.masked_fill(self.mask[..., :T, :T] == 0, float('-inf')) + # normalize weights, all -inf -> 0 after softmax + normalized_weights = F.softmax(weights, dim=-1) + + # attention (B, N, T, D) + attention = self.att_drop(normalized_weights @ v) + + # gather heads and project (B, N, T, D) -> (B, T, N*D) + attention = attention.transpose(1, 2).contiguous().view(B, T, N * D) + + out = self.proj_drop(self.proj_net(attention)) + return out + + +class Block(nn.Module): + """ + Overview: + The implementation of a transformer block in decision transformer. + Interfaces: + ``__init__``, ``forward`` + """ + + def __init__(self, h_dim: int, max_T: int, n_heads: int, drop_p: float) -> None: + """ + Overview: + Initialize the Block Model according to input arguments. + Arguments: + - h_dim (:obj:`int`): The dimension of the hidden layers, such as 128. + - max_T (:obj:`int`): The max context length of the attention, such as 6. + - n_heads (:obj:`int`): The number of heads in calculating attention, such as 8. + - drop_p (:obj:`float`): The drop rate of the drop-out layer, such as 0.1. + """ + super().__init__() + self.attention = MaskedCausalAttention(h_dim, max_T, n_heads, drop_p) self.mlp = nn.Sequential( nn.Linear(h_dim, 4 * h_dim), nn.GELU(), @@ -24,101 +132,282 @@ def __init__(self, h_dim: int, max_T: int, n_heads: int, drop_p: float) -> None: self.ln1 = nn.LayerNorm(h_dim) self.ln2 = nn.LayerNorm(h_dim) - mask = torch.tril(torch.ones((max_T, max_T), dtype=torch.bool)).view(1, 1, max_T, max_T) - # register buffer makes sure mask does not get updated - # during backpropagation - self.register_buffer('mask', mask) - - def forward(self, x: torch.Tensor): + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Overview: + Forward computation graph of the decision transformer block, input a sequence tensor \ + and return a tensor with the same shape. + Arguments: + - x (:obj:`torch.Tensor`): The input tensor. + Returns: + - output (:obj:`torch.Tensor`): Output tensor, the shape is the same as the input. + Examples: + >>> inputs = torch.randn(2, 4, 64) + >>> model = Block(64, 5, 4, 0.1) + >>> outputs = model(inputs) + >>> outputs.shape == torch.Size([2, 4, 64]) + """ # Attention -> LayerNorm -> MLP -> LayerNorm - - x = x + self.att_drop(self.attention(x, self.mask)) # residual + x = x + self.attention(x) # residual x = self.ln1(x) x = x + self.mlp(x) # residual x = self.ln2(x) + # x = x + self.attention(self.ln1(x)) + # x = x + self.mlp(self.ln2(x)) return x -@MODEL_REGISTRY.register('dt') class DecisionTransformer(nn.Module): + """ + Overview: + The implementation of decision transformer. + Interfaces: + ``__init__``, ``forward``, ``configure_optimizers`` + """ def __init__( - self, - state_dim: int, - act_dim: int, - n_blocks: int, - h_dim: int, - context_len: int, - n_heads: int, - drop_p: float, - max_timestep: int = 4096, - continuous: bool = True - ) -> None: + self, + state_dim: Union[int, SequenceType], + act_dim: int, + n_blocks: int, + h_dim: int, + context_len: int, + n_heads: int, + drop_p: float, + max_timestep: int = 4096, + state_encoder: Optional[nn.Module] = None, + continuous: bool = False + ): + """ + Overview: + Initialize the DecisionTransformer Model according to input arguments. + Arguments: + - obs_shape (:obj:`Union[int, SequenceType]`): Dimension of state, such as 128 or (4, 84, 84). + - act_dim (:obj:`int`): The dimension of actions, such as 6. + - n_blocks (:obj:`int`): The number of transformer blocks in the decision transformer, such as 3. + - h_dim (:obj:`int`): The dimension of the hidden layers, such as 128. + - context_len (:obj:`int`): The max context length of the attention, such as 6. + - n_heads (:obj:`int`): The number of heads in calculating attention, such as 8. + - drop_p (:obj:`float`): The drop rate of the drop-out layer, such as 0.1. + - max_timestep (:obj:`int`): The max length of the total sequence, defaults to be 4096. + - state_encoder (:obj:`Optional[nn.Module]`): The encoder to pre-process the given input. If it is set to \ + None, the raw state will be pushed into the transformer. + - continuous (:obj:`bool`): Whether the action space is continuous, defaults to be ``False``. + """ super().__init__() - self.continuous = continuous + self.state_dim = state_dim self.act_dim = act_dim self.h_dim = h_dim # transformer blocks - # we will serially arrange `return`, `state` and `action`, so here the input_seq_len is 3 * context_len input_seq_len = 3 * context_len - blocks = [Block(h_dim, input_seq_len, n_heads, drop_p) for _ in range(n_blocks)] - self.transformer = nn.Sequential(*blocks) # projection heads (project to embedding) self.embed_ln = nn.LayerNorm(h_dim) self.embed_timestep = nn.Embedding(max_timestep, h_dim) - self.embed_rtg = torch.nn.Linear(1, h_dim) - self.embed_state = torch.nn.Linear(state_dim, h_dim) + self.drop = nn.Dropout(drop_p) - if self.continuous: - action_tanh = True # True for continuous actions - self.embed_action = torch.nn.Linear(act_dim, h_dim) + self.pos_emb = nn.Parameter(torch.zeros(1, input_seq_len + 1, self.h_dim)) + self.global_pos_emb = nn.Parameter(torch.zeros(1, max_timestep + 1, self.h_dim)) + if state_encoder is None: + self.state_encoder = None + blocks = [Block(h_dim, input_seq_len, n_heads, drop_p) for _ in range(n_blocks)] + self.embed_rtg = torch.nn.Linear(1, h_dim) + self.embed_state = torch.nn.Linear(state_dim, h_dim) + self.predict_rtg = torch.nn.Linear(h_dim, 1) + self.predict_state = torch.nn.Linear(h_dim, state_dim) + if continuous: + # continuous actions + self.embed_action = torch.nn.Linear(act_dim, h_dim) + use_action_tanh = True # True for continuous actions + else: + # discrete actions + self.embed_action = torch.nn.Embedding(act_dim, h_dim) + use_action_tanh = False # False for discrete actions + self.predict_action = nn.Sequential( + *([nn.Linear(h_dim, act_dim)] + ([nn.Tanh()] if use_action_tanh else [])) + ) else: - action_tanh = False # False for discrete actions - self.embed_action = torch.nn.Linear(act_dim, h_dim) - - # prediction heads - self.predict_rtg = torch.nn.Linear(h_dim, 1) - self.predict_state = torch.nn.Linear(h_dim, state_dim) - self.predict_action = nn.Sequential(*([nn.Linear(h_dim, act_dim)] + ([nn.Tanh()] if action_tanh else []))) + blocks = [Block(h_dim, input_seq_len + 1, n_heads, drop_p) for _ in range(n_blocks)] + self.state_encoder = state_encoder + self.embed_rtg = nn.Sequential(nn.Linear(1, h_dim), nn.Tanh()) + self.head = nn.Linear(h_dim, act_dim, bias=False) + self.embed_action = nn.Sequential(nn.Embedding(act_dim, h_dim), nn.Tanh()) + self.transformer = nn.Sequential(*blocks) def forward( - self, timesteps: torch.Tensor, states: torch.Tensor, actions: torch.Tensor, returns_to_go: torch.Tensor + self, + timesteps: torch.Tensor, + states: torch.Tensor, + actions: torch.Tensor, + returns_to_go: torch.Tensor, + tar: Optional[int] = None ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Overview: + Forward computation graph of the decision transformer, input a sequence tensor \ + and return a tensor with the same shape. + Arguments: + - timesteps (:obj:`torch.Tensor`): The timestep for input sequence. + - states (:obj:`torch.Tensor`): The sequence of states. + - actions (:obj:`torch.Tensor`): The sequence of actions. + - returns_to_go (:obj:`torch.Tensor`): The sequence of return-to-go. + - tar (:obj:`Optional[int]`): Whether to predict action, regardless of index. + Returns: + - output (:obj:`Tuple[torch.Tensor, torch.Tensor, torch.Tensor]`): Output contains three tensors, \ + they are correspondingly the predicted states, predicted actions and predicted return-to-go. + Examples: + >>> B, T = 4, 6 + >>> state_dim = 3 + >>> act_dim = 2 + >>> DT_model = DecisionTransformer(\ + state_dim=state_dim,\ + act_dim=act_dim,\ + n_blocks=3,\ + h_dim=8,\ + context_len=T,\ + n_heads=2,\ + drop_p=0.1,\ + ) + >>> timesteps = torch.randint(0, 100, [B, 3 * T - 1, 1], dtype=torch.long) # B x T + >>> states = torch.randn([B, T, state_dim]) # B x T x state_dim + >>> actions = torch.randint(0, act_dim, [B, T, 1]) + >>> action_target = torch.randint(0, act_dim, [B, T, 1]) + >>> returns_to_go_sample = torch.tensor([1, 0.8, 0.6, 0.4, 0.2, 0.]).repeat([B, 1]).unsqueeze(-1).float() + >>> traj_mask = torch.ones([B, T], dtype=torch.long) # B x T + >>> actions = actions.squeeze(-1) + >>> state_preds, action_preds, return_preds = DT_model.forward(\ + timesteps=timesteps, states=states, actions=actions, returns_to_go=returns_to_go\ + ) + >>> assert state_preds.shape == torch.Size([B, T, state_dim]) + >>> assert return_preds.shape == torch.Size([B, T, 1]) + >>> assert action_preds.shape == torch.Size([B, T, act_dim]) + """ + B, T = states.shape[0], states.shape[1] + if self.state_encoder is None: + time_embeddings = self.embed_timestep(timesteps) - B, T, _ = states.shape + # time embeddings are treated similar to positional embeddings + state_embeddings = self.embed_state(states) + time_embeddings + action_embeddings = self.embed_action(actions) + time_embeddings + returns_embeddings = self.embed_rtg(returns_to_go) + time_embeddings - time_embeddings = self.embed_timestep(timesteps) # shape: (B,context_len/T,h_dim) + # stack rtg, states and actions and reshape sequence as + # (r_0, s_0, a_0, r_1, s_1, a_1, r_2, s_2, a_2 ...) + t_p = torch.stack((returns_embeddings, state_embeddings, action_embeddings), + dim=1).permute(0, 2, 1, 3).reshape(B, 3 * T, self.h_dim) + h = self.embed_ln(t_p) + # transformer and prediction + h = self.transformer(h) + # get h reshaped such that its size = (B x 3 x T x h_dim) and + # h[:, 0, t] is conditioned on the input sequence r_0, s_0, a_0 ... r_t + # h[:, 1, t] is conditioned on the input sequence r_0, s_0, a_0 ... r_t, s_t + # h[:, 2, t] is conditioned on the input sequence r_0, s_0, a_0 ... r_t, s_t, a_t + # that is, for each timestep (t) we have 3 output embeddings from the transformer, + # each conditioned on all previous timesteps plus + # the 3 input variables at that timestep (r_t, s_t, a_t) in sequence. + h = h.reshape(B, T, 3, self.h_dim).permute(0, 2, 1, 3) - # time embeddings are treated similar to positional embeddings - # shape: (B,context_len,h_dim) - state_embeddings = self.embed_state(states) + time_embeddings - action_embeddings = self.embed_action(actions) + time_embeddings - returns_embeddings = self.embed_rtg(returns_to_go) + time_embeddings + return_preds = self.predict_rtg(h[:, 2]) # predict next rtg given r, s, a + state_preds = self.predict_state(h[:, 2]) # predict next state given r, s, a + action_preds = self.predict_action(h[:, 1]) # predict action given r, s + else: + state_embeddings = self.state_encoder( + states.reshape(-1, *self.state_dim).type(torch.float32).contiguous() + ) # (batch * block_size, h_dim) + state_embeddings = state_embeddings.reshape(B, T, self.h_dim) # (batch, block_size, h_dim) + returns_embeddings = self.embed_rtg(returns_to_go.type(torch.float32)) + action_embeddings = self.embed_action(actions.type(torch.long).squeeze(-1)) # (batch, block_size, h_dim) + + token_embeddings = torch.zeros( + (B, T * 3 - int(tar is None), self.h_dim), dtype=torch.float32, device=state_embeddings.device + ) + token_embeddings[:, ::3, :] = returns_embeddings + token_embeddings[:, 1::3, :] = state_embeddings + token_embeddings[:, 2::3, :] = action_embeddings[:, -T + int(tar is None):, :] - # stack rtg, states and actions and reshape sequence as - # (r1, s1, a1, r2, s2, a2 ...) - # after stack shape: (B, 3, context_len/T, h_dim) - h = torch.stack((returns_embeddings, state_embeddings, action_embeddings), - dim=1).permute(0, 2, 1, 3).reshape(B, 3 * T, self.h_dim) + all_global_pos_emb = torch.repeat_interleave( + self.global_pos_emb, B, dim=0 + ) # batch_size, traj_length, h_dim - h = self.embed_ln(h) + position_embeddings = torch.gather( + all_global_pos_emb, 1, torch.repeat_interleave(timesteps, self.h_dim, dim=-1) + ) + self.pos_emb[:, :token_embeddings.shape[1], :] - # transformer and prediction - h = self.transformer(h) + t_p = token_embeddings + position_embeddings - # get h reshaped such that its size = (B x 3 x T x h_dim) and - # h[:, 0, t] is conditioned on r_0, s_0, a_0 ... r_t - # h[:, 1, t] is conditioned on r_0, s_0, a_0 ... r_t, s_t - # h[:, 2, t] is conditioned on r_0, s_0, a_0 ... r_t, s_t, a_t - h = h.reshape(B, T, 3, self.h_dim) + h = self.drop(t_p) + h = self.transformer(h) + h = self.embed_ln(h) + logits = self.head(h) - # get predictions - return_preds = self.predict_rtg(h[..., 2, :]) # predict next rtg given r, s, a - state_preds = self.predict_state(h[..., 2, :]) # predict next state given r, s, a - action_preds = self.predict_action(h[..., 1, :]) # predict action given r, s + return_preds = None + state_preds = None + action_preds = logits[:, 1::3, :] # only keep predictions from state_embeddings return state_preds, action_preds, return_preds + + def configure_optimizers( + self, weight_decay: float, learning_rate: float, betas: Tuple[float, float] = (0.9, 0.95) + ) -> torch.optim.Optimizer: + """ + Overview: + This function returns an optimizer given the input arguments. \ + We are separating out all parameters of the model into two buckets: those that will experience \ + weight decay for regularization and those that won't (biases, and layernorm/embedding weights). + Arguments: + - weight_decay (:obj:`float`): The weigh decay of the optimizer. + - learning_rate (:obj:`float`): The learning rate of the optimizer. + - betas (:obj:`Tuple[float, float]`): The betas for Adam optimizer. + Outputs: + - optimizer (:obj:`torch.optim.Optimizer`): The desired optimizer. + """ + + # separate out all parameters to those that will and won't experience regularizing weight decay + decay = set() + no_decay = set() + # whitelist_weight_modules = (torch.nn.Linear, ) + whitelist_weight_modules = (torch.nn.Linear, torch.nn.Conv2d) + blacklist_weight_modules = (torch.nn.LayerNorm, torch.nn.Embedding) + for mn, m in self.named_modules(): + for pn, p in m.named_parameters(): + fpn = '%s.%s' % (mn, pn) if mn else pn # full param name + + if pn.endswith('bias'): + # all biases will not be decayed + no_decay.add(fpn) + elif pn.endswith('weight') and isinstance(m, whitelist_weight_modules): + # weights of whitelist modules will be weight decayed + decay.add(fpn) + elif pn.endswith('weight') and isinstance(m, blacklist_weight_modules): + # weights of blacklist modules will NOT be weight decayed + no_decay.add(fpn) + + # special case the position embedding parameter in the root GPT module as not decayed + no_decay.add('pos_emb') + no_decay.add('global_pos_emb') + + # validate that we considered every parameter + param_dict = {pn: p for pn, p in self.named_parameters()} + inter_params = decay & no_decay + union_params = decay | no_decay + assert len(inter_params) == 0, "parameters %s made it into both decay/no_decay sets!" % (str(inter_params), ) + assert len(param_dict.keys() - union_params) == 0,\ + "parameters %s were not separated into either decay/no_decay set!" \ + % (str(param_dict.keys() - union_params), ) + + # create the pytorch optimizer object + optim_groups = [ + { + "params": [param_dict[pn] for pn in sorted(list(decay))], + "weight_decay": weight_decay + }, + { + "params": [param_dict[pn] for pn in sorted(list(no_decay))], + "weight_decay": 0.0 + }, + ] + optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas) + return optimizer diff --git a/ding/model/template/diffusion.py b/ding/model/template/diffusion.py new file mode 100755 index 0000000000..f8b48f3061 --- /dev/null +++ b/ding/model/template/diffusion.py @@ -0,0 +1,645 @@ +from typing import Union, List, Dict +from collections import namedtuple +import numpy as np +import math +import torch +import torch.nn as nn +import torch.nn.functional as F +from ding.utils import list_split, MODEL_REGISTRY, squeeze, SequenceType +from ding.torch_utils.network.diffusion import extract, cosine_beta_schedule, apply_conditioning, \ + DiffusionUNet1d, TemporalValue + +Sample = namedtuple('Sample', 'trajectories values chains') + + +def default_sample_fn(model, x, cond, t): + b, *_, device = *x.shape, x.device + model_mean, _, model_log_variance = model.p_mean_variance( + x=x, + cond=cond, + t=t, + ) + noise = 0.5 * torch.randn_like(x) + # no noise when t == 0 + nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1, ) * (len(x.shape) - 1))) + values = torch.zeros(len(x), device=device) + return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise, values + + +def get_guide_output(guide, x, cond, t): + x.requires_grad_() + y = guide(x, cond, t).squeeze(dim=-1) + grad = torch.autograd.grad([y.sum()], [x])[0] + x.detach() + return y, grad + + +def n_step_guided_p_sample( + model, + x, + cond, + t, + guide, + scale=0.001, + t_stopgrad=0, + n_guide_steps=1, + scale_grad_by_std=True, +): + model_log_variance = extract(model.posterior_log_variance_clipped, t, x.shape) + model_std = torch.exp(0.5 * model_log_variance) + model_var = torch.exp(model_log_variance) + + for _ in range(n_guide_steps): + with torch.enable_grad(): + y, grad = get_guide_output(guide, x, cond, t) + + if scale_grad_by_std: + grad = model_var * grad + + grad[t < t_stopgrad] = 0 + + x = x + scale * grad + x = apply_conditioning(x, cond, model.action_dim) + + model_mean, _, model_log_variance = model.p_mean_variance(x=x, cond=cond, t=t) + + # no noise when t == 0 + noise = torch.randn_like(x) + noise[t == 0] = 0 + + return model_mean + model_std * noise, y + + +class GaussianDiffusion(nn.Module): + """ + Overview: + Gaussian diffusion model + Arguments: + - model (:obj:`str`): type of model + - model_cfg (:obj:'dict') config of model + - horizon (:obj:`int`): horizon of trajectory + - obs_dim (:obj:`int`): Dim of the ovservation + - action_dim (:obj:`int`): Dim of the ation + - n_timesteps (:obj:`int`): Number of timesteps + - predict_epsilon (:obj:'bool'): Whether predict epsilon + - loss_discount (:obj:'float'): discount of loss + - clip_denoised (:obj:'bool'): Whether use clip_denoised + - action_weight (:obj:'float'): weight of action + - loss_weights (:obj:'dict'): weight of loss + """ + + def __init__( + self, + model: str, + model_cfg: dict, + horizon: int, + obs_dim: Union[int, SequenceType], + action_dim: Union[int, SequenceType], + n_timesteps: int = 1000, + predict_epsilon: bool = True, + loss_discount: float = 1.0, + clip_denoised: bool = False, + action_weight: float = 1.0, + loss_weights: dict = None, + ) -> None: + super().__init__() + self.horizon = horizon + self.obs_dim = obs_dim + self.action_dim = action_dim + self.transition_dim = obs_dim + action_dim + if type(model) == str: + model = eval(model) + self.model = model(**model_cfg) + self.predict_epsilon = predict_epsilon + self.clip_denoised = clip_denoised + + betas = cosine_beta_schedule(n_timesteps) + alphas = 1. - betas + alphas_cumprod = torch.cumprod(alphas, axis=0) + alphas_cumprod_prev = torch.cat([torch.ones(1), alphas_cumprod[:-1]]) + self.n_timesteps = int(n_timesteps) + + self.register_buffer('betas', betas) + self.register_buffer('alphas_cumprod', alphas_cumprod) + self.register_buffer('alphas_cumprod_prev', alphas_cumprod_prev) + + # calculations for diffusion q(x_t | x_{t-1}) and others + self.register_buffer('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod)) + self.register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1. - alphas_cumprod)) + self.register_buffer('log_one_minus_alphas_cumprod', torch.log(1. - alphas_cumprod)) + self.register_buffer('sqrt_recip_alphas_cumprod', torch.sqrt(1. / alphas_cumprod)) + self.register_buffer('sqrt_recipm1_alphas_cumprod', torch.sqrt(1. / alphas_cumprod - 1)) + + # calculations for posterior q(x_{t-1} | x_t, x_0) + posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod) + self.register_buffer('posterior_variance', posterior_variance) + + # log calculation clipped because the posterior variance + # is 0 at the beginning of the diffusion chain + self.register_buffer('posterior_log_variance_clipped', torch.log(torch.clamp(posterior_variance, min=1e-20))) + self.register_buffer('posterior_mean_coef1', betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod)) + self.register_buffer( + 'posterior_mean_coef2', (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod) + ) + + self.loss_weights = self.get_loss_weights(action_weight, loss_discount, loss_weights) + + def get_loss_weights(self, action_weight: float, discount: float, weights_dict: dict): + """ + Overview: + sets loss coefficients for trajectory + Arguments: + - action_weight (:obj:'float') coefficient on first action loss + - discount (:obj:'float') multiplies t^th timestep of trajectory loss by discount**t + - weights_dict (:obj:'dict') { i: c } multiplies dimension i of observation loss by c + """ + self.action_weight = action_weight + dim_weights = torch.ones(self.transition_dim, dtype=torch.float32) + + # set loss coefficients for dimensions of observation + if weights_dict is None: + weights_dict = {} + for ind, w in weights_dict.items(): + dim_weights[self.action_dim + ind] *= w + + # decay loss with trajectory timestep: discount**t + discounts = discount ** torch.arange(self.horizon, dtype=torch.float) + discounts = discounts / discounts.mean() + loss_weights = torch.einsum('h,t->ht', discounts, dim_weights) + + # manually set a0 weight + loss_weights[0, :self.action_dim] = action_weight + return loss_weights + + def predict_start_from_noise(self, x_t, t, noise): + """ + if self.predict_epsilon, model output is (scaled) noise; + otherwise, model predicts x0 directly + """ + if self.predict_epsilon: + return ( + extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - + extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise + ) + else: + return noise + + def q_posterior(self, x_start, x_t, t): + """ + Overview: + give noise and step, compute mean, variance. + Arguments: + x_start (:obj:'tensor') noise trajectory in timestep 0 + x_t (:obj:'tuple') noise trajectory in timestep t + t (:obj:'int') timestep of diffusion step + """ + posterior_mean = ( + extract(self.posterior_mean_coef1, t, x_t.shape) * x_start + + extract(self.posterior_mean_coef2, t, x_t.shape) * x_t + ) + posterior_variance = extract(self.posterior_variance, t, x_t.shape) + posterior_log_variance_clipped = extract(self.posterior_log_variance_clipped, t, x_t.shape) + return posterior_mean, posterior_variance, posterior_log_variance_clipped + + def p_mean_variance(self, x, cond, t): + x_recon = self.predict_start_from_noise(x, t=t, noise=self.model(x, cond, t)) + + if self.clip_denoised: + x_recon.clamp_(-1., 1.) + else: + assert RuntimeError() + + model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t) + return model_mean, posterior_variance, posterior_log_variance + + @torch.no_grad() + def p_sample_loop(self, shape, cond, return_chain=False, sample_fn=default_sample_fn, plan_size=1, **sample_kwargs): + device = self.betas.device + + batch_size = shape[0] + x = torch.randn(shape, device=device) + x = apply_conditioning(x, cond, self.action_dim) + + chain = [x] if return_chain else None + + for i in reversed(range(0, self.n_timesteps)): + t = torch.full((batch_size, ), i, device=device, dtype=torch.long) + x, values = sample_fn(self, x, cond, t, **sample_kwargs) + x = apply_conditioning(x, cond, self.action_dim) + + if return_chain: + chain.append(x) + values = values.reshape(-1, plan_size, *values.shape[1:]) + x = x.reshape(-1, plan_size, *x.shape[1:]) + if plan_size > 1: + inds = torch.argsort(values, dim=1, descending=True) + x = x[torch.arange(x.size(0)).unsqueeze(1), inds] + values = values[torch.arange(values.size(0)).unsqueeze(1), inds] + if return_chain: + chain = torch.stack(chain, dim=1) + return Sample(x, values, chain) + + @torch.no_grad() + def conditional_sample(self, cond, horizon=None, **sample_kwargs): + """ + conditions : [ (time, state), ... ] + """ + device = self.betas.device + batch_size = len(cond[0]) + horizon = horizon or self.horizon + shape = (batch_size, horizon, self.transition_dim) + + return self.p_sample_loop(shape, cond, **sample_kwargs) + + def q_sample(self, x_start, t, noise=None): + """ + Arguments: + conditions (:obj:'tuple') [ (time, state), ... ] conditions of diffusion + t (:obj:'int') timestep of diffusion + noise (:obj:'tensor.float') timestep's noise of diffusion + """ + if noise is None: + noise = torch.randn_like(x_start) + + sample = ( + extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start + + extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise + ) + + return sample + + def p_losses(self, x_start, cond, t): + noise = torch.randn_like(x_start) + + x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise) + x_noisy = apply_conditioning(x_noisy, cond, self.action_dim) + + x_recon = self.model(x_noisy, cond, t) + x_recon = apply_conditioning(x_recon, cond, self.action_dim) + + assert noise.shape == x_recon.shape + + if self.predict_epsilon: + loss = F.mse_loss(x_recon, noise, reduction='none') + a0_loss = (loss[:, 0, :self.action_dim] / self.loss_weights[0, :self.action_dim].to(loss.device)).mean() + loss = (loss * self.loss_weights.to(loss.device)).mean() + else: + loss = F.mse_loss(x_recon, x_start, reduction='none') + a0_loss = (loss[:, 0, :self.action_dim] / self.loss_weights[0, :self.action_dim].to(loss.device)).mean() + loss = (loss * self.loss_weights.to(loss.device)).mean() + return loss, a0_loss + + def forward(self, cond, *args, **kwargs): + return self.conditional_sample(cond, *args, **kwargs) + + +class ValueDiffusion(GaussianDiffusion): + """ + Overview: + Gaussian diffusion model for value function. + """ + + def p_losses(self, x_start, cond, target, t): + noise = torch.randn_like(x_start) + x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise) + x_noisy = apply_conditioning(x_noisy, cond, self.action_dim) + + pred = self.model(x_noisy, cond, t) + loss = F.mse_loss(pred, target, reduction='none').mean() + log = { + 'mean_pred': pred.mean().item(), + 'max_pred': pred.max().item(), + 'min_pred': pred.min().item(), + } + + return loss, log + + def forward(self, x, cond, t): + return self.model(x, cond, t) + + +@MODEL_REGISTRY.register('pd') +class PlanDiffuser(nn.Module): + """ + Overview: + Diffuser model for plan. + Arguments: + - diffuser_model (:obj:`str`): type of plan model + - diffuser_model_cfg (:obj:'dict') config of diffuser_model + - value_model (:obj:`str`): type of value model, if haven't use, set it as None + - value_model_cfg (:obj:`int`): config of value_model + - sample_kwargs : config of sample function + """ + + def __init__( + self, diffuser_model: str, diffuser_model_cfg: dict, value_model: str, value_model_cfg: dict, **sample_kwargs + ): + super().__init__() + diffuser_model = eval(diffuser_model) + self.diffuser = diffuser_model(**diffuser_model_cfg) + self.value = None + if value_model: + value_model = eval(value_model) + self.value = value_model(**value_model_cfg) + self.sample_kwargs = sample_kwargs + + def diffuser_loss(self, x_start, cond, t): + return self.diffuser.p_losses(x_start, cond, t) + + def value_loss(self, x_start, cond, target, t): + return self.value.p_losses(x_start, cond, target, t) + + def get_eval(self, cond, batch_size=1): + cond = self.repeat_cond(cond, batch_size) + if self.value: + samples = self.diffuser( + cond, sample_fn=n_step_guided_p_sample, plan_size=batch_size, guide=self.value, **self.sample_kwargs + ) + # extract action [eval_num, batch_size, horizon, transition_dim] + actions = samples.trajectories[:, :, :, :self.diffuser.action_dim] + action = actions[:, 0, 0] + return action + else: + samples = self.diffuser(cond, plan_size=batch_size) + return samples.trajectories[:, :, :, self.diffuser.action_dim:].squeeze(1) + + def repeat_cond(self, cond, batch_size): + for k, v in cond.items(): + cond[k] = v.repeat_interleave(batch_size, dim=0) + return cond + + +@MODEL_REGISTRY.register('dd') +class GaussianInvDynDiffusion(nn.Module): + """ + Overview: + Gaussian diffusion model with Invdyn action model. + Arguments: + - model (:obj:`str`): type of model + - model_cfg (:obj:'dict') config of model + - horizon (:obj:`int`): horizon of trajectory + - obs_dim (:obj:`int`): Dim of the ovservation + - action_dim (:obj:`int`): Dim of the ation + - n_timesteps (:obj:`int`): Number of timesteps + - hidden_dim (:obj:'int'): hidden dim of inv_model + - returns_condition (:obj:'bool'): Whether use returns condition + - ar_inv (:obj:'bool'): Whether use inverse action learning + - train_only_inv (:obj:'bool'): Whether train inverse action model only + - predict_epsilon (:obj:'bool'): Whether predict epsilon + - condition_guidance_w (:obj:'float'): weight of condition guidance + - loss_discount (:obj:'float'): discount of loss + """ + + def __init__( + self, + model: str, + model_cfg: dict, + horizon: int, + obs_dim: Union[int, SequenceType], + action_dim: Union[int, SequenceType], + n_timesteps: int = 1000, + hidden_dim: int = 256, + returns_condition: bool = False, + ar_inv: bool = False, + train_only_inv: bool = False, + predict_epsilon: bool = True, + condition_guidance_w: float = 0.1, + loss_discount: float = 1.0, + clip_denoised: bool = False, + ) -> None: + super().__init__() + self.horizon = horizon + self.obs_dim = obs_dim + self.action_dim = action_dim + self.transition_dim = obs_dim + action_dim + if type(model) == str: + model = eval(model) + self.model = model(**model_cfg) + self.ar_inv = ar_inv + self.train_only_inv = train_only_inv + self.predict_epsilon = predict_epsilon + self.condition_guidance_w = condition_guidance_w + + self.inv_model = nn.Sequential( + nn.Linear(2 * self.obs_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, self.action_dim), + ) + + self.returns_condition = returns_condition + self.clip_denoised = clip_denoised + + betas = cosine_beta_schedule(n_timesteps) + alphas = 1. - betas + alphas_cumprod = torch.cumprod(alphas, axis=0) + alphas_cumprod_prev = torch.cat([torch.ones(1), alphas_cumprod[:-1]]) + self.n_timesteps = int(n_timesteps) + + self.register_buffer('betas', betas) + self.register_buffer('alphas_cumprod', alphas_cumprod) + self.register_buffer('alphas_cumprod_prev', alphas_cumprod_prev) + + # calculations for diffusion q(x_t | x_{t-1}) and others + self.register_buffer('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod)) + self.register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1. - alphas_cumprod)) + self.register_buffer('log_one_minus_alphas_cumprod', torch.log(1. - alphas_cumprod)) + self.register_buffer('sqrt_recip_alphas_cumprod', torch.sqrt(1. / alphas_cumprod)) + self.register_buffer('sqrt_recipm1_alphas_cumprod', torch.sqrt(1. / alphas_cumprod - 1)) + + # calculations for posterior q(x_{t-1} | x_t, x_0) + posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod) + self.register_buffer('posterior_variance', posterior_variance) + + # log calculation clipped because the posterior variance + # is 0 at the beginning of the diffusion chain + self.register_buffer('posterior_log_variance_clipped', torch.log(torch.clamp(posterior_variance, min=1e-20))) + self.register_buffer('posterior_mean_coef1', betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod)) + self.register_buffer( + 'posterior_mean_coef2', (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod) + ) + + self.loss_weights = self.get_loss_weights(loss_discount) + + def get_loss_weights(self, discount: int): + self.action_weight = 1 + dim_weights = torch.ones(self.obs_dim, dtype=torch.float32) + + # decay loss with trajectory timestep: discount**t + discounts = discount ** torch.arange(self.horizon, dtype=torch.float) + discounts = discounts / discounts.mean() + loss_weights = torch.einsum('h,t->ht', discounts, dim_weights) + # Cause things are conditioned on t=0 + if self.predict_epsilon: + loss_weights[0, :] = 0 + + return loss_weights + + def predict_start_from_noise(self, x_t, t, noise): + """ + if self.predict_epsilon, model output is (scaled) noise; + otherwise, model predicts x0 directly + """ + if self.predict_epsilon: + return ( + extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - + extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise + ) + else: + return noise + + def q_posterior(self, x_start, x_t, t): + """ + Arguments: + x_start (:obj:'tensor') noise trajectory in timestep 0 + x_t (:obj:'tuple') noise trajectory in timestep t + t (:obj:'int') timestep of diffusion step + """ + posterior_mean = ( + extract(self.posterior_mean_coef1, t, x_t.shape) * x_start + + extract(self.posterior_mean_coef2, t, x_t.shape) * x_t + ) + posterior_variance = extract(self.posterior_variance, t, x_t.shape) + posterior_log_variance_clipped = extract(self.posterior_log_variance_clipped, t, x_t.shape) + return posterior_mean, posterior_variance, posterior_log_variance_clipped + + def p_mean_variance(self, x, cond, t, returns=None): + """ + Arguments: + x (:obj:'tensor') noise trajectory in timestep t + cond (:obj:'tuple') [ (time, state), ... ] state is init state of env, time = 0 + t (:obj:'int') timestep of diffusion step + returns (:obj:'tensor') condition returns of trajectory, returns is normal return + returns: + model_mean (:obj:'tensor.float') + posterior_variance (:obj:'float') + posterior_log_variance (:obj:'float') + """ + if self.returns_condition: + # epsilon could be epsilon or x0 itself + epsilon_cond = self.model(x, cond, t, returns, use_dropout=False) + epsilon_uncond = self.model(x, cond, t, returns, force_dropout=True) + epsilon = epsilon_uncond + self.condition_guidance_w * (epsilon_cond - epsilon_uncond) + else: + epsilon = self.model(x, cond, t) + + t = t.detach().to(torch.int64) + x_recon = self.predict_start_from_noise(x, t=t, noise=epsilon) + + if self.clip_denoised: + x_recon.clamp_(-1., 1.) + else: + assert RuntimeError() + + model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t) + return model_mean, posterior_variance, posterior_log_variance + + @torch.no_grad() + def p_sample(self, x, cond, t, returns=None): + """ + Arguments: + x (:obj:'tensor') noise trajectory in timestep t + cond (:obj:'tuple') [ (time, state), ... ] state is init state of env, time = 0 + t (:obj:'int') timestep of diffusion step + returns (:obj:'tensor') condition returns of trajectory, returns is normal return + """ + b, *_, device = *x.shape, x.device + model_mean, _, model_log_variance = self.p_mean_variance(x=x, cond=cond, t=t, returns=returns) + noise = 0.5 * torch.randn_like(x) + # no noise when t == 0 + nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1, ) * (len(x.shape) - 1))) + return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise + + @torch.no_grad() + def p_sample_loop(self, shape, cond, returns=None, verbose=True, return_diffusion=False): + """ + Arguments: + shape (:obj:'tuple') (batch_size, horizon, self.obs_dim) + cond (:obj:'tuple') [ (time, state), ... ] state is init state of env, time = 0 + returns (:obj:'tensor') condition returns of trajectory, returns is normal return + horizon (:obj:'int') horizon of trajectory + verbose (:obj:'bool') whether log diffusion progress + return_diffusion (:obj:'bool') whether use return diffusion + """ + device = self.betas.device + + batch_size = shape[0] + x = 0.5 * torch.randn(shape, device=device) + # In this model, init state must be given by the env and without noise. + x = apply_conditioning(x, cond, 0) + + if return_diffusion: + diffusion = [x] + + for i in reversed(range(0, self.n_timesteps)): + timesteps = torch.full((batch_size, ), i, device=device, dtype=torch.long) + x = self.p_sample(x, cond, timesteps, returns) + x = apply_conditioning(x, cond, 0) + + if return_diffusion: + diffusion.append(x) + + if return_diffusion: + return x, torch.stack(diffusion, dim=1) + else: + return x + + @torch.no_grad() + def conditional_sample(self, cond, returns=None, horizon=None, *args, **kwargs): + """ + Arguments: + conditions (:obj:'tuple') [ (time, state), ... ] state is init state of env, time is timestep of trajectory + returns (:obj:'tensor') condition returns of trajectory, returns is normal return + horizon (:obj:'int') horizon of trajectory + returns: + x (:obj:'tensor') tarjctory of env + """ + device = self.betas.device + batch_size = len(cond[0]) + horizon = horizon or self.horizon + shape = (batch_size, horizon, self.obs_dim) + + return self.p_sample_loop(shape, cond, returns, *args, **kwargs) + + def q_sample(self, x_start, t, noise=None): + """ + Arguments: + conditions (:obj:'tuple') [ (time, state), ... ] conditions of diffusion + t (:obj:'int') timestep of diffusion + noise (:obj:'tensor.float') timestep's noise of diffusion + """ + if noise is None: + noise = torch.randn_like(x_start) + + sample = ( + extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start + + extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise + ) + + return sample + + def p_losses(self, x_start, cond, t, returns=None): + noise = torch.randn_like(x_start) + + x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise) + x_noisy = apply_conditioning(x_noisy, cond, 0) + + x_recon = self.model(x_noisy, cond, t, returns) + + if not self.predict_epsilon: + x_recon = apply_conditioning(x_recon, cond, 0) + + assert noise.shape == x_recon.shape + + if self.predict_epsilon: + loss = F.mse_loss(x_recon, noise, reduction='none') + loss = (loss * self.loss_weights.to(loss.device)).mean() + else: + loss = F.mse_loss(x_recon, x_start, reduction='none') + loss = (loss * self.loss_weights.to(loss.device)).mean() + + return loss + + def forward(self, cond, *args, **kwargs): + return self.conditional_sample(cond=cond, *args, **kwargs) diff --git a/ding/model/template/ebm.py b/ding/model/template/ebm.py index fe1c073b17..4b91fd1b6d 100644 --- a/ding/model/template/ebm.py +++ b/ding/model/template/ebm.py @@ -15,10 +15,17 @@ from ding.utils import MODEL_REGISTRY, STOCHASTIC_OPTIMIZER_REGISTRY from ding.torch_utils import unsqueeze_repeat from ding.model.wrapper import IModelWrapper -from ..common import RegressionHead +from ding.model.common import RegressionHead def create_stochastic_optimizer(device: str, stochastic_optimizer_config: dict): + """ + Overview: + Create stochastic optimizer. + Arguments: + - device (:obj:`str`): Device. + - stochastic_optimizer_config (:obj:`dict`): Stochastic optimizer config. + """ return STOCHASTIC_OPTIMIZER_REGISTRY.build( stochastic_optimizer_config.pop("type"), device=device, **stochastic_optimizer_config ) @@ -45,20 +52,34 @@ def wrapper(*args, **kwargs): class StochasticOptimizer(ABC): + """ + Overview: + Base class for stochastic optimizers. + Interface: + ``__init__``, ``_sample``, ``_get_best_action_sample``, ``set_action_bounds``, ``sample``, ``infer`` + """ def _sample(self, obs: torch.Tensor, num_samples: int) -> Tuple[torch.Tensor, torch.Tensor]: """ Overview: - Helper method for drawing action samples from the uniform random distribution \ + Drawing action samples from the uniform random distribution \ and tiling observations to the same shape as action samples. - Arguments: - - obs (:obj:`torch.Tensor`): Observation of shape (B, O). - - num_samples (:obj:`int`): The number of negative samples (N). - + - obs (:obj:`torch.Tensor`): Observation. + - num_samples (:obj:`int`): The number of negative samples. Returns: - - tiled_obs (:obj:`torch.Tensor`): Observation of shape (B, N, O). - - action (:obj:`torch.Tensor`): Action of shape (B, N, A). + - tiled_obs (:obj:`torch.Tensor`): Observations tiled. + - action (:obj:`torch.Tensor`): Action sampled. + Shapes: + - obs (:obj:`torch.Tensor`): :math:`(B, O)`. + - num_samples (:obj:`int`): :math:`N`. + - tiled_obs (:obj:`torch.Tensor`): :math:`(B, N, O)`. + - action (:obj:`torch.Tensor`): :math:`(B, N, A)`. + Examples: + >>> obs = torch.randn(2, 4) + >>> opt = StochasticOptimizer() + >>> opt.set_action_bounds(np.stack([np.zeros(5), np.ones(5)], axis=0)) + >>> tiled_obs, action = opt._sample(obs, 8) """ size = (obs.shape[0], num_samples, self.action_bounds.shape[1]) low, high = self.action_bounds[0, :], self.action_bounds[1, :] @@ -72,13 +93,22 @@ def _get_best_action_sample(obs: torch.Tensor, action_samples: torch.Tensor, ebm """ Overview: Return one action for each batch with highest probability (lowest energy). - Arguments: - - obs (:obj:`torch.Tensor`): Observation of shape (B, N, O). - - action_samples (:obj:`torch.Tensor`): Action of shape (B, N, A). - + - obs (:obj:`torch.Tensor`): Observation. + - action_samples (:obj:`torch.Tensor`): Action from uniform distributions. Returns: - - best_action_samples (:obj:`torch.Tensor`): Action of shape (B, A). + - best_action_samples (:obj:`torch.Tensor`): Best action. + Shapes: + - obs (:obj:`torch.Tensor`): :math:`(B, O)`. + - action_samples (:obj:`torch.Tensor`): :math:`(B, N, A)`. + - best_action_samples (:obj:`torch.Tensor`): :math:`(B, A)`. + Examples: + >>> obs = torch.randn(2, 4) + >>> action_samples = torch.randn(2, 8, 5) + >>> ebm = EBM(4, 5) + >>> opt = StochasticOptimizer() + >>> opt.set_action_bounds(np.stack([np.zeros(5), np.ones(5)], axis=0)) + >>> best_action_samples = opt._get_best_action_sample(obs, action_samples, ebm) """ # (B, N) energies = ebm.forward(obs, action_samples) @@ -91,10 +121,17 @@ def set_action_bounds(self, action_bounds: np.ndarray): """ Overview: Set action bounds calculated from the dataset statistics. - Arguments: - action_bounds (:obj:`np.ndarray`): Array of shape (2, A), \ where action_bounds[0] is lower bound and action_bounds[1] is upper bound. + Returns: + - action_bounds (:obj:`torch.Tensor`): Action bounds. + Shapes: + - action_bounds (:obj:`np.ndarray`): :math:`(2, A)`. + - action_bounds (:obj:`torch.Tensor`): :math:`(2, A)`. + Examples: + >>> opt = StochasticOptimizer() + >>> opt.set_action_bounds(np.stack([np.zeros(5), np.ones(5)], axis=0)) """ self.action_bounds = torch.as_tensor(action_bounds, dtype=torch.float32).to(self.device) @@ -103,14 +140,17 @@ def sample(self, obs: torch.Tensor, ebm: nn.Module) -> Tuple[torch.Tensor, torch """ Overview: Create tiled observations and sample counter-negatives for InfoNCE loss. - Arguments: - - obs (:obj:`torch.Tensor`): Observation of shape (B, O). + - obs (:obj:`torch.Tensor`): Observations. - ebm (:obj:`torch.nn.Module`): Energy based model. - Returns: - - tiled_obs (:obj:`torch.Tensor`): Observation of shape (B, N, O). - - action (:obj:`torch.Tensor`): Action of shape (B, N, A). + - tiled_obs (:obj:`torch.Tensor`): Tiled observations. + - action (:obj:`torch.Tensor`): Actions. + Shapes: + - obs (:obj:`torch.Tensor`): :math:`(B, O)`. + - ebm (:obj:`torch.nn.Module`): :math:`(B, N, O)`. + - tiled_obs (:obj:`torch.Tensor`): :math:`(B, N, O)`. + - action (:obj:`torch.Tensor`): :math:`(B, N, A)`. .. note:: In the case of derivative-free optimization, this function will simply call _sample. """ @@ -122,16 +162,27 @@ def infer(self, obs: torch.Tensor, ebm: nn.Module) -> torch.Tensor: Overview: Optimize for the best action conditioned on the current observation. Arguments: - - obs (:obj:`torch.Tensor`): Observation of shape (B, O). + - obs (:obj:`torch.Tensor`): Observations. - ebm (:obj:`torch.nn.Module`): Energy based model. Returns: - - best_action_samples (:obj:`torch.Tensor`): Action of shape (B, A). + - best_action_samples (:obj:`torch.Tensor`): Best actions. + Shapes: + - obs (:obj:`torch.Tensor`): :math:`(B, O)`. + - ebm (:obj:`torch.nn.Module`): :math:`(B, N, O)`. + - best_action_samples (:obj:`torch.Tensor`): :math:`(B, A)`. """ raise NotImplementedError @STOCHASTIC_OPTIMIZER_REGISTRY.register('dfo') class DFO(StochasticOptimizer): + """ + Overview: + Derivative-Free Optimizer in paper Implicit Behavioral Cloning. + https://arxiv.org/abs/2109.00137 + Interface: + ``init``, ``sample``, ``infer`` + """ def __init__( self, @@ -142,6 +193,17 @@ def __init__( inference_samples: int = 16384, device: str = 'cpu', ): + """ + Overview: + Initialize the Derivative-Free Optimizer + Arguments: + - noise_scale (:obj:`float`): Initial noise scale. + - noise_shrink (:obj:`float`): Noise scale shrink rate. + - iters (:obj:`int`): Number of iterations. + - train_samples (:obj:`int`): Number of samples for training. + - inference_samples (:obj:`int`): Number of samples for inference. + - device (:obj:`str`): Device. + """ self.action_bounds = None self.noise_scale = noise_scale self.noise_shrink = noise_shrink @@ -151,10 +213,51 @@ def __init__( self.device = device def sample(self, obs: torch.Tensor, ebm: nn.Module) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Overview: + Drawing action samples from the uniform random distribution \ + and tiling observations to the same shape as action samples. + Arguments: + - obs (:obj:`torch.Tensor`): Observations. + - ebm (:obj:`torch.nn.Module`): Energy based model. + Returns: + - tiled_obs (:obj:`torch.Tensor`): Tiled observation. + - action_samples (:obj:`torch.Tensor`): Action samples. + Shapes: + - obs (:obj:`torch.Tensor`): :math:`(B, O)`. + - ebm (:obj:`torch.nn.Module`): :math:`(B, N, O)`. + - tiled_obs (:obj:`torch.Tensor`): :math:`(B, N, O)`. + - action_samples (:obj:`torch.Tensor`): :math:`(B, N, A)`. + Examples: + >>> obs = torch.randn(2, 4) + >>> ebm = EBM(4, 5) + >>> opt = DFO() + >>> opt.set_action_bounds(np.stack([np.zeros(5), np.ones(5)], axis=0)) + >>> tiled_obs, action_samples = opt.sample(obs, ebm) + """ return self._sample(obs, self.train_samples) @torch.no_grad() def infer(self, obs: torch.Tensor, ebm: nn.Module) -> torch.Tensor: + """ + Overview: + Optimize for the best action conditioned on the current observation. + Arguments: + - obs (:obj:`torch.Tensor`): Observations. + - ebm (:obj:`torch.nn.Module`): Energy based model. + Returns: + - best_action_samples (:obj:`torch.Tensor`): Actions. + Shapes: + - obs (:obj:`torch.Tensor`): :math:`(B, O)`. + - ebm (:obj:`torch.nn.Module`): :math:`(B, N, O)`. + - best_action_samples (:obj:`torch.Tensor`): :math:`(B, A)`. + Examples: + >>> obs = torch.randn(2, 4) + >>> ebm = EBM(4, 5) + >>> opt = DFO() + >>> opt.set_action_bounds(np.stack([np.zeros(5), np.ones(5)], axis=0)) + >>> best_action_samples = opt.infer(obs, ebm) + """ noise_scale = self.noise_scale # (B, N, O), (B, N, A) @@ -181,6 +284,13 @@ def infer(self, obs: torch.Tensor, ebm: nn.Module) -> torch.Tensor: @STOCHASTIC_OPTIMIZER_REGISTRY.register('ardfo') class AutoRegressiveDFO(DFO): + """ + Overview: + AutoRegressive Derivative-Free Optimizer in paper Implicit Behavioral Cloning. + https://arxiv.org/abs/2109.00137 + Interface: + ``__init__``, ``infer`` + """ def __init__( self, @@ -191,10 +301,40 @@ def __init__( inference_samples: int = 4096, device: str = 'cpu', ): + """ + Overview: + Initialize the AutoRegressive Derivative-Free Optimizer + Arguments: + - noise_scale (:obj:`float`): Initial noise scale. + - noise_shrink (:obj:`float`): Noise scale shrink rate. + - iters (:obj:`int`): Number of iterations. + - train_samples (:obj:`int`): Number of samples for training. + - inference_samples (:obj:`int`): Number of samples for inference. + - device (:obj:`str`): Device. + """ super().__init__(noise_scale, noise_shrink, iters, train_samples, inference_samples, device) @torch.no_grad() def infer(self, obs: torch.Tensor, ebm: nn.Module) -> torch.Tensor: + """ + Overview: + Optimize for the best action conditioned on the current observation. + Arguments: + - obs (:obj:`torch.Tensor`): Observations. + - ebm (:obj:`torch.nn.Module`): Energy based model. + Returns: + - best_action_samples (:obj:`torch.Tensor`): Actions. + Shapes: + - obs (:obj:`torch.Tensor`): :math:`(B, O)`. + - ebm (:obj:`torch.nn.Module`): :math:`(B, N, O)`. + - best_action_samples (:obj:`torch.Tensor`): :math:`(B, A)`. + Examples: + >>> obs = torch.randn(2, 4) + >>> ebm = EBM(4, 5) + >>> opt = AutoRegressiveDFO() + >>> opt.set_action_bounds(np.stack([np.zeros(5), np.ones(5)], axis=0)) + >>> best_action_samples = opt.infer(obs, ebm) + """ noise_scale = self.noise_scale # (B, N, O), (B, N, A) @@ -230,38 +370,91 @@ def infer(self, obs: torch.Tensor, ebm: nn.Module) -> torch.Tensor: @STOCHASTIC_OPTIMIZER_REGISTRY.register('mcmc') class MCMC(StochasticOptimizer): + """ + Overview: + MCMC method as stochastic optimizers in paper Implicit Behavioral Cloning. + https://arxiv.org/abs/2109.00137 + Interface: + ``__init__``, ``sample``, ``infer``, ``grad_penalty`` + """ class BaseScheduler(ABC): + """ + Overview: + Base class for learning rate scheduler. + Interface: + ``get_rate`` + """ @abstractmethod def get_rate(self, index): + """ + Overview: + Abstract method for getting learning rate. + """ raise NotImplementedError class ExponentialScheduler: - """Exponential learning rate schedule for Langevin sampler.""" + """ + Overview: + Exponential learning rate schedule for Langevin sampler. + Interface: + ``__init__``, ``get_rate`` + """ def __init__(self, init, decay): + """ + Overview: + Initialize the ExponentialScheduler. + Arguments: + - init (:obj:`float`): Initial learning rate. + - decay (:obj:`float`): Decay rate. + """ self._decay = decay self._latest_lr = init def get_rate(self, index): - """Get learning rate. Assumes calling sequentially.""" + """ + Overview: + Get learning rate. Assumes calling sequentially. + Arguments: + - index (:obj:`int`): Current iteration. + """ del index lr = self._latest_lr self._latest_lr *= self._decay return lr class PolynomialScheduler: - """Polynomial learning rate schedule for Langevin sampler.""" + """ + Overview: + Polynomial learning rate schedule for Langevin sampler. + Interface: + ``__init__``, ``get_rate`` + """ def __init__(self, init, final, power, num_steps): + """ + Overview: + Initialize the PolynomialScheduler. + Arguments: + - init (:obj:`float`): Initial learning rate. + - final (:obj:`float`): Final learning rate. + - power (:obj:`float`): Power of polynomial. + - num_steps (:obj:`int`): Number of steps. + """ self._init = init self._final = final self._power = power self._num_steps = num_steps def get_rate(self, index): - """Get learning rate for index.""" + """ + Overview: + Get learning rate for index. + Arguments: + - index (:obj:`int`): Current iteration. + """ if index == -1: return self._init return ( @@ -298,6 +491,26 @@ def __init__( grad_loss_weight: float = 1.0, **kwargs, ): + """ + Overview: + Initialize the MCMC. + Arguments: + - iters (:obj:`int`): Number of iterations. + - use_langevin_negative_samples (:obj:`bool`): Whether to use Langevin sampler. + - train_samples (:obj:`int`): Number of samples for training. + - inference_samples (:obj:`int`): Number of samples for inference. + - stepsize_scheduler (:obj:`dict`): Step size scheduler for Langevin sampler. + - optimize_again (:obj:`bool`): Whether to run a second optimization. + - again_stepsize_scheduler (:obj:`dict`): Step size scheduler for the second optimization. + - device (:obj:`str`): Device. + - noise_scale (:obj:`float`): Initial noise scale. + - grad_clip (:obj:`float`): Gradient clip. + - delta_action_clip (:obj:`float`): Action clip. + - add_grad_penalty (:obj:`bool`): Whether to add gradient penalty. + - grad_norm_type (:obj:`str`): Gradient norm type. + - grad_margin (:obj:`float`): Gradient margin. + - grad_loss_weight (:obj:`float`): Gradient loss weight. + """ self.iters = iters self.use_langevin_negative_samples = use_langevin_negative_samples self.train_samples = train_samples @@ -323,9 +536,20 @@ def _gradient_wrt_act( create_graph: bool = False, ) -> torch.Tensor: """ - Calculate gradient w.r.t action. - obs: (B, N, O), action: (B, N, A). - return: (B, N, A). + Overview: + Calculate gradient w.r.t action. + Arguments: + - obs (:obj:`torch.Tensor`): Observations. + - action (:obj:`torch.Tensor`): Actions. + - ebm (:obj:`torch.nn.Module`): Energy based model. + - create_graph (:obj:`bool`): Whether to create graph. + Returns: + - grad (:obj:`torch.Tensor`): Gradient w.r.t action. + Shapes: + - obs (:obj:`torch.Tensor`): :math:`(B, N, O)`. + - action (:obj:`torch.Tensor`): :math:`(B, N, A)`. + - ebm (:obj:`torch.nn.Module`): :math:`(B, N, O)`. + - grad (:obj:`torch.Tensor`): :math:`(B, N, A)`. """ action.requires_grad_(True) energy = ebm.forward(obs, action).sum() @@ -337,9 +561,19 @@ def _gradient_wrt_act( def grad_penalty(self, obs: torch.Tensor, action: torch.Tensor, ebm: nn.Module) -> torch.Tensor: """ - Calculate gradient penalty. - obs: (B, N+1, O), action: (B, N+1, A). - return: loss. + Overview: + Calculate gradient penalty. + Arguments: + - obs (:obj:`torch.Tensor`): Observations. + - action (:obj:`torch.Tensor`): Actions. + - ebm (:obj:`torch.nn.Module`): Energy based model. + Returns: + - loss (:obj:`torch.Tensor`): Gradient penalty. + Shapes: + - obs (:obj:`torch.Tensor`): :math:`(B, N+1, O)`. + - action (:obj:`torch.Tensor`): :math:`(B, N+1, A)`. + - ebm (:obj:`torch.nn.Module`): :math:`(B, N+1, O)`. + - loss (:obj:`torch.Tensor`): :math:`(B, )`. """ if not self.add_grad_penalty: return 0. @@ -371,9 +605,20 @@ def compute_grad_norm(grad_norm_type, de_dact) -> torch.Tensor: @no_ebm_grad() def _langevin_step(self, obs: torch.Tensor, action: torch.Tensor, stepsize: float, ebm: nn.Module) -> torch.Tensor: """ - Run one langevin MCMC step. - obs: (B, N, O), action: (B, N, A) - return: (B, N, A). + Overview: + Run one langevin MCMC step. + Arguments: + - obs (:obj:`torch.Tensor`): Observations. + - action (:obj:`torch.Tensor`): Actions. + - stepsize (:obj:`float`): Step size. + - ebm (:obj:`torch.nn.Module`): Energy based model. + Returns: + - action (:obj:`torch.Tensor`): Actions. + Shapes: + - obs (:obj:`torch.Tensor`): :math:`(B, N, O)`. + - action (:obj:`torch.Tensor`): :math:`(B, N, A)`. + - stepsize (:obj:`float`): :math:`(B, )`. + - ebm (:obj:`torch.nn.Module`): :math:`(B, N, O)`. """ l_lambda = 1.0 de_dact = MCMC._gradient_wrt_act(obs, action, ebm) @@ -402,9 +647,19 @@ def _langevin_action_given_obs( scheduler: BaseScheduler = None ) -> torch.Tensor: """ - Run langevin MCMC for `self.iters` steps. - obs: (B, N, O), action: (B, N, A) - return: (B, N, A) + Overview: + Run langevin MCMC for `self.iters` steps. + Arguments: + - obs (:obj:`torch.Tensor`): Observations. + - action (:obj:`torch.Tensor`): Actions. + - ebm (:obj:`torch.nn.Module`): Energy based model. + - scheduler (:obj:`BaseScheduler`): Learning rate scheduler. + Returns: + - action (:obj:`torch.Tensor`): Actions. + Shapes: + - obs (:obj:`torch.Tensor`): :math:`(B, N, O)`. + - action (:obj:`torch.Tensor`): :math:`(B, N, A)`. + - ebm (:obj:`torch.nn.Module`): :math:`(B, N, O)`. """ if not scheduler: self.stepsize_scheduler['num_steps'] = self.iters @@ -417,6 +672,27 @@ def _langevin_action_given_obs( @no_ebm_grad() def sample(self, obs: torch.Tensor, ebm: nn.Module) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Overview: + Create tiled observations and sample counter-negatives for InfoNCE loss. + Arguments: + - obs (:obj:`torch.Tensor`): Observations. + - ebm (:obj:`torch.nn.Module`): Energy based model. + Returns: + - tiled_obs (:obj:`torch.Tensor`): Tiled observations. + - action_samples (:obj:`torch.Tensor`): Action samples. + Shapes: + - obs (:obj:`torch.Tensor`): :math:`(B, O)`. + - ebm (:obj:`torch.nn.Module`): :math:`(B, N, O)`. + - tiled_obs (:obj:`torch.Tensor`): :math:`(B, N, O)`. + - action_samples (:obj:`torch.Tensor`): :math:`(B, N, A)`. + Examples: + >>> obs = torch.randn(2, 4) + >>> ebm = EBM(4, 5) + >>> opt = MCMC() + >>> opt.set_action_bounds(np.stack([np.zeros(5), np.ones(5)], axis=0)) + >>> tiled_obs, action_samples = opt.sample(obs, ebm) + """ obs, uniform_action_samples = self._sample(obs, self.train_samples) if not self.use_langevin_negative_samples: return obs, uniform_action_samples @@ -425,6 +701,25 @@ def sample(self, obs: torch.Tensor, ebm: nn.Module) -> Tuple[torch.Tensor, torch @no_ebm_grad() def infer(self, obs: torch.Tensor, ebm: nn.Module) -> torch.Tensor: + """ + Overview: + Optimize for the best action conditioned on the current observation. + Arguments: + - obs (:obj:`torch.Tensor`): Observations. + - ebm (:obj:`torch.nn.Module`): Energy based model. + Returns: + - best_action_samples (:obj:`torch.Tensor`): Actions. + Shapes: + - obs (:obj:`torch.Tensor`): :math:`(B, O)`. + - ebm (:obj:`torch.nn.Module`): :math:`(B, N, O)`. + - best_action_samples (:obj:`torch.Tensor`): :math:`(B, A)`. + Examples: + >>> obs = torch.randn(2, 4) + >>> ebm = EBM(4, 5) + >>> opt = MCMC() + >>> opt.set_action_bounds(np.stack([np.zeros(5), np.ones(5)], axis=0)) + >>> best_action_samples = opt.infer(obs, ebm) + """ # (B, N, O), (B, N, A) obs, uniform_action_samples = self._sample(obs, self.inference_samples) action_samples = self._langevin_action_given_obs( @@ -449,6 +744,12 @@ def infer(self, obs: torch.Tensor, ebm: nn.Module) -> torch.Tensor: @MODEL_REGISTRY.register('ebm') class EBM(nn.Module): + """ + Overview: + Energy based model. + Interface: + ``__init__``, ``forward`` + """ def __init__( self, @@ -458,6 +759,15 @@ def __init__( hidden_layer_num: int = 4, **kwargs, ): + """ + Overview: + Initialize the EBM. + Arguments: + - obs_shape (:obj:`int`): Observation shape. + - action_shape (:obj:`int`): Action shape. + - hidden_size (:obj:`int`): Hidden size. + - hidden_layer_num (:obj:`int`): Number of hidden layers. + """ super().__init__() input_size = obs_shape + action_shape self.net = nn.Sequential( @@ -471,9 +781,20 @@ def __init__( ) def forward(self, obs, action): - # obs: (B, N, O) - # action: (B, N, A) - # return: (B, N) + """ + Overview: + Forward computation graph of EBM. + Arguments: + - obs (:obj:`torch.Tensor`): Observation of shape (B, N, O). + - action (:obj:`torch.Tensor`): Action of shape (B, N, A). + Returns: + - pred (:obj:`torch.Tensor`): Energy of shape (B, N). + Examples: + >>> obs = torch.randn(2, 3, 4) + >>> action = torch.randn(2, 3, 5) + >>> ebm = EBM(4, 5) + >>> pred = ebm(obs, action) + """ x = torch.cat([obs, action], -1) x = self.net(x) return x['pred'] @@ -481,6 +802,12 @@ def forward(self, obs, action): @MODEL_REGISTRY.register('arebm') class AutoregressiveEBM(nn.Module): + """ + Overview: + Autoregressive energy based model. + Interface: + ``__init__``, ``forward`` + """ def __init__( self, @@ -488,21 +815,37 @@ def __init__( action_shape: int, hidden_size: int = 512, hidden_layer_num: int = 4, - **kwargs, ): + """ + Overview: + Initialize the AutoregressiveEBM. + Arguments: + - obs_shape (:obj:`int`): Observation shape. + - action_shape (:obj:`int`): Action shape. + - hidden_size (:obj:`int`): Hidden size. + - hidden_layer_num (:obj:`int`): Number of hidden layers. + """ super().__init__() self.ebm_list = nn.ModuleList() for i in range(action_shape): self.ebm_list.append(EBM(obs_shape, i + 1, hidden_size, hidden_layer_num)) def forward(self, obs, action): - # obs: (B, N, O) - # action: (B, N, A) - # return: (B, N, A) - - # (B, N) + """ + Overview: + Forward computation graph of AutoregressiveEBM. + Arguments: + - obs (:obj:`torch.Tensor`): Observation of shape (B, N, O). + - action (:obj:`torch.Tensor`): Action of shape (B, N, A). + Returns: + - pred (:obj:`torch.Tensor`): Energy of shape (B, N, A). + Examples: + >>> obs = torch.randn(2, 3, 4) + >>> action = torch.randn(2, 3, 5) + >>> arebm = AutoregressiveEBM(4, 5) + >>> pred = arebm(obs, action) + """ output_list = [] for i, ebm in enumerate(self.ebm_list): output_list.append(ebm(obs, action[..., :i + 1])) - # (B, N, A) return torch.stack(output_list, axis=-1) diff --git a/ding/model/template/edac.py b/ding/model/template/edac.py new file mode 100755 index 0000000000..397ba69763 --- /dev/null +++ b/ding/model/template/edac.py @@ -0,0 +1,182 @@ +from typing import Union, Optional, Dict +from easydict import EasyDict + +import torch +import torch.nn as nn +from ding.model.common import ReparameterizationHead, EnsembleHead +from ding.utils import SequenceType, squeeze + +from ding.utils import MODEL_REGISTRY + + +@MODEL_REGISTRY.register('edac') +class EDAC(nn.Module): + """ + Overview: + The Q-value Actor-Critic network with the ensemble mechanism, which is used in EDAC. + Interfaces: + ``__init__``, ``forward``, ``compute_actor``, ``compute_critic`` + """ + mode = ['compute_actor', 'compute_critic'] + + def __init__( + self, + obs_shape: Union[int, SequenceType], + action_shape: Union[int, SequenceType, EasyDict], + ensemble_num: int = 2, + actor_head_hidden_size: int = 64, + actor_head_layer_num: int = 1, + critic_head_hidden_size: int = 64, + critic_head_layer_num: int = 1, + activation: Optional[nn.Module] = nn.ReLU(), + norm_type: Optional[str] = None, + **kwargs + ) -> None: + """ + Overview: + Initailize the EDAC Model according to input arguments. + Arguments: + - obs_shape (:obj:`Union[int, SequenceType]`): Observation's shape, such as 128, (156, ). + - action_shape (:obj:`Union[int, SequenceType, EasyDict]`): Action's shape, such as 4, (3, ), \ + EasyDict({'action_type_shape': 3, 'action_args_shape': 4}). + - ensemble_num (:obj:`int`): Q-net number. + - actor_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to actor head. + - actor_head_layer_num (:obj:`int`): The num of layers used in the network to compute Q value output \ + for actor head. + - critic_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to critic head. + - critic_head_layer_num (:obj:`int`): The num of layers used in the network to compute Q value output \ + for critic head. + - activation (:obj:`Optional[nn.Module]`): The type of activation function to use in ``MLP`` \ + after each FC layer, if ``None`` then default set to ``nn.ReLU()``. + - norm_type (:obj:`Optional[str]`): The type of normalization to after network layer (FC, Conv), \ + see ``ding.torch_utils.network`` for more details. + """ + super(EDAC, self).__init__() + obs_shape: int = squeeze(obs_shape) + action_shape = squeeze(action_shape) + self.action_shape = action_shape + self.ensemble_num = ensemble_num + self.actor = nn.Sequential( + nn.Linear(obs_shape, actor_head_hidden_size), activation, + ReparameterizationHead( + actor_head_hidden_size, + action_shape, + actor_head_layer_num, + sigma_type='conditioned', + activation=activation, + norm_type=norm_type + ) + ) + + critic_input_size = obs_shape + action_shape + self.critic = EnsembleHead( + critic_input_size, + 1, + critic_head_hidden_size, + critic_head_layer_num, + self.ensemble_num, + activation=activation, + norm_type=norm_type + ) + + def forward(self, inputs: Union[torch.Tensor, Dict[str, torch.Tensor]], mode: str) -> Dict[str, torch.Tensor]: + """ + Overview: + The unique execution (forward) method of EDAC method, and one can indicate different modes to implement \ + different computation graph, including ``compute_actor`` and ``compute_critic`` in EDAC. + Mode compute_actor: + Arguments: + - inputs (:obj:`torch.Tensor`): Observation data, defaults to tensor. + Returns: + - output (:obj:`Dict`): Output dict data, including differnet key-values among distinct action_space. + Mode compute_critic: + Arguments: + - inputs (:obj:`Dict`): Input dict data, including obs and action tensor. + Returns: + - output (:obj:`Dict`): Output dict data, including q_value tensor. + + .. note:: + For specific examples, one can refer to API doc of ``compute_actor`` and ``compute_critic`` respectively. + """ + assert mode in self.mode, "not support forward mode: {}/{}".format(mode, self.mode) + return getattr(self, mode)(inputs) + + def compute_actor(self, obs: torch.Tensor) -> Dict[str, Union[torch.Tensor, Dict[str, torch.Tensor]]]: + """ + Overview: + The forward computation graph of compute_actor mode, uses observation tensor to produce actor output, + such as ``action``, ``logit`` and so on. + Arguments: + - obs (:obj:`torch.Tensor`): Observation tensor data, now supports a batch of 1-dim vector data, \ + i.e. ``(B, obs_shape)``. + Returns: + - outputs (:obj:`Dict[str, Union[torch.Tensor, Dict[str, torch.Tensor]]]`): Actor output varying \ + from action_space: ``reparameterization``. + ReturnsKeys (either): + - logit (:obj:`Dict[str, torch.Tensor]`): Reparameterization logit, usually in SAC. + - mu (:obj:`torch.Tensor`): Mean of parameterization gaussion distribution. + - sigma (:obj:`torch.Tensor`): Standard variation of parameterization gaussion distribution. + Shapes: + - obs (:obj:`torch.Tensor`): :math:`(B, N0)`, B is batch size and N0 corresponds to ``obs_shape``. + - action (:obj:`torch.Tensor`): :math:`(B, N1)`, B is batch size and N1 corresponds to ``action_shape``. + - logit.mu (:obj:`torch.Tensor`): :math:`(B, N1)`, B is batch size and N1 corresponds to ``action_shape``. + - logit.sigma (:obj:`torch.Tensor`): :math:`(B, N1)`, B is batch size. + - logit (:obj:`torch.Tensor`): :math:`(B, N2)`, B is batch size and N2 corresponds to \ + ``action_shape.action_type_shape``. + - action_args (:obj:`torch.Tensor`): :math:`(B, N3)`, B is batch size and N3 corresponds to \ + ``action_shape.action_args_shape``. + Examples: + >>> model = EDAC(64, 64,) + >>> obs = torch.randn(4, 64) + >>> actor_outputs = model(obs,'compute_actor') + >>> assert actor_outputs['logit'][0].shape == torch.Size([4, 64]) # mu + >>> actor_outputs['logit'][1].shape == torch.Size([4, 64]) # sigma + """ + x = self.actor(obs) + return {'logit': [x['mu'], x['sigma']]} + + def compute_critic(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: + """ + Overview: + The forward computation graph of compute_critic mode, uses observation and action tensor to produce critic + output, such as ``q_value``. + Arguments: + - inputs (:obj:`Dict[str, torch.Tensor]`): Dict strcture of input data, including ``obs`` and \ + ``action`` tensor + Returns: + - outputs (:obj:`Dict[str, torch.Tensor]`): Critic output, such as ``q_value``. + ArgumentsKeys: + - obs: (:obj:`torch.Tensor`): Observation tensor data, now supports a batch of 1-dim vector data. + - action (:obj:`Union[torch.Tensor, Dict]`): Continuous action with same size as ``action_shape``. + ReturnKeys: + - q_value (:obj:`torch.Tensor`): Q value tensor with same size as batch size. + Shapes: + - obs (:obj:`torch.Tensor`): :math:`(B, N1)` or '(Ensemble_num, B, N1)', where B is batch size and N1 is \ + ``obs_shape``. + - action (:obj:`torch.Tensor`): :math:`(B, N2)` or '(Ensemble_num, B, N2)', where B is batch size and N4 \ + is ``action_shape``. + - q_value (:obj:`torch.Tensor`): :math:`(Ensemble_num, B)`, where B is batch size. + Examples: + >>> inputs = {'obs': torch.randn(4, 8), 'action': torch.randn(4, 1)} + >>> model = EDAC(obs_shape=(8, ),action_shape=1) + >>> model(inputs, mode='compute_critic')['q_value'] # q value + ... tensor([0.0773, 0.1639, 0.0917, 0.0370], grad_fn=) + """ + + obs, action = inputs['obs'], inputs['action'] + if len(action.shape) == 1: # (B, ) -> (B, 1) + action = action.unsqueeze(1) + x = torch.cat([obs, action], dim=-1) + if len(obs.shape) < 3: + # [batch_size,dim] -> [batch_size,Ensemble_num * dim,1] + x = x.repeat(1, self.ensemble_num).unsqueeze(-1) + else: + # [Ensemble_num,batch_size,dim] -> [batch_size,Ensemble_num,dim] -> [batch_size,Ensemble_num * dim, 1] + x = x.transpose(0, 1) + batch_size = obs.shape[1] + x = x.reshape(batch_size, -1, 1) + # [Ensemble_num,batch_size,1] + x = self.critic(x)['pred'] + # [batch_size,1*Ensemble_num] -> [Ensemble_num,batch_size] + x = x.permute(1, 0) + return {'q_value': x} diff --git a/ding/model/template/havac.py b/ding/model/template/havac.py new file mode 100644 index 0000000000..77489ed517 --- /dev/null +++ b/ding/model/template/havac.py @@ -0,0 +1,500 @@ +from typing import Union, Dict, Optional +import torch +import torch.nn as nn + +from ding.torch_utils import get_lstm +from ding.utils import SequenceType, squeeze, MODEL_REGISTRY +from ding.model.template.q_learning import parallel_wrapper +from ..common import ReparameterizationHead, RegressionHead, DiscreteHead, \ + FCEncoder, ConvEncoder + + +class RNNLayer(nn.Module): + + def __init__(self, lstm_type, input_size, hidden_size, res_link: bool = False): + super(RNNLayer, self).__init__() + self.rnn = get_lstm(lstm_type, input_size=input_size, hidden_size=hidden_size) + self.res_link = res_link + + def forward(self, x, prev_state, inference: bool = False): + """ + Forward pass of the RNN layer. + If inference is True, sequence length of input is set to 1. + If res_link is True, a residual link is added to the output. + """ + # x: obs_embedding + if self.res_link: + a = x + if inference: + x = x.unsqueeze(0) # for rnn input, put the seq_len of x as 1 instead of none. + # prev_state: DataType: List[Tuple[torch.Tensor]]; Initially, it is a list of None + x, next_state = self.rnn(x, prev_state) + x = x.squeeze(0) # to delete the seq_len dim to match head network input + if self.res_link: + x = x + a + return {'output': x, 'next_state': next_state} + else: + # lstm_embedding stores all hidden_state + lstm_embedding = [] + hidden_state_list = [] + for t in range(x.shape[0]): # T timesteps + # use x[t:t+1] but not x[t] can keep original dimension + output, prev_state = self.rnn(x[t:t + 1], prev_state) # output: (1,B, head_hidden_size) + lstm_embedding.append(output) + hidden_state = [p['h'] for p in prev_state] + # only keep ht, {list: x.shape[0]{Tensor:(1, batch_size, head_hidden_size)}} + hidden_state_list.append(torch.cat(hidden_state, dim=1)) + x = torch.cat(lstm_embedding, 0) # (T, B, head_hidden_size) + if self.res_link: + x = x + a + all_hidden_state = torch.cat(hidden_state_list, dim=0) + return {'output': x, 'next_state': prev_state, 'hidden_state': all_hidden_state} + + +@MODEL_REGISTRY.register('havac') +class HAVAC(nn.Module): + """ + Overview: + The HAVAC model of each agent for HAPPO. + Interfaces: + ``__init__``, ``forward`` + """ + mode = ['compute_actor', 'compute_critic', 'compute_actor_critic'] + + def __init__( + self, + agent_obs_shape: Union[int, SequenceType], + global_obs_shape: Union[int, SequenceType], + action_shape: Union[int, SequenceType], + agent_num: int, + use_lstm: bool = False, + lstm_type: str = 'gru', + encoder_hidden_size_list: SequenceType = [128, 128, 64], + actor_head_hidden_size: int = 64, + actor_head_layer_num: int = 2, + critic_head_hidden_size: int = 64, + critic_head_layer_num: int = 1, + action_space: str = 'discrete', + activation: Optional[nn.Module] = nn.ReLU(), + norm_type: Optional[str] = None, + sigma_type: Optional[str] = 'independent', + bound_type: Optional[str] = None, + res_link: bool = False, + ) -> None: + r""" + Overview: + Init the VAC Model for HAPPO according to arguments. + Arguments: + - agent_obs_shape (:obj:`Union[int, SequenceType]`): Observation's space for single agent. + - global_obs_shape (:obj:`Union[int, SequenceType]`): Observation's space for global agent + - action_shape (:obj:`Union[int, SequenceType]`): Action's space. + - agent_num (:obj:`int`): Number of agents. + - lstm_type (:obj:`str`): use lstm or gru, default to gru + - encoder_hidden_size_list (:obj:`SequenceType`): Collection of ``hidden_size`` to pass to ``Encoder`` + - actor_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to actor-nn's ``Head``. + - actor_head_layer_num (:obj:`int`): + The num of layers used in the network to compute Q value output for actor's nn. + - critic_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to critic-nn's ``Head``. + - critic_head_layer_num (:obj:`int`): + The num of layers used in the network to compute Q value output for critic's nn. + - activation (:obj:`Optional[nn.Module]`): + The type of activation function to use in ``MLP`` the after ``layer_fn``, + if ``None`` then default set to ``nn.ReLU()`` + - norm_type (:obj:`Optional[str]`): + The type of normalization to use, see ``ding.torch_utils.fc_block`` for more details` + - res_link (:obj:`bool`): use the residual link or not, default to False + """ + super(HAVAC, self).__init__() + self.agent_num = agent_num + self.agent_models = nn.ModuleList( + [ + HAVACAgent( + agent_obs_shape=agent_obs_shape, + global_obs_shape=global_obs_shape, + action_shape=action_shape, + use_lstm=use_lstm, + action_space=action_space, + ) for _ in range(agent_num) + ] + ) + + def forward(self, agent_idx, input_data, mode): + selected_agent_model = self.agent_models[agent_idx] + output = selected_agent_model(input_data, mode) + return output + + +class HAVACAgent(nn.Module): + """ + Overview: + The HAVAC model of each agent for HAPPO. + Interfaces: + ``__init__``, ``forward``, ``compute_actor``, ``compute_critic``, ``compute_actor_critic`` + """ + mode = ['compute_actor', 'compute_critic', 'compute_actor_critic'] + + def __init__( + self, + agent_obs_shape: Union[int, SequenceType], + global_obs_shape: Union[int, SequenceType], + action_shape: Union[int, SequenceType], + use_lstm: bool = False, + lstm_type: str = 'gru', + encoder_hidden_size_list: SequenceType = [128, 128, 64], + actor_head_hidden_size: int = 64, + actor_head_layer_num: int = 2, + critic_head_hidden_size: int = 64, + critic_head_layer_num: int = 1, + action_space: str = 'discrete', + activation: Optional[nn.Module] = nn.ReLU(), + norm_type: Optional[str] = None, + sigma_type: Optional[str] = 'happo', + bound_type: Optional[str] = None, + res_link: bool = False, + ) -> None: + r""" + Overview: + Init the VAC Model for HAPPO according to arguments. + Arguments: + - agent_obs_shape (:obj:`Union[int, SequenceType]`): Observation's space for single agent. + - global_obs_shape (:obj:`Union[int, SequenceType]`): Observation's space for global agent + - action_shape (:obj:`Union[int, SequenceType]`): Action's space. + - lstm_type (:obj:`str`): use lstm or gru, default to gru + - encoder_hidden_size_list (:obj:`SequenceType`): Collection of ``hidden_size`` to pass to ``Encoder`` + - actor_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to actor-nn's ``Head``. + - actor_head_layer_num (:obj:`int`): + The num of layers used in the network to compute Q value output for actor's nn. + - critic_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to critic-nn's ``Head``. + - critic_head_layer_num (:obj:`int`): + The num of layers used in the network to compute Q value output for critic's nn. + - activation (:obj:`Optional[nn.Module]`): + The type of activation function to use in ``MLP`` the after ``layer_fn``, + if ``None`` then default set to ``nn.ReLU()`` + - norm_type (:obj:`Optional[str]`): + The type of normalization to use, see ``ding.torch_utils.fc_block`` for more details` + - res_link (:obj:`bool`): use the residual link or not, default to False + """ + super(HAVACAgent, self).__init__() + agent_obs_shape: int = squeeze(agent_obs_shape) + global_obs_shape: int = squeeze(global_obs_shape) + action_shape: int = squeeze(action_shape) + self.global_obs_shape, self.agent_obs_shape, self.action_shape = global_obs_shape, agent_obs_shape, action_shape + self.action_space = action_space + # Encoder Type + if isinstance(agent_obs_shape, int) or len(agent_obs_shape) == 1: + actor_encoder_cls = FCEncoder + elif len(agent_obs_shape) == 3: + actor_encoder_cls = ConvEncoder + else: + raise RuntimeError( + "not support obs_shape for pre-defined encoder: {}, please customize your own VAC". + format(agent_obs_shape) + ) + if isinstance(global_obs_shape, int) or len(global_obs_shape) == 1: + critic_encoder_cls = FCEncoder + elif len(global_obs_shape) == 3: + critic_encoder_cls = ConvEncoder + else: + raise RuntimeError( + "not support obs_shape for pre-defined encoder: {}, please customize your own VAC". + format(global_obs_shape) + ) + + # We directly connect the Head after a Liner layer instead of using the 3-layer FCEncoder. + # In SMAC task it can obviously improve the performance. + # Users can change the model according to their own needs. + self.actor_encoder = actor_encoder_cls( + obs_shape=agent_obs_shape, + hidden_size_list=encoder_hidden_size_list, + activation=activation, + norm_type=norm_type + ) + self.critic_encoder = critic_encoder_cls( + obs_shape=global_obs_shape, + hidden_size_list=encoder_hidden_size_list, + activation=activation, + norm_type=norm_type + ) + # RNN part + self.use_lstm = use_lstm + if self.use_lstm: + self.actor_rnn = RNNLayer( + lstm_type, + input_size=encoder_hidden_size_list[-1], + hidden_size=actor_head_hidden_size, + res_link=res_link + ) + self.critic_rnn = RNNLayer( + lstm_type, + input_size=encoder_hidden_size_list[-1], + hidden_size=critic_head_hidden_size, + res_link=res_link + ) + # Head Type + self.critic_head = RegressionHead( + critic_head_hidden_size, 1, critic_head_layer_num, activation=activation, norm_type=norm_type + ) + assert self.action_space in ['discrete', 'continuous'], self.action_space + if self.action_space == 'discrete': + self.actor_head = DiscreteHead( + actor_head_hidden_size, action_shape, actor_head_layer_num, activation=activation, norm_type=norm_type + ) + elif self.action_space == 'continuous': + self.actor_head = ReparameterizationHead( + actor_head_hidden_size, + action_shape, + actor_head_layer_num, + sigma_type=sigma_type, + activation=activation, + norm_type=norm_type, + bound_type=bound_type + ) + # must use list, not nn.ModuleList + self.actor = [self.actor_encoder, self.actor_rnn, self.actor_head] if self.use_lstm \ + else [self.actor_encoder, self.actor_head] + self.critic = [self.critic_encoder, self.critic_rnn, self.critic_head] if self.use_lstm \ + else [self.critic_encoder, self.critic_head] + # for convenience of call some apis(such as: self.critic.parameters()), but may cause + # misunderstanding when print(self) + self.actor = nn.ModuleList(self.actor) + self.critic = nn.ModuleList(self.critic) + + def forward(self, inputs: Union[torch.Tensor, Dict], mode: str) -> Dict: + r""" + Overview: + Use encoded embedding tensor to predict output. + Parameter updates with VAC's MLPs forward setup. + Arguments: + Forward with ``'compute_actor'`` or ``'compute_critic'``: + - inputs (:obj:`torch.Tensor`): + The encoded embedding tensor, determined with given ``hidden_size``, i.e. ``(B, N=hidden_size)``. + Whether ``actor_head_hidden_size`` or ``critic_head_hidden_size`` depend on ``mode``. + Returns: + - outputs (:obj:`Dict`): + Run with encoder and head. + + Forward with ``'compute_actor'``, Necessary Keys: + - logit (:obj:`torch.Tensor`): Logit encoding tensor, with same size as input ``x``. + + Forward with ``'compute_critic'``, Necessary Keys: + - value (:obj:`torch.Tensor`): Q value tensor with same size as batch size. + Shapes: + - inputs (:obj:`torch.Tensor`): :math:`(B, N)`, where B is batch size and N corresponding ``hidden_size`` + - logit (:obj:`torch.FloatTensor`): :math:`(B, N)`, where B is batch size and N is ``action_shape`` + - value (:obj:`torch.FloatTensor`): :math:`(B, )`, where B is batch size. + + Actor Examples: + >>> model = VAC(64,128) + >>> inputs = torch.randn(4, 64) + >>> actor_outputs = model(inputs,'compute_actor') + >>> assert actor_outputs['logit'].shape == torch.Size([4, 128]) + + Critic Examples: + >>> model = VAC(64,64) + >>> inputs = torch.randn(4, 64) + >>> critic_outputs = model(inputs,'compute_critic') + >>> critic_outputs['value'] + tensor([0.0252, 0.0235, 0.0201, 0.0072], grad_fn=) + + Actor-Critic Examples: + >>> model = VAC(64,64) + >>> inputs = torch.randn(4, 64) + >>> outputs = model(inputs,'compute_actor_critic') + >>> outputs['value'] + tensor([0.0252, 0.0235, 0.0201, 0.0072], grad_fn=) + >>> assert outputs['logit'].shape == torch.Size([4, 64]) + + """ + assert mode in self.mode, "not support forward mode: {}/{}".format(mode, self.mode) + return getattr(self, mode)(inputs) + + def compute_actor(self, inputs: Dict, inference: bool = False) -> Dict: + r""" + Overview: + Execute parameter updates with ``'compute_actor'`` mode + Use encoded embedding tensor to predict output. + Arguments: + - inputs (:obj:`torch.Tensor`): + input data dict with keys ['obs'(with keys ['agent_state', 'global_state', 'action_mask']), + 'actor_prev_state'] + Returns: + - outputs (:obj:`Dict`): + Run with encoder RNN(optional) and head. + + ReturnsKeys: + - logit (:obj:`torch.Tensor`): Logit encoding tensor. + - actor_next_state: + - hidden_state + Shapes: + - logit (:obj:`torch.FloatTensor`): :math:`(B, N)`, where B is batch size and N is ``action_shape`` + - actor_next_state: (B,) + - hidden_state: + + Examples: + >>> model = HAVAC( + agent_obs_shape=obs_dim, + global_obs_shape=global_obs_dim, + action_shape=action_dim, + use_lstm = True, + ) + >>> inputs = { + 'obs': { + 'agent_state': torch.randn(T, bs, obs_dim), + 'global_state': torch.randn(T, bs, global_obs_dim), + 'action_mask': torch.randint(0, 2, size=(T, bs, action_dim)) + }, + 'actor_prev_state': [None for _ in range(bs)], + } + >>> actor_outputs = model(inputs,'compute_actor') + >>> assert actor_outputs['logit'].shape == (T, bs, action_dim) + """ + x = inputs['obs']['agent_state'] + output = {} + if self.use_lstm: + rnn_actor_prev_state = inputs['actor_prev_state'] + if inference: + x = self.actor_encoder(x) + rnn_output = self.actor_rnn(x, rnn_actor_prev_state, inference) + x = rnn_output['output'] + x = self.actor_head(x) + output['next_state'] = rnn_output['next_state'] + # output: 'logit'/'next_state' + else: + assert len(x.shape) in [3, 5], x.shape + x = parallel_wrapper(self.actor_encoder)(x) # (T, B, N) + rnn_output = self.actor_rnn(x, rnn_actor_prev_state, inference) + x = rnn_output['output'] + x = parallel_wrapper(self.actor_head)(x) + output['actor_next_state'] = rnn_output['next_state'] + output['actor_hidden_state'] = rnn_output['hidden_state'] + # output: 'logit'/'actor_next_state'/'hidden_state' + else: + x = self.actor_encoder(x) + x = self.actor_head(x) + # output: 'logit' + + if self.action_space == 'discrete': + action_mask = inputs['obs']['action_mask'] + logit = x['logit'] + logit[action_mask == 0.0] = -99999999 + elif self.action_space == 'continuous': + logit = x + output['logit'] = logit + return output + + def compute_critic(self, inputs: Dict, inference: bool = False) -> Dict: + r""" + Overview: + Execute parameter updates with ``'compute_critic'`` mode + Use encoded embedding tensor to predict output. + Arguments: + - inputs (:obj:`Dict`): + input data dict with keys ['obs'(with keys ['agent_state', 'global_state', 'action_mask']), + 'critic_prev_state'(when you are using rnn)] + Returns: + - outputs (:obj:`Dict`): + Run with encoder [rnn] and head. + + Necessary Keys: + - value (:obj:`torch.Tensor`): Q value tensor with same size as batch size. + - logits + Shapes: + - value (:obj:`torch.FloatTensor`): :math:`(B, )`, where B is batch size. + - logits + + Examples: + >>> model = HAVAC( + agent_obs_shape=obs_dim, + global_obs_shape=global_obs_dim, + action_shape=action_dim, + use_lstm = True, + ) + >>> inputs = { + 'obs': { + 'agent_state': torch.randn(T, bs, obs_dim), + 'global_state': torch.randn(T, bs, global_obs_dim), + 'action_mask': torch.randint(0, 2, size=(T, bs, action_dim)) + }, + 'critic_prev_state': [None for _ in range(bs)], + } + >>> critic_outputs = model(inputs,'compute_critic') + >>> assert critic_outputs['value'].shape == (T, bs)) + """ + global_obs = inputs['obs']['global_state'] + output = {} + if self.use_lstm: + rnn_critic_prev_state = inputs['critic_prev_state'] + if inference: + x = self.critic_encoder(global_obs) + rnn_output = self.critic_rnn(x, rnn_critic_prev_state, inference) + x = rnn_output['output'] + x = self.critic_head(x) + output['next_state'] = rnn_output['next_state'] + # output: 'value'/'next_state' + else: + assert len(global_obs.shape) in [3, 5], global_obs.shape + x = parallel_wrapper(self.critic_encoder)(global_obs) # (T, B, N) + rnn_output = self.critic_rnn(x, rnn_critic_prev_state, inference) + x = rnn_output['output'] + x = parallel_wrapper(self.critic_head)(x) + output['critic_next_state'] = rnn_output['next_state'] + output['critic_hidden_state'] = rnn_output['hidden_state'] + # output: 'value'/'critic_next_state'/'hidden_state' + else: + x = self.critic_encoder(global_obs) + x = self.critic_head(x) + # output: 'value' + output['value'] = x['pred'] + return output + + def compute_actor_critic(self, inputs: Dict, inference: bool = False) -> Dict: + r""" + Overview: + Execute parameter updates with ``'compute_actor_critic'`` mode + Use encoded embedding tensor to predict output. + Arguments: + - inputs (:dict): input data dict with keys + ['obs'(with keys ['agent_state', 'global_state', 'action_mask']), + 'actor_prev_state', 'critic_prev_state'(when you are using rnn)] + + Returns: + - outputs (:obj:`Dict`): + Run with encoder and head. + + ReturnsKeys: + - logit (:obj:`torch.Tensor`): Logit encoding tensor, with same size as input ``x``. + - value (:obj:`torch.Tensor`): Q value tensor with same size as batch size. + Shapes: + - logit (:obj:`torch.FloatTensor`): :math:`(B, N)`, where B is batch size and N is ``action_shape`` + - value (:obj:`torch.FloatTensor`): :math:`(B, )`, where B is batch size. + + Examples: + >>> model = VAC(64,64) + >>> inputs = torch.randn(4, 64) + >>> outputs = model(inputs,'compute_actor_critic') + >>> outputs['value'] + tensor([0.0252, 0.0235, 0.0201, 0.0072], grad_fn=) + >>> assert outputs['logit'].shape == torch.Size([4, 64]) + + + .. note:: + ``compute_actor_critic`` interface aims to save computation when shares encoder. + Returning the combination dictionry. + + """ + actor_output = self.compute_actor(inputs, inference) + critic_output = self.compute_critic(inputs, inference) + if self.use_lstm: + return { + 'logit': actor_output['logit'], + 'value': critic_output['value'], + 'actor_next_state': actor_output['actor_next_state'], + 'actor_hidden_state': actor_output['actor_hidden_state'], + 'critic_next_state': critic_output['critic_next_state'], + 'critic_hidden_state': critic_output['critic_hidden_state'], + } + else: + return { + 'logit': actor_output['logit'], + 'value': critic_output['value'], + } diff --git a/ding/model/template/hpt.py b/ding/model/template/hpt.py new file mode 100644 index 0000000000..6ca1209c12 --- /dev/null +++ b/ding/model/template/hpt.py @@ -0,0 +1,207 @@ +from typing import Union, Optional, Dict, Callable, List +from einops import rearrange, repeat +import torch +import torch.nn as nn +from ding.model.common.head import DuelingHead +from ding.utils import MODEL_REGISTRY, squeeze + + +@MODEL_REGISTRY.register('hpt') +class HPT(nn.Module): + """ + Overview: + The HPT model for reinforcement learning, which consists of a Policy Stem and a Dueling Head. + The Policy Stem utilizes cross-attention to process input data, + and the Dueling Head computes Q-values for discrete action spaces. + + Interfaces: + ``__init__``, ``forward`` + + GitHub: [https://github.com/liruiw/HPT/blob/main/hpt/models/policy_stem.py] + + """ + + def __init__(self, state_dim: int, action_dim: int): + """ + Overview: + Initialize the HPT model, including the Policy Stem and the Dueling Head. + + Arguments: + - state_dim (:obj:`int`): The dimension of the input state. + - action_dim (:obj:`int`): The dimension of the action space. + + .. note:: + The Policy Stem is initialized with cross-attention, + and the Dueling Head is set to process the resulting tokens. + """ + super(HPT, self).__init__() + # Initialise Policy Stem + self.policy_stem = PolicyStem(state_dim, 128) + self.policy_stem.init_cross_attn() + + action_dim = squeeze(action_dim) + # Dueling Head, input is 16*128, output is action dimension + self.head = DuelingHead(hidden_size=16 * 128, output_size=action_dim) + + def forward(self, x: torch.Tensor): + """ + Overview: + Forward pass of the HPT model. + Computes latent tokens from the input state and passes them through the Dueling Head. + + Arguments: + - x (:obj:`torch.Tensor`): The input tensor representing the state. + + Returns: + - q_values (:obj:`torch.Tensor`): The predicted Q-values for each action. + """ + # Policy Stem Outputs [B, 16, 128] + tokens = self.policy_stem.compute_latent(x) + # Flatten Operation + tokens_flattened = tokens.view(tokens.size(0), -1) # [B, 16*128] + # Enter to Dueling Head + q_values = self.head(tokens_flattened) + return q_values + + +class PolicyStem(nn.Module): + """ + Overview: + The Policy Stem module is responsible for processing input features + and generating latent tokens using a cross-attention mechanism. + It extracts features from the input and then applies cross-attention + to generate a set of latent tokens. + + Interfaces: + ``__init__``, ``init_cross_attn``, ``compute_latent``, ``forward`` + + .. note:: + This module is inspired by the implementation in the Perceiver IO model + and uses attention mechanisms for feature extraction. + """ + INIT_CONST = 0.02 + + def __init__(self, feature_dim: int = 8, token_dim: int = 128): + """ + Overview: + Initialize the Policy Stem module with a feature extractor and cross-attention mechanism. + + Arguments: + - feature_dim (:obj:`int`): The dimension of the input features. + - token_dim (:obj:`int`): The dimension of the latent tokens generated + by the attention mechanism. + """ + super().__init__() + # Initialise the feature extraction module + self.feature_extractor = nn.Linear(feature_dim, token_dim) + # Initialise CrossAttention + self.init_cross_attn() + + def init_cross_attn(self): + """Initialize cross-attention module and learnable tokens.""" + token_num = 16 + self.tokens = nn.Parameter(torch.randn(1, token_num, 128) * self.INIT_CONST) + self.cross_attention = CrossAttention(128, heads=8, dim_head=64, dropout=0.1) + + def compute_latent(self, x: torch.Tensor) -> torch.Tensor: + """ + Overview: + Compute latent representations of the input data using + the feature extractor and cross-attention. + + Arguments: + - x (:obj:`torch.Tensor`): Input tensor with shape [B, T, ..., F]. + + Returns: + - stem_tokens (:obj:`torch.Tensor`): Latent tokens with shape [B, 16, 128]. + """ + # Using the Feature Extractor + stem_feat = self.feature_extractor(x) + stem_feat = stem_feat.reshape(stem_feat.shape[0], -1, stem_feat.shape[-1]) # (B, N, 128) + # Calculating latent tokens using CrossAttention + stem_tokens = self.tokens.repeat(len(stem_feat), 1, 1) # (B, 16, 128) + stem_tokens = self.cross_attention(stem_tokens, stem_feat) # (B, 16, 128) + return stem_tokens + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Overview: + Forward pass to compute latent tokens. + + Arguments: + - x (:obj:`torch.Tensor`): Input tensor. + + Returns: + - torch.Tensor: Latent tokens tensor. + """ + return self.compute_latent(x) + + @property + def device(self) -> torch.device: + """Returns the device on which the model parameters are located.""" + return next(self.parameters()).device + + +class CrossAttention(nn.Module): + + def __init__(self, query_dim: int, heads: int = 8, dim_head: int = 64, dropout: float = 0.0): + """ + Overview: + CrossAttention module used in the Perceiver IO model. + It computes the attention between the query and context tensors, + and returns the output tensor after applying attention. + + Arguments: + - query_dim (:obj:`int`): The dimension of the query input. + - heads (:obj:`int`, optional): The number of attention heads. Defaults to 8. + - dim_head (:obj:`int`, optional): The dimension of each attention head. Defaults to 64. + - dropout (:obj:`float`, optional): The dropout probability. Defaults to 0.0. + """ + super().__init__() + inner_dim = dim_head * heads + context_dim = query_dim + # Scaling factor for the attention logits to ensure stable gradients. + self.scale = dim_head ** -0.5 + self.heads = heads + + self.to_q = nn.Linear(query_dim, inner_dim, bias=False) + self.to_kv = nn.Linear(context_dim, inner_dim * 2, bias=False) + self.to_out = nn.Linear(inner_dim, query_dim) + + self.dropout = nn.Dropout(dropout) + + def forward(self, x: torch.Tensor, context: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor: + """ + Overview: + Forward pass of the CrossAttention module. + Computes the attention between the query and context tensors. + + Arguments: + - x (:obj:`torch.Tensor`): The query input tensor. + - context (:obj:`torch.Tensor`): The context input tensor. + - mask (:obj:`Optional[torch.Tensor]`): The attention mask tensor. Defaults to None. + + Returns: + - torch.Tensor: The output tensor after applying attention. + """ + h = self.heads + q = self.to_q(x) + k, v = self.to_kv(context).chunk(2, dim=-1) + q, k, v = map(lambda t: rearrange(t, "b n (h d) -> (b h) n d", h=h), (q, k, v)) + sim = torch.einsum("b i d, b j d -> b i j", q, k) * self.scale + + if mask is not None: + # fill in the masks with negative values + mask = rearrange(mask, "b ... -> b (...)") + max_neg_value = -torch.finfo(sim.dtype).max + mask = repeat(mask, "b j -> (b h) () j", h=h) + sim.masked_fill_(~mask, max_neg_value) + + # attention, what we cannot get enough of + attn = sim.softmax(dim=-1) + + # dropout + attn = self.dropout(attn) + out = torch.einsum("b i j, b j d -> b i d", attn, v) + out = rearrange(out, "(b h) n d -> b n (h d)", h=h) + return self.to_out(out) diff --git a/ding/model/template/language_transformer.py b/ding/model/template/language_transformer.py new file mode 100644 index 0000000000..796b2ba0a7 --- /dev/null +++ b/ding/model/template/language_transformer.py @@ -0,0 +1,129 @@ +from typing import List, Dict, Optional +import torch +from torch import nn + +try: + from transformers import AutoTokenizer, AutoModelForTokenClassification +except ImportError: + from ditk import logging + logging.warning("not found transformer, please install it using: pip install transformers") +from ding.utils import MODEL_REGISTRY + + +@MODEL_REGISTRY.register('language_transformer') +class LanguageTransformer(nn.Module): + """ + Overview: + The LanguageTransformer network. Download a pre-trained language model and add head on it. + In the default case, we use BERT model as the text encoder, whose bi-directional character is good + for obtaining the embedding of the whole sentence. + Interfaces: + ``__init__``, ``forward`` + """ + mode = ['compute_actor', 'compute_critic', 'compute_actor_critic'] + + def __init__( + self, + model_name: str = "bert-base-uncased", + add_linear: bool = False, + embedding_size: int = 128, + freeze_encoder: bool = True, + hidden_dim: int = 768, + norm_embedding: bool = False + ) -> None: + """ + Overview: + Init the LanguageTransformer Model according to input arguments. + Arguments: + - model_name (:obj:`str`): The base language model name in huggingface, such as "bert-base-uncased". + - add_linear (:obj:`bool`): Whether to add a linear layer on the top of language model, defaults to be \ + ``False``. + - embedding_size (:obj:`int`): The embedding size of the added linear layer, such as 128. + - freeze_encoder (:obj:`bool`): Whether to freeze the encoder language model while training, \ + defaults to be ``True``. + - hidden_dim (:obj:`int`): The embedding dimension of the encoding model (e.g. BERT). This value should \ + correspond to the model you use. For bert-base-uncased, this value is 768. + - norm_embedding (:obj:`bool`): Whether to normalize the embedding vectors. Default to be ``False``. + """ + super().__init__() + self.tokenizer = AutoTokenizer.from_pretrained(model_name) + self.model = AutoModelForTokenClassification.from_pretrained(model_name) + in_channel = hidden_dim if not add_linear else embedding_size + self.value_head = nn.Linear(in_channel, 1) + self.norm = nn.Identity() if not norm_embedding else nn.LayerNorm( + normalized_shape=in_channel, elementwise_affine=False + ) + + # Freeze transformer encoder and only train the linear layer + if freeze_encoder: + for param in self.model.parameters(): + param.requires_grad = False + + if add_linear: + # Add a small, adjustable linear layer on top of language model tuned through RL + self.embedding_size = embedding_size + self.linear = nn.Linear(self.model.config.hidden_size, embedding_size) + else: + self.linear = None + + def _calc_embedding(self, x: list) -> torch.Tensor: + # ``truncation=True`` means that if the length of the prompt exceed the ``max_length`` of the tokenizer, + # the exceeded part will be truncated. ``padding=True`` means that if the length of the prompt does not reach + # the ``max_length``, the latter part will be padded. These settings ensure the length of encoded tokens is + # exactly ``max_length``, which can enable batch-wise computing. + input = self.tokenizer(x, truncation=True, padding=True, return_tensors="pt").to(self.model.device) + output = self.model(**input, output_hidden_states=True) + # Get last layer hidden states + last_hidden_states = output.hidden_states[-1] + # Get [CLS] hidden states + sentence_embedding = last_hidden_states[:, 0, :] # len(input_list) x hidden_size + sentence_embedding = self.norm(sentence_embedding) + + if self.linear: + sentence_embedding = self.linear(sentence_embedding) # len(input_list) x embedding_size + + return sentence_embedding + + def forward( + self, + train_samples: List[str], + candidate_samples: Optional[List[str]] = None, + mode: str = 'compute_actor' + ) -> Dict: + """ + Overview: + LanguageTransformer forward computation graph, input two lists of strings and predict their matching scores. + Different ``mode`` will forward with different network modules to get different outputs. + Arguments: + - train_samples (:obj:`List[str]`): One list of strings. + - candidate_samples (:obj:`Optional[List[str]]`): The other list of strings to calculate matching scores. + - - mode (:obj:`str`): The forward mode, all the modes are defined in the beginning of this class. + Returns: + - output (:obj:`Dict`): Output dict data, including the logit of matching scores and the \ + corresponding ``torch.distributions.Categorical`` object. + + Examples: + >>> test_pids = [1] + >>> cand_pids = [0, 2, 4] + >>> problems = [ \ + "This is problem 0", "This is the first question", "Second problem is here", "Another problem", \ + "This is the last problem" \ + ] + >>> ctxt_list = [problems[pid] for pid in test_pids] + >>> cands_list = [problems[pid] for pid in cand_pids] + >>> model = LanguageTransformer(model_name="bert-base-uncased", add_linear=True, embedding_size=256) + >>> scores = model(ctxt_list, cands_list) + >>> assert scores.shape == (1, 3) + """ + assert mode in self.mode + prompt_embedding = self._calc_embedding(train_samples) + + res_dict = {} + if mode in ['compute_actor', 'compute_actor_critic']: + cands_embedding = self._calc_embedding(candidate_samples) + scores = torch.mm(prompt_embedding, cands_embedding.t()) + res_dict.update({'dist': torch.distributions.Categorical(logits=scores), 'logit': scores}) + if mode in ['compute_critic', 'compute_actor_critic']: + value = self.value_head(prompt_embedding) + res_dict.update({'value': value}) + return res_dict diff --git a/ding/model/template/madqn.py b/ding/model/template/madqn.py new file mode 100644 index 0000000000..4cab2b1e98 --- /dev/null +++ b/ding/model/template/madqn.py @@ -0,0 +1,54 @@ +import torch.nn as nn +from ding.utils import MODEL_REGISTRY +from .qmix import QMix + + +@MODEL_REGISTRY.register('madqn') +class MADQN(nn.Module): + + def __init__( + self, + agent_num: int, + obs_shape: int, + action_shape: int, + hidden_size_list: list, + global_obs_shape: int = None, + mixer: bool = False, + global_cooperation: bool = True, + lstm_type: str = 'gru', + dueling: bool = False + ) -> None: + super(MADQN, self).__init__() + self.current = QMix( + agent_num=agent_num, + obs_shape=obs_shape, + action_shape=action_shape, + hidden_size_list=hidden_size_list, + global_obs_shape=global_obs_shape, + mixer=mixer, + lstm_type=lstm_type, + dueling=dueling + ) + self.global_cooperation = global_cooperation + if self.global_cooperation: + cooperation_obs_shape = global_obs_shape + else: + cooperation_obs_shape = obs_shape + self.cooperation = QMix( + agent_num=agent_num, + obs_shape=cooperation_obs_shape, + action_shape=action_shape, + hidden_size_list=hidden_size_list, + global_obs_shape=global_obs_shape, + mixer=mixer, + lstm_type=lstm_type, + dueling=dueling + ) + + def forward(self, data: dict, cooperation: bool = False, single_step: bool = True) -> dict: + if cooperation: + if self.global_cooperation: + data['obs']['agent_state'] = data['obs']['global_state'] + return self.cooperation(data, single_step=single_step) + else: + return self.current(data, single_step=single_step) diff --git a/ding/model/template/maqac.py b/ding/model/template/maqac.py index e6ddd996dd..2d72e43d53 100644 --- a/ding/model/template/maqac.py +++ b/ding/model/template/maqac.py @@ -1,6 +1,5 @@ from typing import Union, Dict, Optional from easydict import EasyDict -import numpy as np import torch import torch.nn as nn @@ -9,11 +8,14 @@ FCEncoder, ConvEncoder -@MODEL_REGISTRY.register('maqac') -class MAQAC(nn.Module): - r""" +@MODEL_REGISTRY.register('discrete_maqac') +class DiscreteMAQAC(nn.Module): + """ Overview: - The MAQAC model. + The neural network and computation graph of algorithms related to discrete action Multi-Agent Q-value \ + Actor-CritiC (MAQAC) model. The model is composed of actor and critic, where actor is a MLP network and \ + critic is a MLP network. The actor network is used to predict the action probability distribution, and the \ + critic network is used to predict the Q value of the state-action pair. Interfaces: ``__init__``, ``forward``, ``compute_actor``, ``compute_critic`` """ @@ -32,9 +34,9 @@ def __init__( activation: Optional[nn.Module] = nn.ReLU(), norm_type: Optional[str] = None, ) -> None: - r""" + """ Overview: - Init the MAQAC Model according to arguments. + Initialize the DiscreteMAQAC Model according to arguments. Arguments: - agent_obs_shape (:obj:`Union[int, SequenceType]`): Agent's observation's space. - global_obs_shape (:obj:`Union[int, SequenceType]`): Global observation's space. @@ -42,18 +44,17 @@ def __init__( - action_shape (:obj:`Union[int, SequenceType]`): Action's space. - twin_critic (:obj:`bool`): Whether include twin critic. - actor_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to actor-nn's ``Head``. - - actor_head_layer_num (:obj:`int`): - The num of layers used in the network to compute Q value output for actor's nn. + - actor_head_layer_num (:obj:`int`): The num of layers used in the network to compute Q value output \ + for actor's nn. - critic_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to critic-nn's ``Head``. - - critic_head_layer_num (:obj:`int`): - The num of layers used in the network to compute Q value output for critic's nn. - - activation (:obj:`Optional[nn.Module]`): - The type of activation function to use in ``MLP`` the after ``layer_fn``, - if ``None`` then default set to ``nn.ReLU()`` - - norm_type (:obj:`Optional[str]`): - The type of normalization to use, see ``ding.torch_utils.fc_block`` for more details. + - critic_head_layer_num (:obj:`int`): The num of layers used in the network to compute Q value output \ + for critic's nn. + - activation (:obj:`Optional[nn.Module]`): The type of activation function to use in ``MLP`` the after \ + ``layer_fn``, if ``None`` then default set to ``nn.ReLU()`` + - norm_type (:obj:`Optional[str]`): The type of normalization to use, see ``ding.torch_utils.fc_block`` \ + for more details. """ - super(MAQAC, self).__init__() + super(DiscreteMAQAC, self).__init__() agent_obs_shape: int = squeeze(agent_obs_shape) action_shape: int = squeeze(action_shape) self.actor = nn.Sequential( @@ -92,93 +93,125 @@ def __init__( ) def forward(self, inputs: Union[torch.Tensor, Dict], mode: str) -> Dict: - r""" + """ Overview: - Use observation and action tensor to predict output. - Parameter updates with QAC's MLPs forward setup. + Use observation tensor to predict output, with ``compute_actor`` or ``compute_critic`` mode. Arguments: - Forward with ``'compute_actor'``: - - inputs (:obj:`torch.Tensor`): - The encoded embedding tensor, determined with given ``hidden_size``, i.e. ``(B, N=hidden_size)``. - Whether ``actor_head_hidden_size`` or ``critic_head_hidden_size`` depend on ``mode``. - Forward with ``'compute_critic'``, inputs (`Dict`) Necessary Keys: - - ``obs``, ``action`` encoded tensors. - - mode (:obj:`str`): Name of the forward mode. + - inputs (:obj:`Dict[str, torch.Tensor]`): The input dict tensor data, has keys: + - ``obs`` (:obj:`Dict[str, torch.Tensor]`): The input dict tensor data, has keys: + - ``agent_state`` (:obj:`torch.Tensor`): The agent's observation tensor data, \ + with shape :math:`(B, A, N0)`, where B is batch size and A is agent num. \ + N0 corresponds to ``agent_obs_shape``. + - ``global_state`` (:obj:`torch.Tensor`): The global observation tensor data, \ + with shape :math:`(B, A, N1)`, where B is batch size and A is agent num. \ + N1 corresponds to ``global_obs_shape``. + - ``action_mask`` (:obj:`torch.Tensor`): The action mask tensor data, \ + with shape :math:`(B, A, N2)`, where B is batch size and A is agent num. \ + N2 corresponds to ``action_shape``. + + - mode (:obj:`str`): The forward mode, all the modes are defined in the beginning of this class. Returns: - - outputs (:obj:`Dict`): Outputs of network forward. - Forward with ``'compute_actor'``, Necessary Keys (either): - - action (:obj:`torch.Tensor`): Action tensor with same size as input ``x``. - - logit (:obj:`torch.Tensor`): Action's probabilities. - Forward with ``'compute_critic'``, Necessary Keys: - - q_value (:obj:`torch.Tensor`): Q value tensor with same size as batch size. - Actor Shapes: - - inputs (:obj:`torch.Tensor`): :math:`(B, N0)`, B is batch size and N0 corresponds to ``hidden_size`` - - action (:obj:`torch.Tensor`): :math:`(B, N0)` - - q_value (:obj:`torch.FloatTensor`): :math:`(B, )`, where B is batch size. - Critic Shapes: - - obs (:obj:`torch.Tensor`): :math:`(B, N1)`, where B is batch size and N1 is ``global_obs_shape`` - - logit (:obj:`torch.FloatTensor`): :math:`(B, N2)`, where B is batch size and N2 is ``action_shape`` + - output (:obj:`Dict[str, torch.Tensor]`): The output dict of DiscreteMAQAC forward computation graph, \ + whose key-values vary in different forward modes. + Examples: + >>> B = 32 + >>> agent_obs_shape = 216 + >>> global_obs_shape = 264 + >>> agent_num = 8 + >>> action_shape = 14 + >>> data = { + >>> 'obs': { + >>> 'agent_state': torch.randn(B, agent_num, agent_obs_shape), + >>> 'global_state': torch.randn(B, agent_num, global_obs_shape), + >>> 'action_mask': torch.randint(0, 2, size=(B, agent_num, action_shape)) + >>> } + >>> } + >>> model = DiscreteMAQAC(agent_obs_shape, global_obs_shape, action_shape, twin_critic=True) + >>> logit = model(data, mode='compute_actor')['logit'] + >>> value = model(data, mode='compute_critic')['q_value'] """ assert mode in self.mode, "not support forward mode: {}/{}".format(mode, self.mode) return getattr(self, mode)(inputs) def compute_actor(self, inputs: Dict) -> Dict: - r""" + """ Overview: - Use encoded embedding tensor to predict output. - Execute parameter updates with ``'compute_actor'`` mode - Use encoded embedding tensor to predict output. + Use observation tensor to predict action logits. Arguments: - - inputs (:obj:`torch.Tensor`): - The encoded embedding tensor, determined with given ``hidden_size``, i.e. ``(B, N=hidden_size)``. - ``hidden_size = actor_head_hidden_size`` - - mode (:obj:`str`): Name of the forward mode. + - inputs (:obj:`Dict[str, torch.Tensor]`): The input dict tensor data, has keys: + - ``obs`` (:obj:`Dict[str, torch.Tensor]`): The input dict tensor data, has keys: + - ``agent_state`` (:obj:`torch.Tensor`): The agent's observation tensor data, \ + with shape :math:`(B, A, N0)`, where B is batch size and A is agent num. \ + N0 corresponds to ``agent_obs_shape``. + - ``global_state`` (:obj:`torch.Tensor`): The global observation tensor data, \ + with shape :math:`(B, A, N1)`, where B is batch size and A is agent num. \ + N1 corresponds to ``global_obs_shape``. + - ``action_mask`` (:obj:`torch.Tensor`): The action mask tensor data, \ + with shape :math:`(B, A, N2)`, where B is batch size and A is agent num. \ + N2 corresponds to ``action_shape``. Returns: - - outputs (:obj:`Dict`): Outputs of forward pass encoder and head. - ReturnsKeys (either): - - action (:obj:`torch.Tensor`): Continuous action tensor with same size as ``action_shape``. - - logit (:obj:`torch.Tensor`): - Logit tensor encoding ``mu`` and ``sigma``, both with same size as input ``x``. - Shapes: - - inputs (:obj:`torch.Tensor`): :math:`(B, N0)`, B is batch size and N0 corresponds to ``hidden_size`` - - action (:obj:`torch.Tensor`): :math:`(B, N0)` - - logit (:obj:`list`): 2 elements, mu and sigma, each is the shape of :math:`(B, N0)`. - - q_value (:obj:`torch.FloatTensor`): :math:`(B, )`, B is batch size. + - output (:obj:`Dict[str, torch.Tensor]`): The output dict of DiscreteMAQAC forward computation graph, \ + whose key-values vary in different forward modes. + - logit (:obj:`torch.Tensor`): Action's output logit (real value range), whose shape is \ + :math:`(B, A, N2)`, where N2 corresponds to ``action_shape``. + - action_mask (:obj:`torch.Tensor`): Action mask tensor with same size as ``action_shape``. Examples: - >>> # Regression mode - >>> model = QAC(64, 64, 'regression') - >>> inputs = torch.randn(4, 64) - >>> actor_outputs = model(inputs,'compute_actor') - >>> assert actor_outputs['action'].shape == torch.Size([4, 64]) - >>> # Reparameterization Mode - >>> model = QAC(64, 64, 'reparameterization') - >>> inputs = torch.randn(4, 64) - >>> actor_outputs = model(inputs,'compute_actor') - >>> actor_outputs['logit'][0].shape # mu - >>> torch.Size([4, 64]) - >>> actor_outputs['logit'][1].shape # sigma - >>> torch.Size([4, 64]) + >>> B = 32 + >>> agent_obs_shape = 216 + >>> global_obs_shape = 264 + >>> agent_num = 8 + >>> action_shape = 14 + >>> data = { + >>> 'obs': { + >>> 'agent_state': torch.randn(B, agent_num, agent_obs_shape), + >>> 'global_state': torch.randn(B, agent_num, global_obs_shape), + >>> 'action_mask': torch.randint(0, 2, size=(B, agent_num, action_shape)) + >>> } + >>> } + >>> model = DiscreteMAQAC(agent_obs_shape, global_obs_shape, action_shape, twin_critic=True) + >>> logit = model.compute_actor(data)['logit'] """ action_mask = inputs['obs']['action_mask'] x = self.actor(inputs['obs']['agent_state']) return {'logit': x['logit'], 'action_mask': action_mask} def compute_critic(self, inputs: Dict) -> Dict: - r""" + """ Overview: - Execute parameter updates with ``'compute_critic'`` mode - Use encoded embedding tensor to predict output. + use observation tensor to predict Q value. Arguments: - - ``obs``, ``action`` encoded tensors. - - mode (:obj:`str`): Name of the forward mode. + - inputs (:obj:`Dict[str, torch.Tensor]`): The input dict tensor data, has keys: + - ``obs`` (:obj:`Dict[str, torch.Tensor]`): The input dict tensor data, has keys: + - ``agent_state`` (:obj:`torch.Tensor`): The agent's observation tensor data, \ + with shape :math:`(B, A, N0)`, where B is batch size and A is agent num. \ + N0 corresponds to ``agent_obs_shape``. + - ``global_state`` (:obj:`torch.Tensor`): The global observation tensor data, \ + with shape :math:`(B, A, N1)`, where B is batch size and A is agent num. \ + N1 corresponds to ``global_obs_shape``. + - ``action_mask`` (:obj:`torch.Tensor`): The action mask tensor data, \ + with shape :math:`(B, A, N2)`, where B is batch size and A is agent num. \ + N2 corresponds to ``action_shape``. Returns: - - outputs (:obj:`Dict`): Q-value output. - ReturnKeys: - - q_value (:obj:`torch.Tensor`): Q value tensor with same size as batch size. - Shapes: - - obs (:obj:`torch.Tensor`): :math:`(B, N1)`, where B is batch size and N1 is ``obs_shape`` - - action (:obj:`torch.Tensor`): :math:`(B, N2)`, where B is batch size and N2 is ``action_shape`` - - q_value (:obj:`torch.FloatTensor`): :math:`(B, )`, where B is batch size. + - output (:obj:`Dict[str, torch.Tensor]`): The output dict of DiscreteMAQAC forward computation graph, \ + whose key-values vary in different values of ``twin_critic``. + - q_value (:obj:`list`): If ``twin_critic=True``, q_value should be 2 elements, each is the shape of \ + :math:`(B, A, N2)`, where B is batch size and A is agent num. N2 corresponds to ``action_shape``. \ + Otherwise, q_value should be ``torch.Tensor``. + Examples: + >>> B = 32 + >>> agent_obs_shape = 216 + >>> global_obs_shape = 264 + >>> agent_num = 8 + >>> action_shape = 14 + >>> data = { + >>> 'obs': { + >>> 'agent_state': torch.randn(B, agent_num, agent_obs_shape), + >>> 'global_state': torch.randn(B, agent_num, global_obs_shape), + >>> 'action_mask': torch.randint(0, 2, size=(B, agent_num, action_shape)) + >>> } + >>> } + >>> model = DiscreteMAQAC(agent_obs_shape, global_obs_shape, action_shape, twin_critic=True) + >>> value = model.compute_critic(data)['q_value'] """ if self.twin_critic: @@ -188,11 +221,14 @@ def compute_critic(self, inputs: Dict) -> Dict: return {'q_value': x} -@MODEL_REGISTRY.register('maqac_continuous') +@MODEL_REGISTRY.register('continuous_maqac') class ContinuousMAQAC(nn.Module): - r""" + """ Overview: - The Continuous MAQAC model. + The neural network and computation graph of algorithms related to continuous action Multi-Agent Q-value \ + Actor-CritiC (MAQAC) model. The model is composed of actor and critic, where actor is a MLP network and \ + critic is a MLP network. The actor network is used to predict the action probability distribution, and the \ + critic network is used to predict the Q value of the state-action pair. Interfaces: ``__init__``, ``forward``, ``compute_actor``, ``compute_critic`` """ @@ -212,25 +248,24 @@ def __init__( activation: Optional[nn.Module] = nn.ReLU(), norm_type: Optional[str] = None, ) -> None: - r""" + """ Overview: - Init the QAC Model according to arguments. + Initialize the QAC Model according to arguments. Arguments: - obs_shape (:obj:`Union[int, SequenceType]`): Observation's space. - action_shape (:obj:`Union[int, SequenceType, EasyDict]`): Action's space, such as 4, (3, ) - action_space (:obj:`str`): Whether choose ``regression`` or ``reparameterization``. - twin_critic (:obj:`bool`): Whether include twin critic. - actor_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to actor-nn's ``Head``. - - actor_head_layer_num (:obj:`int`): - The num of layers used in the network to compute Q value output for actor's nn. + - actor_head_layer_num (:obj:`int`): The num of layers used in the network to compute Q value output \ + for actor's nn. - critic_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to critic-nn's ``Head``. - - critic_head_layer_num (:obj:`int`): - The num of layers used in the network to compute Q value output for critic's nn. - - activation (:obj:`Optional[nn.Module]`): - The type of activation function to use in ``MLP`` the after ``layer_fn``, - if ``None`` then default set to ``nn.ReLU()`` - - norm_type (:obj:`Optional[str]`): - The type of normalization to use, see ``ding.torch_utils.fc_block`` for more details. + - critic_head_layer_num (:obj:`int`): The num of layers used in the network to compute Q value output \ + for critic's nn. + - activation (:obj:`Optional[nn.Module]`): The type of activation function to use in ``MLP`` the after \ + ``layer_fn``, if ``None`` then default set to ``nn.ReLU()`` + - norm_type (:obj:`Optional[str]`): The type of normalization to use, see ``ding.torch_utils.fc_block`` \ + for more details. """ super(ContinuousMAQAC, self).__init__() obs_shape: int = squeeze(agent_obs_shape) @@ -238,7 +273,7 @@ def __init__( action_shape = squeeze(action_shape) self.action_shape = action_shape self.action_space = action_space - assert self.action_space in ['regression', 'reparameterization'] + assert self.action_space in ['regression', 'reparameterization'], self.action_space if self.action_space == 'regression': # DDPG, TD3 self.actor = nn.Sequential( nn.Linear(obs_shape, actor_head_hidden_size), activation, @@ -295,101 +330,86 @@ def __init__( ) def forward(self, inputs: Union[torch.Tensor, Dict], mode: str) -> Dict: - r""" + """ Overview: - Use observation and action tensor to predict output. - Parameter updates with QAC's MLPs forward setup. + Use observation and action tensor to predict output in ``compute_actor`` or ``compute_critic`` mode. Arguments: - Forward with ``'compute_actor'``: - - inputs (:obj:`torch.Tensor`): - The encoded embedding tensor, determined with given ``hidden_size``, i.e. ``(B, N=hidden_size)``. - Whether ``actor_head_hidden_size`` or ``critic_head_hidden_size`` depend on ``mode``. - - Forward with ``'compute_critic'``, inputs (`Dict`) Necessary Keys: - - ``obs``, ``action`` encoded tensors. + - inputs (:obj:`Dict[str, torch.Tensor]`): The input dict tensor data, has keys: + - ``obs`` (:obj:`Dict[str, torch.Tensor]`): The input dict tensor data, has keys: + - ``agent_state`` (:obj:`torch.Tensor`): The agent's observation tensor data, \ + with shape :math:`(B, A, N0)`, where B is batch size and A is agent num. \ + N0 corresponds to ``agent_obs_shape``. + - ``global_state`` (:obj:`torch.Tensor`): The global observation tensor data, \ + with shape :math:`(B, A, N1)`, where B is batch size and A is agent num. \ + N1 corresponds to ``global_obs_shape``. + - ``action_mask`` (:obj:`torch.Tensor`): The action mask tensor data, \ + with shape :math:`(B, A, N2)`, where B is batch size and A is agent num. \ + N2 corresponds to ``action_shape``. + - ``action`` (:obj:`torch.Tensor`): The action tensor data, \ + with shape :math:`(B, A, N3)`, where B is batch size and A is agent num. \ + N3 corresponds to ``action_shape``. - mode (:obj:`str`): Name of the forward mode. Returns: - - outputs (:obj:`Dict`): Outputs of network forward. - - Forward with ``'compute_actor'``, Necessary Keys (either): - - action (:obj:`torch.Tensor`): Action tensor with same size as input ``x``. - - logit (:obj:`torch.Tensor`): - Logit tensor encoding ``mu`` and ``sigma``, both with same size as input ``x``. - - Forward with ``'compute_critic'``, Necessary Keys: - - q_value (:obj:`torch.Tensor`): Q value tensor with same size as batch size. - Actor Shapes: - - inputs (:obj:`torch.Tensor`): :math:`(B, N0)`, B is batch size and N0 corresponds to ``hidden_size`` - - action (:obj:`torch.Tensor`): :math:`(B, N0)` - - q_value (:obj:`torch.FloatTensor`): :math:`(B, )`, where B is batch size. - - Critic Shapes: - - obs (:obj:`torch.Tensor`): :math:`(B, N1)`, where B is batch size and N1 is ``obs_shape`` - - action (:obj:`torch.Tensor`): :math:`(B, N2)`, where B is batch size and N2 is``action_shape`` - - logit (:obj:`torch.FloatTensor`): :math:`(B, N2)`, where B is batch size and N3 is ``action_shape`` - - Actor Examples: - >>> # Regression mode - >>> model = QAC(64, 64, 'regression') - >>> inputs = torch.randn(4, 64) - >>> actor_outputs = model(inputs,'compute_actor') - >>> assert actor_outputs['action'].shape == torch.Size([4, 64]) - >>> # Reparameterization Mode - >>> model = QAC(64, 64, 'reparameterization') - >>> inputs = torch.randn(4, 64) - >>> actor_outputs = model(inputs,'compute_actor') - >>> actor_outputs['logit'][0].shape # mu - >>> torch.Size([4, 64]) - >>> actor_outputs['logit'][1].shape # sigma - >>> torch.Size([4, 64]) - + - outputs (:obj:`Dict`): Outputs of network forward, whose key-values will be different for different \ + ``mode``, ``twin_critic``, ``action_space``. + Examples: + >>> B = 32 + >>> agent_obs_shape = 216 + >>> global_obs_shape = 264 + >>> agent_num = 8 + >>> action_shape = 14 + >>> act_space = 'reparameterization' # regression + >>> data = { + >>> 'obs': { + >>> 'agent_state': torch.randn(B, agent_num, agent_obs_shape), + >>> 'global_state': torch.randn(B, agent_num, global_obs_shape), + >>> 'action_mask': torch.randint(0, 2, size=(B, agent_num, action_shape)) + >>> }, + >>> 'action': torch.randn(B, agent_num, squeeze(action_shape)) + >>> } + >>> model = ContinuousMAQAC(agent_obs_shape, global_obs_shape, action_shape, act_space, twin_critic=False) + >>> if action_space == 'regression': + >>> action = model(data['obs'], mode='compute_actor')['action'] + >>> elif action_space == 'reparameterization': + >>> (mu, sigma) = model(data['obs'], mode='compute_actor')['logit'] + >>> value = model(data, mode='compute_critic')['q_value'] """ assert mode in self.mode, "not support forward mode: {}/{}".format(mode, self.mode) return getattr(self, mode)(inputs) def compute_actor(self, inputs: Dict) -> Dict: - r""" + """ Overview: - Use encoded embedding tensor to predict output. - Execute parameter updates with ``'compute_actor'`` mode - Use encoded embedding tensor to predict output. + Use observation tensor to predict action logits. Arguments: - - inputs (:obj:`torch.Tensor`): - The encoded embedding tensor, determined with given ``hidden_size``, i.e. ``(B, N=hidden_size)``. - ``hidden_size = actor_head_hidden_size`` - - mode (:obj:`str`): Name of the forward mode. - Returns: - - outputs (:obj:`Dict`): Outputs of forward pass encoder and head. + - inputs (:obj:`Dict[str, torch.Tensor]`): The input dict tensor data, has keys: + - ``agent_state`` (:obj:`torch.Tensor`): The agent's observation tensor data, \ + with shape :math:`(B, A, N0)`, where B is batch size and A is agent num. \ + N0 corresponds to ``agent_obs_shape``. - ReturnsKeys (either): - - action (:obj:`torch.Tensor`): Continuous action tensor with same size as ``action_shape``. - - logit (:obj:`torch.Tensor`): - Logit tensor encoding ``mu`` and ``sigma``, both with same size as input ``x``. - - logit + action_args - Shapes: - - inputs (:obj:`torch.Tensor`): :math:`(B, N0)`, B is batch size and N0 corresponds to ``hidden_size`` - - action (:obj:`torch.Tensor`): :math:`(B, N0)` - - logit (:obj:`Union[list, torch.Tensor]`): - - case1(continuous space, list): 2 elements, mu and sigma, each is the shape of :math:`(B, N0)`. - - case2(hybrid space, torch.Tensor): :math:`(B, N1)`, where N1 is action_type_shape - - q_value (:obj:`torch.FloatTensor`): :math:`(B, )`, B is batch size. - - action_args (:obj:`torch.FloatTensor`): :math:`(B, N2)`, where N2 is action_args_shape - (action_args are continuous real value) + Returns: + - outputs (:obj:`Dict`): Outputs of network forward. + ReturnKeys (``action_space == 'regression'``): + - action (:obj:`torch.Tensor`): Action tensor with same size as ``action_shape``. + ReturnKeys (``action_space == 'reparameterization'``): + - logit (:obj:`list`): 2 elements, each is the shape of :math:`(B, A, N3)`, where B is batch size and \ + A is agent num. N3 corresponds to ``action_shape``. Examples: - >>> # Regression mode - >>> model = QAC(64, 64, 'regression') - >>> inputs = torch.randn(4, 64) - >>> actor_outputs = model(inputs,'compute_actor') - >>> assert actor_outputs['action'].shape == torch.Size([4, 64]) - >>> # Reparameterization Mode - >>> model = QAC(64, 64, 'reparameterization') - >>> inputs = torch.randn(4, 64) - >>> actor_outputs = model(inputs,'compute_actor') - >>> actor_outputs['logit'][0].shape # mu - >>> torch.Size([4, 64]) - >>> actor_outputs['logit'][1].shape # sigma - >>> torch.Size([4, 64]) + >>> B = 32 + >>> agent_obs_shape = 216 + >>> global_obs_shape = 264 + >>> agent_num = 8 + >>> action_shape = 14 + >>> act_space = 'reparameterization' # 'regression' + >>> data = { + >>> 'agent_state': torch.randn(B, agent_num, agent_obs_shape), + >>> } + >>> model = ContinuousMAQAC(agent_obs_shape, global_obs_shape, action_shape, act_space, twin_critic=False) + >>> if action_space == 'regression': + >>> action = model.compute_actor(data)['action'] + >>> elif action_space == 'reparameterization': + >>> (mu, sigma) = model.compute_actor(data)['logit'] """ inputs = inputs['agent_state'] if self.action_space == 'regression': @@ -400,28 +420,50 @@ def compute_actor(self, inputs: Dict) -> Dict: return {'logit': [x['mu'], x['sigma']]} def compute_critic(self, inputs: Dict) -> Dict: - r""" + """ Overview: - Execute parameter updates with ``'compute_critic'`` mode - Use encoded embedding tensor to predict output. + Use observation tensor and action tensor to predict Q value. Arguments: - - inputs (:obj: `Dict`): ``obs``, ``action`` and ``logit` tensors. - - mode (:obj:`str`): Name of the forward mode. - Returns: - - outputs (:obj:`Dict`): Q-value output. + - inputs (:obj:`Dict[str, torch.Tensor]`): The input dict tensor data, has keys: + - ``obs`` (:obj:`Dict[str, torch.Tensor]`): The input dict tensor data, has keys: + - ``agent_state`` (:obj:`torch.Tensor`): The agent's observation tensor data, \ + with shape :math:`(B, A, N0)`, where B is batch size and A is agent num. \ + N0 corresponds to ``agent_obs_shape``. + - ``global_state`` (:obj:`torch.Tensor`): The global observation tensor data, \ + with shape :math:`(B, A, N1)`, where B is batch size and A is agent num. \ + N1 corresponds to ``global_obs_shape``. + - ``action_mask`` (:obj:`torch.Tensor`): The action mask tensor data, \ + with shape :math:`(B, A, N2)`, where B is batch size and A is agent num. \ + N2 corresponds to ``action_shape``. - ArgumentsKeys: - - necessary: - - obs: (:obj:`torch.Tensor`): 2-dim vector observation - - action (:obj:`Union[torch.Tensor, Dict]`): action from actor - - optional: - - logit (:obj:`torch.Tensor`): discrete action logit - ReturnKeys: - - q_value (:obj:`torch.Tensor`): Q value tensor with same size as batch size. - Shapes: - - obs (:obj:`torch.Tensor`): :math:`(B, N1)`, where B is batch size and N1 is ``obs_shape`` - - action (:obj:`torch.Tensor`): :math:`(B, N2)`, where B is batch size and N2 is ``action_shape`` - - q_value (:obj:`torch.FloatTensor`): :math:`(B, )`, where B is batch size. + - ``action`` (:obj:`torch.Tensor`): The action tensor data, \ + with shape :math:`(B, A, N3)`, where B is batch size and A is agent num. \ + N3 corresponds to ``action_shape``. + + Returns: + - outputs (:obj:`Dict`): Outputs of network forward. + ReturnKeys (``twin_critic=True``): + - q_value (:obj:`list`): 2 elements, each is the shape of :math:`(B, A)`, where B is batch size and \ + A is agent num. + ReturnKeys (``twin_critic=False``): + - q_value (:obj:`torch.Tensor`): :math:`(B, A)`, where B is batch size and A is agent num. + Examples: + >>> B = 32 + >>> agent_obs_shape = 216 + >>> global_obs_shape = 264 + >>> agent_num = 8 + >>> action_shape = 14 + >>> act_space = 'reparameterization' # 'regression' + >>> data = { + >>> 'obs': { + >>> 'agent_state': torch.randn(B, agent_num, agent_obs_shape), + >>> 'global_state': torch.randn(B, agent_num, global_obs_shape), + >>> 'action_mask': torch.randint(0, 2, size=(B, agent_num, action_shape)) + >>> }, + >>> 'action': torch.randn(B, agent_num, squeeze(action_shape)) + >>> } + >>> model = ContinuousMAQAC(agent_obs_shape, global_obs_shape, action_shape, act_space, twin_critic=False) + >>> value = model.compute_critic(data)['q_value'] """ obs, action = inputs['obs']['global_state'], inputs['action'] diff --git a/ding/model/template/mavac.py b/ding/model/template/mavac.py index 47a1171dbf..9018aa9f04 100644 --- a/ding/model/template/mavac.py +++ b/ding/model/template/mavac.py @@ -1,19 +1,22 @@ -from typing import Union, Dict, Optional +from typing import Union, Dict, Tuple, Optional import torch import torch.nn as nn from ding.utils import SequenceType, squeeze, MODEL_REGISTRY -from ..common import ReparameterizationHead, RegressionHead, DiscreteHead, MultiHead, \ - FCEncoder, ConvEncoder +from ..common import ReparameterizationHead, RegressionHead, DiscreteHead @MODEL_REGISTRY.register('mavac') class MAVAC(nn.Module): - r""" + """ Overview: - The MAVAC model. + The neural network and computation graph of algorithms related to (state) Value Actor-Critic (VAC) for \ + multi-agent, such as MAPPO(https://arxiv.org/abs/2103.01955). This model now supports discrete and \ + continuous action space. The MAVAC is composed of four parts: ``actor_encoder``, ``critic_encoder``, \ + ``actor_head`` and ``critic_head``. Encoders are used to extract the feature from various observation. \ + Heads are used to predict corresponding value or action logit. Interfaces: - ``__init__``, ``forward``, ``compute_actor``, ``compute_critic`` + ``__init__``, ``forward``, ``compute_actor``, ``compute_critic``, ``compute_actor_critic``. """ mode = ['compute_actor', 'compute_critic', 'compute_actor_critic'] @@ -32,27 +35,41 @@ def __init__( norm_type: Optional[str] = None, sigma_type: Optional[str] = 'independent', bound_type: Optional[str] = None, + encoder: Optional[Tuple[torch.nn.Module, torch.nn.Module]] = None, ) -> None: - r""" + """ Overview: - Init the VAC Model according to arguments. + Init the MAVAC Model according to arguments. Arguments: - - obs_shape (:obj:`Union[int, SequenceType]`): Observation's space. - - action_shape (:obj:`Union[int, SequenceType]`): Action's space. - - share_encoder (:obj:`bool`): Whether share encoder. - - continuous (:obj:`bool`): Whether collect continuously. - - encoder_hidden_size_list (:obj:`SequenceType`): Collection of ``hidden_size`` to pass to ``Encoder`` - - actor_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to actor-nn's ``Head``. - - actor_head_layer_num (:obj:`int`): - The num of layers used in the network to compute Q value output for actor's nn. - - critic_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to critic-nn's ``Head``. - - critic_head_layer_num (:obj:`int`): - The num of layers used in the network to compute Q value output for critic's nn. - - activation (:obj:`Optional[nn.Module]`): - The type of activation function to use in ``MLP`` the after ``layer_fn``, - if ``None`` then default set to ``nn.ReLU()`` - - norm_type (:obj:`Optional[str]`): - The type of normalization to use, see ``ding.torch_utils.fc_block`` for more details` + - agent_obs_shape (:obj:`Union[int, SequenceType]`): Observation's space for single agent, \ + such as 8 or [4, 84, 84]. + - global_obs_shape (:obj:`Union[int, SequenceType]`): Global observation's space, such as 8 or [4, 84, 84]. + - action_shape (:obj:`Union[int, SequenceType]`): Action space shape for single agent, such as 6 \ + or [2, 3, 3]. + - agent_num (:obj:`int`): This parameter is temporarily reserved. This parameter may be required for \ + subsequent changes to the model + - actor_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` of ``actor_head`` network, defaults \ + to 256, it must match the last element of ``agent_obs_shape``. + - actor_head_layer_num (:obj:`int`): The num of layers used in the ``actor_head`` network to compute action. + - critic_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` of ``critic_head`` network, defaults \ + to 512, it must match the last element of ``global_obs_shape``. + - critic_head_layer_num (:obj:`int`): The num of layers used in the network to compute Q value output for \ + critic's nn. + - action_space (:obj:`Union[int, SequenceType]`): The type of different action spaces, including \ + ['discrete', 'continuous'], then will instantiate corresponding head, including ``DiscreteHead`` \ + and ``ReparameterizationHead``. + - activation (:obj:`Optional[nn.Module]`): The type of activation function to use in ``MLP`` the after \ + ``layer_fn``, if ``None`` then default set to ``nn.ReLU()``. + - norm_type (:obj:`Optional[str]`): The type of normalization in networks, see \ + ``ding.torch_utils.fc_block`` for more details. you can choose one of ['BN', 'IN', 'SyncBN', 'LN']. + - sigma_type (:obj:`Optional[str]`): The type of sigma in continuous action space, see \ + ``ding.torch_utils.network.dreamer.ReparameterizationHead`` for more details, in MAPPO, it defaults \ + to ``independent``, which means state-independent sigma parameters. + - bound_type (:obj:`Optional[str]`): The type of action bound methods in continuous action space, defaults \ + to ``None``, which means no bound. + - encoder (:obj:`Optional[Tuple[torch.nn.Module, torch.nn.Module]]`): The encoder module list, defaults \ + to ``None``, you can define your own actor and critic encoder module and pass it into MAVAC to \ + deal with different observation space. """ super(MAVAC, self).__init__() agent_obs_shape: int = squeeze(agent_obs_shape) @@ -61,61 +78,38 @@ def __init__( self.global_obs_shape, self.agent_obs_shape, self.action_shape = global_obs_shape, agent_obs_shape, action_shape self.action_space = action_space # Encoder Type - if isinstance(agent_obs_shape, int) or len(agent_obs_shape) == 1: - encoder_cls = FCEncoder - elif len(agent_obs_shape) == 3: - encoder_cls = ConvEncoder + if encoder: + self.actor_encoder, self.critic_encoder = encoder else: - raise RuntimeError( - "not support obs_shape for pre-defined encoder: {}, please customize your own DQN". - format(agent_obs_shape) + # We directly connect the Head after a Liner layer instead of using the 3-layer FCEncoder. + # In SMAC task it can obviously improve the performance. + # Users can change the model according to their own needs. + self.actor_encoder = nn.Sequential( + nn.Linear(agent_obs_shape, actor_head_hidden_size), + activation, ) - if isinstance(global_obs_shape, int) or len(global_obs_shape) == 1: - global_encoder_cls = FCEncoder - elif len(global_obs_shape) == 3: - global_encoder_cls = ConvEncoder - else: - raise RuntimeError( - "not support obs_shape for pre-defined encoder: {}, please customize your own DQN". - format(global_obs_shape) + self.critic_encoder = nn.Sequential( + nn.Linear(global_obs_shape, critic_head_hidden_size), + activation, ) - - # We directly connect the Head after a Liner layer instead of using the 3-layer FCEncoder. - # In SMAC task it can obviously improve the performance. - # Users can change the model according to their own needs. - self.actor_encoder = nn.Identity() - self.critic_encoder = nn.Identity() # Head Type - self.critic_head = nn.Sequential( - nn.Linear(global_obs_shape, critic_head_hidden_size), activation, - RegressionHead( - critic_head_hidden_size, 1, critic_head_layer_num, activation=activation, norm_type=norm_type - ) + self.critic_head = RegressionHead( + critic_head_hidden_size, 1, critic_head_layer_num, activation=activation, norm_type=norm_type ) assert self.action_space in ['discrete', 'continuous'], self.action_space if self.action_space == 'discrete': - self.actor_head = nn.Sequential( - nn.Linear(agent_obs_shape, actor_head_hidden_size), activation, - DiscreteHead( - actor_head_hidden_size, - action_shape, - actor_head_layer_num, - activation=activation, - norm_type=norm_type - ) + self.actor_head = DiscreteHead( + actor_head_hidden_size, action_shape, actor_head_layer_num, activation=activation, norm_type=norm_type ) elif self.action_space == 'continuous': - self.actor_head = nn.Sequential( - nn.Linear(agent_obs_shape, actor_head_hidden_size), activation, - ReparameterizationHead( - actor_head_hidden_size, - action_shape, - actor_head_layer_num, - sigma_type=sigma_type, - activation=activation, - norm_type=norm_type, - bound_type=bound_type - ) + self.actor_head = ReparameterizationHead( + actor_head_hidden_size, + action_shape, + actor_head_layer_num, + sigma_type=sigma_type, + activation=activation, + norm_type=norm_type, + bound_type=bound_type ) # must use list, not nn.ModuleList self.actor = [self.actor_encoder, self.actor_head] @@ -126,77 +120,86 @@ def __init__( self.critic = nn.ModuleList(self.critic) def forward(self, inputs: Union[torch.Tensor, Dict], mode: str) -> Dict: - r""" + """ Overview: - Use encoded embedding tensor to predict output. - Parameter updates with VAC's MLPs forward setup. + MAVAC forward computation graph, input observation tensor to predict state value or action logit. \ + ``mode`` includes ``compute_actor``, ``compute_critic``, ``compute_actor_critic``. + Different ``mode`` will forward with different network modules to get different outputs and save \ + computation. Arguments: - Forward with ``'compute_actor'`` or ``'compute_critic'``: - - inputs (:obj:`torch.Tensor`): - The encoded embedding tensor, determined with given ``hidden_size``, i.e. ``(B, N=hidden_size)``. - Whether ``actor_head_hidden_size`` or ``critic_head_hidden_size`` depend on ``mode``. + - inputs (:obj:`Dict`): The input dict including observation and related info, \ + whose key-values vary from different ``mode``. + - mode (:obj:`str`): The forward mode, all the modes are defined in the beginning of this class. Returns: - - outputs (:obj:`Dict`): - Run with encoder and head. + - outputs (:obj:`Dict`): The output dict of MAVAC's forward computation graph, whose key-values vary from \ + different ``mode``. - Forward with ``'compute_actor'``, Necessary Keys: - - logit (:obj:`torch.Tensor`): Logit encoding tensor, with same size as input ``x``. - - Forward with ``'compute_critic'``, Necessary Keys: - - value (:obj:`torch.Tensor`): Q value tensor with same size as batch size. - Shapes: - - inputs (:obj:`torch.Tensor`): :math:`(B, N)`, where B is batch size and N corresponding ``hidden_size`` - - logit (:obj:`torch.FloatTensor`): :math:`(B, N)`, where B is batch size and N is ``action_shape`` - - value (:obj:`torch.FloatTensor`): :math:`(B, )`, where B is batch size. - - Actor Examples: - >>> model = VAC(64,128) - >>> inputs = torch.randn(4, 64) + Examples (Actor): + >>> model = MAVAC(agent_obs_shape=64, global_obs_shape=128, action_shape=14) + >>> inputs = { + 'agent_state': torch.randn(10, 8, 64), + 'global_state': torch.randn(10, 8, 128), + 'action_mask': torch.randint(0, 2, size=(10, 8, 14)) + } >>> actor_outputs = model(inputs,'compute_actor') - >>> assert actor_outputs['logit'].shape == torch.Size([4, 128]) + >>> assert actor_outputs['logit'].shape == torch.Size([10, 8, 14]) - Critic Examples: - >>> model = VAC(64,64) - >>> inputs = torch.randn(4, 64) + Examples (Critic): + >>> model = MAVAC(agent_obs_shape=64, global_obs_shape=128, action_shape=14) + >>> inputs = { + 'agent_state': torch.randn(10, 8, 64), + 'global_state': torch.randn(10, 8, 128), + 'action_mask': torch.randint(0, 2, size=(10, 8, 14)) + } >>> critic_outputs = model(inputs,'compute_critic') - >>> critic_outputs['value'] - tensor([0.0252, 0.0235, 0.0201, 0.0072], grad_fn=) + >>> assert actor_outputs['value'].shape == torch.Size([10, 8]) - Actor-Critic Examples: - >>> model = VAC(64,64) - >>> inputs = torch.randn(4, 64) + Examples (Actor-Critic): + >>> model = MAVAC(64, 64) + >>> inputs = { + 'agent_state': torch.randn(10, 8, 64), + 'global_state': torch.randn(10, 8, 128), + 'action_mask': torch.randint(0, 2, size=(10, 8, 14)) + } >>> outputs = model(inputs,'compute_actor_critic') - >>> outputs['value'] - tensor([0.0252, 0.0235, 0.0201, 0.0072], grad_fn=) - >>> assert outputs['logit'].shape == torch.Size([4, 64]) + >>> assert outputs['value'].shape == torch.Size([10, 8, 14]) + >>> assert outputs['logit'].shape == torch.Size([10, 8]) """ assert mode in self.mode, "not support forward mode: {}/{}".format(mode, self.mode) return getattr(self, mode)(inputs) - def compute_actor(self, x: torch.Tensor) -> Dict: - r""" + def compute_actor(self, x: Dict) -> Dict: + """ Overview: - Execute parameter updates with ``'compute_actor'`` mode - Use encoded embedding tensor to predict output. + MAVAC forward computation graph for actor part, \ + predicting action logit with agent observation tensor in ``x``. Arguments: - - inputs (:obj:`torch.Tensor`): - The encoded embedding tensor, determined with given ``hidden_size``, i.e. ``(B, N=hidden_size)``. - ``hidden_size = actor_head_hidden_size`` + - x (:obj:`Dict`): Input data dict with keys ['agent_state', 'action_mask'(optional)]. + - agent_state: (:obj:`torch.Tensor`): Each agent local state(obs). + - action_mask(optional): (:obj:`torch.Tensor`): When ``action_space`` is discrete, action_mask needs \ + to be provided to mask illegal actions. Returns: - - outputs (:obj:`Dict`): - Run with encoder and head. - + - outputs (:obj:`Dict`): The output dict of the forward computation graph for actor, including ``logit``. ReturnsKeys: - - logit (:obj:`torch.Tensor`): Logit encoding tensor, with same size as input ``x``. + - logit (:obj:`torch.Tensor`): The predicted action logit tensor, for discrete action space, it will be \ + the same dimension real-value ranged tensor of possible action choices, and for continuous action \ + space, it will be the mu and sigma of the Gaussian distribution, and the number of mu and sigma is the \ + same as the number of continuous actions. Shapes: - - logit (:obj:`torch.FloatTensor`): :math:`(B, N)`, where B is batch size and N is ``action_shape`` + - logit (:obj:`torch.FloatTensor`): :math:`(B, M, N)`, where B is batch size and N is ``action_shape`` \ + and M is ``agent_num``. Examples: - >>> model = VAC(64,64) - >>> inputs = torch.randn(4, 64) + >>> model = MAVAC(agent_obs_shape=64, global_obs_shape=128, action_shape=14) + >>> inputs = { + 'agent_state': torch.randn(10, 8, 64), + 'global_state': torch.randn(10, 8, 128), + 'action_mask': torch.randint(0, 2, size=(10, 8, 14)) + } >>> actor_outputs = model(inputs,'compute_actor') - >>> assert actor_outputs['action'].shape == torch.Size([4, 64]) + >>> assert actor_outputs['logit'].shape == torch.Size([10, 8, 14]) + """ if self.action_space == 'discrete': action_mask = x['action_mask'] @@ -213,29 +216,30 @@ def compute_actor(self, x: torch.Tensor) -> Dict: return {'logit': logit} def compute_critic(self, x: Dict) -> Dict: - r""" + """ Overview: - Execute parameter updates with ``'compute_critic'`` mode - Use encoded embedding tensor to predict output. + MAVAC forward computation graph for critic part. \ + Predict state value with global observation tensor in ``x``. Arguments: - - inputs (:obj:`Dict`): - The encoded embedding tensor, determined with given ``hidden_size``, i.e. ``(B, N=hidden_size)``. - ``hidden_size = critic_head_hidden_size`` + - x (:obj:`Dict`): Input data dict with keys ['global_state']. + - global_state: (:obj:`torch.Tensor`): Global state(obs). Returns: - - outputs (:obj:`Dict`): - Run with encoder and head. - - Necessary Keys: - - value (:obj:`torch.Tensor`): Q value tensor with same size as batch size. + - outputs (:obj:`Dict`): The output dict of MAVAC's forward computation graph for critic, \ + including ``value``. + ReturnsKeys: + - value (:obj:`torch.Tensor`): The predicted state value tensor. Shapes: - - value (:obj:`torch.FloatTensor`): :math:`(B, )`, where B is batch size. + - value (:obj:`torch.FloatTensor`): :math:`(B, M)`, where B is batch size and M is ``agent_num``. Examples: - >>> model = VAC(64,64) - >>> inputs = torch.randn(4, 64) + >>> model = MAVAC(agent_obs_shape=64, global_obs_shape=128, action_shape=14) + >>> inputs = { + 'agent_state': torch.randn(10, 8, 64), + 'global_state': torch.randn(10, 8, 128), + 'action_mask': torch.randint(0, 2, size=(10, 8, 14)) + } >>> critic_outputs = model(inputs,'compute_critic') - >>> critic_outputs['value'] - tensor([0.0252, 0.0235, 0.0201, 0.0072], grad_fn=) + >>> assert critic_outputs['value'].shape == torch.Size([10, 8]) """ x = self.critic_encoder(x['global_state']) @@ -243,38 +247,44 @@ def compute_critic(self, x: Dict) -> Dict: return {'value': x['pred']} def compute_actor_critic(self, x: Dict) -> Dict: - r""" + """ Overview: - Execute parameter updates with ``'compute_actor_critic'`` mode - Use encoded embedding tensor to predict output. + MAVAC forward computation graph for both actor and critic part, input observation to predict action \ + logit and state value. Arguments: - - inputs (:obj:`torch.Tensor`): The encoded embedding tensor. - + - x (:obj:`Dict`): The input dict contains ``agent_state``, ``global_state`` and other related info. Returns: - - outputs (:obj:`Dict`): - Run with encoder and head. - + - outputs (:obj:`Dict`): The output dict of MAVAC's forward computation graph for both actor and critic, \ + including ``logit`` and ``value``. ReturnsKeys: - logit (:obj:`torch.Tensor`): Logit encoding tensor, with same size as input ``x``. - value (:obj:`torch.Tensor`): Q value tensor with same size as batch size. Shapes: - - logit (:obj:`torch.FloatTensor`): :math:`(B, N)`, where B is batch size and N is ``action_shape`` - - value (:obj:`torch.FloatTensor`): :math:`(B, )`, where B is batch size. + - logit (:obj:`torch.FloatTensor`): :math:`(B, M, N)`, where B is batch size and N is ``action_shape`` \ + and M is ``agent_num``. + - value (:obj:`torch.FloatTensor`): :math:`(B, M)`, where B is batch sizeand M is ``agent_num``. Examples: - >>> model = VAC(64,64) - >>> inputs = torch.randn(4, 64) + >>> model = MAVAC(64, 64) + >>> inputs = { + 'agent_state': torch.randn(10, 8, 64), + 'global_state': torch.randn(10, 8, 128), + 'action_mask': torch.randint(0, 2, size=(10, 8, 14)) + } >>> outputs = model(inputs,'compute_actor_critic') - >>> outputs['value'] - tensor([0.0252, 0.0235, 0.0201, 0.0072], grad_fn=) - >>> assert outputs['logit'].shape == torch.Size([4, 64]) - - - .. note:: - ``compute_actor_critic`` interface aims to save computation when shares encoder. - Returning the combination dictionry. - + >>> assert outputs['value'].shape == torch.Size([10, 8]) + >>> assert outputs['logit'].shape == torch.Size([10, 8, 14]) """ - logit = self.compute_actor(x)['logit'] - value = self.compute_critic(x)['value'] + x_actor = self.actor_encoder(x['agent_state']) + x_critic = self.critic_encoder(x['global_state']) + + if self.action_space == 'discrete': + action_mask = x['action_mask'] + x = self.actor_head(x_actor) + logit = x['logit'] + logit[action_mask == 0.0] = -99999999 + elif self.action_space == 'continuous': + x = self.actor_head(x_actor) + logit = x + value = self.critic_head(x_critic)['pred'] return {'logit': logit, 'value': value} diff --git a/ding/model/template/ngu.py b/ding/model/template/ngu.py index eb605982bd..caa3c14760 100644 --- a/ding/model/template/ngu.py +++ b/ding/model/template/ngu.py @@ -10,9 +10,9 @@ def parallel_wrapper(forward_fn: Callable) -> Callable: - r""" + """ Overview: - Process timestep T and batch_size B at the same time, in other words, treat different timestep data as + Process timestep T and batch_size B at the same time, in other words, treat different timestep data as \ different trajectories in a batch. Arguments: - forward_fn (:obj:`Callable`): Normal ``nn.Module`` 's forward function. @@ -44,9 +44,12 @@ def reshape(d): class NGU(nn.Module): """ Overview: - The recurrent Q model for NGU policy, modified from the class DRQN in q_leaning.py - input: x_t, a_{t-1}, r_e_{t-1}, r_i_{t-1}, beta - output: + The recurrent Q model for NGU(https://arxiv.org/pdf/2002.06038.pdf) policy, modified from the class DRQN in \ + q_leaning.py. The implementation mentioned in the original paper is 'adapt the R2D2 agent that uses the \ + dueling network architecture with an LSTM layer after a convolutional neural network'. The NGU network \ + includes encoder, LSTM core(rnn) and head. + Interface: + ``__init__``, ``forward``. """ def __init__( @@ -62,20 +65,26 @@ def __init__( activation: Optional[nn.Module] = nn.ReLU(), norm_type: Optional[str] = None ) -> None: - r""" + """ Overview: - Init the DRQN Model according to arguments. + Init the DRQN Model for NGU according to arguments. Arguments: - - obs_shape (:obj:`Union[int, SequenceType]`): Observation's space. - - action_shape (:obj:`Union[int, SequenceType]`): Action's space. - - encoder_hidden_size_list (:obj:`SequenceType`): Collection of ``hidden_size`` to pass to ``Encoder`` - - head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to ``Head``. - - lstm_type (:obj:`Optional[str]`): Version of rnn cell, now support ['normal', 'pytorch', 'hpc', 'gru'] + - obs_shape (:obj:`Union[int, SequenceType]`): Observation's space, such as 8 or [4, 84, 84]. + - action_shape (:obj:`Union[int, SequenceType]`): Action's space, such as 6 or [2, 3, 3]. + - encoder_hidden_size_list (:obj:`SequenceType`): Collection of ``hidden_size`` to pass to ``Encoder``. + - collector_env_num (:obj:`Optional[int]`): The number of environments used to collect data simultaneously. + - dueling (:obj:`bool`): Whether choose ``DuelingHead`` (True) or ``DiscreteHead (False)``, \ + default to True. + - head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to ``Head``, should match the \ + last element of ``encoder_hidden_size_list``. + - head_layer_num (:obj:`int`): The number of layers in head network. + - lstm_type (:obj:`Optional[str]`): Version of rnn cell, now support ['normal', 'pytorch', 'hpc', 'gru'], \ + default is 'normal'. - activation (:obj:`Optional[nn.Module]`): - The type of activation function to use in ``MLP`` the after ``layer_fn``, - if ``None`` then default set to ``nn.ReLU()`` + The type of activation function to use in ``MLP`` the after ``layer_fn``, \ + if ``None`` then default set to ``nn.ReLU()``. - norm_type (:obj:`Optional[str]`): - The type of normalization to use, see ``ding.torch_utils.fc_block`` for more details` + The type of normalization to use, see ``ding.torch_utils.fc_block`` for more details`. """ super(NGU, self).__init__() # For compatibility: 1, (1, ), [4, H, H] @@ -120,34 +129,31 @@ def __init__( ) def forward(self, inputs: Dict, inference: bool = False, saved_state_timesteps: Optional[list] = None) -> Dict: - r""" + """ Overview: - Use observation, prev_action prev_reward_extrinsic to predict NGU Q output. - Parameter updates with NGU's MLPs forward setup. + Forward computation graph of NGU R2D2 network. Input observation, prev_action prev_reward_extrinsic \ + to predict NGU Q output. Parameter updates with NGU's MLPs forward setup. Arguments: - inputs (:obj:`Dict`): - - inference: (:obj:'bool'): if inference is True, we unroll the one timestep transition, + - obs (:obj:`torch.Tensor`): Encoded observation. + - prev_state (:obj:`list`): Previous state's tensor of size ``(B, N)``. + - inference: (:obj:'bool'): If inference is True, we unroll the one timestep transition, \ if inference is False, we unroll the sequence transitions. - - saved_state_timesteps: (:obj:'Optional[list]'): when inference is False, - we unroll the sequence transitions, then we would save rnn hidden states at timesteps + - saved_state_timesteps: (:obj:'Optional[list]'): When inference is False, \ + we unroll the sequence transitions, then we would save rnn hidden states at timesteps \ that are listed in list saved_state_timesteps. - - ArgumentsKeys: - - obs (:obj:`torch.Tensor`): Encoded observation - - prev_state (:obj:`list`): Previous state's tensor of size ``(B, N)`` - Returns: - outputs (:obj:`Dict`): Run ``MLP`` with ``DRQN`` setups and return the result prediction dictionary. ReturnsKeys: - logit (:obj:`torch.Tensor`): Logit tensor with same size as input ``obs``. - - next_state (:obj:`list`): Next state's tensor of size ``(B, N)`` + - next_state (:obj:`list`): Next state's tensor of size ``(B, N)``. Shapes: - obs (:obj:`torch.Tensor`): :math:`(B, N=obs_space)`, where B is batch size. - - prev_state(:obj:`torch.FloatTensor list`): :math:`[(B, N)]` - - logit (:obj:`torch.FloatTensor`): :math:`(B, N)` - - next_state(:obj:`torch.FloatTensor list`): :math:`[(B, N)]` + - prev_state(:obj:`torch.FloatTensor list`): :math:`[(B, N)]`. + - logit (:obj:`torch.FloatTensor`): :math:`(B, N)`. + - next_state(:obj:`torch.FloatTensor list`): :math:`[(B, N)]`. """ x, prev_state = inputs['obs'], inputs['prev_state'] if 'prev_action' in inputs.keys(): diff --git a/ding/model/template/pdqn.py b/ding/model/template/pdqn.py index 7d24d1d5af..ec94cb3fe1 100644 --- a/ding/model/template/pdqn.py +++ b/ding/model/template/pdqn.py @@ -11,6 +11,17 @@ @MODEL_REGISTRY.register('pdqn') class PDQN(nn.Module): + """ + Overview: + The neural network and computation graph of PDQN(https://arxiv.org/abs/1810.06394v1) and \ + MPDQN(https://arxiv.org/abs/1905.04388) algorithms for parameterized action space. \ + This model supports parameterized action space with discrete ``action_type`` and continuous ``action_arg``. \ + In principle, PDQN consists of x network (continuous action parameter network) and Q network (discrete \ + action type network). But for simplicity, the code is split into ``encoder`` and ``actor_head``, which \ + contain the encoder and head of the above two networks respectively. + Interface: + ``__init__``, ``forward``, ``compute_discrete``, ``compute_continuous``. + """ mode = ['compute_discrete', 'compute_continuous'] def __init__( @@ -26,7 +37,7 @@ def __init__( multi_pass: Optional[bool] = False, action_mask: Optional[list] = None ) -> None: - r""" + """ Overview: Init the PDQN (encoder + head) Model according to input arguments. Arguments: @@ -37,17 +48,17 @@ def __init__( the last element must match ``head_hidden_size``. - dueling (:obj:`dueling`): Whether choose ``DuelingHead`` or ``DiscreteHead(default)``. - head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` of head network. - - head_layer_num (:obj:`int`): The number of layers used in the head network to compute Q value output + - head_layer_num (:obj:`int`): The number of layers used in the head network to compute Q value output. - activation (:obj:`Optional[nn.Module]`): The type of activation function in networks \ - if ``None`` then default set it to ``nn.ReLU()`` + if ``None`` then default set it to ``nn.ReLU()``. - norm_type (:obj:`Optional[str]`): The type of normalization in networks, see \ ``ding.torch_utils.fc_block`` for more details. - multi_pass (:obj:`Optional[bool]`): Whether to use multi pass version. - - action_mask: (:obj:`Optional[list]`): An action mask indicating how action args are - associated to each discrete action. For example, if there are 3 discrete action, - 4 continous action args, and the first discrete action associates with the first - continuous action args, the second discrete action associates with the second continuous - action args, and the third discrete action associates with the remaining 2 action args, + - action_mask: (:obj:`Optional[list]`): An action mask indicating how action args are \ + associated to each discrete action. For example, if there are 3 discrete action, \ + 4 continous action args, and the first discrete action associates with the first \ + continuous action args, the second discrete action associates with the second continuous \ + action args, and the third discrete action associates with the remaining 2 action args, \ the action mask will be like: [[1,0,0,0],[0,1,0,0],[0,0,1,1]] with shape 3*4. """ super(PDQN, self).__init__() @@ -120,33 +131,42 @@ def __init__( self.actor_head = nn.ModuleList([self.dis_head, self.cont_head]) # self.encoder = nn.ModuleList([self.dis_encoder, self.cont_encoder]) + # To speed up the training process, the X network and the Q network share the encoder for the state self.encoder = nn.ModuleList([self.cont_encoder, self.cont_encoder]) def forward(self, inputs: Union[torch.Tensor, Dict, EasyDict], mode: str) -> Dict: - r""" + """ Overview: PDQN forward computation graph, input observation tensor to predict q_value for \ - discrete actions and values for continuous action_args + discrete actions and values for continuous action_args. Arguments: - - inputs (:obj:`torch.Tensor`): Observation inputs + - inputs (:obj:`Union[torch.Tensor, Dict, EasyDict]`): Inputs including observation and \ + other info according to `mode`. - mode (:obj:`str`): Name of the forward mode. Shapes: - - inputs (:obj:`torch.Tensor`): :math:`(B, N)`, where B is batch size and N is ``obs_shape`` + - inputs (:obj:`torch.Tensor`): :math:`(B, N)`, where B is batch size and N is ``obs_shape``. """ assert mode in self.mode, "not support forward mode: {}/{}".format(mode, self.mode) return getattr(self, mode)(inputs) def compute_continuous(self, inputs: torch.Tensor) -> Dict: - r""" + """ Overview: Use observation tensor to predict continuous action args. Arguments: - - inputs (:obj:`torch.Tensor`): Observation inputs - Shapes: - - inputs (:obj:`torch.Tensor`): :math:`(B, N)`, where B is batch size and N is ``obs_shape`` + - inputs (:obj:`torch.Tensor`): Observation inputs. Returns: - - outputs (:obj:`Dict`): A dict with key 'action_args' - - 'action_args': the continuous action args + - outputs (:obj:`Dict`): A dict with key 'action_args'. + - 'action_args' (:obj:`torch.Tensor`): The continuous action args. + Shapes: + - inputs (:obj:`torch.Tensor`): :math:`(B, N)`, where B is batch size and N is ``obs_shape``. + - action_args (:obj:`torch.Tensor`): :math:`(B, M)`, where M is ``action_args_shape``. + Examples: + >>> act_shape = EasyDict({'action_type_shape': (3, ), 'action_args_shape': (5, )}) + >>> model = PDQN(4, act_shape) + >>> inputs = torch.randn(64, 4) + >>> outputs = model.forward(inputs, mode='compute_continuous') + >>> assert outputs['action_args'].shape == torch.Size([64, 5]) """ cont_x = self.encoder[1](inputs) # size (B, encoded_state_shape) action_args = self.actor_head[1](cont_x)['pred'] # size (B, action_args_shape) @@ -154,15 +174,25 @@ def compute_continuous(self, inputs: torch.Tensor) -> Dict: return outputs def compute_discrete(self, inputs: Union[Dict, EasyDict]) -> Dict: - r""" + """ Overview: Use observation tensor and continuous action args to predict discrete action types. Arguments: - - inputs (:obj:`torch.Tensor`): A dict with keys 'state', 'action_args' + - inputs (:obj:`Union[Dict, EasyDict]`): A dict with keys 'state', 'action_args'. + - state (:obj:`torch.Tensor`): Observation inputs. + - action_args (:obj:`torch.Tensor`): Action parameters are used to concatenate with the observation \ + and serve as input to the discrete action type network. Returns: - - outputs (:obj:`Dict`): A dict with keys 'logit', 'action_args' - - 'logit': the logit value for each discrete action, - - 'action_args': the continuous action args(same as the inputs['action_args']) for later usage + - outputs (:obj:`Dict`): A dict with keys 'logit', 'action_args'. + - 'logit': The logit value for each discrete action. + - 'action_args': The continuous action args(same as the inputs['action_args']) for later usage. + Examples: + >>> act_shape = EasyDict({'action_type_shape': (3, ), 'action_args_shape': (5, )}) + >>> model = PDQN(4, act_shape) + >>> inputs = {'state': torch.randn(64, 4), 'action_args': torch.randn(64, 5)} + >>> outputs = model.forward(inputs, mode='compute_discrete') + >>> assert outputs['logit'].shape == torch.Size([64, 3]) + >>> assert outputs['action_args'].shape == torch.Size([64, 5]) """ dis_x = self.encoder[0](inputs['state']) # size (B, encoded_state_shape) action_args = inputs['action_args'] # size (B, action_args_shape) @@ -190,6 +220,8 @@ def compute_discrete(self, inputs: Union[Dict, EasyDict]) -> Dict: logit = torch.diagonal(logit, dim1=-2, dim2=-1) # (B, K) else: # pdqn # size (B, encoded_state_shape + action_args_shape) + if len(action_args.shape) == 1: # (B, ) -> (B, 1) + action_args = action_args.unsqueeze(1) state_action_cat = torch.cat((dis_x, action_args), dim=-1) logit = self.actor_head[0](state_action_cat)['logit'] # size (B, K) where K is action_type_shape diff --git a/ding/model/template/pg.py b/ding/model/template/pg.py new file mode 100644 index 0000000000..6059642dd3 --- /dev/null +++ b/ding/model/template/pg.py @@ -0,0 +1,111 @@ +from typing import Union, Optional, Dict, Callable, List +import torch +import torch.nn as nn +from easydict import EasyDict + +from ding.torch_utils import get_lstm +from ding.utils import MODEL_REGISTRY, SequenceType, squeeze +from ..common import FCEncoder, ConvEncoder, DiscreteHead, DuelingHead, \ + MultiHead, RegressionHead, ReparameterizationHead, independent_normal_dist + + +@MODEL_REGISTRY.register('pg') +class PG(nn.Module): + """ + Overview: + The neural network and computation graph of algorithms related to Policy Gradient(PG) \ + (https://proceedings.neurips.cc/paper/1999/file/464d828b85b0bed98e80ade0a5c43b0f-Paper.pdf). \ + The PG model is composed of two parts: encoder and head. Encoders are used to extract the feature \ + from various observation. Heads are used to predict corresponding action logit. + Interface: + ``__init__``, ``forward``. + """ + + def __init__( + self, + obs_shape: Union[int, SequenceType], + action_shape: Union[int, SequenceType], + action_space: str = 'discrete', + encoder_hidden_size_list: SequenceType = [128, 128, 64], + head_hidden_size: Optional[int] = None, + head_layer_num: int = 1, + activation: Optional[nn.Module] = nn.ReLU(), + norm_type: Optional[str] = None + ) -> None: + """ + Overview: + Initialize the PG model according to corresponding input arguments. + Arguments: + - obs_shape (:obj:`Union[int, SequenceType]`): Observation space shape, such as 8 or [4, 84, 84]. + - action_shape (:obj:`Union[int, SequenceType]`): Action space shape, such as 6 or [2, 3, 3]. + - action_space (:obj:`str`): The type of different action spaces, including ['discrete', 'continuous'], \ + then will instantiate corresponding head, including ``DiscreteHead`` and ``ReparameterizationHead``. + - encoder_hidden_size_list (:obj:`SequenceType`): Collection of ``hidden_size`` to pass to ``Encoder``, \ + the last element must match ``head_hidden_size``. + - head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` of ``head`` network, defaults \ + to None, it must match the last element of ``encoder_hidden_size_list``. + - head_layer_num (:obj:`int`): The num of layers used in the ``head`` network to compute action. + - activation (:obj:`Optional[nn.Module]`): The type of activation function in networks \ + if ``None`` then default set it to ``nn.ReLU()``. + - norm_type (:obj:`Optional[str]`): The type of normalization in networks, see \ + ``ding.torch_utils.fc_block`` for more details. you can choose one of ['BN', 'IN', 'SyncBN', 'LN'] + Examples: + >>> model = PG((4, 84, 84), 5) + >>> inputs = torch.randn(8, 4, 84, 84) + >>> outputs = model(inputs) + >>> assert isinstance(outputs, dict) + >>> assert outputs['logit'].shape == (8, 5) + >>> assert outputs['dist'].sample().shape == (8, ) + """ + super(PG, self).__init__() + # For compatibility: 1, (1, ), [4, 32, 32] + obs_shape, action_shape = squeeze(obs_shape), squeeze(action_shape) + if head_hidden_size is None: + head_hidden_size = encoder_hidden_size_list[-1] + # FC Encoder + if isinstance(obs_shape, int) or len(obs_shape) == 1: + self.encoder = FCEncoder(obs_shape, encoder_hidden_size_list, activation=activation, norm_type=norm_type) + # Conv Encoder + elif len(obs_shape) == 3: + self.encoder = ConvEncoder(obs_shape, encoder_hidden_size_list, activation=activation, norm_type=norm_type) + else: + raise RuntimeError( + "not support obs_shape for pre-defined encoder: {}, please customize your own BC".format(obs_shape) + ) + self.action_space = action_space + # Head + if self.action_space == 'discrete': + self.head = DiscreteHead( + head_hidden_size, action_shape, head_layer_num, activation=activation, norm_type=norm_type + ) + elif self.action_space == 'continuous': + self.head = ReparameterizationHead( + head_hidden_size, + action_shape, + head_layer_num, + activation=activation, + norm_type=norm_type, + sigma_type='independent' + ) + else: + raise KeyError("not support action space: {}".format(self.action_space)) + + def forward(self, x: torch.Tensor) -> Dict: + """ + Overview: + PG forward computation graph, input observation tensor to predict policy distribution. + Arguments: + - x (:obj:`torch.Tensor`): The input observation tensor data. + Returns: + - outputs (:obj:`torch.distributions`): The output policy distribution. If action space is \ + discrete, the output is Categorical distribution; if action space is continuous, the output is Normal \ + distribution. + """ + x = self.encoder(x) + x = self.head(x) + if self.action_space == 'discrete': + x['dist'] = torch.distributions.Categorical(logits=x['logit']) + elif self.action_space == 'continuous': + x = {'logit': {'mu': x['mu'], 'sigma': x['sigma']}} + x['dist'] = independent_normal_dist(x['logit']) + return x diff --git a/ding/model/template/ppg.py b/ding/model/template/ppg.py index f9dd64d21e..76df579e71 100644 --- a/ding/model/template/ppg.py +++ b/ding/model/template/ppg.py @@ -8,6 +8,15 @@ @MODEL_REGISTRY.register('ppg') class PPG(nn.Module): + """ + Overview: + Phasic Policy Gradient (PPG) model from paper `Phasic Policy Gradient` + https://arxiv.org/abs/2009.04416 \ + This module contains VAC module and an auxiliary critic module. + Interfaces: + ``forward``, ``compute_actor``, ``compute_critic``, ``compute_actor_critic`` + """ + mode = ['compute_actor', 'compute_critic', 'compute_actor_critic'] def __init__( @@ -25,6 +34,27 @@ def __init__( norm_type: Optional[str] = None, impala_cnn_encoder: bool = False, ) -> None: + """ + Overview: + Initailize the PPG Model according to input arguments. + Arguments: + - obs_shape (:obj:`Union[int, SequenceType]`): Observation's shape, such as 128, (156, ). + - action_shape (:obj:`Union[int, SequenceType]`): Action's shape, such as 4, (3, ). + - action_space (:obj:`str`): The action space type, such as 'discrete', 'continuous'. + - share_encoder (:obj:`bool`): Whether to share encoder. + - encoder_hidden_size_list (:obj:`SequenceType`): The hidden size list of encoder. + - actor_head_hidden_size (:obj:`int`): The ``hidden_size`` to pass to actor head. + - actor_head_layer_num (:obj:`int`): The num of layers used in the network to compute Q value output \ + for actor head. + - critic_head_hidden_size (:obj:`int`): The ``hidden_size`` to pass to critic head. + - critic_head_layer_num (:obj:`int`): The num of layers used in the network to compute Q value output \ + for critic head. + - activation (:obj:`Optional[nn.Module]`): The type of activation function to use in ``MLP`` \ + after each FC layer, if ``None`` then default set to ``nn.ReLU()``. + - norm_type (:obj:`Optional[str]`): The type of normalization to after network layer (FC, Conv), \ + see ``ding.torch_utils.network`` for more details. + - impala_cnn_encoder (:obj:`bool`): Whether to use impala cnn encoder. + """ super(PPG, self).__init__() self.actor_critic = VAC( obs_shape, @@ -43,20 +73,53 @@ def __init__( self.aux_critic = copy.deepcopy(self.actor_critic.critic) def forward(self, inputs: Union[torch.Tensor, Dict], mode: str) -> Dict: + """ + Overview: + Compute action logits or value according to mode being ``compute_actor``, ``compute_critic`` or \ + ``compute_actor_critic``. + Arguments: + - x (:obj:`torch.Tensor`): The input observation tensor data. + - mode (:obj:`str`): The forward mode, all the modes are defined in the beginning of this class. + Returns: + - outputs (:obj:`Dict`): The output dict of PPG's forward computation graph, whose key-values vary from \ + different ``mode``. + """ assert mode in self.mode, "not support forward mode: {}/{}".format(mode, self.mode) return getattr(self, mode)(inputs) def compute_actor(self, x: torch.Tensor) -> Dict: """ + Overview: + Use actor to compute action logits. + Arguments: + - x (:obj:`torch.Tensor`): The input observation tensor data. + Returns: + - output (:obj:`Dict`): The output data containing action logits. ReturnsKeys: - - necessary: ``logit`` + - logit (:obj:`torch.Tensor`): The predicted action logit tensor, for discrete action space, it will be \ + the same dimension real-value ranged tensor of possible action choices, and for continuous action \ + space, it will be the mu and sigma of the Gaussian distribution, and the number of mu and sigma is the \ + same as the number of continuous actions. Hybrid action space is a kind of combination of discrete \ + and continuous action space, so the logit will be a dict with ``action_type`` and ``action_args``. + Shapes: + - x (:obj:`torch.Tensor`): :math:`(B, N)`, where B is batch size and N is the input feature size. + - output (:obj:`Dict`): ``logit``: :math:`(B, A)`, where B is batch size and A is the action space size. """ return self.actor_critic(x, mode='compute_actor') def compute_critic(self, x: torch.Tensor) -> Dict: """ + Overview: + Use critic to compute value. + Arguments: + - x (:obj:`torch.Tensor`): The input observation tensor data. + Returns: + - output (:obj:`Dict`): The output dict of VAC's forward computation graph for critic, including ``value``. ReturnsKeys: - necessary: ``value`` + Shapes: + - x (:obj:`torch.Tensor`): :math:`(B, N)`, where B is batch size and N is the input feature size. + - output (:obj:`Dict`): ``value``: :math:`(B, 1)`, where B is batch size. """ x = self.aux_critic[0](x) # encoder x = self.aux_critic[1](x) # head @@ -64,10 +127,26 @@ def compute_critic(self, x: torch.Tensor) -> Dict: def compute_actor_critic(self, x: torch.Tensor) -> Dict: """ + Overview: + Use actor and critic to compute action logits and value. + Arguments: + - x (:obj:`torch.Tensor`): The input observation tensor data. + Returns: + - outputs (:obj:`Dict`): The output dict of PPG's forward computation graph for both actor and critic, \ + including ``logit`` and ``value``. + ReturnsKeys: + - logit (:obj:`torch.Tensor`): The predicted action logit tensor, for discrete action space, it will be \ + the same dimension real-value ranged tensor of possible action choices, and for continuous action \ + space, it will be the mu and sigma of the Gaussian distribution, and the number of mu and sigma is the \ + same as the number of continuous actions. Hybrid action space is a kind of combination of discrete \ + and continuous action space, so the logit will be a dict with ``action_type`` and ``action_args``. + - value (:obj:`torch.Tensor`): The predicted state value tensor. + Shapes: + - x (:obj:`torch.Tensor`): :math:`(B, N)`, where B is batch size and N is the input feature size. + - output (:obj:`Dict`): ``value``: :math:`(B, 1)`, where B is batch size. + - output (:obj:`Dict`): ``logit``: :math:`(B, A)`, where B is batch size and A is the action space size. + .. note:: ``compute_actor_critic`` interface aims to save computation when shares encoder. - - ReturnsKeys: - - necessary: ``value``, ``logit`` """ return self.actor_critic(x, mode='compute_actor_critic') diff --git a/ding/model/template/procedure_cloning.py b/ding/model/template/procedure_cloning.py new file mode 100644 index 0000000000..4f03c8a4bf --- /dev/null +++ b/ding/model/template/procedure_cloning.py @@ -0,0 +1,327 @@ +from typing import Optional, Tuple, Union, Dict + +import torch +import torch.nn as nn + +from ding.utils import MODEL_REGISTRY, SequenceType +from ding.torch_utils.network.transformer import Attention +from ding.torch_utils.network.nn_module import fc_block, build_normalization +from ..common import FCEncoder, ConvEncoder + + +class PCTransformer(nn.Module): + """ + Overview: + The transformer block for neural network of algorithms related to Procedure cloning (PC). + Interfaces: + ``__init__``, ``forward``. + """ + + def __init__( + self, cnn_hidden: int, att_hidden: int, att_heads: int, drop_p: float, max_T: int, n_att: int, + feedforward_hidden: int, n_feedforward: int + ) -> None: + """ + Overview: + Initialize the procedure cloning transformer model according to corresponding input arguments. + Arguments: + - cnn_hidden (:obj:`int`): The last channel dimension of CNN encoder, such as 32. + - att_hidden (:obj:`int`): The dimension of attention blocks, such as 32. + - att_heads (:obj:`int`): The number of heads in attention blocks, such as 4. + - drop_p (:obj:`float`): The drop out rate of attention, such as 0.5. + - max_T (:obj:`int`): The sequence length of procedure cloning, such as 4. + - n_attn (:obj:`int`): The number of attention layers, such as 4. + - feedforward_hidden (:obj:`int`):The dimension of feedforward layers, such as 32. + - n_feedforward (:obj:`int`): The number of feedforward layers, such as 4. + """ + super().__init__() + self.n_att = n_att + self.n_feedforward = n_feedforward + self.attention_layer = [] + + self.norm_layer = [nn.LayerNorm(att_hidden)] * n_att + self.attention_layer.append(Attention(cnn_hidden, att_hidden, att_hidden, att_heads, nn.Dropout(drop_p))) + for i in range(n_att - 1): + self.attention_layer.append(Attention(att_hidden, att_hidden, att_hidden, att_heads, nn.Dropout(drop_p))) + + self.att_drop = nn.Dropout(drop_p) + + self.fc_blocks = [] + self.fc_blocks.append(fc_block(att_hidden, feedforward_hidden, activation=nn.ReLU())) + for i in range(n_feedforward - 1): + self.fc_blocks.append(fc_block(feedforward_hidden, feedforward_hidden, activation=nn.ReLU())) + self.norm_layer.extend([nn.LayerNorm(feedforward_hidden)] * n_feedforward) + self.mask = torch.tril(torch.ones((max_T, max_T), dtype=torch.bool)).view(1, 1, max_T, max_T) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Overview: + The unique execution (forward) method of PCTransformer. + Arguments: + - x (:obj:`torch.Tensor`): Sequential data of several hidden states. + Returns: + - output (:obj:`torch.Tensor`): A tensor with the same shape as the input. + Examples: + >>> model = PCTransformer(128, 128, 8, 0, 16, 2, 128, 2) + >>> h = torch.randn((2, 16, 128)) + >>> h = model(h) + >>> assert h.shape == torch.Size([2, 16, 128]) + """ + for i in range(self.n_att): + x = self.att_drop(self.attention_layer[i](x, self.mask)) + x = self.norm_layer[i](x) + for i in range(self.n_feedforward): + x = self.fc_blocks[i](x) + x = self.norm_layer[i + self.n_att](x) + return x + + +@MODEL_REGISTRY.register('pc_mcts') +class ProcedureCloningMCTS(nn.Module): + """ + Overview: + The neural network of algorithms related to Procedure cloning (PC). + Interfaces: + ``__init__``, ``forward``. + """ + + def __init__( + self, + obs_shape: SequenceType, + action_dim: int, + cnn_hidden_list: SequenceType = [128, 128, 256, 256, 256], + cnn_activation: nn.Module = nn.ReLU(), + cnn_kernel_size: SequenceType = [3, 3, 3, 3, 3], + cnn_stride: SequenceType = [1, 1, 1, 1, 1], + cnn_padding: SequenceType = [1, 1, 1, 1, 1], + mlp_hidden_list: SequenceType = [256, 256], + mlp_activation: nn.Module = nn.ReLU(), + att_heads: int = 8, + att_hidden: int = 128, + n_att: int = 4, + n_feedforward: int = 2, + feedforward_hidden: int = 256, + drop_p: float = 0.5, + max_T: int = 17 + ) -> None: + """ + Overview: + Initialize the MCTS procedure cloning model according to corresponding input arguments. + Arguments: + - obs_shape (:obj:`SequenceType`): Observation space shape, such as [4, 84, 84]. + - action_dim (:obj:`int`): Action space shape, such as 6. + - cnn_hidden_list (:obj:`SequenceType`): The cnn channel dims for each block, such as\ + [128, 128, 256, 256, 256]. + - cnn_activation (:obj:`nn.Module`): The activation function for cnn blocks, such as ``nn.ReLU()``. + - cnn_kernel_size (:obj:`SequenceType`): The kernel size for each cnn block, such as [3, 3, 3, 3, 3]. + - cnn_stride (:obj:`SequenceType`): The stride for each cnn block, such as [1, 1, 1, 1, 1]. + - cnn_padding (:obj:`SequenceType`): The padding for each cnn block, such as [1, 1, 1, 1, 1]. + - mlp_hidden_list (:obj:`SequenceType`): The last dim for this must match the last dim of \ + ``cnn_hidden_list``, such as [256, 256]. + - mlp_activation (:obj:`nn.Module`): The activation function for mlp layers, such as ``nn.ReLU()``. + - att_heads (:obj:`int`): The number of attention heads in transformer, such as 8. + - att_hidden (:obj:`int`): The number of attention dimension in transformer, such as 128. + - n_att (:obj:`int`): The number of attention blocks in transformer, such as 4. + - n_feedforward (:obj:`int`): The number of feedforward layers in transformer, such as 2. + - drop_p (:obj:`float`): The drop out rate of attention, such as 0.5. + - max_T (:obj:`int`): The sequence length of procedure cloning, such as 17. + """ + super().__init__() + + # Conv Encoder + self.embed_state = ConvEncoder( + obs_shape, cnn_hidden_list, cnn_activation, cnn_kernel_size, cnn_stride, cnn_padding + ) + self.embed_action = FCEncoder(action_dim, mlp_hidden_list, activation=mlp_activation) + + self.cnn_hidden_list = cnn_hidden_list + + assert cnn_hidden_list[-1] == mlp_hidden_list[-1] + layers = [] + for i in range(n_att): + if i == 0: + layers.append(Attention(cnn_hidden_list[-1], att_hidden, att_hidden, att_heads, nn.Dropout(drop_p))) + else: + layers.append(Attention(att_hidden, att_hidden, att_hidden, att_heads, nn.Dropout(drop_p))) + layers.append(build_normalization('LN')(att_hidden)) + for i in range(n_feedforward): + if i == 0: + layers.append(fc_block(att_hidden, feedforward_hidden, activation=nn.ReLU())) + else: + layers.append(fc_block(feedforward_hidden, feedforward_hidden, activation=nn.ReLU())) + self.layernorm2 = build_normalization('LN')(feedforward_hidden) + + self.transformer = PCTransformer( + cnn_hidden_list[-1], att_hidden, att_heads, drop_p, max_T, n_att, feedforward_hidden, n_feedforward + ) + + self.predict_goal = torch.nn.Linear(cnn_hidden_list[-1], cnn_hidden_list[-1]) + self.predict_action = torch.nn.Linear(cnn_hidden_list[-1], action_dim) + + def forward(self, states: torch.Tensor, goals: torch.Tensor, + actions: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Overview: + ProcedureCloningMCTS forward computation graph, input states tensor and goals tensor, \ + calculate the predicted states and actions. + Arguments: + - states (:obj:`torch.Tensor`): The observation of current time. + - goals (:obj:`torch.Tensor`): The target observation after a period. + - actions (:obj:`torch.Tensor`): The actions executed during the period. + Returns: + - outputs (:obj:`Tuple[torch.Tensor, torch.Tensor]`): Predicted states and actions. + Examples: + >>> inputs = { \ + 'states': torch.randn(2, 3, 64, 64), \ + 'goals': torch.randn(2, 3, 64, 64), \ + 'actions': torch.randn(2, 15, 9) \ + } + >>> model = ProcedureCloningMCTS(obs_shape=(3, 64, 64), action_dim=9) + >>> goal_preds, action_preds = model(inputs['states'], inputs['goals'], inputs['actions']) + >>> assert goal_preds.shape == (2, 256) + >>> assert action_preds.shape == (2, 16, 9) + """ + B, T, _ = actions.shape + + # shape: (B, h_dim) + state_embeddings = self.embed_state(states).reshape(B, 1, self.cnn_hidden_list[-1]) + goal_embeddings = self.embed_state(goals).reshape(B, 1, self.cnn_hidden_list[-1]) + # shape: (B, context_len, h_dim) + actions_embeddings = self.embed_action(actions) + + h = torch.cat((state_embeddings, goal_embeddings, actions_embeddings), dim=1) + h = self.transformer(h) + h = h.reshape(B, T + 2, self.cnn_hidden_list[-1]) + + goal_preds = self.predict_goal(h[:, 0, :]) + action_preds = self.predict_action(h[:, 1:, :]) + + return goal_preds, action_preds + + +class BFSConvEncoder(nn.Module): + """ + Overview: + The ``BFSConvolution Encoder`` used to encode raw 3-dim observations. And output a feature map with the + same height and width as input. Interfaces: ``__init__``, ``forward``. + """ + + def __init__( + self, + obs_shape: SequenceType, + hidden_size_list: SequenceType = [32, 64, 64, 128], + activation: Optional[nn.Module] = nn.ReLU(), + kernel_size: SequenceType = [8, 4, 3], + stride: SequenceType = [4, 2, 1], + padding: Optional[SequenceType] = None, + ) -> None: + """ + Overview: + Init the ``BFSConvolution Encoder`` according to the provided arguments. + Arguments: + - obs_shape (:obj:`SequenceType`): Sequence of ``in_channel``, plus one or more ``input size``. + - hidden_size_list (:obj:`SequenceType`): Sequence of ``hidden_size`` of subsequent conv layers \ + and the final dense layer. + - activation (:obj:`nn.Module`): Type of activation to use in the conv ``layers`` and ``ResBlock``. \ + Default is ``nn.ReLU()``. + - kernel_size (:obj:`SequenceType`): Sequence of ``kernel_size`` of subsequent conv layers. + - stride (:obj:`SequenceType`): Sequence of ``stride`` of subsequent conv layers. + - padding (:obj:`SequenceType`): Padding added to all four sides of the input for each conv layer. \ + See ``nn.Conv2d`` for more details. Default is ``None``. + """ + super(BFSConvEncoder, self).__init__() + self.obs_shape = obs_shape + self.act = activation + self.hidden_size_list = hidden_size_list + if padding is None: + padding = [0 for _ in range(len(kernel_size))] + + layers = [] + input_size = obs_shape[0] # in_channel + for i in range(len(kernel_size)): + layers.append(nn.Conv2d(input_size, hidden_size_list[i], kernel_size[i], stride[i], padding[i])) + layers.append(self.act) + input_size = hidden_size_list[i] + layers = layers[:-1] + self.main = nn.Sequential(*layers) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Overview: + Return output tensor of the env observation. + Arguments: + - x (:obj:`torch.Tensor`): Env raw observation. + Returns: + - outputs (:obj:`torch.Tensor`): Output embedding tensor. + Examples: + >>> model = BFSConvEncoder([3, 16, 16], [32, 32, 4], kernel_size=[3, 3, 3], stride=[1, 1, 1]\ + , padding=[1, 1, 1]) + >>> inputs = torch.randn(3, 16, 16).unsqueeze(0) + >>> outputs = model(inputs) + >>> assert outputs['logit'].shape == torch.Size([4, 16, 16]) + """ + return self.main(x) + + +@MODEL_REGISTRY.register('pc_bfs') +class ProcedureCloningBFS(nn.Module): + """ + Overview: + The neural network introduced in procedure cloning (PC) to process 3-dim observations.\ + Given an input, this model will perform several 3x3 convolutions and output a feature map with \ + the same height and width of input. The channel number of output will be the ``action_shape``. + Interfaces: + ``__init__``, ``forward``. + """ + + def __init__( + self, + obs_shape: SequenceType, + action_shape: int, + encoder_hidden_size_list: SequenceType = [128, 128, 256, 256], + ): + """ + Overview: + Init the ``BFSConvolution Encoder`` according to the provided arguments. + Arguments: + - obs_shape (:obj:`SequenceType`): Sequence of ``in_channel``, plus one or more ``input size``,\ + such as [4, 84, 84]. + - action_dim (:obj:`int`): Action space shape, such as 6. + - cnn_hidden_list (:obj:`SequenceType`): The cnn channel dims for each block, such as [128, 128, 256, 256]. + """ + super().__init__() + num_layers = len(encoder_hidden_size_list) + + kernel_sizes = (3, ) * (num_layers + 1) + stride_sizes = (1, ) * (num_layers + 1) + padding_sizes = (1, ) * (num_layers + 1) + # The output channel equals to action_shape + 1 + encoder_hidden_size_list.append(action_shape + 1) + + self._encoder = BFSConvEncoder( + obs_shape=obs_shape, + hidden_size_list=encoder_hidden_size_list, + kernel_size=kernel_sizes, + stride=stride_sizes, + padding=padding_sizes, + ) + + def forward(self, x: torch.Tensor) -> Dict: + """ + Overview: + The computation graph. Given a 3-dim observation, this function will return a tensor with the same \ + height and width. The channel number of output will be the ``action_shape``. + Arguments: + - x (:obj:`torch.Tensor`): The input observation tensor data. + Returns: + - outputs (:obj:`Dict`): The output dict of model's forward computation graph, \ + only contains a single key ``logit``. + Examples: + >>> model = ProcedureCloningBFS([3, 16, 16], 4) + >>> inputs = torch.randn(16, 16, 3).unsqueeze(0) + >>> outputs = model(inputs) + >>> assert outputs['logit'].shape == torch.Size([16, 16, 4]) + """ + x = x.permute(0, 3, 1, 2) + x = self._encoder(x) + return {'logit': x.permute(0, 2, 3, 1)} diff --git a/ding/model/template/q_learning.py b/ding/model/template/q_learning.py index 54df388446..d7c1f25a18 100644 --- a/ding/model/template/q_learning.py +++ b/ding/model/template/q_learning.py @@ -5,50 +5,75 @@ from ding.torch_utils import get_lstm from ding.utils import MODEL_REGISTRY, SequenceType, squeeze from ..common import FCEncoder, ConvEncoder, DiscreteHead, DuelingHead, MultiHead, RainbowHead, \ - QuantileHead, FQFHead, QRDQNHead, DistributionHead + QuantileHead, FQFHead, QRDQNHead, DistributionHead, BranchingHead from ding.torch_utils.network.gtrxl import GTrXL @MODEL_REGISTRY.register('dqn') class DQN(nn.Module): + """ + Overview: + The neural nework structure and computation graph of Deep Q Network (DQN) algorithm, which is the most classic \ + value-based RL algorithm for discrete action. The DQN is composed of two parts: ``encoder`` and ``head``. \ + The ``encoder`` is used to extract the feature from various observation, and the ``head`` is used to compute \ + the Q value of each action dimension. + Interfaces: + ``__init__``, ``forward``. + + .. note:: + Current ``DQN`` supports two types of encoder: ``FCEncoder`` and ``ConvEncoder``, two types of head: \ + ``DiscreteHead`` and ``DuelingHead``. You can customize your own encoder or head by inheriting this class. + """ def __init__( - self, - obs_shape: Union[int, SequenceType], - action_shape: Union[int, SequenceType], - encoder_hidden_size_list: SequenceType = [128, 128, 64], - dueling: bool = True, - head_hidden_size: Optional[int] = None, - head_layer_num: int = 1, - activation: Optional[nn.Module] = nn.ReLU(), - norm_type: Optional[str] = None + self, + obs_shape: Union[int, SequenceType], + action_shape: Union[int, SequenceType], + encoder_hidden_size_list: SequenceType = [128, 128, 64], + dueling: bool = True, + head_hidden_size: Optional[int] = None, + head_layer_num: int = 1, + activation: Optional[nn.Module] = nn.ReLU(), + norm_type: Optional[str] = None, + dropout: Optional[float] = None, + init_bias: Optional[float] = None, + noise: bool = False, ) -> None: """ Overview: - Init the DQN (encoder + head) Model according to input arguments. + initialize the DQN (encoder + head) Model according to corresponding input arguments. Arguments: - obs_shape (:obj:`Union[int, SequenceType]`): Observation space shape, such as 8 or [4, 84, 84]. - action_shape (:obj:`Union[int, SequenceType]`): Action space shape, such as 6 or [2, 3, 3]. - encoder_hidden_size_list (:obj:`SequenceType`): Collection of ``hidden_size`` to pass to ``Encoder``, \ the last element must match ``head_hidden_size``. - - dueling (:obj:`dueling`): Whether choose ``DuelingHead`` or ``DiscreteHead(default)``. - - head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` of head network. - - head_layer_num (:obj:`int`): The number of layers used in the head network to compute Q value output + - dueling (:obj:`Optional[bool]`): Whether choose ``DuelingHead`` or ``DiscreteHead (default)``. + - head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` of head network, defaults to None, \ + then it will be set to the last element of ``encoder_hidden_size_list``. + - head_layer_num (:obj:`int`): The number of layers used in the head network to compute Q value output. - activation (:obj:`Optional[nn.Module]`): The type of activation function in networks \ - if ``None`` then default set it to ``nn.ReLU()`` + if ``None`` then default set it to ``nn.ReLU()``. - norm_type (:obj:`Optional[str]`): The type of normalization in networks, see \ - ``ding.torch_utils.fc_block`` for more details. + ``ding.torch_utils.fc_block`` for more details. you can choose one of ['BN', 'IN', 'SyncBN', 'LN'] + - dropout (:obj:`Optional[float]`): The dropout rate of the dropout layer. \ + if ``None`` then default disable dropout layer. + - init_bias (:obj:`Optional[float]`): The initial value of the last layer bias in the head network. \ + - noise (:obj:`bool`): Whether to use ``NoiseLinearLayer`` as ``layer_fn`` to boost exploration in \ + Q networks' MLP. Default to ``False``. """ super(DQN, self).__init__() - # For compatibility: 1, (1, ), [4, 32, 32] + # Squeeze data from tuple, list or dict to single object. For example, from (4, ) to 4 obs_shape, action_shape = squeeze(obs_shape), squeeze(action_shape) if head_hidden_size is None: head_hidden_size = encoder_hidden_size_list[-1] # FC Encoder if isinstance(obs_shape, int) or len(obs_shape) == 1: - self.encoder = FCEncoder(obs_shape, encoder_hidden_size_list, activation=activation, norm_type=norm_type) + self.encoder = FCEncoder( + obs_shape, encoder_hidden_size_list, activation=activation, norm_type=norm_type, dropout=dropout + ) # Conv Encoder elif len(obs_shape) == 3: + assert dropout is None, "dropout is not supported in ConvEncoder" self.encoder = ConvEncoder(obs_shape, encoder_hidden_size_list, activation=activation, norm_type=norm_type) else: raise RuntimeError( @@ -67,39 +92,161 @@ def __init__( action_shape, layer_num=head_layer_num, activation=activation, - norm_type=norm_type + norm_type=norm_type, + dropout=dropout, + noise=noise, ) else: self.head = head_cls( - head_hidden_size, action_shape, head_layer_num, activation=activation, norm_type=norm_type + head_hidden_size, + action_shape, + head_layer_num, + activation=activation, + norm_type=norm_type, + dropout=dropout, + noise=noise, ) + if init_bias is not None and head_cls == DuelingHead: + # Zero the last layer bias of advantage head + self.head.A[-1][0].bias.data.fill_(init_bias) def forward(self, x: torch.Tensor) -> Dict: - r""" + """ Overview: DQN forward computation graph, input observation tensor to predict q_value. Arguments: - - x (:obj:`torch.Tensor`): Observation inputs + - x (:obj:`torch.Tensor`): The input observation tensor data. Returns: - - outputs (:obj:`Dict`): DQN forward outputs, such as q_value. + - outputs (:obj:`Dict`): The output of DQN's forward, including q_value. ReturnsKeys: - - logit (:obj:`torch.Tensor`): Discrete Q-value output of each action dimension. + - logit (:obj:`torch.Tensor`): Discrete Q-value output of each possible action dimension. Shapes: - x (:obj:`torch.Tensor`): :math:`(B, N)`, where B is batch size and N is ``obs_shape`` - - logit (:obj:`torch.FloatTensor`): :math:`(B, M)`, where B is batch size and M is ``action_shape`` + - logit (:obj:`torch.Tensor`): :math:`(B, M)`, where B is batch size and M is ``action_shape`` Examples: >>> model = DQN(32, 6) # arguments: 'obs_shape' and 'action_shape' >>> inputs = torch.randn(4, 32) >>> outputs = model(inputs) >>> assert isinstance(outputs, dict) and outputs['logit'].shape == torch.Size([4, 6]) + + .. note:: + For consistency and compatibility, we name all the outputs of the network which are related to action \ + selections as ``logit``. """ x = self.encoder(x) x = self.head(x) return x +@MODEL_REGISTRY.register('bdq') +class BDQ(nn.Module): + + def __init__( + self, + obs_shape: Union[int, SequenceType], + num_branches: int = 0, + action_bins_per_branch: int = 2, + layer_num: int = 3, + a_layer_num: Optional[int] = None, + v_layer_num: Optional[int] = None, + encoder_hidden_size_list: SequenceType = [128, 128, 64], + head_hidden_size: Optional[int] = None, + norm_type: Optional[nn.Module] = None, + activation: Optional[nn.Module] = nn.ReLU(), + ) -> None: + """ + Overview: + Init the BDQ (encoder + head) Model according to input arguments. \ + referenced paper Action Branching Architectures for Deep Reinforcement Learning \ + + Arguments: + - obs_shape (:obj:`Union[int, SequenceType]`): Observation space shape, such as 8 or [4, 84, 84]. + - num_branches (:obj:`int`): The number of branches, which is equivalent to the action dimension, \ + such as 6 in mujoco's halfcheetah environment. + - action_bins_per_branch (:obj:`int`): The number of actions in each dimension. + - layer_num (:obj:`int`): The number of layers used in the network to compute Advantage and Value output. + - a_layer_num (:obj:`int`): The number of layers used in the network to compute Advantage output. + - v_layer_num (:obj:`int`): The number of layers used in the network to compute Value output. + - encoder_hidden_size_list (:obj:`SequenceType`): Collection of ``hidden_size`` to pass to ``Encoder``, \ + the last element must match ``head_hidden_size``. + - head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` of head network. + - norm_type (:obj:`Optional[str]`): The type of normalization in networks, see \ + ``ding.torch_utils.fc_block`` for more details. + - activation (:obj:`Optional[nn.Module]`): The type of activation function in networks \ + if ``None`` then default set it to ``nn.ReLU()`` + """ + super(BDQ, self).__init__() + # For compatibility: 1, (1, ), [4, 32, 32] + obs_shape, num_branches = squeeze(obs_shape), squeeze(num_branches) + if head_hidden_size is None: + head_hidden_size = encoder_hidden_size_list[-1] + + # backbone + # FC Encoder + if isinstance(obs_shape, int) or len(obs_shape) == 1: + self.encoder = FCEncoder(obs_shape, encoder_hidden_size_list, activation=activation, norm_type=norm_type) + # Conv Encoder + elif len(obs_shape) == 3: + self.encoder = ConvEncoder(obs_shape, encoder_hidden_size_list, activation=activation, norm_type=norm_type) + else: + raise RuntimeError( + "not support obs_shape for pre-defined encoder: {}, please customize your own DQN".format(obs_shape) + ) + + self.num_branches = num_branches + self.action_bins_per_branch = action_bins_per_branch + + # head + self.head = BranchingHead( + head_hidden_size, + num_branches=self.num_branches, + action_bins_per_branch=self.action_bins_per_branch, + layer_num=layer_num, + a_layer_num=a_layer_num, + v_layer_num=v_layer_num, + activation=activation, + norm_type=norm_type + ) + + def forward(self, x: torch.Tensor) -> Dict: + """ + Overview: + BDQ forward computation graph, input observation tensor to predict q_value. + Arguments: + - x (:obj:`torch.Tensor`): Observation inputs + Returns: + - outputs (:obj:`Dict`): BDQ forward outputs, such as q_value. + ReturnsKeys: + - logit (:obj:`torch.Tensor`): Discrete Q-value output of each action dimension. + Shapes: + - x (:obj:`torch.Tensor`): :math:`(B, N)`, where B is batch size and N is ``obs_shape`` + - logit (:obj:`torch.FloatTensor`): :math:`(B, M)`, where B is batch size and M is + ``num_branches * action_bins_per_branch`` + Examples: + >>> model = BDQ(8, 5, 2) # arguments: 'obs_shape', 'num_branches' and 'action_bins_per_branch'. + >>> inputs = torch.randn(4, 8) + >>> outputs = model(inputs) + >>> assert isinstance(outputs, dict) and outputs['logit'].shape == torch.Size([4, 5, 2]) + """ + x = self.encoder(x) / (self.num_branches + 1) # corresponds to the "Gradient Rescaling" in the paper + x = self.head(x) + return x + + @MODEL_REGISTRY.register('c51dqn') class C51DQN(nn.Module): + """ + Overview: + The neural network structure and computation graph of C51DQN, which combines distributional RL and DQN. \ + You can refer to https://arxiv.org/pdf/1707.06887.pdf for more details. The C51DQN is composed of \ + ``encoder`` and ``head``. ``encoder`` is used to extract the feature of observation, and ``head`` is \ + used to compute the distribution of Q-value. + Interfaces: + ``__init__``, ``forward`` + + .. note:: + Current C51DQN supports two types of encoder: ``FCEncoder`` and ``ConvEncoder``. + """ def __init__( self, @@ -114,21 +261,27 @@ def __init__( v_max: Optional[float] = 10, n_atom: Optional[int] = 51, ) -> None: - r""" + """ Overview: - Init the C51 Model according to input arguments. + initialize the C51 Model according to corresponding input arguments. Arguments: - - obs_shape (:obj:`Union[int, SequenceType]`): Observation's space. - - action_shape (:obj:`Union[int, SequenceType]`): Action's space. - - encoder_hidden_size_list (:obj:`SequenceType`): Collection of ``hidden_size`` to pass to ``Encoder`` - - head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to ``Head``. - - head_layer_num (:obj:`int`): The num of layers used in the network to compute Q value output - - activation (:obj:`Optional[nn.Module]`): - The type of activation function to use in ``MLP`` the after ``layer_fn``, - if ``None`` then default set to ``nn.ReLU()`` - - norm_type (:obj:`Optional[str]`): - The type of normalization to use, see ``ding.torch_utils.fc_block`` for more details` - - n_atom (:obj:`Optional[int]`): Number of atoms in the prediction distribution. + - obs_shape (:obj:`Union[int, SequenceType]`): Observation space shape, such as 8 or [4, 84, 84]. + - action_shape (:obj:`Union[int, SequenceType]`): Action space shape, such as 6 or [2, 3, 3]. + - encoder_hidden_size_list (:obj:`SequenceType`): Collection of ``hidden_size`` to pass to ``Encoder``, \ + the last element must match ``head_hidden_size``. + - head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` of head network, defaults to None, \ + then it will be set to the last element of ``encoder_hidden_size_list``. + - head_layer_num (:obj:`int`): The number of layers used in the head network to compute Q value output. + - activation (:obj:`Optional[nn.Module]`): The type of activation function in networks \ + if ``None`` then default set it to ``nn.ReLU()``. + - norm_type (:obj:`Optional[str]`): The type of normalization in networks, see \ + ``ding.torch_utils.fc_block`` for more details. you can choose one of ['BN', 'IN', 'SyncBN', 'LN'] + - v_min (:obj:`Optional[float]`): The minimum value of the support of the distribution, which is related \ + to the value (discounted sum of reward) scale of the specific environment. Defaults to -10. + - v_max (:obj:`Optional[float]`): The maximum value of the support of the distribution, which is related \ + to the value (discounted sum of reward) scale of the specific environment. Defaults to 10. + - n_atom (:obj:`Optional[int]`): The number of atoms in the prediction distribution, 51 is the default \ + value in the paper, you can also try other values such as 301. """ super(C51DQN, self).__init__() # For compatibility: 1, (1, ), [4, 32, 32] @@ -172,25 +325,21 @@ def __init__( ) def forward(self, x: torch.Tensor) -> Dict: - r""" + """ Overview: - Use observation tensor to predict C51DQN's output. - Parameter updates with C51DQN's MLPs forward setup. + C51DQN forward computation graph, input observation tensor to predict q_value and its distribution. Arguments: - - x (:obj:`torch.Tensor`): - The encoded embedding tensor w/ ``(B, N=head_hidden_size)``. + - x (:obj:`torch.Tensor`): The input observation tensor data. Returns: - - outputs (:obj:`Dict`): - Run with encoder and head. Return the result prediction dictionary. - + - outputs (:obj:`Dict`): The output of DQN's forward, including q_value, and distribution. ReturnsKeys: - - logit (:obj:`torch.Tensor`): Logit tensor with same size as input ``x``. - - distribution (:obj:`torch.Tensor`): Distribution tensor of size ``(B, N, n_atom)`` + - logit (:obj:`torch.Tensor`): Discrete Q-value output of each possible action dimension. + - distribution (:obj:`torch.Tensor`): Q-Value discretized distribution, i.e., probability of each \ + uniformly spaced atom Q-value, such as dividing [-10, 10] into 51 uniform spaces. Shapes: - x (:obj:`torch.Tensor`): :math:`(B, N)`, where B is batch size and N is head_hidden_size. - - logit (:obj:`torch.FloatTensor`): :math:`(B, M)`, where M is action_shape. - - distribution(:obj:`torch.FloatTensor`): :math:`(B, M, P)`, where P is n_atom. - + - logit (:obj:`torch.Tensor`): :math:`(B, M)`, where M is action_shape. + - distribution(:obj:`torch.Tensor`): :math:`(B, M, P)`, where P is n_atom. Examples: >>> model = C51DQN(128, 64) # arguments: 'obs_shape' and 'action_shape' >>> inputs = torch.randn(4, 128) @@ -200,6 +349,14 @@ def forward(self, x: torch.Tensor) -> Dict: >>> assert outputs['logit'].shape == torch.Size([4, 64]) >>> # default n_atom: int = 51 >>> assert outputs['distribution'].shape == torch.Size([4, 64, 51]) + + .. note:: + For consistency and compatibility, we name all the outputs of the network which are related to action \ + selections as ``logit``. + + .. note:: + For convenience, we recommend that the number of atoms should be odd, so that the middle atom is exactly \ + the value of the Q-value. """ x = self.encoder(x) x = self.head(x) @@ -208,6 +365,14 @@ def forward(self, x: torch.Tensor) -> Dict: @MODEL_REGISTRY.register('qrdqn') class QRDQN(nn.Module): + """ + Overview: + The neural network structure and computation graph of QRDQN, which combines distributional RL and DQN. \ + You can refer to Distributional Reinforcement Learning with Quantile Regression \ + https://arxiv.org/pdf/1710.10044.pdf for more details. + Interfaces: + ``__init__``, ``forward`` + """ def __init__( self, @@ -220,9 +385,9 @@ def __init__( activation: Optional[nn.Module] = nn.ReLU(), norm_type: Optional[str] = None, ) -> None: - r""" + """ Overview: - Init the QRDQN Model according to input arguments. + Initialize the QRDQN Model according to input arguments. Arguments: - obs_shape (:obj:`Union[int, SequenceType]`): Observation's space. - action_shape (:obj:`Union[int, SequenceType]`): Action's space. @@ -274,7 +439,7 @@ def __init__( ) def forward(self, x: torch.Tensor) -> Dict: - r""" + """ Overview: Use observation tensor to predict QRDQN's output. Parameter updates with QRDQN's MLPs forward setup. @@ -284,7 +449,6 @@ def forward(self, x: torch.Tensor) -> Dict: Returns: - outputs (:obj:`Dict`): Run with encoder and head. Return the result prediction dictionary. - ReturnsKeys: - logit (:obj:`torch.Tensor`): Logit tensor with same size as input ``x``. - q (:obj:`torch.Tensor`): Q valye tensor tensor of size ``(B, N, num_quantiles)`` @@ -293,7 +457,6 @@ def forward(self, x: torch.Tensor) -> Dict: - x (:obj:`torch.Tensor`): :math:`(B, N)`, where B is batch size and N is head_hidden_size. - logit (:obj:`torch.FloatTensor`): :math:`(B, M)`, where M is action_shape. - tau (:obj:`torch.Tensor`): :math:`(B, M, 1)` - Examples: >>> model = QRDQN(64, 64) >>> inputs = torch.randn(4, 64) @@ -311,6 +474,14 @@ def forward(self, x: torch.Tensor) -> Dict: @MODEL_REGISTRY.register('iqn') class IQN(nn.Module): + """ + Overview: + The neural network structure and computation graph of IQN, which combines distributional RL and DQN. \ + You can refer to paper Implicit Quantile Networks for Distributional Reinforcement Learning \ + https://arxiv.org/pdf/1806.06923.pdf for more details. + Interfaces: + ``__init__``, ``forward`` + """ def __init__( self, @@ -324,9 +495,9 @@ def __init__( activation: Optional[nn.Module] = nn.ReLU(), norm_type: Optional[str] = None ) -> None: - r""" + """ Overview: - Init the IQN Model according to input arguments. + Initialize the IQN Model according to input arguments. Arguments: - obs_shape (:obj:`Union[int, SequenceType]`): Observation space shape. - action_shape (:obj:`Union[int, SequenceType]`): Action space shape. @@ -381,7 +552,7 @@ def __init__( ) def forward(self, x: torch.Tensor) -> Dict: - r""" + """ Overview: Use encoded embedding tensor to predict IQN's output. Parameter updates with IQN's MLPs forward setup. @@ -391,7 +562,6 @@ def forward(self, x: torch.Tensor) -> Dict: Returns: - outputs (:obj:`Dict`): Run with encoder and head. Return the result prediction dictionary. - ReturnsKeys: - logit (:obj:`torch.Tensor`): Logit tensor with same size as input ``x``. - q (:obj:`torch.Tensor`): Q valye tensor tensor of size ``(num_quantiles, N, B)`` @@ -418,6 +588,14 @@ def forward(self, x: torch.Tensor) -> Dict: @MODEL_REGISTRY.register('fqf') class FQF(nn.Module): + """ + Overview: + The neural network structure and computation graph of FQF, which combines distributional RL and DQN. \ + You can refer to paper Fully Parameterized Quantile Function for Distributional Reinforcement Learning \ + https://arxiv.org/pdf/1911.02140.pdf for more details. + Interface: + ``__init__``, ``forward`` + """ def __init__( self, @@ -431,9 +609,9 @@ def __init__( activation: Optional[nn.Module] = nn.ReLU(), norm_type: Optional[str] = None ) -> None: - r""" + """ Overview: - Init the FQF Model according to input arguments. + Initialize the FQF Model according to input arguments. Arguments: - obs_shape (:obj:`Union[int, SequenceType]`): Observation space shape. - action_shape (:obj:`Union[int, SequenceType]`): Action space shape. @@ -488,7 +666,7 @@ def __init__( ) def forward(self, x: torch.Tensor) -> Dict: - r""" + """ Overview: Use encoded embedding tensor to predict FQF's output. Parameter updates with FQF's MLPs forward setup. @@ -530,10 +708,14 @@ def forward(self, x: torch.Tensor) -> Dict: class RainbowDQN(nn.Module): """ Overview: - RainbowDQN network (C51 + Dueling + Noisy Block) + The neural network structure and computation graph of RainbowDQN, which combines distributional RL and DQN. \ + You can refer to paper Rainbow: Combining Improvements in Deep Reinforcement Learning \ + https://arxiv.org/pdf/1710.02298.pdf for more details. + Interfaces: + ``__init__``, ``forward`` .. note:: - RainbowDQN contains dueling architecture by default + RainbowDQN contains dueling architecture by default. """ def __init__( @@ -607,7 +789,7 @@ def __init__( ) def forward(self, x: torch.Tensor) -> Dict: - r""" + """ Overview: Use observation tensor to predict Rainbow output. Parameter updates with Rainbow's MLPs forward setup. @@ -617,7 +799,6 @@ def forward(self, x: torch.Tensor) -> Dict: Returns: - outputs (:obj:`Dict`): Run ``MLP`` with ``RainbowHead`` setups and return the result prediction dictionary. - ReturnsKeys: - logit (:obj:`torch.Tensor`): Logit tensor with same size as input ``x``. - distribution (:obj:`torch.Tensor`): Distribution tensor of size ``(B, N, n_atom)`` @@ -625,7 +806,6 @@ def forward(self, x: torch.Tensor) -> Dict: - x (:obj:`torch.Tensor`): :math:`(B, N)`, where B is batch size and N is head_hidden_size. - logit (:obj:`torch.FloatTensor`): :math:`(B, M)`, where M is action_shape. - distribution(:obj:`torch.FloatTensor`): :math:`(B, M, P)`, where P is n_atom. - Examples: >>> model = RainbowDQN(64, 64) # arguments: 'obs_shape' and 'action_shape' >>> inputs = torch.randn(4, 64) @@ -641,7 +821,7 @@ def forward(self, x: torch.Tensor) -> Dict: def parallel_wrapper(forward_fn: Callable) -> Callable: - r""" + """ Overview: Process timestep T and batch_size B at the same time, in other words, treat different timestep data as different trajectories in a batch. @@ -680,7 +860,22 @@ def reshape(d): class DRQN(nn.Module): """ Overview: - DQN + RNN = DRQN + The DRQN (Deep Recurrent Q-Network) is a neural network model combining DQN with RNN to handle sequential + data and partially observable environments. It consists of three main components: ``encoder``, ``rnn``, + and ``head``. + - **Encoder**: Extracts features from various observation inputs. + - **RNN**: Processes sequential observations and other data. + - **Head**: Computes Q-values for each action dimension. + + Interfaces: + ``__init__``, ``forward``. + + .. note:: + The current implementation supports: + - Two encoder types: ``FCEncoder`` and ``ConvEncoder``. + - Two head types: ``DiscreteHead`` and ``DuelingHead``. + - Three RNN types: ``normal (LSTM with LayerNorm)``, ``pytorch`` (PyTorch's native LSTM), and ``gru``. + You can extend the model by customizing your own encoder, RNN, or head by inheriting this class. """ def __init__( @@ -696,41 +891,50 @@ def __init__( norm_type: Optional[str] = None, res_link: bool = False ) -> None: - r""" + """ Overview: - Init the DRQN Model according to arguments. + Initialize the DRQN model with specified parameters. Arguments: - - obs_shape (:obj:`Union[int, SequenceType]`): Observation's space. - - action_shape (:obj:`Union[int, SequenceType]`): Action's space. - - encoder_hidden_size_list (:obj:`SequenceType`): Collection of ``hidden_size`` to pass to ``Encoder`` - - head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to ``Head``. - - lstm_type (:obj:`Optional[str]`): Version of rnn cell, now support ['normal', 'pytorch', 'hpc', 'gru'] - - activation (:obj:`Optional[nn.Module]`): - The type of activation function to use in ``MLP`` the after ``layer_fn``, - if ``None`` then default set to ``nn.ReLU()`` - - norm_type (:obj:`Optional[str]`): - The type of normalization to use, see ``ding.torch_utils.fc_block`` for more details` - - res_link (:obj:`bool`): use the residual link or not, default to False + - obs_shape (:obj:`Union[int, SequenceType]`): Shape of the observation space, e.g., 8 or [4, 84, 84]. + - action_shape (:obj:`Union[int, SequenceType]`): Shape of the action space, e.g., 6 or [2, 3, 3]. + - encoder_hidden_size_list (:obj:`SequenceType`): List of hidden sizes for the encoder. The last element \ + must match ``head_hidden_size``. + - dueling (:obj:`Optional[bool]`): Use ``DuelingHead`` if True, otherwise use ``DiscreteHead``. + - head_hidden_size (:obj:`Optional[int]`): Hidden size for the head network. Defaults to the last \ + element of ``encoder_hidden_size_list`` if None. + - head_layer_num (:obj:`int`): Number of layers in the head network to compute Q-value outputs. + - lstm_type (:obj:`Optional[str]`): Type of RNN module. Supported types are ``normal``, ``pytorch``, \ + and ``gru``. + - activation (:obj:`Optional[nn.Module]`): Activation function used in the network. Defaults to \ + ``nn.ReLU()``. + - norm_type (:obj:`Optional[str]`): Normalization type for the networks. Supported types are: \ + ['BN', 'IN', 'SyncBN', 'LN']. See ``ding.torch_utils.fc_block`` for more details. + - res_link (:obj:`bool`): Enables residual connections between single-frame data and sequential data. \ + Defaults to False. """ super(DRQN, self).__init__() - # For compatibility: 1, (1, ), [4, 32, 32] + # Compatibility for obs_shape/action_shape: Handles scalar, tuple, or multi-dimensional inputs. obs_shape, action_shape = squeeze(obs_shape), squeeze(action_shape) if head_hidden_size is None: head_hidden_size = encoder_hidden_size_list[-1] - # FC Encoder + + # Encoder: Determines the encoder type based on the observation shape. if isinstance(obs_shape, int) or len(obs_shape) == 1: + # FC Encoder self.encoder = FCEncoder(obs_shape, encoder_hidden_size_list, activation=activation, norm_type=norm_type) - # Conv Encoder elif len(obs_shape) == 3: + # Conv Encoder self.encoder = ConvEncoder(obs_shape, encoder_hidden_size_list, activation=activation, norm_type=norm_type) else: raise RuntimeError( - "not support obs_shape for pre-defined encoder: {}, please customize your own DRQN".format(obs_shape) + f"Unsupported obs_shape for pre-defined encoder: {obs_shape}. Please customize your own DRQN." ) - # LSTM Type + + # RNN: Initializes the RNN module based on the specified lstm_type. self.rnn = get_lstm(lstm_type, input_size=head_hidden_size, hidden_size=head_hidden_size) self.res_link = res_link - # Head Type + + # Head: Determines the head type (Dueling or Discrete) and its configuration. if dueling: head_cls = DuelingHead else: @@ -751,42 +955,34 @@ def __init__( ) def forward(self, inputs: Dict, inference: bool = False, saved_state_timesteps: Optional[list] = None) -> Dict: - r""" + """ Overview: - Use observation tensor to predict DRQN output. - Parameter updates with DRQN's MLPs forward setup. + Defines the forward pass of the DRQN model. Takes observation and previous RNN states as inputs \ + and predicts Q-values. Arguments: - - inputs (:obj:`Dict`): - - inference: (:obj:'bool'): if inference is True, we unroll the one timestep transition, - if inference is False, we unroll the sequence transitions. - - saved_state_timesteps: (:obj:'Optional[list]'): when inference is False, - we unroll the sequence transitions, then we would save rnn hidden states at timesteps - that are listed in list saved_state_timesteps. - - ArgumentsKeys: - - obs (:obj:`torch.Tensor`): Encoded observation - - prev_state (:obj:`list`): Previous state's tensor of size ``(B, N)`` - + - inputs (:obj:`Dict`): Input data dictionary containing observation and previous RNN state. + - inference (:obj:`bool`): If True, unrolls one timestep (used during evaluation). If False, unrolls \ + the entire sequence (used during training). + - saved_state_timesteps (:obj:`Optional[list]`): When inference is False, specifies the timesteps \ + whose hidden states are saved and returned. + ArgumentsKeys: + - obs (:obj:`torch.Tensor`): Raw observation tensor. + - prev_state (:obj:`list`): Previous RNN state tensor, structure depends on ``lstm_type``. Returns: - - outputs (:obj:`Dict`): - Run ``MLP`` with ``DRQN`` setups and return the result prediction dictionary. - + - outputs (:obj:`Dict`): The output of DRQN's forward, including logit (q_value) and next state. ReturnsKeys: - - logit (:obj:`torch.Tensor`): Logit tensor with same size as input ``obs``. - - next_state (:obj:`list`): Next state's tensor of size ``(B, N)`` + - logit (:obj:`torch.Tensor`): Discrete Q-value output for each action dimension. + - next_state (:obj:`list`): Next RNN state tensor. Shapes: - - obs (:obj:`torch.Tensor`): :math:`(B, N=obs_space)`, where B is batch size. - - prev_state(:obj:`torch.FloatTensor list`): :math:`[(B, N)]` - - logit (:obj:`torch.FloatTensor`): :math:`(B, N)` - - next_state(:obj:`torch.FloatTensor list`): :math:`[(B, N)]` - + - obs (:obj:`torch.Tensor`): :math:`(B, N)` where B is batch size and N is ``obs_shape``. + - logit (:obj:`torch.Tensor`): :math:`(B, M)` where B is batch size and M is ``action_shape``. Examples: - >>> # Init input's Keys: + >>> # Initialize input keys >>> prev_state = [[torch.randn(1, 1, 64) for __ in range(2)] for _ in range(4)] # B=4 >>> obs = torch.randn(4,64) >>> model = DRQN(64, 64) # arguments: 'obs_shape' and 'action_shape' >>> outputs = model({'obs': inputs, 'prev_state': prev_state}, inference=True) - >>> # Check outputs's Keys + >>> # Validate output keys and shapes >>> assert isinstance(outputs, dict) >>> assert outputs['logit'].shape == (4, 64) >>> assert len(outputs['next_state']) == 4 @@ -795,9 +991,9 @@ def forward(self, inputs: Dict, inference: bool = False, saved_state_timesteps: """ x, prev_state = inputs['obs'], inputs['prev_state'] - # for both inference and other cases, the network structure is encoder -> rnn network -> head - # the difference is inference take the data with seq_len=1 (or T = 1) - # NOTE(rjy): in most situations, set inference=True when evaluate and inference=False when training + # Forward pass: Encoder -> RNN -> Head + # in most situations, set inference=True when evaluate and inference=False when training + # Inference mode: Processes one timestep (seq_len=1). if inference: x = self.encoder(x) if self.res_link: @@ -811,6 +1007,7 @@ def forward(self, inputs: Dict, inference: bool = False, saved_state_timesteps: x = self.head(x) x['next_state'] = next_state return x + # Training mode: Processes the entire sequence. else: # In order to better explain why rnn needs saved_state and which states need to be stored, # let's take r2d2 as an example @@ -818,20 +1015,20 @@ def forward(self, inputs: Dict, inference: bool = False, saved_state_timesteps: # 1) data['burnin_nstep_obs'] = data['obs'][:bs + self._nstep] # 2) data['main_obs'] = data['obs'][bs:-self._nstep] # 3) data['target_obs'] = data['obs'][bs + self._nstep:] - # NOTE(rjy): (T, B, N) or (T, B, C, H, W) - assert len(x.shape) in [3, 5], x.shape + assert len(x.shape) in [3, 5], f"Expected shape (T, B, N) or (T, B, C, H, W), got {x.shape}" x = parallel_wrapper(self.encoder)(x) # (T, B, N) if self.res_link: a = x - # NOTE(rjy) lstm_embedding stores all hidden_state + # lstm_embedding stores all hidden_state lstm_embedding = [] # TODO(nyz) how to deal with hidden_size key-value hidden_state_list = [] + if saved_state_timesteps is not None: saved_state = [] - for t in range(x.shape[0]): # T timesteps - # NOTE(rjy) use x[t:t+1] but not x[t] can keep original dimension - output, prev_state = self.rnn(x[t:t + 1], prev_state) # output: (1,B, head_hidden_size) + for t in range(x.shape[0]): # Iterate over timesteps (T). + # use x[t:t+1] but not x[t] can keep the original dimension + output, prev_state = self.rnn(x[t:t + 1], prev_state) # RNN step output: (1, B, hidden_size) if saved_state_timesteps is not None and t + 1 in saved_state_timesteps: saved_state.append(prev_state) lstm_embedding.append(output) @@ -842,7 +1039,7 @@ def forward(self, inputs: Dict, inference: bool = False, saved_state_timesteps: if self.res_link: x = x + a x = parallel_wrapper(self.head)(x) # (T, B, action_shape) - # NOTE(rjy): x['next_state'] is the hidden state of the last timestep inputted to lstm + # x['next_state'] is the hidden state of the last timestep inputted to lstm # the last timestep state including the hidden state (h) and the cell state (c) # shape: {list: B{dict: 2{Tensor:(1, 1, head_hidden_size}}} x['next_state'] = prev_state @@ -851,18 +1048,24 @@ def forward(self, inputs: Dict, inference: bool = False, saved_state_timesteps: x['hidden_state'] = torch.cat(hidden_state_list, dim=0) if saved_state_timesteps is not None: # the selected saved hidden states, including the hidden state (h) and the cell state (c) - # in r2d2, set 'saved_hidden_​​state_timesteps=[self._burnin_step, self._burnin_step + self._nstep]', + # in r2d2, set 'saved_hidden_state_timesteps=[self._burnin_step, self._burnin_step + self._nstep]', # then saved_state will record the hidden_state for main_obs and target_obs to # initialize their lstm (h c) x['saved_state'] = saved_state return x -@MODEL_REGISTRY.register('gtrxl_discrete') -class GTrXLDiscreteHead(nn.Module): +@MODEL_REGISTRY.register('gtrxldqn') +class GTrXLDQN(nn.Module): """ Overview: - Add a discrete head on top of the GTrXL module. + The neural network structure and computation graph of Gated Transformer-XL DQN algorithm, which is the \ + enhanced version of DRQN, using Transformer-XL to improve long-term sequential modelling ability. The \ + GTrXL-DQN is composed of three parts: ``encoder``, ``head`` and ``core``. The ``encoder`` is used to extract \ + the feature from various observation, the ``core`` is used to process the sequential observation and other \ + data, and the ``head`` is used to compute the Q value of each action dimension. + Interfaces: + ``__init__``, ``forward``, ``reset_memory``, ``get_memory`` . """ def __init__( @@ -885,11 +1088,15 @@ def __init__( encoder_hidden_size_list: SequenceType = [128, 128, 256], encoder_norm_type: Optional[str] = None, ) -> None: - r""" + """ Overview: - Init the model according to arguments. + Initialize the GTrXLDQN model accoding to corresponding input arguments. + + .. tip:: + You can refer to GTrXl class in ``ding.torch_utils.network.gtrxl`` for more details about the input \ + arguments. + Arguments: - Refer to GTrXl class in `ding.torch_utils.network.gtrxl` for more details about the input arguments. - obs_shape (:obj:`Union[int, SequenceType]`): Used by Transformer. Observation's space. - action_shape (:obj:Union[int, SequenceType]): Used by Head. Action's space. - head_layer_num (:obj:`int`): Used by Head. Number of layers. @@ -899,20 +1106,20 @@ def __init__( - att_mlp_num (:obj:`int`): Used by Transformer. - att_layer_num (:obj:`int`): Used by Transformer. - memory_len (:obj:`int`): Used by Transformer. - - activation (:obj:`Optional[nn.Module]`): Used by Transformer and Head. if ``None`` then default set to - ``nn.ReLU()``. - - head_norm_type (:obj:`Optional[str]`): Used by Head. The type of normalization to use, see - ``ding.torch_utils.fc_block`` for more details`. + - activation (:obj:`Optional[nn.Module]`): Used by Transformer and Head. if ``None`` then default set to \ + ``nn.ReLU()``. + - head_norm_type (:obj:`Optional[str]`): Used by Head. The type of normalization to use, see \ + ``ding.torch_utils.fc_block`` for more details`. - dropout (:obj:`bool`): Used by Transformer. - gru_gating (:obj:`bool`): Used by Transformer. - gru_bias (:obj:`float`): Used by Transformer. - dueling (:obj:`bool`): Used by Head. Make the head dueling. - - encoder_hidden_size_list(:obj:`SequenceType`): Used by Encoder. The collection of ``hidden_size`` if using - a custom convolutional encoder. - - encoder_norm_type (:obj:`Optional[str]`): Used by Encoder. The type of normalization to use, see + - encoder_hidden_size_list(:obj:`SequenceType`): Used by Encoder. The collection of ``hidden_size`` if \ + using a custom convolutional encoder. + - encoder_norm_type (:obj:`Optional[str]`): Used by Encoder. The type of normalization to use, see \ ``ding.torch_utils.fc_block`` for more details`. """ - super(GTrXLDiscreteHead, self).__init__() + super(GTrXLDQN, self).__init__() self.core = GTrXL( input_dim=obs_shape, head_dim=att_head_dim, @@ -927,6 +1134,7 @@ def __init__( gru_bias=gru_bias, ) + # for vector obs, use Identity Encoder, i.e. pass if isinstance(obs_shape, int) or len(obs_shape) == 1: pass # replace the embedding layer of Transformer with Conv Encoder @@ -962,24 +1170,22 @@ def __init__( ) def forward(self, x: torch.Tensor) -> Dict: - r""" + """ Overview: Let input tensor go through GTrXl and the Head sequentially. Arguments: - x (:obj:`torch.Tensor`): input tensor of shape (seq_len, bs, obs_shape). Returns: - out (:obj:`Dict`): run ``GTrXL`` with ``DiscreteHead`` setups and return the result prediction dictionary. - Necessary Keys: - - logit (:obj:`torch.Tensor`): discrete Q-value output of each action dimension. - Shape is (bs, action_space) - - memory (:obj:`torch.Tensor`): - memory tensor of size ``(bs x layer_num+1 x memory_len x embedding_dim)`` - - transformer_out (:obj:`torch.Tensor`): output tensor of transformer with same size as input ``x``. + ReturnKeys: + - logit (:obj:`torch.Tensor`): discrete Q-value output of each action dimension, shape is (B, action_space). + - memory (:obj:`torch.Tensor`): memory tensor of size ``(bs x layer_num+1 x memory_len x embedding_dim)``. + - transformer_out (:obj:`torch.Tensor`): output tensor of transformer with same size as input ``x``. Examples: >>> # Init input's Keys: >>> obs_dim, seq_len, bs, action_dim = 128, 64, 32, 4 >>> obs = torch.rand(seq_len, bs, obs_dim) - >>> model = GTrXLDiscreteHead(obs_dim, action_dim) + >>> model = GTrXLDQN(obs_dim, action_dim) >>> outputs = model(obs) >>> assert isinstance(outputs, dict) """ @@ -995,27 +1201,23 @@ def forward(self, x: torch.Tensor) -> Dict: out['transformer_out'] = o1['logit'] # output of gtrxl, out['logit'] is final output return out - def reset_memory(self, batch_size: Optional[int] = None, state: Optional[torch.Tensor] = None): - r""" + def reset_memory(self, batch_size: Optional[int] = None, state: Optional[torch.Tensor] = None) -> None: + """ Overview: - Clear or set the memory of GTrXL. - Arguments: - - batch_size (:obj:`Optional[int]`): batch size - - state (:obj:`Optional[torch.Tensor]`): input memory. - Shape is (layer_num, memory_len, bs, embedding_dim). + Clear or reset the memory of GTrXL. + Arguments: + - batch_size (:obj:`Optional[int]`): The number of samples in a training batch. + - state (:obj:`Optional[torch.Tensor]`): The input memory data, whose shape is \ + (layer_num, memory_len, bs, embedding_dim). """ self.core.reset_memory(batch_size, state) def get_memory(self) -> Optional[torch.Tensor]: - r""" + """ Overview: Return the memory of GTrXL. Returns: - - memory: (:obj:`Optional[torch.Tensor]`): output memory or None if memory has not been initialized. - Shape is (layer_num, memory_len, bs, embedding_dim). + - memory: (:obj:`Optional[torch.Tensor]`): output memory or None if memory has not been initialized, \ + whose shape is (layer_num, memory_len, bs, embedding_dim). """ return self.core.get_memory() - - -class GeneralQNetwork(nn.Module): - pass diff --git a/ding/model/template/qac.py b/ding/model/template/qac.py old mode 100644 new mode 100755 index 67cfd9fa2b..6034a4d74c --- a/ding/model/template/qac.py +++ b/ding/model/template/qac.py @@ -9,59 +9,103 @@ FCEncoder, ConvEncoder -@MODEL_REGISTRY.register('qac') -class QAC(nn.Module): - r""" +@MODEL_REGISTRY.register('continuous_qac') +class ContinuousQAC(nn.Module): + """ Overview: - The QAC network, which is used in DDPG/TD3/SAC. + The neural network and computation graph of algorithms related to Q-value Actor-Critic (QAC), such as \ + DDPG/TD3/SAC. This model now supports continuous and hybrid action space. The ContinuousQAC is composed of \ + four parts: ``actor_encoder``, ``critic_encoder``, ``actor_head`` and ``critic_head``. Encoders are used to \ + extract the feature from various observation. Heads are used to predict corresponding Q-value or action logit. \ + In high-dimensional observation space like 2D image, we often use a shared encoder for both ``actor_encoder`` \ + and ``critic_encoder``. In low-dimensional observation space like 1D vector, we often use different encoders. Interfaces: ``__init__``, ``forward``, ``compute_actor``, ``compute_critic`` """ mode = ['compute_actor', 'compute_critic'] def __init__( - self, - obs_shape: Union[int, SequenceType], - action_shape: Union[int, SequenceType, EasyDict], - action_space: str, - twin_critic: bool = False, - actor_head_hidden_size: int = 64, - actor_head_layer_num: int = 1, - critic_head_hidden_size: int = 64, - critic_head_layer_num: int = 1, - activation: Optional[nn.Module] = nn.ReLU(), - norm_type: Optional[str] = None, + self, + obs_shape: Union[int, SequenceType], + action_shape: Union[int, SequenceType, EasyDict], + action_space: str, + twin_critic: bool = False, + actor_head_hidden_size: int = 64, + actor_head_layer_num: int = 1, + critic_head_hidden_size: int = 64, + critic_head_layer_num: int = 1, + activation: Optional[nn.Module] = nn.ReLU(), + norm_type: Optional[str] = None, + encoder_hidden_size_list: Optional[SequenceType] = None, + share_encoder: Optional[bool] = False, ) -> None: """ Overview: - Initailize the QAC Model according to input arguments. + Initailize the ContinuousQAC Model according to input arguments. Arguments: - obs_shape (:obj:`Union[int, SequenceType]`): Observation's shape, such as 128, (156, ). - action_shape (:obj:`Union[int, SequenceType, EasyDict]`): Action's shape, such as 4, (3, ), \ EasyDict({'action_type_shape': 3, 'action_args_shape': 4}). - - action_space (:obj:`str`): The type of action space, \ - including [``regression``, ``reparameterization``, ``hybrid``]. + - action_space (:obj:`str`): The type of action space, including [``regression``, ``reparameterization``, \ + ``hybrid``], ``regression`` is used for DDPG/TD3, ``reparameterization`` is used for SAC and \ + ``hybrid`` for PADDPG. - twin_critic (:obj:`bool`): Whether to use twin critic, one of tricks in TD3. - actor_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to actor head. - - actor_head_layer_num (:obj:`int`): The num of layers used in the network to compute Q value output \ - for actor head. + - actor_head_layer_num (:obj:`int`): The num of layers used in the actor network to compute action. - critic_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to critic head. - - critic_head_layer_num (:obj:`int`): The num of layers used in the network to compute Q value output \ - for critic head. + - critic_head_layer_num (:obj:`int`): The num of layers used in the critic network to compute Q-value. - activation (:obj:`Optional[nn.Module]`): The type of activation function to use in ``MLP`` \ after each FC layer, if ``None`` then default set to ``nn.ReLU()``. - norm_type (:obj:`Optional[str]`): The type of normalization to after network layer (FC, Conv), \ see ``ding.torch_utils.network`` for more details. + - encoder_hidden_size_list (:obj:`SequenceType`): Collection of ``hidden_size`` to pass to ``Encoder``, \ + the last element must match ``head_hidden_size``, this argument is only used in image observation. + - share_encoder (:obj:`Optional[bool]`): Whether to share encoder between actor and critic. """ - super(QAC, self).__init__() + super(ContinuousQAC, self).__init__() obs_shape: int = squeeze(obs_shape) action_shape = squeeze(action_shape) self.action_shape = action_shape self.action_space = action_space - assert self.action_space in ['regression', 'reparameterization', 'hybrid'] + assert self.action_space in ['regression', 'reparameterization', 'hybrid'], self.action_space + + # encoder + self.share_encoder = share_encoder + if np.isscalar(obs_shape) or len(obs_shape) == 1: + assert not self.share_encoder, "Vector observation doesn't need share encoder." + assert encoder_hidden_size_list is None, "Vector obs encoder only uses one layer nn.Linear" + # Because there is already a layer nn.Linear in the head, so we use nn.Identity here to keep + # compatible with the image observation and avoid adding an extra layer nn.Linear. + self.actor_encoder = nn.Identity() + self.critic_encoder = nn.Identity() + encoder_output_size = obs_shape + elif len(obs_shape) == 3: + + def setup_conv_encoder(): + kernel_size = [3 for _ in range(len(encoder_hidden_size_list))] + stride = [2] + [1 for _ in range(len(encoder_hidden_size_list) - 1)] + return ConvEncoder( + obs_shape, + encoder_hidden_size_list, + activation=activation, + norm_type=norm_type, + kernel_size=kernel_size, + stride=stride + ) + + if self.share_encoder: + encoder = setup_conv_encoder() + self.actor_encoder = self.critic_encoder = encoder + else: + self.actor_encoder = setup_conv_encoder() + self.critic_encoder = setup_conv_encoder() + encoder_output_size = self.actor_encoder.output_size + else: + raise RuntimeError("not support observation shape: {}".format(obs_shape)) + # head if self.action_space == 'regression': # DDPG, TD3 - self.actor = nn.Sequential( - nn.Linear(obs_shape, actor_head_hidden_size), activation, + self.actor_head = nn.Sequential( + nn.Linear(encoder_output_size, actor_head_hidden_size), activation, RegressionHead( actor_head_hidden_size, action_shape, @@ -72,8 +116,8 @@ def __init__( ) ) elif self.action_space == 'reparameterization': # SAC - self.actor = nn.Sequential( - nn.Linear(obs_shape, actor_head_hidden_size), activation, + self.actor_head = nn.Sequential( + nn.Linear(encoder_output_size, actor_head_hidden_size), activation, ReparameterizationHead( actor_head_hidden_size, action_shape, @@ -89,7 +133,7 @@ def __init__( action_shape.action_args_shape = squeeze(action_shape.action_args_shape) action_shape.action_type_shape = squeeze(action_shape.action_type_shape) actor_action_args = nn.Sequential( - nn.Linear(obs_shape, actor_head_hidden_size), activation, + nn.Linear(encoder_output_size, actor_head_hidden_size), activation, RegressionHead( actor_head_hidden_size, action_shape.action_args_shape, @@ -100,7 +144,7 @@ def __init__( ) ) actor_action_type = nn.Sequential( - nn.Linear(obs_shape, actor_head_hidden_size), activation, + nn.Linear(encoder_output_size, actor_head_hidden_size), activation, DiscreteHead( actor_head_hidden_size, action_shape.action_type_shape, @@ -109,17 +153,17 @@ def __init__( norm_type=norm_type, ) ) - self.actor = nn.ModuleList([actor_action_type, actor_action_args]) + self.actor_head = nn.ModuleList([actor_action_type, actor_action_args]) self.twin_critic = twin_critic if self.action_space == 'hybrid': - critic_input_size = obs_shape + action_shape.action_type_shape + action_shape.action_args_shape + critic_input_size = encoder_output_size + action_shape.action_type_shape + action_shape.action_args_shape else: - critic_input_size = obs_shape + action_shape + critic_input_size = encoder_output_size + action_shape if self.twin_critic: - self.critic = nn.ModuleList() + self.critic_head = nn.ModuleList() for _ in range(2): - self.critic.append( + self.critic_head.append( nn.Sequential( nn.Linear(critic_input_size, critic_head_hidden_size), activation, RegressionHead( @@ -133,7 +177,7 @@ def __init__( ) ) else: - self.critic = nn.Sequential( + self.critic_head = nn.Sequential( nn.Linear(critic_input_size, critic_head_hidden_size), activation, RegressionHead( critic_head_hidden_size, @@ -145,24 +189,41 @@ def __init__( ) ) + # Convenient for calling some apis (e.g. self.critic.parameters()), + # but may cause misunderstanding when `print(self)` + self.actor = nn.ModuleList([self.actor_encoder, self.actor_head]) + self.critic = nn.ModuleList([self.critic_encoder, self.critic_head]) + def forward(self, inputs: Union[torch.Tensor, Dict[str, torch.Tensor]], mode: str) -> Dict[str, torch.Tensor]: """ Overview: - The unique execution (forward) method of QAC method, and one can indicate different modes to implement \ - different computation graph, including ``compute_actor`` and ``compute_critic`` in QAC. - Mode compute_actor: - Arguments: - - inputs (:obj:`torch.Tensor`): Observation data, defaults to tensor. - Returns: - - output (:obj:`Dict`): Output dict data, including differnet key-values among distinct action_space. - Mode compute_critic: - Arguments: - - inputs (:obj:`Dict`): Input dict data, including obs and action tensor. - Returns: - - output (:obj:`Dict`): Output dict data, including q_value tensor. + QAC forward computation graph, input observation tensor to predict Q-value or action logit. Different \ + ``mode`` will forward with different network modules to get different outputs and save computation. + Arguments: + - inputs (:obj:`Union[torch.Tensor, Dict[str, torch.Tensor]]`): The input data for forward computation \ + graph, for ``compute_actor``, it is the observation tensor, for ``compute_critic``, it is the \ + dict data including obs and action tensor. + - mode (:obj:`str`): The forward mode, all the modes are defined in the beginning of this class. + Returns: + - output (:obj:`Dict[str, torch.Tensor]`): The output dict of QAC forward computation graph, whose \ + key-values vary in different forward modes. + Examples (Actor): + >>> # Regression mode + >>> model = ContinuousQAC(64, 6, 'regression') + >>> obs = torch.randn(4, 64) + >>> actor_outputs = model(obs,'compute_actor') + >>> assert actor_outputs['action'].shape == torch.Size([4, 6]) + >>> # Reparameterization Mode + >>> model = ContinuousQAC(64, 6, 'reparameterization') + >>> obs = torch.randn(4, 64) + >>> actor_outputs = model(obs,'compute_actor') + >>> assert actor_outputs['logit'][0].shape == torch.Size([4, 6]) # mu + >>> actor_outputs['logit'][1].shape == torch.Size([4, 6]) # sigma - .. note:: - For specific examples, one can refer to API doc of ``compute_actor`` and ``compute_critic`` respectively. + Examples (Critic): + >>> inputs = {'obs': torch.randn(4, 8), 'action': torch.randn(4, 1)} + >>> model = ContinuousQAC(obs_shape=(8, ),action_shape=1, action_space='regression') + >>> assert model(inputs, mode='compute_critic')['q_value'].shape == (4, ) # q value """ assert mode in self.mode, "not support forward mode: {}/{}".format(mode, self.mode) return getattr(self, mode)(inputs) @@ -170,26 +231,22 @@ def forward(self, inputs: Union[torch.Tensor, Dict[str, torch.Tensor]], mode: st def compute_actor(self, obs: torch.Tensor) -> Dict[str, Union[torch.Tensor, Dict[str, torch.Tensor]]]: """ Overview: - The forward computation graph of compute_actor mode, uses observation tensor to produce actor output, - such as ``action``, ``logit`` and so on. + QAC forward computation graph for actor part, input observation tensor to predict action or action logit. Arguments: - - obs (:obj:`torch.Tensor`): Observation tensor data, now supports a batch of 1-dim vector data, \ - i.e. ``(B, obs_shape)``. + - x (:obj:`torch.Tensor`): The input observation tensor data. Returns: - - outputs (:obj:`Dict[str, Union[torch.Tensor, Dict[str, torch.Tensor]]]`): Actor output varying \ + - outputs (:obj:`Dict[str, Union[torch.Tensor, Dict[str, torch.Tensor]]]`): Actor output dict varying \ from action_space: ``regression``, ``reparameterization``, ``hybrid``. - - ReturnsKeys (either): - - regression action_space - - action (:obj:`torch.Tensor`): Continuous action with same size as ``action_shape``, usually in DDPG. - - reparameterization action_space - - logit (:obj:`Dict[str, torch.Tensor]`): Reparameterization logit, usually in SAC. - - - mu (:obj:`torch.Tensor`): Mean of parameterization gaussion distribution. - - sigma (:obj:`torch.Tensor`): Standard variation of parameterization gaussion distribution. - - hybrid action_space - - logit (:obj:`torch.Tensor`): Discrete action type logit. - - action_args (:obj:`torch.Tensor`): Continuous action arguments. + ReturnsKeys (regression): + - action (:obj:`torch.Tensor`): Continuous action with same size as ``action_shape``, usually in DDPG/TD3. + ReturnsKeys (reparameterization): + - logit (:obj:`Dict[str, torch.Tensor]`): The predictd reparameterization action logit, usually in SAC. \ + It is a list containing two tensors: ``mu`` and ``sigma``. The former is the mean of the gaussian \ + distribution, the latter is the standard deviation of the gaussian distribution. + ReturnsKeys (hybrid): + - logit (:obj:`torch.Tensor`): The predicted discrete action type logit, it will be the same dimension \ + as ``action_type_shape``, i.e., all the possible discrete action types. + - action_args (:obj:`torch.Tensor`): Continuous action arguments with same size as ``action_args_shape``. Shapes: - obs (:obj:`torch.Tensor`): :math:`(B, N0)`, B is batch size and N0 corresponds to ``obs_shape``. - action (:obj:`torch.Tensor`): :math:`(B, N1)`, B is batch size and N1 corresponds to ``action_shape``. @@ -201,44 +258,44 @@ def compute_actor(self, obs: torch.Tensor) -> Dict[str, Union[torch.Tensor, Dict ``action_shape.action_args_shape``. Examples: >>> # Regression mode - >>> model = QAC(64, 64, 'regression') + >>> model = ContinuousQAC(64, 6, 'regression') >>> obs = torch.randn(4, 64) >>> actor_outputs = model(obs,'compute_actor') - >>> assert actor_outputs['action'].shape == torch.Size([4, 64]) + >>> assert actor_outputs['action'].shape == torch.Size([4, 6]) >>> # Reparameterization Mode - >>> model = QAC(64, 64, 'reparameterization') + >>> model = ContinuousQAC(64, 6, 'reparameterization') >>> obs = torch.randn(4, 64) >>> actor_outputs = model(obs,'compute_actor') - >>> assert actor_outputs['logit'][0].shape == torch.Size([4, 64]) # mu - >>> actor_outputs['logit'][1].shape == torch.Size([4, 64]) # sigma + >>> assert actor_outputs['logit'][0].shape == torch.Size([4, 6]) # mu + >>> actor_outputs['logit'][1].shape == torch.Size([4, 6]) # sigma """ + obs = self.actor_encoder(obs) if self.action_space == 'regression': - x = self.actor(obs) + x = self.actor_head(obs) return {'action': x['pred']} elif self.action_space == 'reparameterization': - x = self.actor(obs) + x = self.actor_head(obs) return {'logit': [x['mu'], x['sigma']]} elif self.action_space == 'hybrid': - logit = self.actor[0](obs) - action_args = self.actor[1](obs) + logit = self.actor_head[0](obs) + action_args = self.actor_head[1](obs) return {'logit': logit['logit'], 'action_args': action_args['pred']} def compute_critic(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: """ Overview: - The forward computation graph of compute_critic mode, uses observation and action tensor to produce critic - output, such as ``q_value``. + QAC forward computation graph for critic part, input observation and action tensor to predict Q-value. Arguments: - - inputs (:obj:`Dict[str, torch.Tensor]`): Dict strcture of input data, including ``obs`` and ``action`` \ - tensor, also contains ``logit`` tensor in hybrid action_space. - Returns: - - outputs (:obj:`Dict[str, torch.Tensor]`): Critic output, such as ``q_value``. - + - inputs (:obj:`Dict[str, torch.Tensor]`): The dict of input data, including ``obs`` and ``action`` \ + tensor, also contains ``logit`` and ``action_args`` tensor in hybrid action_space. ArgumentsKeys: - obs: (:obj:`torch.Tensor`): Observation tensor data, now supports a batch of 1-dim vector data. - action (:obj:`Union[torch.Tensor, Dict]`): Continuous action with same size as ``action_shape``. - logit (:obj:`torch.Tensor`): Discrete action logit, only in hybrid action_space. - action_args (:obj:`torch.Tensor`): Continuous action arguments, only in hybrid action_space. + Returns: + - outputs (:obj:`Dict[str, torch.Tensor]`): The output dict of QAC's forward computation graph for critic, \ + including ``q_value``. ReturnKeys: - q_value (:obj:`torch.Tensor`): Q value tensor with same size as batch size. Shapes: @@ -252,12 +309,12 @@ def compute_critic(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Ten Examples: >>> inputs = {'obs': torch.randn(4, 8), 'action': torch.randn(4, 1)} - >>> model = QAC(obs_shape=(8, ),action_shape=1, action_space='regression') - >>> model(inputs, mode='compute_critic')['q_value'] # q value - ... tensor([0.0773, 0.1639, 0.0917, 0.0370], grad_fn=) + >>> model = ContinuousQAC(obs_shape=(8, ),action_shape=1, action_space='regression') + >>> assert model(inputs, mode='compute_critic')['q_value'].shape == (4, ) # q value """ obs, action = inputs['obs'], inputs['action'] + obs = self.critic_encoder(obs) assert len(obs.shape) == 2 if self.action_space == 'hybrid': action_type_logit = inputs['logit'] @@ -271,80 +328,100 @@ def compute_critic(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Ten action = action.unsqueeze(1) x = torch.cat([obs, action], dim=1) if self.twin_critic: - x = [m(x)['pred'] for m in self.critic] + x = [m(x)['pred'] for m in self.critic_head] else: - x = self.critic(x)['pred'] + x = self.critic_head(x)['pred'] return {'q_value': x} @MODEL_REGISTRY.register('discrete_qac') class DiscreteQAC(nn.Module): - r""" + """ Overview: - The Discrete QAC model, used in DiscreteSAC. + The neural network and computation graph of algorithms related to discrete action Q-value Actor-Critic (QAC), \ + such as DiscreteSAC. This model now supports only discrete action space. The DiscreteQAC is composed of \ + four parts: ``actor_encoder``, ``critic_encoder``, ``actor_head`` and ``critic_head``. Encoders are used to \ + extract the feature from various observation. Heads are used to predict corresponding Q-value or action logit. \ + In high-dimensional observation space like 2D image, we often use a shared encoder for both ``actor_encoder`` \ + and ``critic_encoder``. In low-dimensional observation space like 1D vector, we often use different encoders. Interfaces: ``__init__``, ``forward``, ``compute_actor``, ``compute_critic`` """ mode = ['compute_actor', 'compute_critic'] def __init__( - self, - agent_obs_shape: Union[int, SequenceType], - global_obs_shape: Union[int, SequenceType], - action_shape: Union[int, SequenceType], - encoder_hidden_size_list: SequenceType = [64], - twin_critic: bool = False, - actor_head_hidden_size: int = 64, - actor_head_layer_num: int = 1, - critic_head_hidden_size: int = 64, - critic_head_layer_num: int = 1, - activation: Optional[nn.Module] = nn.ReLU(), - norm_type: Optional[str] = None, + self, + obs_shape: Union[int, SequenceType], + action_shape: Union[int, SequenceType], + twin_critic: bool = False, + actor_head_hidden_size: int = 64, + actor_head_layer_num: int = 1, + critic_head_hidden_size: int = 64, + critic_head_layer_num: int = 1, + activation: Optional[nn.Module] = nn.ReLU(), + norm_type: Optional[str] = None, + encoder_hidden_size_list: SequenceType = None, + share_encoder: Optional[bool] = False, ) -> None: - r""" + """ Overview: - Init the QAC Model according to arguments. + Initailize the DiscreteQAC Model according to input arguments. Arguments: - - obs_shape (:obj:`Union[int, SequenceType]`): Observation's space. - - action_shape (:obj:`Union[int, SequenceType]`): Action's space. - - twin_critic (:obj:`bool`): Whether include twin critic. - - actor_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to actor-nn's ``Head``. - - actor_head_layer_num (:obj:`int`): - The num of layers used in the network to compute Q value output for actor's nn. - - critic_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to critic-nn's ``Head``. - - critic_head_layer_num (:obj:`int`): - The num of layers used in the network to compute Q value output for critic's nn. - - activation (:obj:`Optional[nn.Module]`): - The type of activation function to use in ``MLP`` the after ``layer_fn``, - if ``None`` then default set to ``nn.ReLU()`` - - norm_type (:obj:`Optional[str]`): - The type of normalization to use, see ``ding.torch_utils.fc_block`` for more details. + - obs_shape (:obj:`Union[int, SequenceType]`): Observation's shape, such as 128, (156, ). + - action_shape (:obj:`Union[int, SequenceType, EasyDict]`): Action's shape, such as 4, (3, ). + - twin_critic (:obj:`bool`): Whether to use twin critic. + - actor_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to actor head. + - actor_head_layer_num (:obj:`int`): The num of layers used in the actor network to compute action. + - critic_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to critic head. + - critic_head_layer_num (:obj:`int`): The num of layers used in the critic network to compute Q-value. + - activation (:obj:`Optional[nn.Module]`): The type of activation function to use in ``MLP`` \ + after each FC layer, if ``None`` then default set to ``nn.ReLU()``. + - norm_type (:obj:`Optional[str]`): The type of normalization to after network layer (FC, Conv), \ + see ``ding.torch_utils.network`` for more details. + - encoder_hidden_size_list (:obj:`SequenceType`): Collection of ``hidden_size`` to pass to ``Encoder``, \ + the last element must match ``head_hidden_size``, this argument is only used in image observation. + - share_encoder (:obj:`Optional[bool]`): Whether to share encoder between actor and critic. """ super(DiscreteQAC, self).__init__() - agent_obs_shape: int = squeeze(agent_obs_shape) + obs_shape: int = squeeze(obs_shape) action_shape: int = squeeze(action_shape) + # encoder + self.share_encoder = share_encoder + if np.isscalar(obs_shape) or len(obs_shape) == 1: + assert not self.share_encoder, "Vector observation doesn't need share encoder." + assert encoder_hidden_size_list is None, "Vector obs encoder only uses one layer nn.Linear" + # Because there is already a layer nn.Linear in the head, so we use nn.Identity here to keep + # compatible with the image observation and avoid adding an extra layer nn.Linear. + self.actor_encoder = nn.Identity() + self.critic_encoder = nn.Identity() + encoder_output_size = obs_shape + elif len(obs_shape) == 3: + + def setup_conv_encoder(): + kernel_size = [3 for _ in range(len(encoder_hidden_size_list))] + stride = [2] + [1 for _ in range(len(encoder_hidden_size_list) - 1)] + return ConvEncoder( + obs_shape, + encoder_hidden_size_list, + activation=activation, + norm_type=norm_type, + kernel_size=kernel_size, + stride=stride + ) - if isinstance(agent_obs_shape, int) or len(agent_obs_shape) == 1: - encoder_cls = FCEncoder - elif len(agent_obs_shape) == 3: - encoder_cls = ConvEncoder - else: - raise RuntimeError( - "not support obs_shape for pre-defined encoder: {}, please customize your own DQN". - format(agent_obs_shape) - ) - if isinstance(global_obs_shape, int) or len(global_obs_shape) == 1: - global_encoder_cls = FCEncoder - elif len(global_obs_shape) == 3: - global_encoder_cls = ConvEncoder + if self.share_encoder: + encoder = setup_conv_encoder() + self.actor_encoder = self.critic_encoder = encoder + else: + self.actor_encoder = setup_conv_encoder() + self.critic_encoder = setup_conv_encoder() + encoder_output_size = self.actor_encoder.output_size else: - raise RuntimeError( - "not support obs_shape for pre-defined encoder: {}, please customize your own DQN". - format(global_obs_shape) - ) + raise RuntimeError("not support observation shape: {}".format(obs_shape)) - self.actor = nn.Sequential( - encoder_cls(agent_obs_shape, encoder_hidden_size_list, activation=activation, norm_type=norm_type), + # head + self.actor_head = nn.Sequential( + nn.Linear(encoder_output_size, actor_head_hidden_size), activation, DiscreteHead( actor_head_hidden_size, action_shape, actor_head_layer_num, activation=activation, norm_type=norm_type ) @@ -352,13 +429,11 @@ def __init__( self.twin_critic = twin_critic if self.twin_critic: - self.critic = nn.ModuleList() + self.critic_head = nn.ModuleList() for _ in range(2): - self.critic.append( + self.critic_head.append( nn.Sequential( - global_encoder_cls( - agent_obs_shape, encoder_hidden_size_list, activation=activation, norm_type=norm_type - ), + nn.Linear(encoder_output_size, critic_head_hidden_size), activation, DiscreteHead( critic_head_hidden_size, action_shape, @@ -369,10 +444,8 @@ def __init__( ) ) else: - self.critic = nn.Sequential( - global_encoder_cls( - agent_obs_shape, encoder_hidden_size_list, activation=activation, norm_type=norm_type - ), + self.critic_head = nn.Sequential( + nn.Linear(encoder_output_size, critic_head_hidden_size), activation, DiscreteHead( critic_head_hidden_size, action_shape, @@ -381,136 +454,88 @@ def __init__( norm_type=norm_type ) ) + # Convenient for calling some apis (e.g. self.critic.parameters()), + # but may cause misunderstanding when `print(self)` + self.actor = nn.ModuleList([self.actor_encoder, self.actor_head]) + self.critic = nn.ModuleList([self.critic_encoder, self.critic_head]) - def forward(self, inputs: Union[torch.Tensor, Dict], mode: str) -> Dict: - r""" + def forward(self, inputs: torch.Tensor, mode: str) -> Dict[str, torch.Tensor]: + """ Overview: - Use bbservation and action tensor to predict output. - Parameter updates with QAC's MLPs forward setup. + QAC forward computation graph, input observation tensor to predict Q-value or action logit. Different \ + ``mode`` will forward with different network modules to get different outputs and save computation. Arguments: - Forward with ``'compute_actor'``: - - inputs (:obj:`torch.Tensor`): - The encoded embedding tensor, determined with given ``hidden_size``, i.e. ``(B, N=hidden_size)``. - Whether ``actor_head_hidden_size`` or ``critic_head_hidden_size`` depend on ``mode``. - - Forward with ``'compute_critic'``, inputs (`Dict`) Necessary Keys: - - ``obs``, ``action`` encoded tensors. - - - mode (:obj:`str`): Name of the forward mode. + - inputs (:obj:`torch.Tensor`): The input observation tensor data. + - mode (:obj:`str`): The forward mode, all the modes are defined in the beginning of this class. Returns: - - outputs (:obj:`Dict`): Outputs of network forward. - - Forward with ``'compute_actor'``, Necessary Keys (either): - - action (:obj:`torch.Tensor`): Action tensor with same size as input ``x``. - - logit (:obj:`torch.Tensor`): - Logit tensor encoding ``mu`` and ``sigma``, both with same size as input ``x``. - - Forward with ``'compute_critic'``, Necessary Keys: - - q_value (:obj:`torch.Tensor`): Q value tensor with same size as batch size. - Actor Shapes: - - inputs (:obj:`torch.Tensor`): :math:`(B, N0)`, B is batch size and N0 corresponds to ``hidden_size`` - - action (:obj:`torch.Tensor`): :math:`(B, N0)` - - q_value (:obj:`torch.FloatTensor`): :math:`(B, )`, where B is batch size. - - Critic Shapes: - - obs (:obj:`torch.Tensor`): :math:`(B, N1)`, where B is batch size and N1 is ``obs_shape`` - - action (:obj:`torch.Tensor`): :math:`(B, N2)`, where B is batch size and N2 is``action_shape`` - - logit (:obj:`torch.FloatTensor`): :math:`(B, N2)`, where B is batch size and N3 is ``action_shape`` - - Actor Examples: - >>> # Regression mode - >>> model = QAC(64, 64, 'regression') - >>> inputs = torch.randn(4, 64) - >>> actor_outputs = model(inputs,'compute_actor') - >>> assert actor_outputs['action'].shape == torch.Size([4, 64]) - >>> # Reparameterization Mode - >>> model = QAC(64, 64, 'reparameterization') - >>> inputs = torch.randn(4, 64) - >>> actor_outputs = model(inputs,'compute_actor') - >>> actor_outputs['logit'][0].shape # mu - >>> torch.Size([4, 64]) - >>> actor_outputs['logit'][1].shape # sigma - >>> torch.Size([4, 64]) - - Critic Examples: - >>> inputs = {'obs': torch.randn(4,N), 'action': torch.randn(4,1)} - >>> model = QAC(obs_shape=(N, ), action_shape=1, action_space='regression') - >>> model(inputs, mode='compute_critic')['q_value'] # q value - tensor([0.0773, 0.1639, 0.0917, 0.0370], grad_fn=) + - output (:obj:`Dict[str, torch.Tensor]`): The output dict of QAC forward computation graph, whose \ + key-values vary in different forward modes. + Examples (Actor): + >>> model = DiscreteQAC(64, 6) + >>> obs = torch.randn(4, 64) + >>> actor_outputs = model(obs,'compute_actor') + >>> assert actor_outputs['logit'].shape == torch.Size([4, 6]) + Examples(Critic): + >>> model = DiscreteQAC(64, 6, twin_critic=False) + >>> obs = torch.randn(4, 64) + >>> actor_outputs = model(obs,'compute_critic') + >>> assert actor_outputs['q_value'].shape == torch.Size([4, 6]) """ assert mode in self.mode, "not support forward mode: {}/{}".format(mode, self.mode) return getattr(self, mode)(inputs) - def compute_actor(self, inputs: torch.Tensor) -> Dict: - r""" + def compute_actor(self, inputs: torch.Tensor) -> Dict[str, torch.Tensor]: + """ Overview: - Use encoded embedding tensor to predict output. - Execute parameter updates with ``'compute_actor'`` mode - Use encoded embedding tensor to predict output. + QAC forward computation graph for actor part, input observation tensor to predict action or action logit. Arguments: - - inputs (:obj:`torch.Tensor`): - The encoded embedding tensor, determined with given ``hidden_size``, i.e. ``(B, N=hidden_size)``. - ``hidden_size = actor_head_hidden_size`` - - mode (:obj:`str`): Name of the forward mode. + - inputs (:obj:`torch.Tensor`): The input observation tensor data. Returns: - - outputs (:obj:`Dict`): Outputs of forward pass encoder and head. - - ReturnsKeys (either): - - action (:obj:`torch.Tensor`): Continuous action tensor with same size as ``action_shape``. - - logit (:obj:`torch.Tensor`): - Logit tensor encoding ``mu`` and ``sigma``, both with same size as input ``x``. + - outputs (:obj:`Dict[str, torch.Tensor]`): The output dict of QAC forward computation graph for actor, \ + including discrete action ``logit``. + ReturnsKeys: + - logit (:obj:`torch.Tensor`): The predicted discrete action type logit, it will be the same dimension \ + as ``action_shape``, i.e., all the possible discrete action choices. Shapes: - - inputs (:obj:`torch.Tensor`): :math:`(B, N0)`, B is batch size and N0 corresponds to ``hidden_size`` - - action (:obj:`torch.Tensor`): :math:`(B, N0)` - - logit (:obj:`list`): 2 elements, mu and sigma, each is the shape of :math:`(B, N0)`. - - q_value (:obj:`torch.FloatTensor`): :math:`(B, )`, B is batch size. + - inputs (:obj:`torch.Tensor`): :math:`(B, N0)`, B is batch size and N0 corresponds to ``obs_shape``. + - logit (:obj:`torch.Tensor`): :math:`(B, N2)`, B is batch size and N2 corresponds to \ + ``action_shape``. Examples: - >>> # Regression mode - >>> model = QAC(64, 64, 'regression') - >>> inputs = torch.randn(4, 64) - >>> actor_outputs = model(inputs,'compute_actor') - >>> assert actor_outputs['action'].shape == torch.Size([4, 64]) - >>> # Reparameterization Mode - >>> model = QAC(64, 64, 'reparameterization') - >>> inputs = torch.randn(4, 64) - >>> actor_outputs = model(inputs,'compute_actor') - >>> actor_outputs['logit'][0].shape # mu - >>> torch.Size([4, 64]) - >>> actor_outputs['logit'][1].shape # sigma - >>> torch.Size([4, 64]) + >>> model = DiscreteQAC(64, 6) + >>> obs = torch.randn(4, 64) + >>> actor_outputs = model(obs,'compute_actor') + >>> assert actor_outputs['logit'].shape == torch.Size([4, 6]) """ - x = self.actor(inputs['obs']) + x = self.actor_encoder(inputs) + x = self.actor_head(x) return {'logit': x['logit']} - def compute_critic(self, inputs: Dict) -> Dict: - r""" + def compute_critic(self, inputs: torch.Tensor) -> Dict[str, torch.Tensor]: + """ Overview: - Execute parameter updates with ``'compute_critic'`` mode - Use encoded embedding tensor to predict output. + QAC forward computation graph for critic part, input observation to predict Q-value for each possible \ + discrete action choices. Arguments: - - ``obs``, ``action`` encoded tensors. - - mode (:obj:`str`): Name of the forward mode. + - inputs (:obj:`torch.Tensor`): The input observation tensor data. Returns: - - outputs (:obj:`Dict`): Q-value output. - + - outputs (:obj:`Dict[str, torch.Tensor]`): The output dict of QAC forward computation graph for critic, \ + including ``q_value`` for each possible discrete action choices. ReturnKeys: - - q_value (:obj:`torch.Tensor`): Q value tensor with same size as batch size. + - q_value (:obj:`torch.Tensor`): The predicted Q-value for each possible discrete action choices, it will \ + be the same dimension as ``action_shape`` and used to calculate the loss. Shapes: - - obs (:obj:`torch.Tensor`): :math:`(B, N1)`, where B is batch size and N1 is ``obs_shape`` - - action (:obj:`torch.Tensor`): :math:`(B, N2)`, where B is batch size and N2 is ``action_shape`` - - q_value (:obj:`torch.FloatTensor`): :math:`(B, )`, where B is batch size. - + - obs (:obj:`torch.Tensor`): :math:`(B, N1)`, where B is batch size and N1 is ``obs_shape``. + - q_value (:obj:`torch.Tensor`): :math:`(B, N2)`, where B is batch size and N2 is ``action_shape``. Examples: - >>> inputs = {'obs': torch.randn(4, N), 'action': torch.randn(4, 1)} - >>> model = QAC(obs_shape=(N, ),action_shape=1, action_space='regression') - >>> model(inputs, mode='compute_critic')['q_value'] # q value - tensor([0.0773, 0.1639, 0.0917, 0.0370], grad_fn=) - + >>> model = DiscreteQAC(64, 6, twin_critic=False) + >>> obs = torch.randn(4, 64) + >>> actor_outputs = model(obs,'compute_critic') + >>> assert actor_outputs['q_value'].shape == torch.Size([4, 6]) """ - + inputs = self.critic_encoder(inputs) if self.twin_critic: - x = [m(inputs['obs'])['logit'] for m in self.critic] + x = [m(inputs)['logit'] for m in self.critic_head] else: - x = self.critic(inputs['obs'])['logit'] + x = self.critic_head(inputs)['logit'] return {'q_value': x} diff --git a/ding/model/template/qac_dist.py b/ding/model/template/qac_dist.py index e2ce65f34c..d9390cb06e 100644 --- a/ding/model/template/qac_dist.py +++ b/ding/model/template/qac_dist.py @@ -8,7 +8,7 @@ @MODEL_REGISTRY.register('qac_dist') class QACDIST(nn.Module): - r""" + """ Overview: The QAC model with distributional Q-value. Interfaces: @@ -32,7 +32,7 @@ def __init__( v_max: Optional[float] = 10, n_atom: Optional[int] = 51, ) -> None: - r""" + """ Overview: Init the QAC Distributional Model according to arguments. Arguments: @@ -102,7 +102,7 @@ def __init__( ) def forward(self, inputs: Union[torch.Tensor, Dict], mode: str) -> Dict: - r""" + """ Overview: Use observation and action tensor to predict output. Parameter updates with QACDIST's MLPs forward setup. @@ -166,7 +166,7 @@ def forward(self, inputs: Union[torch.Tensor, Dict], mode: str) -> Dict: return getattr(self, mode)(inputs) def compute_actor(self, inputs: torch.Tensor) -> Dict: - r""" + """ Overview: Use encoded embedding tensor to predict output. Execute parameter updates with ``'compute_actor'`` mode @@ -210,7 +210,7 @@ def compute_actor(self, inputs: torch.Tensor) -> Dict: return {'logit': [x['mu'], x['sigma']]} def compute_critic(self, inputs: Dict) -> Dict: - r""" + """ Overview: Execute parameter updates with ``'compute_critic'`` mode Use encoded embedding tensor to predict output. diff --git a/ding/model/template/qgpo.py b/ding/model/template/qgpo.py new file mode 100644 index 0000000000..dd384846a7 --- /dev/null +++ b/ding/model/template/qgpo.py @@ -0,0 +1,465 @@ +############################################################# +# This QGPO model is a modification implementation from https://github.com/ChenDRAG/CEP-energy-guided-diffusion +############################################################# + +from easydict import EasyDict +import torch +import torch.nn as nn +import torch.nn.functional as F +import copy +from ding.torch_utils import MLP +from ding.torch_utils.diffusion_SDE.dpm_solver_pytorch import DPM_Solver, NoiseScheduleVP +from ding.model.common.encoder import GaussianFourierProjectionTimeEncoder +from ding.torch_utils.network.res_block import TemporalSpatialResBlock + + +def marginal_prob_std(t, device): + """ + Overview: + Compute the mean and standard deviation of $p_{0t}(x(t) | x(0))$. + Arguments: + - t (:obj:`torch.Tensor`): The input time. + - device (:obj:`torch.device`): The device to use. + """ + + t = torch.tensor(t, device=device) + beta_1 = 20.0 + beta_0 = 0.1 + log_mean_coeff = -0.25 * t ** 2 * (beta_1 - beta_0) - 0.5 * t * beta_0 + alpha_t = torch.exp(log_mean_coeff) + std = torch.sqrt(1. - torch.exp(2. * log_mean_coeff)) + return alpha_t, std + + +class TwinQ(nn.Module): + """ + Overview: + Twin Q network for QGPO, which has two Q networks. + Interfaces: + ``__init__``, ``forward``, ``both`` + """ + + def __init__(self, action_dim, state_dim): + """ + Overview: + Initialization of Twin Q. + Arguments: + - action_dim (:obj:`int`): The dimension of action. + - state_dim (:obj:`int`): The dimension of state. + """ + super().__init__() + self.q1 = MLP( + in_channels=state_dim + action_dim, + hidden_channels=256, + out_channels=1, + activation=nn.ReLU(), + layer_num=4, + output_activation=False + ) + self.q2 = MLP( + in_channels=state_dim + action_dim, + hidden_channels=256, + out_channels=1, + activation=nn.ReLU(), + layer_num=4, + output_activation=False + ) + + def both(self, action, condition=None): + """ + Overview: + Return the output of two Q networks. + Arguments: + - action (:obj:`torch.Tensor`): The input action. + - condition (:obj:`torch.Tensor`): The input condition. + """ + as_ = torch.cat([action, condition], -1) if condition is not None else action + return self.q1(as_), self.q2(as_) + + def forward(self, action, condition=None): + """ + Overview: + Return the minimum output of two Q networks. + Arguments: + - action (:obj:`torch.Tensor`): The input action. + - condition (:obj:`torch.Tensor`): The input condition. + """ + return torch.min(*self.both(action, condition)) + + +class GuidanceQt(nn.Module): + """ + Overview: + Energy Guidance Qt network for QGPO. \ + In the origin paper, the energy guidance is trained by CEP method. + Interfaces: + ``__init__``, ``forward`` + """ + + def __init__(self, action_dim, state_dim, time_embed_dim=32): + """ + Overview: + Initialization of Guidance Qt. + Arguments: + - action_dim (:obj:`int`): The dimension of action. + - state_dim (:obj:`int`): The dimension of state. + - time_embed_dim (:obj:`int`): The dimension of time embedding. \ + The time embedding is a Gaussian Fourier Feature tensor. + """ + super().__init__() + self.qt = MLP( + in_channels=action_dim + time_embed_dim + state_dim, + hidden_channels=256, + out_channels=1, + activation=torch.nn.SiLU(), + layer_num=4, + output_activation=False + ) + self.embed = nn.Sequential( + GaussianFourierProjectionTimeEncoder(embed_dim=time_embed_dim), nn.Linear(time_embed_dim, time_embed_dim) + ) + + def forward(self, action, t, condition=None): + """ + Overview: + Return the output of Guidance Qt. + Arguments: + - action (:obj:`torch.Tensor`): The input action. + - t (:obj:`torch.Tensor`): The input time. + - condition (:obj:`torch.Tensor`): The input condition. + """ + embed = self.embed(t) + ats = torch.cat([action, embed, condition], -1) if condition is not None else torch.cat([action, embed], -1) + return self.qt(ats) + + +class QGPOCritic(nn.Module): + """ + Overview: + QGPO critic network. + Interfaces: + ``__init__``, ``forward``, ``calculateQ``, ``calculate_guidance`` + """ + + def __init__(self, device, cfg, action_dim, state_dim) -> None: + """ + Overview: + Initialization of QGPO critic. + Arguments: + - device (:obj:`torch.device`): The device to use. + - cfg (:obj:`EasyDict`): The config dict. + - action_dim (:obj:`int`): The dimension of action. + - state_dim (:obj:`int`): The dimension of state. + """ + + super().__init__() + # is state_dim is 0 means unconditional guidance + assert state_dim > 0 + # only apply to conditional sampling here + self.device = device + self.q0 = TwinQ(action_dim, state_dim).to(self.device) + self.q0_target = copy.deepcopy(self.q0).requires_grad_(False).to(self.device) + self.qt = GuidanceQt(action_dim, state_dim).to(self.device) + + self.alpha = cfg.alpha + self.q_alpha = cfg.q_alpha + + def calculate_guidance(self, a, t, condition=None, guidance_scale=1.0): + """ + Overview: + Calculate the guidance for conditional sampling. + Arguments: + - a (:obj:`torch.Tensor`): The input action. + - t (:obj:`torch.Tensor`): The input time. + - condition (:obj:`torch.Tensor`): The input condition. + - guidance_scale (:obj:`float`): The scale of guidance. + """ + + with torch.enable_grad(): + a.requires_grad_(True) + Q_t = self.qt(a, t, condition) + guidance = guidance_scale * torch.autograd.grad(torch.sum(Q_t), a)[0] + return guidance.detach() + + def forward(self, a, condition=None): + """ + Overview: + Return the output of QGPO critic. + Arguments: + - a (:obj:`torch.Tensor`): The input action. + - condition (:obj:`torch.Tensor`): The input condition. + """ + + return self.q0(a, condition) + + def calculateQ(self, a, condition=None): + """ + Overview: + Return the output of QGPO critic. + Arguments: + - a (:obj:`torch.Tensor`): The input action. + - condition (:obj:`torch.Tensor`): The input condition. + """ + + return self(a, condition) + + +class ScoreNet(nn.Module): + """ + Overview: + Score-based generative model for QGPO. + Interfaces: + ``__init__``, ``forward`` + """ + + def __init__(self, device, input_dim, output_dim, embed_dim=32): + """ + Overview: + Initialization of ScoreNet. + Arguments: + - device (:obj:`torch.device`): The device to use. + - input_dim (:obj:`int`): The dimension of input. + - output_dim (:obj:`int`): The dimension of output. + - embed_dim (:obj:`int`): The dimension of time embedding. \ + The time embedding is a Gaussian Fourier Feature tensor. + """ + + super().__init__() + + # origin score base + self.output_dim = output_dim + self.embed = nn.Sequential( + GaussianFourierProjectionTimeEncoder(embed_dim=embed_dim), nn.Linear(embed_dim, embed_dim) + ) + + self.device = device + self.pre_sort_condition = nn.Sequential(nn.Linear(input_dim - output_dim, 32), torch.nn.SiLU()) + self.sort_t = nn.Sequential( + nn.Linear(64, 128), + torch.nn.SiLU(), + nn.Linear(128, 128), + ) + self.down_block1 = TemporalSpatialResBlock(output_dim, 512) + self.down_block2 = TemporalSpatialResBlock(512, 256) + self.down_block3 = TemporalSpatialResBlock(256, 128) + self.middle1 = TemporalSpatialResBlock(128, 128) + self.up_block3 = TemporalSpatialResBlock(256, 256) + self.up_block2 = TemporalSpatialResBlock(512, 512) + self.last = nn.Linear(1024, output_dim) + + def forward(self, x, t, condition): + """ + Overview: + Return the output of ScoreNet. + Arguments: + - x (:obj:`torch.Tensor`): The input tensor. + - t (:obj:`torch.Tensor`): The input time. + - condition (:obj:`torch.Tensor`): The input condition. + """ + + embed = self.embed(t) + embed = torch.cat([self.pre_sort_condition(condition), embed], dim=-1) + embed = self.sort_t(embed) + d1 = self.down_block1(x, embed) + d2 = self.down_block2(d1, embed) + d3 = self.down_block3(d2, embed) + u3 = self.middle1(d3, embed) + u2 = self.up_block3(torch.cat([d3, u3], dim=-1), embed) + u1 = self.up_block2(torch.cat([d2, u2], dim=-1), embed) + u0 = torch.cat([d1, u1], dim=-1) + h = self.last(u0) + self.h = h + # Normalize output + return h / marginal_prob_std(t, device=self.device)[1][..., None] + + +class QGPO(nn.Module): + """ + Overview: + Model of QGPO algorithm. + Interfaces: + ``__init__``, ``calculateQ``, ``select_actions``, ``sample``, ``score_model_loss_fn``, ``q_loss_fn``, \ + ``qt_loss_fn`` + """ + + def __init__(self, cfg: EasyDict) -> None: + """ + Overview: + Initialization of QGPO. + Arguments: + - cfg (:obj:`EasyDict`): The config dict. + """ + + super(QGPO, self).__init__() + self.device = cfg.device + self.obs_dim = cfg.obs_dim + self.action_dim = cfg.action_dim + + self.noise_schedule = NoiseScheduleVP(schedule='linear') + + self.score_model = ScoreNet( + device=self.device, + input_dim=self.obs_dim + self.action_dim, + output_dim=self.action_dim, + ) + + self.q = QGPOCritic(self.device, cfg.qgpo_critic, action_dim=self.action_dim, state_dim=self.obs_dim) + + def calculateQ(self, s, a): + """ + Overview: + Calculate the Q value. + Arguments: + - s (:obj:`torch.Tensor`): The input state. + - a (:obj:`torch.Tensor`): The input action. + """ + + return self.q(a, s) + + def select_actions(self, states, diffusion_steps=15, guidance_scale=1.0): + """ + Overview: + Select actions for conditional sampling. + Arguments: + - states (:obj:`list`): The input states. + - diffusion_steps (:obj:`int`): The diffusion steps. + - guidance_scale (:obj:`float`): The scale of guidance. + """ + + def forward_dpm_wrapper_fn(x, t): + score = self.score_model(x, t, condition=states) + result = -(score + + self.q.calculate_guidance(x, t, states, guidance_scale=guidance_scale)) * marginal_prob_std( + t, device=self.device + )[1][..., None] + return result + + self.eval() + multiple_input = True + with torch.no_grad(): + states = torch.FloatTensor(states).to(self.device) + if states.dim == 1: + states = states.unsqueeze(0) + multiple_input = False + num_states = states.shape[0] + + init_x = torch.randn(states.shape[0], self.action_dim, device=self.device) + results = DPM_Solver( + forward_dpm_wrapper_fn, self.noise_schedule, predict_x0=True + ).sample( + init_x, steps=diffusion_steps, order=2 + ).cpu().numpy() + + actions = results.reshape(num_states, self.action_dim).copy() # + + out_actions = [actions[i] for i in range(actions.shape[0])] if multiple_input else actions[0] + self.train() + return out_actions + + def sample(self, states, sample_per_state=16, diffusion_steps=15, guidance_scale=1.0): + """ + Overview: + Sample actions for conditional sampling. + Arguments: + - states (:obj:`list`): The input states. + - sample_per_state (:obj:`int`): The number of samples per state. + - diffusion_steps (:obj:`int`): The diffusion steps. + - guidance_scale (:obj:`float`): The scale of guidance. + """ + + def forward_dpm_wrapper_fn(x, t): + score = self.score_model(x, t, condition=states) + result = -(score + self.q.calculate_guidance(x, t, states, guidance_scale=guidance_scale)) \ + * marginal_prob_std(t, device=self.device)[1][..., None] + return result + + self.eval() + num_states = states.shape[0] + with torch.no_grad(): + states = torch.FloatTensor(states).to(self.device) + states = torch.repeat_interleave(states, sample_per_state, dim=0) + + init_x = torch.randn(states.shape[0], self.action_dim, device=self.device) + results = DPM_Solver( + forward_dpm_wrapper_fn, self.noise_schedule, predict_x0=True + ).sample( + init_x, steps=diffusion_steps, order=2 + ).cpu().numpy() + + actions = results[:, :].reshape(num_states, sample_per_state, self.action_dim).copy() + + self.train() + return actions + + def score_model_loss_fn(self, x, s, eps=1e-3): + """ + Overview: + The loss function for training score-based generative models. + Arguments: + model: A PyTorch model instance that represents a \ + time-dependent score-based model. + x: A mini-batch of training data. + eps: A tolerance value for numerical stability. + """ + + random_t = torch.rand(x.shape[0], device=x.device) * (1. - eps) + eps + z = torch.randn_like(x) + alpha_t, std = marginal_prob_std(random_t, device=x.device) + perturbed_x = x * alpha_t[:, None] + z * std[:, None] + score = self.score_model(perturbed_x, random_t, condition=s) + loss = torch.mean(torch.sum((score * std[:, None] + z) ** 2, dim=(1, ))) + return loss + + def q_loss_fn(self, a, s, r, s_, d, fake_a_, discount=0.99): + """ + Overview: + The loss function for training Q function. + Arguments: + - a (:obj:`torch.Tensor`): The input action. + - s (:obj:`torch.Tensor`): The input state. + - r (:obj:`torch.Tensor`): The input reward. + - s\_ (:obj:`torch.Tensor`): The input next state. + - d (:obj:`torch.Tensor`): The input done. + - fake_a (:obj:`torch.Tensor`): The input fake action. + - discount (:obj:`float`): The discount factor. + """ + + with torch.no_grad(): + softmax = nn.Softmax(dim=1) + next_energy = self.q.q0_target(fake_a_, torch.stack([s_] * fake_a_.shape[1], axis=1)).detach().squeeze() + next_v = torch.sum(softmax(self.q.q_alpha * next_energy) * next_energy, dim=-1, keepdim=True) + # Update Q function + targets = r + (1. - d.float()) * discount * next_v.detach() + qs = self.q.q0.both(a, s) + q_loss = sum(F.mse_loss(q, targets) for q in qs) / len(qs) + + return q_loss + + def qt_loss_fn(self, s, fake_a): + """ + Overview: + The loss function for training Guidance Qt. + Arguments: + - s (:obj:`torch.Tensor`): The input state. + - fake_a (:obj:`torch.Tensor`): The input fake action. + """ + + # input many s anction , + energy = self.q.q0_target(fake_a, torch.stack([s] * fake_a.shape[1], axis=1)).detach().squeeze() + + # CEP guidance method, as proposed in the paper + logsoftmax = nn.LogSoftmax(dim=1) + softmax = nn.Softmax(dim=1) + + x0_data_energy = energy * self.q.alpha + random_t = torch.rand((fake_a.shape[0], ), device=self.device) * (1. - 1e-3) + 1e-3 + random_t = torch.stack([random_t] * fake_a.shape[1], dim=1) + z = torch.randn_like(fake_a) + alpha_t, std = marginal_prob_std(random_t, device=self.device) + perturbed_fake_a = fake_a * alpha_t[..., None] + z * std[..., None] + xt_model_energy = self.q.qt(perturbed_fake_a, random_t, torch.stack([s] * fake_a.shape[1], axis=1)).squeeze() + p_label = softmax(x0_data_energy) + + # + qt_loss = -torch.mean(torch.sum(p_label * logsoftmax(xt_model_energy), axis=-1)) + return qt_loss diff --git a/ding/model/template/qmix.py b/ding/model/template/qmix.py index 483b48027c..e9f2475819 100644 --- a/ding/model/template/qmix.py +++ b/ding/model/template/qmix.py @@ -1,30 +1,45 @@ -from typing import Union, List +from functools import reduce +from typing import List, Union + import torch import torch.nn as nn import torch.nn.functional as F -from functools import reduce -from ding.utils import list_split, MODEL_REGISTRY -from ding.torch_utils import fc_block, MLP +from ding.torch_utils import MLP, fc_block +from ding.utils import MODEL_REGISTRY, list_split + +from ..common import ConvEncoder from .q_learning import DRQN class Mixer(nn.Module): """ Overview: - mixer network in QMIX, which mix up the independent q_value of each agent to a total q_value + Mixer network in QMIX, which mix up the independent q_value of each agent to a total q_value. \ + The weights (but not the biases) of the Mixer network are restricted to be non-negative and \ + produced by separate hypernetworks. Each hypernetwork takes the globle state s as input and generates \ + the weights of one layer of the Mixer network. Interface: - __init__, forward + ``__init__``, ``forward``. """ - def __init__(self, agent_num, state_dim, mixing_embed_dim, hypernet_embed=64, activation=nn.ReLU()): + def __init__( + self, + agent_num: int, + state_dim: int, + mixing_embed_dim: int, + hypernet_embed: int = 64, + activation: nn.Module = nn.ReLU() + ): """ Overview: - Initialize mixer network proposed in QMIX. + Initialize mixer network proposed in QMIX according to arguments. Each hypernetwork consists of \ + linear layers, followed by an absolute activation function, to ensure that the Mixer network weights are \ + non-negative. Arguments: - - agent_num (:obj:`int`): the number of agent - - state_dim(:obj:`int`): the dimension of global observation state - - mixing_embed_dim (:obj:`int`): the dimension of mixing state emdedding - - hypernet_embed (:obj:`int`): the dimension of hypernet emdedding, default to 64 + - agent_num (:obj:`int`): The number of agent, such as 8. + - state_dim(:obj:`int`): The dimension of global observation state, such as 16. + - mixing_embed_dim (:obj:`int`): The dimension of mixing state emdedding, such as 128. + - hypernet_embed (:obj:`int`): The dimension of hypernet emdedding, default to 64. - activation (:obj:`nn.Module`): Activation function in network, defaults to nn.ReLU(). """ super(Mixer, self).__init__() @@ -41,7 +56,7 @@ def __init__(self, agent_num, state_dim, mixing_embed_dim, hypernet_embed=64, ac nn.Linear(self.state_dim, hypernet_embed), self.act, nn.Linear(hypernet_embed, self.embed_dim) ) - # State dependent bias for hidden layer + # state dependent bias for hidden layer self.hyper_b_1 = nn.Linear(self.state_dim, self.embed_dim) # V(s) instead of a bias for the last layers @@ -50,16 +65,17 @@ def __init__(self, agent_num, state_dim, mixing_embed_dim, hypernet_embed=64, ac def forward(self, agent_qs, states): """ Overview: - forward computation graph of pymarl mixer network + Forward computation graph of pymarl mixer network. Mix up the input independent q_value of each agent \ + to a total q_value with weights generated by hypernetwork according to global ``states``. Arguments: - - agent_qs (:obj:`torch.FloatTensor`): the independent q_value of each agent - - states (:obj:`torch.FloatTensor`): the emdedding vector of global state + - agent_qs (:obj:`torch.FloatTensor`): The independent q_value of each agent. + - states (:obj:`torch.FloatTensor`): The emdedding vector of global state. Returns: - - q_tot (:obj:`torch.FloatTensor`): the total mixed q_value + - q_tot (:obj:`torch.FloatTensor`): The total mixed q_value. Shapes: - - agent_qs (:obj:`torch.FloatTensor`): :math:`(B, N)`, where B is batch size and N is agent_num - - states (:obj:`torch.FloatTensor`): :math:`(B, M)`, where M is embedding_size - - q_tot (:obj:`torch.FloatTensor`): :math:`(B, )` + - agent_qs (:obj:`torch.FloatTensor`): :math:`(B, N)`, where B is batch size and N is agent_num. + - states (:obj:`torch.FloatTensor`): :math:`(B, M)`, where M is embedding_size. + - q_tot (:obj:`torch.FloatTensor`): :math:`(B, )`. """ bs = agent_qs.shape[:-1] states = states.reshape(-1, self.state_dim) @@ -86,16 +102,19 @@ def forward(self, agent_qs, states): class QMix(nn.Module): """ Overview: - QMIX network + The neural network and computation graph of algorithms related to QMIX(https://arxiv.org/abs/1803.11485). \ + The QMIX is composed of two parts: agent Q network and mixer(optional). The QMIX paper mentions that all \ + agents share local Q network parameters, so only one Q network is initialized here. Then use summation or \ + Mixer network to process the local Q according to the ``mixer`` settings to obtain the global Q. Interface: - __init__, forward, _setup_global_encoder + ``__init__``, ``forward``. """ def __init__( self, agent_num: int, obs_shape: int, - global_obs_shape: int, + global_obs_shape: Union[int, List[int]], action_shape: int, hidden_size_list: list, mixer: bool = True, @@ -105,17 +124,22 @@ def __init__( ) -> None: """ Overview: - Initialize QMIX neural network, i.e. agent Q network and mixer. + Initialize QMIX neural network according to arguments, i.e. agent Q network and mixer. Arguments: - - agent_num (:obj:`int`): the number of agent - - obs_shape (:obj:`int`): the dimension of each agent's observation state - - global_obs_shape (:obj:`int`): the dimension of global observation state - - action_shape (:obj:`int`): the dimension of action shape - - hidden_size_list (:obj:`list`): the list of hidden size - - mixer (:obj:`bool`): use mixer net or not, default to True - - lstm_type (:obj:`str`): use lstm or gru, default to gru - - activation (:obj:`nn.Module`): Activation function in network, defaults to nn.ReLU(). - - dueling (:obj:`bool`): use dueling head or not, default to False. + - agent_num (:obj:`int`): The number of agent, such as 8. + - obs_shape (:obj:`int`): The dimension of each agent's observation state, such as 8 or [4, 84, 84]. + - global_obs_shape (:obj:`int`): The dimension of global observation state, such as 8 or [4, 84, 84]. + - action_shape (:obj:`int`): The dimension of action shape, such as 6 or [2, 3, 3]. + - hidden_size_list (:obj:`list`): The list of hidden size for ``q_network``, \ + the last element must match mixer's ``mixing_embed_dim``. + - mixer (:obj:`bool`): Use mixer net or not, default to True. If it is false, \ + the final local Q is added to obtain the global Q. + - lstm_type (:obj:`str`): The type of RNN module in ``q_network``, now support \ + ['normal', 'pytorch', 'gru'], default to gru. + - activation (:obj:`nn.Module`): The type of activation function to use in ``MLP`` the after \ + ``layer_fn``, if ``None`` then default set to ``nn.ReLU()``. + - dueling (:obj:`bool`): Whether choose ``DuelingHead`` (True) or ``DiscreteHead (False)``, \ + default to False. """ super(QMix, self).__init__() self._act = activation @@ -125,54 +149,90 @@ def __init__( embedding_size = hidden_size_list[-1] self.mixer = mixer if self.mixer: - self._mixer = Mixer(agent_num, global_obs_shape, embedding_size, activation=activation) - self._global_state_encoder = nn.Identity() + global_obs_shape_type = self._get_global_obs_shape_type(global_obs_shape) + + if global_obs_shape_type == "flat": + self._mixer = Mixer(agent_num, global_obs_shape, embedding_size, activation=activation) + self._global_state_encoder = nn.Identity() + elif global_obs_shape_type == "image": + self._mixer = Mixer(agent_num, embedding_size, embedding_size, activation=activation) + self._global_state_encoder = ConvEncoder( + global_obs_shape, hidden_size_list=hidden_size_list, activation=activation, norm_type='BN' + ) + else: + raise ValueError(f"Unsupported global_obs_shape: {global_obs_shape}") + + def _get_global_obs_shape_type(self, global_obs_shape: Union[int, List[int]]) -> str: + """ + Overview: + Determine the type of global observation shape. + Arguments: + - global_obs_shape (:obj:`Union[int, List[int]]`): The global observation state. + Returns: + - obs_shape_type (:obj:`str`): 'flat' for 1D observation or 'image' for 3D observation. + """ + if isinstance(global_obs_shape, int) or (isinstance(global_obs_shape, list) and len(global_obs_shape) == 1): + return "flat" + elif isinstance(global_obs_shape, list) and len(global_obs_shape) == 3: + return "image" + else: + raise ValueError(f"Unsupported global_obs_shape: {global_obs_shape}") def forward(self, data: dict, single_step: bool = True) -> dict: """ Overview: - forward computation graph of qmix network + QMIX forward computation graph, input dict including time series observation and related data to predict \ + total q_value and each agent q_value. Arguments: - - data (:obj:`dict`): input data dict with keys ['obs', 'prev_state', 'action'] - - agent_state (:obj:`torch.Tensor`): each agent local state(obs) - - global_state (:obj:`torch.Tensor`): global state(obs) - - prev_state (:obj:`list`): previous rnn state - - action (:obj:`torch.Tensor` or None): if action is None, use argmax q_value index as action to\ - calculate ``agent_q_act`` - - single_step (:obj:`bool`): whether single_step forward, if so, add timestep dim before forward and\ - remove it after forward + - data (:obj:`dict`): Input data dict with keys ['obs', 'prev_state', 'action']. + - agent_state (:obj:`torch.Tensor`): Time series local observation data of each agents. + - global_state (:obj:`torch.Tensor`): Time series global observation data. + - prev_state (:obj:`list`): Previous rnn state for ``q_network``. + - action (:obj:`torch.Tensor` or None): The actions of each agent given outside the function. \ + If action is None, use argmax q_value index as action to calculate ``agent_q_act``. + - single_step (:obj:`bool`): Whether single_step forward, if so, add timestep dim before forward and\ + remove it after forward. Returns: - - ret (:obj:`dict`): output data dict with keys [``total_q``, ``logit``, ``next_state``] - - total_q (:obj:`torch.Tensor`): total q_value, which is the result of mixer network - - agent_q (:obj:`torch.Tensor`): each agent q_value - - next_state (:obj:`list`): next rnn state + - ret (:obj:`dict`): Output data dict with keys [``total_q``, ``logit``, ``next_state``]. + ReturnsKeys: + - total_q (:obj:`torch.Tensor`): Total q_value, which is the result of mixer network. + - agent_q (:obj:`torch.Tensor`): Each agent q_value. + - next_state (:obj:`list`): Next rnn state for ``q_network``. Shapes: - agent_state (:obj:`torch.Tensor`): :math:`(T, B, A, N)`, where T is timestep, B is batch_size\ - A is agent_num, N is obs_shape - - global_state (:obj:`torch.Tensor`): :math:`(T, B, M)`, where M is global_obs_shape - - prev_state (:obj:`list`): math:`(B, A)`, a list of length B, and each element is a list of length A - - action (:obj:`torch.Tensor`): :math:`(T, B, A)` - - total_q (:obj:`torch.Tensor`): :math:`(T, B)` - - agent_q (:obj:`torch.Tensor`): :math:`(T, B, A, P)`, where P is action_shape - - next_state (:obj:`list`): math:`(B, A)`, a list of length B, and each element is a list of length A + A is agent_num, N is obs_shape. + - global_state (:obj:`torch.Tensor`): :math:`(T, B, M)`, where M is global_obs_shape. + - prev_state (:obj:`list`): math:`(B, A)`, a list of length B, and each element is a list of length A. + - action (:obj:`torch.Tensor`): :math:`(T, B, A)`. + - total_q (:obj:`torch.Tensor`): :math:`(T, B)`. + - agent_q (:obj:`torch.Tensor`): :math:`(T, B, A, P)`, where P is action_shape. + - next_state (:obj:`list`): math:`(B, A)`, a list of length B, and each element is a list of length A. """ agent_state, global_state, prev_state = data['obs']['agent_state'], data['obs']['global_state'], data[ 'prev_state'] action = data.get('action', None) + # If single_step is True, add a new dimension at the front of agent_state + # This is necessary to maintain the expected input shape for the model, + # which requires a time step dimension even when processing a single step. if single_step: - agent_state, global_state = agent_state.unsqueeze(0), global_state.unsqueeze(0) + agent_state = agent_state.unsqueeze(0) + # If single_step is True and global_state has 2 dimensions, add a new dimension at the front of global_state + # This ensures that global_state has the same number of dimensions as agent_state, + # allowing for consistent processing in the forward computation. + if single_step and len(global_state.shape) == 2: + global_state = global_state.unsqueeze(0) T, B, A = agent_state.shape[:3] assert len(prev_state) == B and all( [len(p) == A for p in prev_state] ), '{}-{}-{}-{}'.format([type(p) for p in prev_state], B, A, len(prev_state[0])) prev_state = reduce(lambda x, y: x + y, prev_state) agent_state = agent_state.reshape(T, -1, *agent_state.shape[3:]) - output = self._q_network({'obs': agent_state, 'prev_state': prev_state, 'enable_fast_timestep': True}) + output = self._q_network({'obs': agent_state, 'prev_state': prev_state}) agent_q, next_state = output['logit'], output['next_state'] next_state, _ = list_split(next_state, step=A) agent_q = agent_q.reshape(T, B, A, -1) if action is None: - # For target forward process + # for target forward process if len(data['obs']['action_mask'].shape) == 3: action_mask = data['obs']['action_mask'].unsqueeze(0) else: @@ -182,12 +242,14 @@ def forward(self, data: dict, single_step: bool = True) -> dict: agent_q_act = torch.gather(agent_q, dim=-1, index=action.unsqueeze(-1)) agent_q_act = agent_q_act.squeeze(-1) # T, B, A if self.mixer: - global_state_embedding = self._global_state_encoder(global_state) + global_state_embedding = self._process_global_state(global_state) total_q = self._mixer(agent_q_act, global_state_embedding) else: - total_q = agent_q_act.sum(-1) + total_q = agent_q_act.sum(dim=-1) + if single_step: total_q, agent_q = total_q.squeeze(0), agent_q.squeeze(0) + return { 'total_q': total_q, 'logit': agent_q, @@ -195,14 +257,23 @@ def forward(self, data: dict, single_step: bool = True) -> dict: 'action_mask': data['obs']['action_mask'] } - def _setup_global_encoder(self, global_obs_shape: int, embedding_size: int) -> torch.nn.Module: + def _process_global_state(self, global_state: torch.Tensor) -> torch.Tensor: """ Overview: - Used to encoder global observation. + Process the global state to obtain an embedding. Arguments: - - global_obs_shape (:obj:`int`): the dimension of global observation state - - embedding_size (:obj:`int`): the dimension of state emdedding - Return: - - outputs (:obj:`torch.nn.Module`): Global observation encoding network + - global_state (:obj:`torch.Tensor`): The global state tensor. + + Returns: + - global_state_embedding (:obj:`torch.Tensor`): The processed global state embedding. """ - return MLP(global_obs_shape, embedding_size, embedding_size, 2, activation=self._act) + # If global_state has 5 dimensions, it's likely in the form [batch_size, time_steps, C, H, W] + if global_state.dim() == 5: + # Reshape and apply the global state encoder + batch_time_shape = global_state.shape[:2] # [batch_size, time_steps] + reshaped_state = global_state.view(-1, *global_state.shape[-3:]) # Collapse batch and time dims + encoded_state = self._global_state_encoder(reshaped_state) + return encoded_state.view(*batch_time_shape, -1) # Reshape back to [batch_size, time_steps, embedding_dim] + else: + # For lower-dimensional states, apply the encoder directly + return self._global_state_encoder(global_state) diff --git a/ding/model/template/qtran.py b/ding/model/template/qtran.py index 6e627f1d15..9114245f13 100644 --- a/ding/model/template/qtran.py +++ b/ding/model/template/qtran.py @@ -100,7 +100,7 @@ def forward(self, data: dict, single_step: bool = True) -> dict: ), '{}-{}-{}-{}'.format([type(p) for p in prev_state], B, A, len(prev_state[0])) prev_state = reduce(lambda x, y: x + y, prev_state) agent_state = agent_state.reshape(T, -1, *agent_state.shape[3:]) - output = self._q_network({'obs': agent_state, 'prev_state': prev_state, 'enable_fast_timestep': True}) + output = self._q_network({'obs': agent_state, 'prev_state': prev_state}) agent_q, next_state = output['logit'], output['next_state'] next_state, _ = list_split(next_state, step=A) agent_q = agent_q.reshape(T, B, A, -1) diff --git a/ding/model/template/qvac.py b/ding/model/template/qvac.py new file mode 100644 index 0000000000..c805051c78 --- /dev/null +++ b/ding/model/template/qvac.py @@ -0,0 +1,361 @@ +from typing import Union, Dict, Optional +from easydict import EasyDict +import numpy as np +import torch +import torch.nn as nn + +from ding.utils import SequenceType, squeeze, MODEL_REGISTRY +from ..common import RegressionHead, ReparameterizationHead, DiscreteHead, MultiHead, \ + FCEncoder, ConvEncoder + + +@MODEL_REGISTRY.register('continuous_qvac') +class ContinuousQVAC(nn.Module): + """ + Overview: + The neural network and computation graph of algorithms related to Actor-Critic that have both Q-value and \ + V-value critic, such as IQL. This model now supports continuous and hybrid action space. The ContinuousQVAC is \ + composed of four parts: ``actor_encoder``, ``critic_encoder``, ``actor_head`` and ``critic_head``. Encoders \ + are used to extract the feature. Heads are used to predict corresponding value or action logit. + In high-dimensional observation space like 2D image, we often use a shared encoder for both ``actor_encoder`` \ + and ``critic_encoder``. In low-dimensional observation space like 1D vector, we often use different encoders. + Interfaces: + ``__init__``, ``forward``, ``compute_actor``, ``compute_critic`` + """ + mode = ['compute_actor', 'compute_critic'] + + def __init__( + self, + obs_shape: Union[int, SequenceType], + action_shape: Union[int, SequenceType, EasyDict], + action_space: str, + twin_critic: bool = False, + actor_head_hidden_size: int = 64, + actor_head_layer_num: int = 1, + critic_head_hidden_size: int = 64, + critic_head_layer_num: int = 1, + activation: Optional[nn.Module] = nn.SiLU(), + norm_type: Optional[str] = None, + encoder_hidden_size_list: Optional[SequenceType] = None, + share_encoder: Optional[bool] = False, + ) -> None: + """ + Overview: + Initailize the ContinuousQVAC Model according to input arguments. + Arguments: + - obs_shape (:obj:`Union[int, SequenceType]`): Observation's shape, such as 128, (156, ). + - action_shape (:obj:`Union[int, SequenceType, EasyDict]`): Action's shape, such as 4, (3, ), \ + EasyDict({'action_type_shape': 3, 'action_args_shape': 4}). + - action_space (:obj:`str`): The type of action space, including [``regression``, ``reparameterization``, \ + ``hybrid``], ``regression`` is used for DDPG/TD3, ``reparameterization`` is used for SAC and \ + ``hybrid`` for PADDPG. + - twin_critic (:obj:`bool`): Whether to use twin critic, one of tricks in TD3. + - actor_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to actor head. + - actor_head_layer_num (:obj:`int`): The num of layers used in the actor network to compute action. + - critic_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to critic head. + - critic_head_layer_num (:obj:`int`): The num of layers used in the critic network to compute Q-value. + - activation (:obj:`Optional[nn.Module]`): The type of activation function to use in ``MLP`` \ + after each FC layer, if ``None`` then default set to ``nn.ReLU()``. + - norm_type (:obj:`Optional[str]`): The type of normalization to after network layer (FC, Conv), \ + see ``ding.torch_utils.network`` for more details. + - encoder_hidden_size_list (:obj:`SequenceType`): Collection of ``hidden_size`` to pass to ``Encoder``, \ + the last element must match ``head_hidden_size``, this argument is only used in image observation. + - share_encoder (:obj:`Optional[bool]`): Whether to share encoder between actor and critic. + """ + super(ContinuousQVAC, self).__init__() + obs_shape: int = squeeze(obs_shape) + action_shape = squeeze(action_shape) + self.action_shape = action_shape + self.action_space = action_space + assert self.action_space in ['regression', 'reparameterization', 'hybrid'], self.action_space + + # encoder + self.share_encoder = share_encoder + if np.isscalar(obs_shape) or len(obs_shape) == 1: + assert not self.share_encoder, "Vector observation doesn't need share encoder." + assert encoder_hidden_size_list is None, "Vector obs encoder only uses one layer nn.Linear" + # Because there is already a layer nn.Linear in the head, so we use nn.Identity here to keep + # compatible with the image observation and avoid adding an extra layer nn.Linear. + self.actor_encoder = nn.Identity() + self.critic_encoder = nn.Identity() + encoder_output_size = obs_shape + elif len(obs_shape) == 3: + + def setup_conv_encoder(): + kernel_size = [3 for _ in range(len(encoder_hidden_size_list))] + stride = [2] + [1 for _ in range(len(encoder_hidden_size_list) - 1)] + return ConvEncoder( + obs_shape, + encoder_hidden_size_list, + activation=activation, + norm_type=norm_type, + kernel_size=kernel_size, + stride=stride + ) + + if self.share_encoder: + encoder = setup_conv_encoder() + self.actor_encoder = self.critic_encoder = encoder + else: + self.actor_encoder = setup_conv_encoder() + self.critic_encoder = setup_conv_encoder() + encoder_output_size = self.actor_encoder.output_size + else: + raise RuntimeError("not support observation shape: {}".format(obs_shape)) + # head + if self.action_space == 'regression': # DDPG, TD3 + self.actor_head = nn.Sequential( + nn.Linear(encoder_output_size, actor_head_hidden_size), activation, + RegressionHead( + actor_head_hidden_size, + action_shape, + actor_head_layer_num, + final_tanh=True, + activation=activation, + norm_type=norm_type + ) + ) + elif self.action_space == 'reparameterization': # SAC + self.actor_head = nn.Sequential( + nn.Linear(encoder_output_size, actor_head_hidden_size), activation, + ReparameterizationHead( + actor_head_hidden_size, + action_shape, + actor_head_layer_num, + sigma_type='conditioned', + activation=activation, + norm_type=norm_type + ) + ) + elif self.action_space == 'hybrid': # PADDPG + # hybrid action space: action_type(discrete) + action_args(continuous), + # such as {'action_type_shape': torch.LongTensor([0]), 'action_args_shape': torch.FloatTensor([0.1, -0.27])} + action_shape.action_args_shape = squeeze(action_shape.action_args_shape) + action_shape.action_type_shape = squeeze(action_shape.action_type_shape) + actor_action_args = nn.Sequential( + nn.Linear(encoder_output_size, actor_head_hidden_size), activation, + RegressionHead( + actor_head_hidden_size, + action_shape.action_args_shape, + actor_head_layer_num, + final_tanh=True, + activation=activation, + norm_type=norm_type + ) + ) + actor_action_type = nn.Sequential( + nn.Linear(encoder_output_size, actor_head_hidden_size), activation, + DiscreteHead( + actor_head_hidden_size, + action_shape.action_type_shape, + actor_head_layer_num, + activation=activation, + norm_type=norm_type, + ) + ) + self.actor_head = nn.ModuleList([actor_action_type, actor_action_args]) + + self.twin_critic = twin_critic + if self.action_space == 'hybrid': + critic_q_input_size = encoder_output_size + action_shape.action_type_shape + action_shape.action_args_shape + critic_v_input_size = encoder_output_size + else: + critic_q_input_size = encoder_output_size + action_shape + critic_v_input_size = encoder_output_size + if self.twin_critic: + self.critic_q_head = nn.ModuleList() + self.critic_v_head = nn.ModuleList() + for _ in range(2): + self.critic_q_head.append( + nn.Sequential( + nn.Linear(critic_q_input_size, critic_head_hidden_size), activation, + RegressionHead( + critic_head_hidden_size, + 1, + critic_head_layer_num, + final_tanh=False, + activation=activation, + norm_type=norm_type + ) + ) + ) + self.critic_v_head = nn.Sequential( + nn.Linear(critic_v_input_size, critic_head_hidden_size), activation, + RegressionHead( + critic_head_hidden_size, + 1, + critic_head_layer_num, + final_tanh=False, + activation=activation, + norm_type=norm_type + ) + ) + else: + self.critic_q_head = nn.Sequential( + nn.Linear(critic_q_input_size, critic_head_hidden_size), activation, + RegressionHead( + critic_head_hidden_size, + 1, + critic_head_layer_num, + final_tanh=False, + activation=activation, + norm_type=norm_type + ) + ) + self.critic_v_head = nn.Sequential( + nn.Linear(critic_v_input_size, critic_head_hidden_size), activation, + RegressionHead( + critic_head_hidden_size, + 1, + critic_head_layer_num, + final_tanh=False, + activation=activation, + norm_type=norm_type + ) + ) + + # Convenient for calling some apis (e.g. self.critic.parameters()), + # but may cause misunderstanding when `print(self)` + self.actor = nn.ModuleList([self.actor_encoder, self.actor_head]) + self.critic = nn.ModuleList([self.critic_encoder, self.critic_q_head, self.critic_v_head]) + + def forward(self, inputs: Union[torch.Tensor, Dict[str, torch.Tensor]], mode: str) -> Dict[str, torch.Tensor]: + """ + Overview: + QVAC forward computation graph, input observation tensor to predict Q-value or action logit. Different \ + ``mode`` will forward with different network modules to get different outputs and save computation. + Arguments: + - inputs (:obj:`Union[torch.Tensor, Dict[str, torch.Tensor]]`): The input data for forward computation \ + graph, for ``compute_actor``, it is the observation tensor, for ``compute_critic``, it is the \ + dict data including obs and action tensor. + - mode (:obj:`str`): The forward mode, all the modes are defined in the beginning of this class. + Returns: + - output (:obj:`Dict[str, torch.Tensor]`): The output dict of QVAC forward computation graph, whose \ + key-values vary in different forward modes. + Examples (Actor): + >>> # Regression mode + >>> model = ContinuousQVAC(64, 6, 'regression') + >>> obs = torch.randn(4, 64) + >>> actor_outputs = model(obs,'compute_actor') + >>> assert actor_outputs['action'].shape == torch.Size([4, 6]) + >>> # Reparameterization Mode + >>> model = ContinuousQVAC(64, 6, 'reparameterization') + >>> obs = torch.randn(4, 64) + >>> actor_outputs = model(obs,'compute_actor') + >>> assert actor_outputs['logit'][0].shape == torch.Size([4, 6]) # mu + >>> actor_outputs['logit'][1].shape == torch.Size([4, 6]) # sigma + + Examples (Critic): + >>> inputs = {'obs': torch.randn(4, 8), 'action': torch.randn(4, 1)} + >>> model = ContinuousQVAC(obs_shape=(8, ),action_shape=1, action_space='regression') + >>> assert model(inputs, mode='compute_critic')['q_value'].shape == (4, ) # q value + """ + assert mode in self.mode, "not support forward mode: {}/{}".format(mode, self.mode) + return getattr(self, mode)(inputs) + + def compute_actor(self, obs: torch.Tensor) -> Dict[str, Union[torch.Tensor, Dict[str, torch.Tensor]]]: + """ + Overview: + QVAC forward computation graph for actor part, input observation tensor to predict action or action logit. + Arguments: + - x (:obj:`torch.Tensor`): The input observation tensor data. + Returns: + - outputs (:obj:`Dict[str, Union[torch.Tensor, Dict[str, torch.Tensor]]]`): Actor output dict varying \ + from action_space: ``regression``, ``reparameterization``, ``hybrid``. + ReturnsKeys (regression): + - action (:obj:`torch.Tensor`): Continuous action with same size as ``action_shape``, usually in DDPG/TD3. + ReturnsKeys (reparameterization): + - logit (:obj:`Dict[str, torch.Tensor]`): The predictd reparameterization action logit, usually in SAC. \ + It is a list containing two tensors: ``mu`` and ``sigma``. The former is the mean of the gaussian \ + distribution, the latter is the standard deviation of the gaussian distribution. + ReturnsKeys (hybrid): + - logit (:obj:`torch.Tensor`): The predicted discrete action type logit, it will be the same dimension \ + as ``action_type_shape``, i.e., all the possible discrete action types. + - action_args (:obj:`torch.Tensor`): Continuous action arguments with same size as ``action_args_shape``. + Shapes: + - obs (:obj:`torch.Tensor`): :math:`(B, N0)`, B is batch size and N0 corresponds to ``obs_shape``. + - action (:obj:`torch.Tensor`): :math:`(B, N1)`, B is batch size and N1 corresponds to ``action_shape``. + - logit.mu (:obj:`torch.Tensor`): :math:`(B, N1)`, B is batch size and N1 corresponds to ``action_shape``. + - logit.sigma (:obj:`torch.Tensor`): :math:`(B, N1)`, B is batch size. + - logit (:obj:`torch.Tensor`): :math:`(B, N2)`, B is batch size and N2 corresponds to \ + ``action_shape.action_type_shape``. + - action_args (:obj:`torch.Tensor`): :math:`(B, N3)`, B is batch size and N3 corresponds to \ + ``action_shape.action_args_shape``. + Examples: + >>> # Regression mode + >>> model = ContinuousQVAC(64, 6, 'regression') + >>> obs = torch.randn(4, 64) + >>> actor_outputs = model(obs,'compute_actor') + >>> assert actor_outputs['action'].shape == torch.Size([4, 6]) + >>> # Reparameterization Mode + >>> model = ContinuousQVAC(64, 6, 'reparameterization') + >>> obs = torch.randn(4, 64) + >>> actor_outputs = model(obs,'compute_actor') + >>> assert actor_outputs['logit'][0].shape == torch.Size([4, 6]) # mu + >>> actor_outputs['logit'][1].shape == torch.Size([4, 6]) # sigma + """ + obs = self.actor_encoder(obs) + if self.action_space == 'regression': + x = self.actor_head(obs) + return {'action': x['pred']} + elif self.action_space == 'reparameterization': + x = self.actor_head(obs) + return {'logit': [x['mu'], x['sigma']]} + elif self.action_space == 'hybrid': + logit = self.actor_head[0](obs) + action_args = self.actor_head[1](obs) + return {'logit': logit['logit'], 'action_args': action_args['pred']} + + def compute_critic(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: + """ + Overview: + QVAC forward computation graph for critic part, input observation and action tensor to predict Q-value. + Arguments: + - inputs (:obj:`Dict[str, torch.Tensor]`): The dict of input data, including ``obs`` and ``action`` \ + tensor, also contains ``logit`` and ``action_args`` tensor in hybrid action_space. + ArgumentsKeys: + - obs: (:obj:`torch.Tensor`): Observation tensor data, now supports a batch of 1-dim vector data. + - action (:obj:`Union[torch.Tensor, Dict]`): Continuous action with same size as ``action_shape``. + - logit (:obj:`torch.Tensor`): Discrete action logit, only in hybrid action_space. + - action_args (:obj:`torch.Tensor`): Continuous action arguments, only in hybrid action_space. + Returns: + - outputs (:obj:`Dict[str, torch.Tensor]`): The output of QVAC's forward computation graph for critic, \ + including ``q_value``. + ReturnKeys: + - q_value (:obj:`torch.Tensor`): Q value tensor with same size as batch size. + Shapes: + - obs (:obj:`torch.Tensor`): :math:`(B, N1)`, where B is batch size and N1 is ``obs_shape``. + - logit (:obj:`torch.Tensor`): :math:`(B, N2)`, B is batch size and N2 corresponds to \ + ``action_shape.action_type_shape``. + - action_args (:obj:`torch.Tensor`): :math:`(B, N3)`, B is batch size and N3 corresponds to \ + ``action_shape.action_args_shape``. + - action (:obj:`torch.Tensor`): :math:`(B, N4)`, where B is batch size and N4 is ``action_shape``. + - q_value (:obj:`torch.Tensor`): :math:`(B, )`, where B is batch size. + + Examples: + >>> inputs = {'obs': torch.randn(4, 8), 'action': torch.randn(4, 1)} + >>> model = ContinuousQVAC(obs_shape=(8, ),action_shape=1, action_space='regression') + >>> assert model(inputs, mode='compute_critic')['q_value'].shape == (4, ) # q value + """ + + obs, action = inputs['obs'], inputs['action'] + obs = self.critic_encoder(obs) + assert len(obs.shape) == 2 + if self.action_space == 'hybrid': + action_type_logit = inputs['logit'] + action_type_logit = torch.softmax(action_type_logit, dim=-1) + action_args = action['action_args'] + if len(action_args.shape) == 1: + action_args = action_args.unsqueeze(1) + x = torch.cat([obs, action_type_logit, action_args], dim=1) + else: + if len(action.shape) == 1: # (B, ) -> (B, 1) + action = action.unsqueeze(1) + x = torch.cat([obs, action], dim=1) + if self.twin_critic: + x = [m(x)['pred'] for m in self.critic_q_head] + y = self.critic_v_head(obs)['pred'] + else: + x = self.critic_q_head(x)['pred'] + y = self.critic_v_head(obs)['pred'] + return {'q_value': x, 'v_value': y} diff --git a/ding/model/template/tests/test_acer.py b/ding/model/template/tests/test_acer.py new file mode 100644 index 0000000000..1c3877335a --- /dev/null +++ b/ding/model/template/tests/test_acer.py @@ -0,0 +1,41 @@ +import torch +import pytest +from itertools import product + +from ding.model.template import ACER +from ding.torch_utils import is_differentiable + +B = 4 +obs_shape = [4, (8, ), (4, 64, 64)] +act_shape = [3, (6, )] +args = list(product(*[obs_shape, act_shape])) + + +@pytest.mark.unittest +class TestACER: + + @pytest.mark.parametrize('obs_shape, act_shape', args) + def test_ACER(self, obs_shape, act_shape): + if isinstance(obs_shape, int): + inputs = torch.randn(B, obs_shape) + else: + inputs = torch.randn(B, *obs_shape) + model = ACER(obs_shape, act_shape) + + outputs_c = model(inputs, mode='compute_critic') + assert isinstance(outputs_c, dict) + if isinstance(act_shape, int): + assert outputs_c['q_value'].shape == (B, act_shape) + elif len(act_shape) == 1: + assert outputs_c['q_value'].shape == (B, *act_shape) + + outputs_a = model(inputs, mode='compute_actor') + assert isinstance(outputs_a, dict) + if isinstance(act_shape, int): + assert outputs_a['logit'].shape == (B, act_shape) + elif len(act_shape) == 1: + assert outputs_a['logit'].shape == (B, *act_shape) + + outputs = {**outputs_a, **outputs_c} + loss = sum([v.sum() for v in outputs.values()]) + is_differentiable(loss, model) diff --git a/ding/model/template/tests/test_bcq.py b/ding/model/template/tests/test_bcq.py new file mode 100644 index 0000000000..101cfd9b9c --- /dev/null +++ b/ding/model/template/tests/test_bcq.py @@ -0,0 +1,75 @@ +import pytest +from itertools import product +import torch +from ding.model.template import BCQ +from ding.torch_utils import is_differentiable + +B = 4 +obs_shape = [4, (8, )] +act_shape = [3, (6, )] +args = list(product(*[obs_shape, act_shape])) + + +@pytest.mark.unittest +class TestBCQ: + + def output_check(self, model, outputs): + if isinstance(outputs, torch.Tensor): + loss = outputs.sum() + elif isinstance(outputs, dict): + loss = sum([v.sum() for v in outputs.values()]) + is_differentiable(loss, model) + + @pytest.mark.parametrize('obs_shape, act_shape', args) + def test_BCQ(self, obs_shape, act_shape): + if isinstance(obs_shape, int): + inputs_obs = torch.randn(B, obs_shape) + else: + inputs_obs = torch.randn(B, *obs_shape) + if isinstance(act_shape, int): + inputs_act = torch.randn(B, act_shape) + else: + inputs_act = torch.randn(B, *act_shape) + inputs = {'obs': inputs_obs, 'action': inputs_act} + model = BCQ(obs_shape, act_shape) + + outputs_c = model(inputs, mode='compute_critic') + assert isinstance(outputs_c, dict) + if isinstance(act_shape, int): + assert torch.stack(outputs_c['q_value']).shape == (2, B) + else: + assert torch.stack(outputs_c['q_value']).shape == (2, B) + self.output_check(model.critic, torch.stack(outputs_c['q_value'])) + + outputs_a = model(inputs, mode='compute_actor') + assert isinstance(outputs_a, dict) + if isinstance(act_shape, int): + assert outputs_a['action'].shape == (B, act_shape) + elif len(act_shape) == 1: + assert outputs_a['action'].shape == (B, *act_shape) + self.output_check(model.actor, outputs_a) + + outputs_vae = model(inputs, mode='compute_vae') + assert isinstance(outputs_vae, dict) + if isinstance(act_shape, int): + assert outputs_vae['recons_action'].shape == (B, act_shape) + assert outputs_vae['mu'].shape == (B, act_shape * 2) + assert outputs_vae['log_var'].shape == (B, act_shape * 2) + assert outputs_vae['z'].shape == (B, act_shape * 2) + elif len(act_shape) == 1: + assert outputs_vae['recons_action'].shape == (B, *act_shape) + assert outputs_vae['mu'].shape == (B, act_shape[0] * 2) + assert outputs_vae['log_var'].shape == (B, act_shape[0] * 2) + assert outputs_vae['z'].shape == (B, act_shape[0] * 2) + if isinstance(obs_shape, int): + assert outputs_vae['prediction_residual'].shape == (B, obs_shape) + else: + assert outputs_vae['prediction_residual'].shape == (B, *obs_shape) + + outputs_eval = model(inputs, mode='compute_eval') + assert isinstance(outputs_eval, dict) + assert isinstance(outputs_eval, dict) + if isinstance(act_shape, int): + assert outputs_eval['action'].shape == (B, act_shape) + elif len(act_shape) == 1: + assert outputs_eval['action'].shape == (B, *act_shape) diff --git a/ding/model/template/tests/test_decision_transformer.py b/ding/model/template/tests/test_decision_transformer.py index 69dfb6738c..71f52da4a9 100644 --- a/ding/model/template/tests/test_decision_transformer.py +++ b/ding/model/template/tests/test_decision_transformer.py @@ -1,33 +1,49 @@ import pytest from itertools import product import torch +import torch.nn as nn import torch.nn.functional as F from ding.model.template import DecisionTransformer -from ding.torch_utils import is_differentiable, one_hot +from ding.torch_utils import is_differentiable -args = ['continuous', 'discrete'] +action_space = ['continuous', 'discrete'] +state_encoder = [None, nn.Sequential(nn.Flatten(), nn.Linear(8, 8), nn.Tanh())] +args = list(product(*[action_space, state_encoder])) +args.pop(1) @pytest.mark.unittest -@pytest.mark.parametrize('action_space', args) -def test_decision_transformer(action_space): +@pytest.mark.parametrize('action_space, state_encoder', args) +def test_decision_transformer(action_space, state_encoder): B, T = 4, 6 - state_dim = 3 + if state_encoder: + state_dim = (2, 2, 2) + else: + state_dim = 3 act_dim = 2 DT_model = DecisionTransformer( state_dim=state_dim, act_dim=act_dim, + state_encoder=state_encoder, n_blocks=3, h_dim=8, context_len=T, n_heads=2, drop_p=0.1, + continuous=(action_space == 'continuous') ) + DT_model.configure_optimizers(1.0, 0.0003) is_continuous = True if action_space == 'continuous' else False - timesteps = torch.randint(0, 100, [B, T], dtype=torch.long) # B x T - states = torch.randn([B, T, state_dim]) # B x T x state_dim + if state_encoder: + timesteps = torch.randint(0, 100, [B, 3 * T - 1, 1], dtype=torch.long) # B x T + else: + timesteps = torch.randint(0, 100, [B, T], dtype=torch.long) # B x T + if isinstance(state_dim, int): + states = torch.randn([B, T, state_dim]) # B x T x state_dim + else: + states = torch.randn([B, T, *state_dim]) # B x T x state_dim if action_space == 'continuous': actions = torch.randn([B, T, act_dim]) # B x T x act_dim action_target = torch.randn([B, T, act_dim]) @@ -40,26 +56,29 @@ def test_decision_transformer(action_space): # all ones since no padding traj_mask = torch.ones([B, T], dtype=torch.long) # B x T - # if discrete - if not is_continuous: - actions = one_hot(actions.squeeze(-1), num=act_dim) - - assert actions.shape == (B, T, act_dim) if is_continuous: assert action_target.shape == (B, T, act_dim) else: assert action_target.shape == (B, T, 1) + actions = actions.squeeze(-1) returns_to_go = returns_to_go.float() state_preds, action_preds, return_preds = DT_model.forward( timesteps=timesteps, states=states, actions=actions, returns_to_go=returns_to_go ) - assert state_preds.shape == (B, T, state_dim) + if state_encoder: + assert state_preds is None + assert return_preds is None + else: + assert state_preds.shape == (B, T, state_dim) + assert return_preds.shape == (B, T, 1) assert action_preds.shape == (B, T, act_dim) - assert return_preds.shape == (B, T, 1) # only consider non padded elements - action_preds = action_preds.view(-1, act_dim)[traj_mask.view(-1, ) > 0] + if state_encoder: + action_preds = action_preds.reshape(-1, act_dim) + else: + action_preds = action_preds.view(-1, act_dim)[traj_mask.view(-1, ) > 0] if is_continuous: action_target = action_target.view(-1, act_dim)[traj_mask.view(-1, ) > 0] @@ -71,11 +90,14 @@ def test_decision_transformer(action_space): else: action_loss = F.cross_entropy(action_preds, action_target) - # print(action_loss) - # is_differentiable(action_loss, DT_model) - is_differentiable( - action_loss, [ - DT_model.transformer, DT_model.embed_action, DT_model.predict_action, DT_model.embed_rtg, - DT_model.embed_state - ] - ) # pass + if state_encoder: + is_differentiable( + action_loss, [DT_model.transformer, DT_model.embed_action, DT_model.embed_rtg, DT_model.state_encoder] + ) + else: + is_differentiable( + action_loss, [ + DT_model.transformer, DT_model.embed_action, DT_model.predict_action, DT_model.embed_rtg, + DT_model.embed_state + ] + ) diff --git a/ding/model/template/tests/test_edac.py b/ding/model/template/tests/test_edac.py new file mode 100644 index 0000000000..76f0cca60a --- /dev/null +++ b/ding/model/template/tests/test_edac.py @@ -0,0 +1,57 @@ +import torch +import pytest +from itertools import product + +from ding.model.template import EDAC +from ding.torch_utils import is_differentiable + +B = 4 +obs_shape = [4, (8, )] +act_shape = [3, (6, )] +args = list(product(*[obs_shape, act_shape])) + + +@pytest.mark.unittest +class TestEDAC: + + def output_check(self, model, outputs): + if isinstance(outputs, torch.Tensor): + loss = outputs.sum() + elif isinstance(outputs, list): + loss = sum([t.sum() for t in outputs]) + elif isinstance(outputs, dict): + loss = sum([v.sum() for v in outputs.values()]) + is_differentiable(loss, model) + + @pytest.mark.parametrize('obs_shape, act_shape', args) + def test_EDAC(self, obs_shape, act_shape): + if isinstance(obs_shape, int): + inputs_obs = torch.randn(B, obs_shape) + else: + inputs_obs = torch.randn(B, *obs_shape) + if isinstance(act_shape, int): + inputs_act = torch.randn(B, act_shape) + else: + inputs_act = torch.randn(B, *act_shape) + inputs = {'obs': inputs_obs, 'action': inputs_act} + model = EDAC(obs_shape, act_shape, ensemble_num=2) + + outputs_c = model(inputs, mode='compute_critic') + assert isinstance(outputs_c, dict) + assert outputs_c['q_value'].shape == (2, B) + self.output_check(model.critic, outputs_c) + + if isinstance(obs_shape, int): + inputs = torch.randn(B, obs_shape) + else: + inputs = torch.randn(B, *obs_shape) + outputs_a = model(inputs, mode='compute_actor') + assert isinstance(outputs_a, dict) + if isinstance(act_shape, int): + assert outputs_a['logit'][0].shape == (B, act_shape) + assert outputs_a['logit'][1].shape == (B, act_shape) + elif len(act_shape) == 1: + assert outputs_a['logit'][0].shape == (B, *act_shape) + assert outputs_a['logit'][1].shape == (B, *act_shape) + outputs = {'mu': outputs_a['logit'][0], 'sigma': outputs_a['logit'][1]} + self.output_check(model.actor, outputs) diff --git a/ding/model/template/tests/test_havac.py b/ding/model/template/tests/test_havac.py new file mode 100644 index 0000000000..42982ec5ae --- /dev/null +++ b/ding/model/template/tests/test_havac.py @@ -0,0 +1,103 @@ +import pytest +import torch +import random +from ding.torch_utils import is_differentiable +from ding.model.template import HAVAC + + +@pytest.mark.unittest +class TestHAVAC: + + def test_havac_rnn_actor(self): + # discrete+rnn + bs, T = 3, 8 + obs_dim, global_obs_dim, action_dim = 32, 32 * 4, 9 + agent_num = 5 + data = { + 'obs': { + 'agent_state': torch.randn(T, bs, obs_dim), + 'global_state': torch.randn(T, bs, global_obs_dim), + 'action_mask': torch.randint(0, 2, size=(T, bs, action_dim)) + }, + 'actor_prev_state': [None for _ in range(bs)], + } + model = HAVAC( + agent_obs_shape=obs_dim, + global_obs_shape=global_obs_dim, + action_shape=action_dim, + agent_num=agent_num, + use_lstm=True, + ) + agent_idx = random.randint(0, agent_num - 1) + output = model(agent_idx, data, mode='compute_actor') + assert set(output.keys()) == set(['logit', 'actor_next_state', 'actor_hidden_state']) + assert output['logit'].shape == (T, bs, action_dim) + assert len(output['actor_next_state']) == bs + print(output['actor_next_state'][0]['h'].shape) + loss = output['logit'].sum() + is_differentiable(loss, model.agent_models[agent_idx].actor) + + def test_havac_rnn_critic(self): + # discrete+rnn + bs, T = 3, 8 + obs_dim, global_obs_dim, action_dim = 32, 32 * 4, 9 + agent_num = 5 + data = { + 'obs': { + 'agent_state': torch.randn(T, bs, obs_dim), + 'global_state': torch.randn(T, bs, global_obs_dim), + 'action_mask': torch.randint(0, 2, size=(T, bs, action_dim)) + }, + 'critic_prev_state': [None for _ in range(bs)], + } + model = HAVAC( + agent_obs_shape=obs_dim, + global_obs_shape=global_obs_dim, + action_shape=action_dim, + agent_num=agent_num, + use_lstm=True, + ) + agent_idx = random.randint(0, agent_num - 1) + output = model(agent_idx, data, mode='compute_critic') + assert set(output.keys()) == set(['value', 'critic_next_state', 'critic_hidden_state']) + assert output['value'].shape == (T, bs) + assert len(output['critic_next_state']) == bs + print(output['critic_next_state'][0]['h'].shape) + loss = output['value'].sum() + is_differentiable(loss, model.agent_models[agent_idx].critic) + + def test_havac_rnn_actor_critic(self): + # discrete+rnn + bs, T = 3, 8 + obs_dim, global_obs_dim, action_dim = 32, 32 * 4, 9 + agent_num = 5 + data = { + 'obs': { + 'agent_state': torch.randn(T, bs, obs_dim), + 'global_state': torch.randn(T, bs, global_obs_dim), + 'action_mask': torch.randint(0, 2, size=(T, bs, action_dim)) + }, + 'actor_prev_state': [None for _ in range(bs)], + 'critic_prev_state': [None for _ in range(bs)], + } + model = HAVAC( + agent_obs_shape=obs_dim, + global_obs_shape=global_obs_dim, + action_shape=action_dim, + agent_num=agent_num, + use_lstm=True, + ) + agent_idx = random.randint(0, agent_num - 1) + output = model(agent_idx, data, mode='compute_actor_critic') + assert set(output.keys()) == set( + ['logit', 'actor_next_state', 'actor_hidden_state', 'value', 'critic_next_state', 'critic_hidden_state'] + ) + assert output['logit'].shape == (T, bs, action_dim) + assert output['value'].shape == (T, bs) + loss = output['logit'].sum() + output['value'].sum() + is_differentiable(loss, model.agent_models[agent_idx]) + + +# test_havac_rnn_actor() +# test_havac_rnn_critic() +# test_havac_rnn_actor_critic() diff --git a/ding/model/template/tests/test_hpt.py b/ding/model/template/tests/test_hpt.py new file mode 100644 index 0000000000..d05681a127 --- /dev/null +++ b/ding/model/template/tests/test_hpt.py @@ -0,0 +1,47 @@ +import pytest +import torch +from itertools import product +from ding.model.template.hpt import HPT +from ding.torch_utils import is_differentiable + +T, B = 3, 4 +obs_shape = [4, (8, )] +act_shape = [3, (6, )] +args = list(product(*[obs_shape, act_shape])) + + +@pytest.mark.unittest +class TestHPT: + + def output_check(self, model, outputs): + if isinstance(outputs, torch.Tensor): + loss = outputs.sum() + elif isinstance(outputs, list): + loss = sum([t.sum() for t in outputs]) + elif isinstance(outputs, dict): + loss = sum([v.sum() for v in outputs.values()]) + is_differentiable(loss, model) + + @pytest.mark.parametrize('obs_shape, act_shape', args) + def test_hpt(self, obs_shape, act_shape): + if isinstance(obs_shape, int): + inputs = torch.randn(B, obs_shape) + state_dim = obs_shape + else: + inputs = torch.randn(B, *obs_shape) + state_dim = obs_shape[0] + + model = HPT(state_dim=state_dim, action_dim=act_shape) + outputs = model(inputs) + + assert isinstance(outputs, dict) + + if isinstance(act_shape, int): + assert outputs['logit'].shape == (B, act_shape) + elif len(act_shape) == 1: + assert outputs['logit'].shape == (B, *act_shape) + else: + for i, s in enumerate(act_shape): + assert outputs['logit'][i].shape == (B, s) + + self.output_check(model, outputs) diff --git a/ding/model/template/tests/test_hybrid_qac.py b/ding/model/template/tests/test_hybrid_qac.py index 018c3f2d36..3a81d55350 100644 --- a/ding/model/template/tests/test_hybrid_qac.py +++ b/ding/model/template/tests/test_hybrid_qac.py @@ -3,7 +3,7 @@ import pytest from itertools import product -from ding.model.template import QAC +from ding.model.template import ContinuousQAC from ding.torch_utils import is_differentiable from ding.utils import squeeze from easydict import EasyDict @@ -21,7 +21,7 @@ @pytest.mark.unittest -class TestHybridQAC: +class TestHybridContinuousQAC: def test_hybrid_qac( self, @@ -39,7 +39,7 @@ def test_hybrid_qac( }, 'logit': torch.randn(B, squeeze(action_shape.action_type_shape)) } - model = QAC( + model = ContinuousQAC( obs_shape=(N, ), action_shape=action_shape, action_space=action_space, @@ -50,8 +50,8 @@ def test_hybrid_qac( # compute_q q = model(inputs, mode='compute_critic')['q_value'] if twin: - is_differentiable(q[0].sum(), model.critic[0]) - is_differentiable(q[1].sum(), model.critic[1]) + is_differentiable(q[0].sum(), model.critic[1][0]) + is_differentiable(q[1].sum(), model.critic[1][1]) else: is_differentiable(q.sum(), model.critic) diff --git a/ding/model/template/tests/test_language_transformer.py b/ding/model/template/tests/test_language_transformer.py new file mode 100644 index 0000000000..eaaaae5a84 --- /dev/null +++ b/ding/model/template/tests/test_language_transformer.py @@ -0,0 +1,49 @@ +import pytest + +from ding.model.template.language_transformer import LanguageTransformer + + +@pytest.mark.unittest +class TestNLPPretrainedModel: + + def check_model(self): + test_pids = [1] + cand_pids = [0, 2, 4] + problems = [ + "This is problem 0", "This is the first question", "Second problem is here", "Another problem", + "This is the last problem" + ] + ctxt_list = [problems[pid] for pid in test_pids] + cands_list = [problems[pid] for pid in cand_pids] + + model = LanguageTransformer(model_name="bert-base-uncased", add_linear=True, embedding_size=256) + output = model(ctxt_list, cands_list, mode='compute_actor') + assert 'dist' in output.keys() and 'logit' in output.keys() and len(output.keys()) == 2 + assert output['logit'].shape == (1, 3) + + output = model(ctxt_list, cands_list, mode='compute_critic') + assert 'value' in output.keys() and len(output.keys()) == 1 + assert output['value'].shape == (1, ) + + output = model(ctxt_list, cands_list, mode='compute_critic') + assert 'value' in output.keys() and 'dist' in output.keys() and 'logit' in output.keys() and len( + output.keys() + ) == 3 + assert output['value'].shape == (1, ) + assert output['logit'].shape == (1, 3) + + model = LanguageTransformer(model_name="bert-base-uncased", add_linear=False, norm_embedding=True) + output = model(ctxt_list, cands_list, mode='compute_actor') + assert 'dist' in output.keys() and 'logit' in output.keys() and len(output.keys()) == 2 + assert output['logit'].shape == (1, 3) + + output = model(ctxt_list, cands_list, mode='compute_critic') + assert 'value' in output.keys() and len(output.keys()) == 1 + assert output['value'].shape == (1, ) + + output = model(ctxt_list, cands_list, mode='compute_critic') + assert 'value' in output.keys() and 'dist' in output.keys() and 'logit' in output.keys() and len( + output.keys() + ) == 3 + assert output['value'].shape == (1, ) + assert output['logit'].shape == (1, 3) diff --git a/ding/model/template/tests/test_madqn.py b/ding/model/template/tests/test_madqn.py new file mode 100644 index 0000000000..c2245c332f --- /dev/null +++ b/ding/model/template/tests/test_madqn.py @@ -0,0 +1,30 @@ +import pytest +import torch +from ding.torch_utils import is_differentiable +from ding.model.template import MADQN + + +@pytest.mark.unittest +def test_madqn(): + agent_num, bs, T = 4, 3, 8 + obs_dim, global_obs_dim, action_dim = 32, 32 * 4, 9 + embedding_dim = 64 + madqn_model = MADQN( + agent_num=agent_num, + obs_shape=obs_dim, + action_shape=action_dim, + hidden_size_list=[embedding_dim, embedding_dim], + global_obs_shape=global_obs_dim + ) + data = { + 'obs': { + 'agent_state': torch.randn(T, bs, agent_num, obs_dim), + 'global_state': torch.randn(T, bs, agent_num, global_obs_dim), + 'action_mask': torch.randint(0, 2, size=(T, bs, agent_num, action_dim)) + }, + 'prev_state': [[None for _ in range(agent_num)] for _ in range(bs)], + 'action': torch.randint(0, action_dim, size=(T, bs, agent_num)) + } + output = madqn_model(data, cooperation=True, single_step=False) + assert output['total_q'].shape == (T, bs) + assert len(output['next_state']) == bs and all([len(n) == agent_num for n in output['next_state']]) diff --git a/ding/model/template/tests/test_maqac.py b/ding/model/template/tests/test_maqac.py index 4b6f40e69a..fa917e7ebc 100644 --- a/ding/model/template/tests/test_maqac.py +++ b/ding/model/template/tests/test_maqac.py @@ -3,7 +3,7 @@ import pytest from itertools import product -from ding.model.template import MAQAC, ContinuousMAQAC +from ding.model.template import DiscreteMAQAC, ContinuousMAQAC from ding.torch_utils import is_differentiable from ding.utils.default_helper import squeeze @@ -17,7 +17,7 @@ @pytest.mark.unittest @pytest.mark.parametrize('agent_obs_shape, global_obs_shape, twin_critic', args) -class TestMAQAC: +class TestDiscreteMAQAC: def output_check(self, model, outputs, action_shape): if isinstance(action_shape, tuple): @@ -34,7 +34,7 @@ def test_maqac(self, agent_obs_shape, global_obs_shape, twin_critic): 'action_mask': torch.randint(0, 2, size=(B, agent_num, action_shape)) } } - model = MAQAC(agent_obs_shape, global_obs_shape, action_shape, twin_critic=twin_critic) + model = DiscreteMAQAC(agent_obs_shape, global_obs_shape, action_shape, twin_critic=twin_critic) logit = model(data, mode='compute_actor')['logit'] value = model(data, mode='compute_critic')['q_value'] diff --git a/ding/model/template/tests/test_mavac.py b/ding/model/template/tests/test_mavac.py index f6c6927373..9e8a3ec27b 100644 --- a/ding/model/template/tests/test_mavac.py +++ b/ding/model/template/tests/test_mavac.py @@ -1,6 +1,7 @@ import pytest import numpy as np import torch +import torch.nn as nn from itertools import product from ding.model import mavac @@ -50,3 +51,39 @@ def test_vac(self, agent_obs_shape, global_obs_shape): value = model(data, mode='compute_critic')['value'] assert value.shape == (B, agent_num) self.output_check(model.critic, value, action_shape) + + def test_vac_with_encoder(self, agent_obs_shape, global_obs_shape): + data = { + 'agent_state': torch.randn(B, agent_num, agent_obs_shape), + 'global_state': torch.randn(B, agent_num, global_obs_shape), + 'action_mask': torch.randint(0, 2, size=(B, agent_num, action_shape)) + } + + actor_size, critic_size = 128, 128 + encoder = [nn.Linear(agent_obs_shape, actor_size), nn.Linear(global_obs_shape, critic_size)] + model = MAVAC( + agent_obs_shape, + global_obs_shape, + action_shape, + agent_num, + encoder=encoder, + actor_head_hidden_size=actor_size, + critic_head_hidden_size=critic_size + ) + + logit = model(data, mode='compute_actor_critic')['logit'] + value = model(data, mode='compute_actor_critic')['value'] + + outputs = value.sum() + logit.sum() + self.output_check(model, outputs, action_shape) + + for p in model.parameters(): + p.grad = None + logit = model(data, mode='compute_actor')['logit'] + self.output_check(model.actor, logit, model.action_shape) + + for p in model.parameters(): + p.grad = None + value = model(data, mode='compute_critic')['value'] + assert value.shape == (B, agent_num) + self.output_check(model.critic, value, action_shape) diff --git a/ding/model/template/tests/test_ngu.py b/ding/model/template/tests/test_ngu.py new file mode 100644 index 0000000000..ed0e86f194 --- /dev/null +++ b/ding/model/template/tests/test_ngu.py @@ -0,0 +1,70 @@ +import pytest +from itertools import product +import torch +from ding.model.template import NGU +from ding.torch_utils import is_differentiable + +B = 4 +H = 4 +obs_shape = [4, (8, ), (4, 64, 64)] +act_shape = [4, (4, )] +args = list(product(*[obs_shape, act_shape])) + + +@pytest.mark.unittest +class TestNGU: + + def output_check(self, model, outputs): + if isinstance(outputs, torch.Tensor): + loss = outputs.sum() + elif isinstance(outputs, list): + loss = sum([t.sum() for t in outputs]) + elif isinstance(outputs, dict): + loss = sum([v.sum() for v in outputs.values()]) + is_differentiable(loss, model) + + @pytest.mark.parametrize('obs_shape, act_shape', args) + def test_ngu(self, obs_shape, act_shape): + if isinstance(obs_shape, int): + inputs_obs = torch.randn(B, H, obs_shape) + else: + inputs_obs = torch.randn(B, H, *obs_shape) + if isinstance(act_shape, int): + inputs_prev_action = torch.ones(B, act_shape).long() + else: + inputs_prev_action = torch.ones(B, *act_shape).long() + inputs_prev_reward_extrinsic = torch.randn(B, H, 1) + inputs_beta = 2 * torch.ones([4, 4], dtype=torch.long) + inputs = { + 'obs': inputs_obs, + 'prev_state': None, + 'prev_action': inputs_prev_action, + 'prev_reward_extrinsic': inputs_prev_reward_extrinsic, + 'beta': inputs_beta + } + + model = NGU(obs_shape, act_shape, collector_env_num=3) + outputs = model(inputs) + assert isinstance(outputs, dict) + if isinstance(act_shape, int): + assert outputs['logit'].shape == (B, act_shape, act_shape) + elif len(act_shape) == 1: + assert outputs['logit'].shape == (B, *act_shape, *act_shape) + self.output_check(model, outputs['logit']) + + inputs = { + 'obs': inputs_obs, + 'prev_state': None, + 'action': inputs_prev_action, + 'reward': inputs_prev_reward_extrinsic, + 'prev_reward_extrinsic': inputs_prev_reward_extrinsic, + 'beta': inputs_beta + } + model = NGU(obs_shape, act_shape, collector_env_num=3) + outputs = model(inputs) + assert isinstance(outputs, dict) + if isinstance(act_shape, int): + assert outputs['logit'].shape == (B, act_shape, act_shape) + elif len(act_shape) == 1: + assert outputs['logit'].shape == (B, *act_shape, *act_shape) + self.output_check(model, outputs['logit']) diff --git a/ding/model/template/tests/test_pdqn.py b/ding/model/template/tests/test_pdqn.py index 10dfa25737..6f9b66f9af 100644 --- a/ding/model/template/tests/test_pdqn.py +++ b/ding/model/template/tests/test_pdqn.py @@ -1,16 +1,21 @@ import pytest -from easydict import EasyDict import torch +from easydict import EasyDict + from ding.model.template import PDQN +action_args_shape_values = [1, 5] + @pytest.mark.unittest class TestPQQN: - def test_dqn(self): + @pytest.mark.unittest + @pytest.mark.parametrize('action_type_shape', action_args_shape_values) + def test_dqn(self, action_type_shape): T, B = 3, 4 obs_shape = (4, ) - act_shape = EasyDict({'action_type_shape': (3, ), 'action_args_shape': (5, )}) + act_shape = EasyDict({'action_type_shape': (3, ), 'action_args_shape': (action_type_shape, )}) if isinstance(obs_shape, int): cont_inputs = torch.randn(B, obs_shape) else: diff --git a/ding/model/template/tests/test_pg.py b/ding/model/template/tests/test_pg.py new file mode 100644 index 0000000000..2eb0dfba90 --- /dev/null +++ b/ding/model/template/tests/test_pg.py @@ -0,0 +1,61 @@ +import torch +import numpy as np +import pytest +from itertools import product + +from ding.model.template import PG +from ding.torch_utils import is_differentiable +from ding.utils import squeeze + +B = 4 + + +@pytest.mark.unittest +class TestDiscretePG: + + def output_check(self, model, outputs): + if isinstance(outputs, torch.Tensor): + loss = outputs.sum() + elif isinstance(outputs, list): + loss = sum([t.sum() for t in outputs]) + elif isinstance(outputs, dict): + loss = sum([v.sum() for v in outputs.values()]) + is_differentiable(loss, model) + + def test_discrete_pg(self): + obs_shape = (4, 84, 84) + action_shape = 5 + model = PG( + obs_shape, + action_shape, + ) + inputs = torch.randn(B, 4, 84, 84) + + outputs = model(inputs) + assert isinstance(outputs, dict) + assert outputs['logit'].shape == (B, action_shape) + assert outputs['dist'].sample().shape == (B, ) + self.output_check(model, outputs['logit']) + + def test_continuous_pg(self): + N = 32 + action_shape = (6, ) + inputs = {'obs': torch.randn(B, N), 'action': torch.randn(B, squeeze(action_shape))} + model = PG( + obs_shape=(N, ), + action_shape=action_shape, + action_space='continuous', + ) + # compute_action + print(model) + outputs = model(inputs['obs']) + assert isinstance(outputs, dict) + dist = outputs['dist'] + action = dist.sample() + assert action.shape == (B, *action_shape) + + logit = outputs['logit'] + mu, sigma = logit['mu'], logit['sigma'] + assert mu.shape == (B, *action_shape) + assert sigma.shape == (B, *action_shape) + is_differentiable(mu.sum() + sigma.sum(), model) diff --git a/ding/model/template/tests/test_procedure_cloning.py b/ding/model/template/tests/test_procedure_cloning.py new file mode 100644 index 0000000000..b2bb197954 --- /dev/null +++ b/ding/model/template/tests/test_procedure_cloning.py @@ -0,0 +1,37 @@ +import pytest +from itertools import product + +import torch + +from ding.model.template import ProcedureCloningMCTS, ProcedureCloningBFS + +B = 4 +T = 15 +obs_shape = [(64, 64, 3)] +action_dim = [9] +obs_embeddings = 256 +args = list(product(*[obs_shape, action_dim])) + + +@pytest.mark.unittest +@pytest.mark.parametrize('obs_shape, action_dim', args) +class TestProcedureCloning: + + def test_procedure_cloning_mcts(self, obs_shape, action_dim): + inputs = { + 'states': torch.randn(B, *obs_shape), + 'goals': torch.randn(B, *obs_shape), + 'actions': torch.randn(B, T, action_dim) + } + model = ProcedureCloningMCTS(obs_shape=obs_shape, action_dim=action_dim) + goal_preds, action_preds = model(inputs['states'], inputs['goals'], inputs['actions']) + assert goal_preds.shape == (B, obs_embeddings) + assert action_preds.shape == (B, T + 1, action_dim) + + def test_procedure_cloning_bfs(self, obs_shape, action_dim): + o_shape = (obs_shape[2], obs_shape[0], obs_shape[1]) + model = ProcedureCloningBFS(obs_shape=o_shape, action_shape=action_dim) + + inputs = torch.randn(B, *obs_shape) + map_preds = model(inputs) + assert map_preds['logit'].shape == (B, obs_shape[0], obs_shape[1], action_dim + 1) diff --git a/ding/model/template/tests/test_q_learning.py b/ding/model/template/tests/test_q_learning.py index b444becdf6..1f0ceba3b9 100644 --- a/ding/model/template/tests/test_q_learning.py +++ b/ding/model/template/tests/test_q_learning.py @@ -1,8 +1,8 @@ import pytest from itertools import product import torch -from ding.model.template import DQN, RainbowDQN, QRDQN, IQN, FQF, DRQN, C51DQN -from ding.torch_utils import is_differentiable +from ding.model.template import DQN, RainbowDQN, QRDQN, IQN, FQF, DRQN, C51DQN, BDQ, GTrXLDQN +from ding.torch_utils import is_differentiable, NoiseLinearLayer T, B = 3, 4 obs_shape = [4, (8, ), (4, 64, 64)] @@ -40,6 +40,25 @@ def test_dqn(self, obs_shape, act_shape): assert outputs['logit'][i].shape == (B, s) self.output_check(model, outputs['logit']) + @pytest.mark.parametrize('obs_shape, act_shape', args) + def test_bdq(self, obs_shape, act_shape): + if isinstance(obs_shape, int): + inputs = torch.randn(B, obs_shape) + else: + inputs = torch.randn(B, *obs_shape) + if not isinstance(act_shape, int) and len(act_shape) > 1: + return + num_branches = act_shape + for action_bins_per_branch in range(1, 10): + model = BDQ(obs_shape, num_branches, action_bins_per_branch) + outputs = model(inputs) + assert isinstance(outputs, dict) + if isinstance(act_shape, int): + assert outputs['logit'].shape == (B, act_shape, action_bins_per_branch) + else: + assert outputs['logit'].shape == (B, *act_shape, action_bins_per_branch) + self.output_check(model, outputs['logit']) + @pytest.mark.parametrize('obs_shape, act_shape', args) def test_rainbowdqn(self, obs_shape, act_shape): if isinstance(obs_shape, int): @@ -47,6 +66,9 @@ def test_rainbowdqn(self, obs_shape, act_shape): else: inputs = torch.randn(B, *obs_shape) model = RainbowDQN(obs_shape, act_shape, n_atom=41) + for k, v in model.named_modules(): + if isinstance(v, NoiseLinearLayer): + v.enable_noise = True outputs = model(inputs) assert isinstance(outputs, dict) if isinstance(act_shape, int): @@ -264,3 +286,11 @@ def test_drqn_inference_res_link(self, obs_shape, act_shape): assert all([len(t) == 2 for t in outputs['next_state']]) assert all([t['h'].shape == (1, 1, 64) for t in outputs['next_state']]) self.output_check(model, outputs['logit']) + + @pytest.mark.tmp + def test_GTrXLDQN(self): + obs_dim, seq_len, bs, action_dim = [4, 64, 64], 64, 32, 4 + obs = torch.rand(seq_len, bs, *obs_dim) + model = GTrXLDQN(obs_dim, action_dim, encoder_hidden_size_list=[16, 16, 16]) + outputs = model(obs) + assert isinstance(outputs, dict) diff --git a/ding/model/template/tests/test_qac.py b/ding/model/template/tests/test_qac.py index ea0e9348d9..7ddbf9d511 100644 --- a/ding/model/template/tests/test_qac.py +++ b/ding/model/template/tests/test_qac.py @@ -3,7 +3,7 @@ import pytest from itertools import product -from ding.model.template import QAC, MAQAC, DiscreteQAC +from ding.model.template import ContinuousQAC, DiscreteMAQAC, DiscreteQAC from ding.torch_utils import is_differentiable from ding.utils import squeeze @@ -18,12 +18,12 @@ @pytest.mark.unittest @pytest.mark.parametrize('action_shape, twin, action_space', args) -class TestQAC: +class TestContinuousQAC: def test_fcqac(self, action_shape, twin, action_space): N = 32 inputs = {'obs': torch.randn(B, N), 'action': torch.randn(B, squeeze(action_shape))} - model = QAC( + model = ContinuousQAC( obs_shape=(N, ), action_shape=action_shape, action_space=action_space, @@ -34,8 +34,8 @@ def test_fcqac(self, action_shape, twin, action_space): # compute_q q = model(inputs, mode='compute_critic')['q_value'] if twin: - is_differentiable(q[0].sum(), model.critic[0]) - is_differentiable(q[1].sum(), model.critic[1]) + is_differentiable(q[0].sum(), model.critic[1][0]) + is_differentiable(q[1].sum(), model.critic[1][1]) else: is_differentiable(q.sum(), model.critic) @@ -56,80 +56,75 @@ def test_fcqac(self, action_shape, twin, action_space): is_differentiable(mu.sum() + sigma.sum(), model.actor) -args = list(product(*[[True, False]])) +args = list(product(*[[True, False], [(13, ), [4, 84, 84]]])) @pytest.mark.unittest -@pytest.mark.parametrize('twin', args) +@pytest.mark.parametrize('twin, obs_shape', args) class TestDiscreteQAC: - def test_discreteqac(self, twin): - N = 32 - A = 6 - inputs = {'obs': torch.randn(B, N)} + def test_discreteqac(self, twin, obs_shape): + action_shape = 6 + inputs = torch.randn(B, *obs_shape) model = DiscreteQAC( - agent_obs_shape=N, - global_obs_shape=N, - action_shape=A, + obs_shape=obs_shape, + action_shape=action_shape, twin_critic=twin, + encoder_hidden_size_list=[32, 32, 64] if len(obs_shape) > 1 else None, ) - # compute_q + # compute_critic q = model(inputs, mode='compute_critic')['q_value'] if twin: - is_differentiable(q[0].sum(), model.critic[0]) - is_differentiable(q[1].sum(), model.critic[1]) + is_differentiable(q[0].sum(), model.critic[1][0]) + # is_differentiable(q[1].sum(), model.critic[1][1]) # backward encoder twice + assert q[0].shape == (B, action_shape) + assert q[1].shape == (B, action_shape) else: - is_differentiable(q.sum(), model.critic) + is_differentiable(q.sum(), model.critic[1]) + assert q.shape == (B, action_shape) - # compute_action + # compute_actor print(model) logit = model(inputs, mode='compute_actor')['logit'] - assert logit.shape[0] == B - assert logit.shape[1] == A + assert logit.shape == (B, action_shape) + is_differentiable(logit.sum(), model.actor) -B = 32 -agent_obs_shape = [216, 265] -global_obs_shape = [264, 324] -agent_num = 8 -action_shape = 14 -args = list(product(*[agent_obs_shape, global_obs_shape])) +B = 4 +embedding_size = 64 +action_shape_args = [(6, ), 1] +args = list(product(*[action_shape_args, [True, False], [True, False]])) @pytest.mark.unittest -@pytest.mark.parametrize('agent_obs_shape, global_obs_shape', args) -class TestMAQAC: - - def output_check(self, model, outputs, action_shape): - if isinstance(action_shape, tuple): - loss = sum([t.sum() for t in outputs]) - elif np.isscalar(action_shape): - loss = outputs.sum() - is_differentiable(loss, model) - - def test_maqac(self, agent_obs_shape, global_obs_shape): - data = { - 'obs': { - 'agent_state': torch.randn(B, agent_num, agent_obs_shape), - 'global_state': torch.randn(B, agent_num, global_obs_shape), - 'action_mask': torch.randint(0, 2, size=(B, agent_num, action_shape)) - } - } - model = MAQAC(agent_obs_shape, global_obs_shape, action_shape) - - logit = model(data, mode='compute_actor')['logit'] - value = model(data, mode='compute_critic')['q_value'] - - outputs = value.sum() + logit.sum() - self.output_check(model, outputs, action_shape) - - for p in model.parameters(): - p.grad = None - logit = model(data, mode='compute_actor')['logit'] - self.output_check(model.actor, logit, action_shape) - - for p in model.parameters(): - p.grad = None - value = model(data, mode='compute_critic')['q_value'] - assert value.shape == (B, agent_num, action_shape) - self.output_check(model.critic, value, action_shape) +@pytest.mark.parametrize('action_shape, twin, share_encoder', args) +class TestContinuousQACPixel: + + def test_qacpixel(self, action_shape, twin, share_encoder): + inputs = {'obs': torch.randn(B, 3, 84, 84), 'action': torch.randn(B, squeeze(action_shape))} + model = ContinuousQAC( + obs_shape=(3, 84, 84), + action_shape=action_shape, + action_space='reparameterization', + critic_head_hidden_size=embedding_size, + actor_head_hidden_size=embedding_size, + twin_critic=twin, + share_encoder=share_encoder, + encoder_hidden_size_list=[32, 32, 64], + ) + # compute_q + q = model(inputs, mode='compute_critic')['q_value'] + if twin: + q = torch.min(q[0], q[1]) + is_differentiable(q.sum(), model.critic) + + # compute_action + print(model) + (mu, sigma) = model(inputs['obs'], mode='compute_actor')['logit'] + action_shape = squeeze(action_shape) + assert mu.shape == (B, action_shape) + assert sigma.shape == (B, action_shape) + if share_encoder: # if share_encoder, actor_encoder's grad is not None + is_differentiable(mu.sum() + sigma.sum(), model.actor_head) + else: + is_differentiable(mu.sum() + sigma.sum(), model.actor) diff --git a/ding/model/template/tests/test_qmix.py b/ding/model/template/tests/test_qmix.py index ce1817b697..74062357e7 100644 --- a/ding/model/template/tests/test_qmix.py +++ b/ding/model/template/tests/test_qmix.py @@ -43,3 +43,34 @@ def test_qmix(): is_differentiable(loss, qmix_model) data.pop('action') output = qmix_model(data, single_step=False) + + +@pytest.mark.unittest +def test_qmix_process_global_state(): + # Test the behavior of the _process_global_state method with different global_obs_shape types + agent_num, obs_dim, global_obs_dim, action_dim = 4, 32, 32 * 4, 9 + embedding_dim = 64 + + # Case 1: Test "flat" type global_obs_shape + global_obs_shape = global_obs_dim # Flat global_obs_shape + qmix_model_flat = QMix(agent_num, obs_dim, global_obs_shape, action_dim, [64, 128, embedding_dim], mixer=True) + + # Simulate input for the "flat" type global_state + batch_size, time_steps = 3, 8 + global_state_flat = torch.randn(batch_size, time_steps, global_obs_dim) + processed_flat = qmix_model_flat._process_global_state(global_state_flat) + + # Ensure the output shape is correct [batch_size, time_steps, embedding_dim] + assert processed_flat.shape == (batch_size, time_steps, global_obs_dim) + + # Case 2: Test "image" type global_obs_shape + global_obs_shape = [3, 64, 64] # Image-shaped global_obs_shape (C, H, W) + qmix_model_image = QMix(agent_num, obs_dim, global_obs_shape, action_dim, [64, 128, embedding_dim], mixer=True) + + # Simulate input for the "image" type global_state + C, H, W = global_obs_shape + global_state_image = torch.randn(batch_size, time_steps, C, H, W) + processed_image = qmix_model_image._process_global_state(global_state_image) + + # Ensure the output shape is correct [batch_size, time_steps, embedding_dim] + assert processed_image.shape == (batch_size, time_steps, embedding_dim) diff --git a/ding/model/template/tests/test_qtran.py b/ding/model/template/tests/test_qtran.py new file mode 100644 index 0000000000..c468ec0f15 --- /dev/null +++ b/ding/model/template/tests/test_qtran.py @@ -0,0 +1,33 @@ +import pytest +from itertools import product +import torch +from ding.model.template import QTran +from ding.torch_utils import is_differentiable + + +@pytest.mark.unittest +def test_qtran(): + agent_num, bs, T = 4, 3, 8 + obs_dim, global_obs_dim, action_dim = 32, 32 * 4, 9 + embedding_dim = 64 + data = { + 'obs': { + 'agent_state': torch.randn(T, bs, agent_num, obs_dim), + 'global_state': torch.randn(T, bs, global_obs_dim), + 'action_mask': torch.randint(0, 2, size=(T, bs, agent_num, action_dim)) + }, + 'prev_state': [[None for _ in range(agent_num)] for _ in range(bs)], + 'action': torch.randint(0, action_dim, size=(T, bs, agent_num)) + } + model = QTran(agent_num, obs_dim, global_obs_dim, action_dim, [32, embedding_dim], embedding_dim) + output = model.forward(data, single_step=False) + assert set(output.keys()) == set(['next_state', 'agent_q_act', 'vs', 'logit', 'action_mask', 'total_q']) + assert output['total_q'].shape == (T, bs) + assert output['logit'].shape == (T, bs, agent_num, action_dim) + assert len(output['next_state']) == bs and all([len(n) == agent_num for n in output['next_state']]) + print(output['next_state'][0][0]['h'].shape) + loss = output['total_q'].sum() + output['agent_q_act'].sum() + output['vs'].sum() + is_differentiable(loss, model) + + data.pop('action') + outputs = model.forward(data, single_step=False) diff --git a/ding/model/template/tests/test_vac.py b/ding/model/template/tests/test_vac.py index 4f31e942b6..85e44e8a4c 100644 --- a/ding/model/template/tests/test_vac.py +++ b/ding/model/template/tests/test_vac.py @@ -3,14 +3,16 @@ import torch from itertools import product -from ding.model import VAC +from ding.model import VAC, DREAMERVAC from ding.torch_utils import is_differentiable from ding.model import ConvEncoder +from easydict import EasyDict +ezD = EasyDict({'action_args_shape': (3, ), 'action_type_shape': 4}) B, C, H, W = 4, 3, 128, 128 obs_shape = [4, (8, ), (4, 64, 64)] -act_args = [[6, 'discrete'], [(3, ), 'continuous'], [[2, 3, 6], 'discrete']] +act_args = [[6, 'discrete'], [(3, ), 'continuous'], [[2, 3, 6], 'discrete'], [ezD, 'hybrid']] # act_args = [[(3, ), True]] args = list(product(*[obs_shape, act_args, [False, True]])) @@ -20,6 +22,8 @@ def output_check(model, outputs, action_shape): loss = sum([t.sum() for t in outputs]) elif np.isscalar(action_shape): loss = outputs.sum() + elif isinstance(action_shape, dict): + loss = outputs.sum() is_differentiable(loss, model) @@ -28,6 +32,9 @@ def model_check(model, inputs): value, logit = outputs['value'], outputs['logit'] if model.action_space == 'continuous': outputs = value.sum() + logit['mu'].sum() + logit['sigma'].sum() + elif model.action_space == 'hybrid': + outputs = value.sum() + logit['action_type'].sum() + logit['action_args']['mu'].sum( + ) + logit['action_args']['sigma'].sum() else: if model.multi_head: outputs = value.sum() + sum([t.sum() for t in logit]) @@ -40,6 +47,8 @@ def model_check(model, inputs): logit = model(inputs, mode='compute_actor')['logit'] if model.action_space == 'continuous': logit = logit['mu'].sum() + logit['sigma'].sum() + elif model.action_space == 'hybrid': + logit = logit['action_type'].sum() + logit['action_args']['mu'].sum() + logit['action_args']['sigma'].sum() output_check(model.actor, logit, model.action_shape) for p in model.parameters(): @@ -49,6 +58,15 @@ def model_check(model, inputs): output_check(model.critic, value, 1) +@pytest.mark.unittest +class TestDREAMERVAC: + + def test_DREAMERVAC(self): + obs_shape = 8 + act_shape = 6 + model = DREAMERVAC(obs_shape, act_shape) + + @pytest.mark.unittest @pytest.mark.parametrize('obs_shape, act_args, share_encoder', args) class TestVACGeneral: diff --git a/ding/model/template/vac.py b/ding/model/template/vac.py index cf026e0192..2ab1cf1f4d 100644 --- a/ding/model/template/vac.py +++ b/ding/model/template/vac.py @@ -6,15 +6,21 @@ from ding.utils import SequenceType, squeeze, MODEL_REGISTRY from ..common import ReparameterizationHead, RegressionHead, DiscreteHead, MultiHead, \ FCEncoder, ConvEncoder, IMPALAConvEncoder +from ding.torch_utils.network.dreamer import ActionHead, DenseHead @MODEL_REGISTRY.register('vac') class VAC(nn.Module): - r""" + """ Overview: - The VAC model. + The neural network and computation graph of algorithms related to (state) Value Actor-Critic (VAC), such as \ + A2C/PPO/IMPALA. This model now supports discrete, continuous and hybrid action space. The VAC is composed of \ + four parts: ``actor_encoder``, ``critic_encoder``, ``actor_head`` and ``critic_head``. Encoders are used to \ + extract the feature from various observation. Heads are used to predict corresponding value or action logit. \ + In high-dimensional observation space like 2D image, we often use a shared encoder for both ``actor_encoder`` \ + and ``critic_encoder``. In low-dimensional observation space like 1D vector, we often use different encoders. Interfaces: - ``__init__``, ``forward``, ``compute_actor``, ``compute_critic`` + ``__init__``, ``forward``, ``compute_actor``, ``compute_critic``, ``compute_actor_critic``. """ mode = ['compute_actor', 'compute_critic', 'compute_actor_critic'] @@ -37,26 +43,37 @@ def __init__( encoder: Optional[torch.nn.Module] = None, impala_cnn_encoder: bool = False, ) -> None: - r""" + """ Overview: - Init the VAC Model according to arguments. + Initialize the VAC model according to corresponding input arguments. Arguments: - - obs_shape (:obj:`Union[int, SequenceType]`): Observation's space. - - action_shape (:obj:`Union[int, SequenceType]`): Action's space. - - action_space (:obj:`str`): Choose action head in ['discrete', 'continuous', 'hybrid'] - - share_encoder (:obj:`bool`): Whether share encoder. - - encoder_hidden_size_list (:obj:`SequenceType`): Collection of ``hidden_size`` to pass to ``Encoder`` - - actor_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to actor-nn's ``Head``. - - actor_head_layer_num (:obj:`int`): - The num of layers used in the network to compute Q value output for actor's nn. - - critic_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to critic-nn's ``Head``. - - critic_head_layer_num (:obj:`int`): - The num of layers used in the network to compute Q value output for critic's nn. - - activation (:obj:`Optional[nn.Module]`): - The type of activation function to use in ``MLP`` the after ``layer_fn``, - if ``None`` then default set to ``nn.ReLU()`` - - norm_type (:obj:`Optional[str]`): - The type of normalization to use, see ``ding.torch_utils.fc_block`` for more details` + - obs_shape (:obj:`Union[int, SequenceType]`): Observation space shape, such as 8 or [4, 84, 84]. + - action_shape (:obj:`Union[int, SequenceType]`): Action space shape, such as 6 or [2, 3, 3]. + - action_space (:obj:`str`): The type of different action spaces, including ['discrete', 'continuous', \ + 'hybrid'], then will instantiate corresponding head, including ``DiscreteHead``, \ + ``ReparameterizationHead``, and hybrid heads. + - share_encoder (:obj:`bool`): Whether to share observation encoders between actor and decoder. + - encoder_hidden_size_list (:obj:`SequenceType`): Collection of ``hidden_size`` to pass to ``Encoder``, \ + the last element is used as the input size of ``actor_head`` and ``critic_head``. + - actor_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` of ``actor_head`` network, defaults \ + to 64, it is the hidden size of the last layer of the ``actor_head`` network. + - actor_head_layer_num (:obj:`int`): The num of layers used in the ``actor_head`` network to compute action. + - critic_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` of ``critic_head`` network, defaults \ + to 64, it is the hidden size of the last layer of the ``critic_head`` network. + - critic_head_layer_num (:obj:`int`): The num of layers used in the ``critic_head`` network. + - activation (:obj:`Optional[nn.Module]`): The type of activation function in networks \ + if ``None`` then default set it to ``nn.ReLU()``. + - norm_type (:obj:`Optional[str]`): The type of normalization in networks, see \ + ``ding.torch_utils.fc_block`` for more details. you can choose one of ['BN', 'IN', 'SyncBN', 'LN'] + - sigma_type (:obj:`Optional[str]`): The type of sigma in continuous action space, see \ + ``ding.torch_utils.network.dreamer.ReparameterizationHead`` for more details, in A2C/PPO, it defaults \ + to ``independent``, which means state-independent sigma parameters. + - fixed_sigma_value (:obj:`Optional[int]`): If ``sigma_type`` is ``fixed``, then use this value as sigma. + - bound_type (:obj:`Optional[str]`): The type of action bound methods in continuous action space, defaults \ + to ``None``, which means no bound. + - encoder (:obj:`Optional[torch.nn.Module]`): The encoder module, defaults to ``None``, you can define \ + your own encoder module and pass it into VAC to deal with different observation space. + - impala_cnn_encoder (:obj:`bool`): Whether to use IMPALA CNN encoder, defaults to ``False``. """ super(VAC, self).__init__() obs_shape: int = squeeze(obs_shape) @@ -66,7 +83,7 @@ def __init__( self.share_encoder = share_encoder # Encoder Type - def new_encoder(outsize): + def new_encoder(outsize, activation): if impala_cnn_encoder: return IMPALAConvEncoder(obs_shape=obs_shape, channels=encoder_hidden_size_list, outsize=outsize) else: @@ -91,15 +108,13 @@ def new_encoder(outsize): ) if self.share_encoder: - assert actor_head_hidden_size == critic_head_hidden_size, \ - "actor and critic network head should have same size." if encoder: if isinstance(encoder, torch.nn.Module): self.encoder = encoder else: raise ValueError("illegal encoder instance.") else: - self.encoder = new_encoder(actor_head_hidden_size) + self.encoder = new_encoder(encoder_hidden_size_list[-1], activation) else: if encoder: if isinstance(encoder, torch.nn.Module): @@ -108,25 +123,31 @@ def new_encoder(outsize): else: raise ValueError("illegal encoder instance.") else: - self.actor_encoder = new_encoder(actor_head_hidden_size) - self.critic_encoder = new_encoder(critic_head_hidden_size) + self.actor_encoder = new_encoder(encoder_hidden_size_list[-1], activation) + self.critic_encoder = new_encoder(encoder_hidden_size_list[-1], activation) # Head Type self.critic_head = RegressionHead( - critic_head_hidden_size, 1, critic_head_layer_num, activation=activation, norm_type=norm_type + encoder_hidden_size_list[-1], + 1, + critic_head_layer_num, + activation=activation, + norm_type=norm_type, + hidden_size=critic_head_hidden_size ) self.action_space = action_space assert self.action_space in ['discrete', 'continuous', 'hybrid'], self.action_space if self.action_space == 'continuous': self.multi_head = False self.actor_head = ReparameterizationHead( - actor_head_hidden_size, + encoder_hidden_size_list[-1], action_shape, actor_head_layer_num, sigma_type=sigma_type, activation=activation, norm_type=norm_type, - bound_type=bound_type + bound_type=bound_type, + hidden_size=actor_head_hidden_size, ) elif self.action_space == 'discrete': actor_head_cls = DiscreteHead @@ -155,7 +176,7 @@ def new_encoder(outsize): action_shape.action_args_shape = squeeze(action_shape.action_args_shape) action_shape.action_type_shape = squeeze(action_shape.action_type_shape) actor_action_args = ReparameterizationHead( - actor_head_hidden_size, + encoder_hidden_size_list[-1], action_shape.action_args_shape, actor_head_layer_num, sigma_type=sigma_type, @@ -163,6 +184,7 @@ def new_encoder(outsize): activation=activation, norm_type=norm_type, bound_type=bound_type, + hidden_size=actor_head_hidden_size, ) actor_action_type = DiscreteHead( actor_head_hidden_size, @@ -173,7 +195,6 @@ def new_encoder(outsize): ) self.actor_head = nn.ModuleList([actor_action_type, actor_action_args]) - # must use list, not nn.ModuleList if self.share_encoder: self.actor = [self.encoder, self.actor_head] self.critic = [self.encoder, self.critic_head] @@ -185,86 +206,80 @@ def new_encoder(outsize): self.actor = nn.ModuleList(self.actor) self.critic = nn.ModuleList(self.critic) - def forward(self, inputs: Union[torch.Tensor, Dict], mode: str) -> Dict: - r""" + def forward(self, x: torch.Tensor, mode: str) -> Dict: + """ Overview: - Use encoded embedding tensor to predict output. - Parameter updates with VAC's MLPs forward setup. + VAC forward computation graph, input observation tensor to predict state value or action logit. Different \ + ``mode`` will forward with different network modules to get different outputs and save computation. Arguments: - Forward with ``'compute_actor'`` or ``'compute_critic'``: - - inputs (:obj:`torch.Tensor`): - The encoded embedding tensor, determined with given ``hidden_size``, i.e. ``(B, N=hidden_size)``. - Whether ``actor_head_hidden_size`` or ``critic_head_hidden_size`` depend on ``mode``. + - x (:obj:`torch.Tensor`): The input observation tensor data. + - mode (:obj:`str`): The forward mode, all the modes are defined in the beginning of this class. Returns: - - outputs (:obj:`Dict`): - Run with encoder and head. + - outputs (:obj:`Dict`): The output dict of VAC's forward computation graph, whose key-values vary from \ + different ``mode``. - Forward with ``'compute_actor'``, Necessary Keys: - - logit (:obj:`torch.Tensor`): Logit encoding tensor, with same size as input ``x``. - - Forward with ``'compute_critic'``, Necessary Keys: - - value (:obj:`torch.Tensor`): Q value tensor with same size as batch size. - Shapes: - - inputs (:obj:`torch.Tensor`): :math:`(B, N)`, where B is batch size and N corresponding ``hidden_size`` - - logit (:obj:`torch.FloatTensor`): :math:`(B, N)`, where B is batch size and N is ``action_shape`` - - value (:obj:`torch.FloatTensor`): :math:`(B, )`, where B is batch size. - - Actor Examples: - >>> model = VAC(64,128) + Examples (Actor): + >>> model = VAC(64, 128) >>> inputs = torch.randn(4, 64) >>> actor_outputs = model(inputs,'compute_actor') >>> assert actor_outputs['logit'].shape == torch.Size([4, 128]) - Critic Examples: - >>> model = VAC(64,64) + Examples (Critic): + >>> model = VAC(64, 64) >>> inputs = torch.randn(4, 64) >>> critic_outputs = model(inputs,'compute_critic') - >>> critic_outputs['value'] - tensor([0.0252, 0.0235, 0.0201, 0.0072], grad_fn=) + >>> assert actor_outputs['logit'].shape == torch.Size([4, 64]) - Actor-Critic Examples: - >>> model = VAC(64,64) + Examples (Actor-Critic): + >>> model = VAC(64, 64) >>> inputs = torch.randn(4, 64) >>> outputs = model(inputs,'compute_actor_critic') - >>> outputs['value'] - tensor([0.0252, 0.0235, 0.0201, 0.0072], grad_fn=) + >>> assert critic_outputs['value'].shape == torch.Size([4]) >>> assert outputs['logit'].shape == torch.Size([4, 64]) """ assert mode in self.mode, "not support forward mode: {}/{}".format(mode, self.mode) - return getattr(self, mode)(inputs) + return getattr(self, mode)(x) - def compute_actor(self, x: torch.Tensor) -> Dict: - r""" + def compute_actor(self, x: Union[torch.Tensor, Dict]) -> Dict: + """ Overview: - Execute parameter updates with ``'compute_actor'`` mode - Use encoded embedding tensor to predict output. + VAC forward computation graph for actor part, input observation tensor to predict action logit. Arguments: - - inputs (:obj:`torch.Tensor`): - The encoded embedding tensor, determined with given ``hidden_size``, i.e. ``(B, N=hidden_size)``. - ``hidden_size = actor_head_hidden_size`` + - x (:obj:`Union[torch.Tensor, Dict]`): The input observation tensor data. If a dictionary is provided, \ + it should contain keys 'observation' and optionally 'action_mask'. Returns: - - outputs (:obj:`Dict`): - Run with encoder and head. - + - outputs (:obj:`Dict`): The output dict of VAC's forward computation graph for actor, including ``logit`` \ + and optionally ``action_mask`` if the input is a dictionary. ReturnsKeys: - - logit (:obj:`torch.Tensor`): Logit encoding tensor, with same size as input ``x``. + - logit (:obj:`torch.Tensor`): The predicted action logit tensor, for discrete action space, it will be \ + the same dimension real-value ranged tensor of possible action choices, and for continuous action \ + space, it will be the mu and sigma of the Gaussian distribution, and the number of mu and sigma is the \ + same as the number of continuous actions. Hybrid action space is a kind of combination of discrete \ + and continuous action space, so the logit will be a dict with ``action_type`` and ``action_args``. + - action_mask (:obj:`Optional[torch.Tensor]`): The action mask tensor, included if the input is a \ + dictionary containing 'action_mask'. Shapes: - - logit (:obj:`torch.FloatTensor`): :math:`(B, N)`, where B is batch size and N is ``action_shape`` + - logit (:obj:`torch.Tensor`): :math:`(B, N)`, where B is batch size and N is ``action_shape`` Examples: - >>> model = VAC(64,64) + >>> model = VAC(64, 64) >>> inputs = torch.randn(4, 64) >>> actor_outputs = model(inputs,'compute_actor') - >>> assert actor_outputs['action'].shape == torch.Size([4, 64]) + >>> assert actor_outputs['logit'].shape == torch.Size([4, 64]) """ - if self.share_encoder: - x = self.encoder(x) + if isinstance(x, dict): + action_mask = x['action_mask'] + x = self.encoder(x['observation']) if self.share_encoder else self.actor_encoder(x['observation']) else: - x = self.actor_encoder(x) + action_mask = None + x = self.encoder(x) if self.share_encoder else self.actor_encoder(x) if self.action_space == 'discrete': - return self.actor_head(x) + result = {'logit': self.actor_head(x)['logit']} + if action_mask is not None: + result['action_mask'] = action_mask + return result elif self.action_space == 'continuous': x = self.actor_head(x) # mu, sigma return {'logit': x} @@ -273,82 +288,92 @@ def compute_actor(self, x: torch.Tensor) -> Dict: action_args = self.actor_head[1](x) return {'logit': {'action_type': action_type['logit'], 'action_args': action_args}} - def compute_critic(self, x: torch.Tensor) -> Dict: - r""" + def compute_critic(self, x: Union[torch.Tensor, Dict]) -> Dict: + """ Overview: - Execute parameter updates with ``'compute_critic'`` mode - Use encoded embedding tensor to predict output. + VAC forward computation graph for critic part, input observation tensor to predict state value. Arguments: - - inputs (:obj:`torch.Tensor`): - The encoded embedding tensor, determined with given ``hidden_size``, i.e. ``(B, N=hidden_size)``. - ``hidden_size = critic_head_hidden_size`` + - x (:obj:`Union[torch.Tensor, Dict]`): The input observation tensor data. If a dictionary is provided, \ + it should contain the key 'observation'. Returns: - - outputs (:obj:`Dict`): - Run with encoder and head. - - Necessary Keys: - - value (:obj:`torch.Tensor`): Q value tensor with same size as batch size. + - outputs (:obj:`Dict`): The output dict of VAC's forward computation graph for critic, including ``value``. + ReturnsKeys: + - value (:obj:`torch.Tensor`): The predicted state value tensor. Shapes: - - value (:obj:`torch.FloatTensor`): :math:`(B, )`, where B is batch size. + - value (:obj:`torch.Tensor`): :math:`(B, )`, where B is batch size, (B, 1) is squeezed to (B, ). Examples: - >>> model = VAC(64,64) + >>> model = VAC(64, 64) >>> inputs = torch.randn(4, 64) >>> critic_outputs = model(inputs,'compute_critic') - >>> critic_outputs['value'] - tensor([0.0252, 0.0235, 0.0201, 0.0072], grad_fn=) + >>> assert critic_outputs['value'].shape == torch.Size([4]) """ - if self.share_encoder: - x = self.encoder(x) + if isinstance(x, dict): + x = self.encoder(x['observation']) if self.share_encoder else self.critic_encoder(x['observation']) else: - x = self.critic_encoder(x) + x = self.encoder(x) if self.share_encoder else self.critic_encoder(x) x = self.critic_head(x) return {'value': x['pred']} - def compute_actor_critic(self, x: torch.Tensor) -> Dict: - r""" + def compute_actor_critic(self, x: Union[torch.Tensor, Dict]) -> Dict: + """ Overview: - Execute parameter updates with ``'compute_actor_critic'`` mode - Use encoded embedding tensor to predict output. + VAC forward computation graph for both actor and critic part, input observation tensor to predict action \ + logit and state value. Arguments: - - inputs (:obj:`torch.Tensor`): The encoded embedding tensor. - + - x (:obj:`Union[torch.Tensor, Dict]`): The input observation tensor data. If a dictionary is provided, \ + it should contain keys 'observation' and optionally 'action_mask'. Returns: - - outputs (:obj:`Dict`): - Run with encoder and head. - + - outputs (:obj:`Dict`): The output dict of VAC's forward computation graph for both actor and critic, \ + including ``logit``, ``value``, and optionally ``action_mask`` if the input is a dictionary. ReturnsKeys: - - logit (:obj:`torch.Tensor`): Logit encoding tensor, with same size as input ``x``. - - value (:obj:`torch.Tensor`): Q value tensor with same size as batch size. + - logit (:obj:`torch.Tensor`): The predicted action logit tensor, for discrete action space, it will be \ + the same dimension real-value ranged tensor of possible action choices, and for continuous action \ + space, it will be the mu and sigma of the Gaussian distribution, and the number of mu and sigma is the \ + same as the number of continuous actions. Hybrid action space is a kind of combination of discrete \ + and continuous action space, so the logit will be a dict with ``action_type`` and ``action_args``. + - value (:obj:`torch.Tensor`): The predicted state value tensor. + - action_mask (:obj:`torch.Tensor`, optional): The action mask tensor, included if the input is a \ + dictionary containing 'action_mask'. Shapes: - - logit (:obj:`torch.FloatTensor`): :math:`(B, N)`, where B is batch size and N is ``action_shape`` - - value (:obj:`torch.FloatTensor`): :math:`(B, )`, where B is batch size. + - logit (:obj:`torch.Tensor`): :math:`(B, N)`, where B is batch size and N is ``action_shape`` + - value (:obj:`torch.Tensor`): :math:`(B, )`, where B is batch size, (B, 1) is squeezed to (B, ). Examples: - >>> model = VAC(64,64) + >>> model = VAC(64, 64) >>> inputs = torch.randn(4, 64) >>> outputs = model(inputs,'compute_actor_critic') - >>> outputs['value'] - tensor([0.0252, 0.0235, 0.0201, 0.0072], grad_fn=) + >>> assert critic_outputs['value'].shape == torch.Size([4]) >>> assert outputs['logit'].shape == torch.Size([4, 64]) .. note:: - ``compute_actor_critic`` interface aims to save computation when shares encoder. - Returning the combination dictionry. - + ``compute_actor_critic`` interface aims to save computation when shares encoder and return the combination \ + dict output. """ - if self.share_encoder: - actor_embedding = critic_embedding = self.encoder(x) + if isinstance(x, dict): + action_mask = x['action_mask'] + if self.share_encoder: + actor_embedding = critic_embedding = self.encoder(x['observation']) + else: + actor_embedding = self.actor_encoder(x['observation']) + critic_embedding = self.critic_encoder(x['observation']) else: - actor_embedding = self.actor_encoder(x) - critic_embedding = self.critic_encoder(x) + action_mask = None + if self.share_encoder: + actor_embedding = critic_embedding = self.encoder(x) + else: + actor_embedding = self.actor_encoder(x) + critic_embedding = self.critic_encoder(x) value = self.critic_head(critic_embedding)['pred'] if self.action_space == 'discrete': logit = self.actor_head(actor_embedding)['logit'] - return {'logit': logit, 'value': value} + result = {'logit': logit, 'value': value} + if action_mask is not None: + result['action_mask'] = action_mask + return result elif self.action_space == 'continuous': x = self.actor_head(actor_embedding) return {'logit': x, 'value': value} @@ -356,3 +381,75 @@ def compute_actor_critic(self, x: torch.Tensor) -> Dict: action_type = self.actor_head[0](actor_embedding) action_args = self.actor_head[1](actor_embedding) return {'logit': {'action_type': action_type['logit'], 'action_args': action_args}, 'value': value} + + +@MODEL_REGISTRY.register('dreamervac') +class DREAMERVAC(nn.Module): + """ + Overview: + The neural network and computation graph of DreamerV3 (state) Value Actor-Critic (VAC). + This model now supports discrete, continuous action space. + Interfaces: + ``__init__``, ``forward``. + """ + mode = ['compute_actor', 'compute_critic', 'compute_actor_critic'] + + def __init__( + self, + action_shape: Union[int, SequenceType, EasyDict], + dyn_stoch=32, + dyn_deter=512, + dyn_discrete=32, + actor_layers=2, + value_layers=2, + units=512, + act='SiLU', + norm='LayerNorm', + actor_dist='normal', + actor_init_std=1.0, + actor_min_std=0.1, + actor_max_std=1.0, + actor_temp=0.1, + action_unimix_ratio=0.01, + ) -> None: + """ + Overview: + Initialize the ``DREAMERVAC`` model according to arguments. + Arguments: + - obs_shape (:obj:`Union[int, SequenceType]`): Observation space shape, such as 8 or [4, 84, 84]. + - action_shape (:obj:`Union[int, SequenceType]`): Action space shape, such as 6 or [2, 3, 3]. + """ + super(DREAMERVAC, self).__init__() + action_shape = squeeze(action_shape) + self.action_shape = action_shape + + if dyn_discrete: + feat_size = dyn_stoch * dyn_discrete + dyn_deter + else: + feat_size = dyn_stoch + dyn_deter + self.actor = ActionHead( + feat_size, # pytorch version + action_shape, + actor_layers, + units, + act, + norm, + actor_dist, + actor_init_std, + actor_min_std, + actor_max_std, + actor_temp, + outscale=1.0, + unimix_ratio=action_unimix_ratio, + ) + self.critic = DenseHead( + feat_size, # pytorch version + (255, ), + value_layers, + units, + 'SiLU', # act + 'LN', # norm + 'twohot_symlog', + outscale=0.0, + device='cuda' if torch.cuda.is_available() else 'cpu', + ) diff --git a/ding/model/template/vae.py b/ding/model/template/vae.py index c8d7546ddc..f3181361c7 100644 --- a/ding/model/template/vae.py +++ b/ding/model/template/vae.py @@ -83,7 +83,7 @@ def encode(self, input: Dict[str, Tensor]) -> Dict[str, Any]: return {'mu': mu, 'log_var': log_var, 'obs_encoding': obs_encoding} def decode(self, z: Tensor, obs_encoding: Tensor) -> Dict[str, Any]: - r""" + """ Overview: Maps the given latent action and obs_encoding onto the original action space. Arguments: @@ -108,7 +108,7 @@ def decode(self, z: Tensor, obs_encoding: Tensor) -> Dict[str, Any]: return {'reconstruction_action': reconstruction_action, 'predition_residual': predition_residual} def decode_with_obs(self, z: Tensor, obs: Tensor) -> Dict[str, Any]: - r""" + """ Overview: Maps the given latent action and obs onto the original action space. Using the method self.encode_obs_head(obs) to get the obs_encoding. @@ -136,7 +136,7 @@ def decode_with_obs(self, z: Tensor, obs: Tensor) -> Dict[str, Any]: return {'reconstruction_action': reconstruction_action, 'predition_residual': predition_residual} def reparameterize(self, mu: Tensor, logvar: Tensor) -> Tensor: - r""" + """ Overview: Reparameterization trick to sample from N(mu, var) from N(0,1). Arguments: @@ -184,22 +184,17 @@ def forward(self, input: Dict[str, Tensor], **kwargs) -> dict: 'z': z } - def loss_function(self, args: Dict[str, Tensor], **kwargs) -> dict: + def loss_function(self, args: Dict[str, Tensor], **kwargs) -> Dict[str, Tensor]: """ Overview: Computes the VAE loss function. - KL(N(\mu, \sigma), N(0, 1)) = \log \frac{1}{\sigma} + \frac{\sigma^2 + \mu^2}{2} - \frac{1}{2} Arguments: - - args (:obj:`Dict`): Dict containing keywords `recons_action` (:obj:`torch.Tensor`) \ - and `prediction_residual` (:obj:`torch.Tensor`), `original_action` (:obj:`torch.Tensor`), \ - `mu` (:obj:`torch.Tensor`), `log_var` (:obj:`torch.Tensor`) and \ - `true_residual` (:obj:`torch.Tensor`). - - kwargs (:obj:`Dict`): Dict containing keywords `kld_weight` (:obj:`torch.Tensor`) \ - and `predict_weight` (:obj:`torch.Tensor`). + - args (:obj:`Dict[str, Tensor]`): Dict containing keywords ``recons_action``, ``prediction_residual`` \ + ``original_action``, ``mu``, ``log_var`` and ``true_residual``. + - kwargs (:obj:`Dict`): Dict containing keywords ``kld_weight`` and ``predict_weight``. Returns: - - outputs (:obj: `Dict`): Dict containing keywords `loss` \ - (`obj`:`torch.Tensor`), `reconstruction_loss` (:obj: `torch.Tensor`), \ - `kld_loss` (:obj: `torch.Tensor`) and `predict_loss` (:obj: `torch.Tensor`). + - outputs (:obj:`Dict[str, Tensor]`): Dict containing different ``loss`` results, including ``loss``, \ + ``reconstruction_loss``, ``kld_loss``, ``predict_loss``. Shapes: - recons_action (:obj:`torch.Tensor`): :math:`(B, A)`, where B is batch size \ and A is ``action dim``. diff --git a/ding/model/template/wqmix.py b/ding/model/template/wqmix.py index c5307d9bbf..251c025f9e 100644 --- a/ding/model/template/wqmix.py +++ b/ding/model/template/wqmix.py @@ -13,21 +13,24 @@ class MixerStar(nn.Module): """ Overview: - mixer network for Q_star in WQMIX , which mix up the independent q_value of - each agent to a total q_value and is diffrent from the Qmix's mixer network, - here the mixing network is a feedforward network with 3 hidden layers of 256 dim. + Mixer network for Q_star in WQMIX(https://arxiv.org/abs/2006.10800), which mix up the independent q_value of \ + each agent to a total q_value and is diffrent from the QMIX's mixer network, \ + here the mixing network is a feedforward network with 3 hidden layers of 256 dim. \ + This Q_star mixing network is not constrained to be monotonic by using non-negative weights and \ + having the state and agent_q be inputs, as opposed to having hypernetworks take the state as input \ + and generate the weights in QMIX. Interface: - __init__, forward + ``__init__``, ``forward``. """ def __init__(self, agent_num: int, state_dim: int, mixing_embed_dim: int) -> None: """ Overview: - initialize the mixer network of Q_star in WQMIX. + Initialize the mixer network of Q_star in WQMIX. Arguments: - - agent_num (:obj:`int`): the number of agent - - state_dim(:obj:`int`): the dimension of global observation state - - mixing_embed_dim (:obj:`int`): the dimension of mixing state emdedding + - agent_num (:obj:`int`): The number of agent, e.g., 8. + - state_dim(:obj:`int`): The dimension of global observation state, e.g., 16. + - mixing_embed_dim (:obj:`int`): The dimension of mixing state emdedding, e.g., 128. """ super(MixerStar, self).__init__() self.agent_num = agent_num @@ -46,17 +49,18 @@ def __init__(self, agent_num: int, state_dim: int, mixing_embed_dim: int) -> Non def forward(self, agent_qs: torch.FloatTensor, states: torch.FloatTensor) -> torch.FloatTensor: """ Overview: - forward computation graph of the mixer network for Q_star in WQMIX. + Forward computation graph of the mixer network for Q_star in WQMIX. This mixer network for \ + is a feed-forward network that takes the state and the appropriate actions' utilities as input. Arguments: - - agent_qs (:obj:`torch.FloatTensor`): the independent q_value of each agent - - states (:obj:`torch.FloatTensor`): the emdedding vector of global state + - agent_qs (:obj:`torch.FloatTensor`): The independent q_value of each agent. + - states (:obj:`torch.FloatTensor`): The emdedding vector of global state. Returns: - - q_tot (:obj:`torch.FloatTensor`): the total mixed q_value + - q_tot (:obj:`torch.FloatTensor`): The total mixed q_value. Shapes: - - agent_qs (:obj:`torch.FloatTensor`): :math:`(T,B, N)`, where T is timestep, - B is batch size, A is agent_num, N is obs_shape - - states (:obj:`torch.FloatTensor`): :math:`(T, B, M)`, where M is global_obs_shape - - q_tot (:obj:`torch.FloatTensor`): :math:`(T, B, )` + - agent_qs (:obj:`torch.FloatTensor`): :math:`(T,B, N)`, where T is timestep, \ + B is batch size, A is agent_num, N is obs_shape. + - states (:obj:`torch.FloatTensor`): :math:`(T, B, M)`, where M is global_obs_shape. + - q_tot (:obj:`torch.FloatTensor`): :math:`(T, B, )`. """ # in below annotations about the shape of the variables, T is timestep, # B is batch_size A is agent_num, N is obs_shape， for example, @@ -77,9 +81,13 @@ def forward(self, agent_qs: torch.FloatTensor, states: torch.FloatTensor) -> tor class WQMix(nn.Module): """ Overview: - WQMIX network, which is same as Qmix network + WQMIX (https://arxiv.org/abs/2006.10800) network, There are two components: \ + 1) Q_tot, which is same as QMIX network and composed of agent Q network and mixer network. \ + 2) An unrestricted joint action Q_star, which is composed of agent Q network and mixer_star network. \ + The QMIX paper mentions that all agents share local Q network parameters, so only one Q network is initialized \ + in Q_tot or Q_star. Interface: - __init__, forward, _setup_global_encoder + ``__init__``, ``forward``. """ def __init__( @@ -94,15 +102,19 @@ def __init__( ) -> None: """ Overview: - initialize Qmix network + Initialize WQMIX neural network according to arguments, i.e. agent Q network and mixer, \ + Q_star network and mixer_star. Arguments: - - agent_num (:obj:`int`): the number of agent - - obs_shape (:obj:`int`): the dimension of each agent's observation state - - global_obs_shape (:obj:`int`): the dimension of global observation state - - action_shape (:obj:`int`): the dimension of action shape - - hidden_size_list (:obj:`list`): the list of hidden size - - lstm_type (:obj:`str`): use lstm or gru, default to gru - - dueling (:obj:`bool`): use dueling head or not, default to False. + - agent_num (:obj:`int`): The number of agent, such as 8. + - obs_shape (:obj:`int`): The dimension of each agent's observation state, such as 8. + - global_obs_shape (:obj:`int`): The dimension of global observation state, such as 8. + - action_shape (:obj:`int`): The dimension of action shape, such as 6. + - hidden_size_list (:obj:`list`): The list of hidden size for ``q_network``, \ + the last element must match mixer's ``mixing_embed_dim``. + - lstm_type (:obj:`str`): The type of RNN module in ``q_network``, now support \ + ['normal', 'pytorch', 'gru'], default to gru. + - dueling (:obj:`bool`): Whether choose ``DuelingHead`` (True) or ``DiscreteHead (False)``, \ + default to False. """ super(WQMix, self).__init__() self._act = nn.ReLU() @@ -118,34 +130,36 @@ def __init__( def forward(self, data: dict, single_step: bool = True, q_star: bool = False) -> dict: """ Overview: - forward computation graph of qmix network + Forward computation graph of qmix network. Input dict including time series observation and \ + related data to predict total q_value and each agent q_value. Determine whether to calculate \ + Q_tot or Q_star based on the ``q_star`` parameter. Arguments: - - data (:obj:`dict`): input data dict with keys ['obs', 'prev_state', 'action'] - - agent_state (:obj:`torch.Tensor`): each agent local state(obs) - - global_state (:obj:`torch.Tensor`): global state(obs) - - prev_state (:obj:`list`): previous rnn state - - action (:obj:`torch.Tensor` or None): if action is None, use argmax q_value index as action to\ - calculate ``agent_q_act`` - - single_step (:obj:`bool`): whether single_step forward, if so, add timestep dim before forward and\ - remove it after forward - - Q_star (:obj:`bool`): whether Q_star network forward. If True, using the Q_star network, where the\ + - data (:obj:`dict`): Input data dict with keys ['obs', 'prev_state', 'action']. + - agent_state (:obj:`torch.Tensor`): Time series local observation data of each agents. + - global_state (:obj:`torch.Tensor`): Time series global observation data. + - prev_state (:obj:`list`): Previous rnn state for ``q_network`` or ``_q_network_star``. + - action (:obj:`torch.Tensor` or None): If action is None, use argmax q_value index as action to\ + calculate ``agent_q_act``. + - single_step (:obj:`bool`): Whether single_step forward, if so, add timestep dim before forward and\ + remove it after forward. + - Q_star (:obj:`bool`): Whether Q_star network forward. If True, using the Q_star network, where the\ agent networks have the same architecture as Q network but do not share parameters and the mixing\ network is a feedforward network with 3 hidden layers of 256 dim; if False, using the Q network,\ same as the Q network in Qmix paper. Returns: - - ret (:obj:`dict`): output data dict with keys [``total_q``, ``logit``, ``next_state``] - - total_q (:obj:`torch.Tensor`): total q_value, which is the result of mixer network - - agent_q (:obj:`torch.Tensor`): each agent q_value - - next_state (:obj:`list`): next rnn state + - ret (:obj:`dict`): Output data dict with keys [``total_q``, ``logit``, ``next_state``]. + - total_q (:obj:`torch.Tensor`): Total q_value, which is the result of mixer network. + - agent_q (:obj:`torch.Tensor`): Each agent q_value. + - next_state (:obj:`list`): Next rnn state. Shapes: - agent_state (:obj:`torch.Tensor`): :math:`(T, B, A, N)`, where T is timestep, B is batch_size\ - A is agent_num, N is obs_shape - - global_state (:obj:`torch.Tensor`): :math:`(T, B, M)`, where M is global_obs_shape - - prev_state (:obj:`list`): math:`(T, B, A)`, a list of length B, and each element is a list of length A - - action (:obj:`torch.Tensor`): :math:`(T, B, A)` - - total_q (:obj:`torch.Tensor`): :math:`(T, B)` - - agent_q (:obj:`torch.Tensor`): :math:`(T, B, A, P)`, where P is action_shape - - next_state (:obj:`list`): math:`(T, B, A)`, a list of length B, and each element is a list of length A + A is agent_num, N is obs_shape. + - global_state (:obj:`torch.Tensor`): :math:`(T, B, M)`, where M is global_obs_shape. + - prev_state (:obj:`list`): math:`(T, B, A)`, a list of length B, and each element is a list of length A. + - action (:obj:`torch.Tensor`): :math:`(T, B, A)`. + - total_q (:obj:`torch.Tensor`): :math:`(T, B)`. + - agent_q (:obj:`torch.Tensor`): :math:`(T, B, A, P)`, where P is action_shape. + - next_state (:obj:`list`): math:`(T, B, A)`, a list of length B, and each element is a list of length A. """ if q_star: # forward using Q_star network agent_state, global_state, prev_state = data['obs']['agent_state'], data['obs']['global_state'], data[ @@ -163,7 +177,6 @@ def forward(self, data: dict, single_step: bool = True, q_star: bool = False) -> { 'obs': agent_state, 'prev_state': prev_state, - 'enable_fast_timestep': True } ) # here is the forward pass of the agent networks of Q_star agent_q, next_state = output['logit'], output['next_state'] @@ -209,7 +222,6 @@ def forward(self, data: dict, single_step: bool = True, q_star: bool = False) -> { 'obs': agent_state, 'prev_state': prev_state, - 'enable_fast_timestep': True } ) # here is the forward pass of the agent networks of Q agent_q, next_state = output['logit'], output['next_state'] @@ -239,15 +251,3 @@ def forward(self, data: dict, single_step: bool = True, q_star: bool = False) -> 'next_state': next_state, 'action_mask': data['obs']['action_mask'] } - - def _setup_global_encoder(self, global_obs_shape: int, embedding_size: int) -> torch.nn.Module: - """ - Overview: - Used to encoder global observation. - Arguments: - - global_obs_shape (:obj:`int`): the dimension of global observation state - - embedding_size (:obj:`int`): the dimension of state emdedding - Return: - - outputs (:obj:`torch.nn.Module`): Global observation encoding network - """ - return MLP(global_obs_shape, embedding_size, embedding_size, 2, activation=self._act) diff --git a/ding/model/wrapper/__init__.py b/ding/model/wrapper/__init__.py index f74d3f2668..24d621e973 100644 --- a/ding/model/wrapper/__init__.py +++ b/ding/model/wrapper/__init__.py @@ -1 +1 @@ -from .model_wrappers import model_wrap, register_wrapper, IModelWrapper, BaseModelWrapper +from .model_wrappers import model_wrap, register_wrapper, IModelWrapper diff --git a/ding/model/wrapper/model_wrappers.py b/ding/model/wrapper/model_wrappers.py index 50efbbeafb..1aa71b6597 100644 --- a/ding/model/wrapper/model_wrappers.py +++ b/ding/model/wrapper/model_wrappers.py @@ -1,40 +1,52 @@ -from typing import Any, Tuple, Callable, Optional, List, Dict +from typing import Any, Tuple, Callable, Optional, List, Dict, Union from abc import ABC import numpy as np import torch -from ding.torch_utils import get_tensor_data -from ding.rl_utils import create_noise_generator +import torch.nn as nn +import torch.nn.functional as F from torch.distributions import Categorical, Independent, Normal +from ding.torch_utils import get_tensor_data, zeros_like +from ding.rl_utils import create_noise_generator from ding.utils.data import default_collate -import torch.nn.functional as F class IModelWrapper(ABC): - r""" + """ Overview: - the base class of Model Wrappers + The basic interface class of model wrappers. Model wrapper is a wrapper class of torch.nn.Module model, which \ + is used to add some extra operations for the wrapped model, such as hidden state maintain for RNN-base model, \ + argmax action selection for discrete action space, etc. Interfaces: - register + ``__init__``, ``__getattr__``, ``info``, ``reset``, ``forward``. """ - def __init__(self, model: Any) -> None: + def __init__(self, model: nn.Module) -> None: + """ + Overview: + Initialize model and other necessary member variabls in the model wrapper. + """ self._model = model def __getattr__(self, key: str) -> Any: - r""" + """ Overview: - Get the attrbute in model. + Get original attrbutes of torch.nn.Module model, such as variables and methods defined in model. Arguments: - - key (:obj:`str`): The key to query. + - key (:obj:`str`): The string key to query. Returns: - ret (:obj:`Any`): The queried attribute. """ return getattr(self._model, key) - def info(self, attr_name): - r""" + def info(self, attr_name: str) -> str: + """ Overview: - get info of attr_name + Get some string information of the indicated ``attr_name``, which is used for debug wrappers. + This method will recursively search for the indicated ``attr_name``. + Arguments: + - attr_name (:obj:`str`): The string key to query information. + Returns: + - info_string (:obj:`str`): The information string of the indicated ``attr_name``. """ if attr_name in dir(self): if isinstance(self._model, IModelWrapper): @@ -50,36 +62,48 @@ def info(self, attr_name): else: return '{}'.format(self._model.__class__.__name__) - -class BaseModelWrapper(IModelWrapper): - r""" - Overview: - the base class of Model Wrappers - Interfaces: - register - """ - - def reset(self, data_id: List[int] = None) -> None: - r""" + def reset(self, data_id: List[int] = None, **kwargs) -> None: + """ Overview - the reset function that the Model Wrappers with states should implement - used to reset the stored states + Basic interface, reset some stateful varaibles in the model wrapper, such as hidden state of RNN. + Here we do nothing and just implement this interface method. + Other derived model wrappers can override this method to add some extra operations. + Arguments: + - data_id (:obj:`List[int]`): The data id list to reset. If None, reset all data. In practice, \ + model wrappers often needs to maintain some stateful variables for each data trajectory, \ + so we leave this ``data_id`` argument to reset the stateful variables of the indicated data. """ - pass + # This is necessary when multiple model wrappers. + if hasattr(self._model, 'reset'): + return self._model.reset(data_id=data_id, **kwargs) + def forward(self, *args, **kwargs) -> Any: + """ + Overview: + Basic interface, call the wrapped model's forward method. Other derived model wrappers can override this \ + method to add some extra operations. + """ + return self._model.forward(*args, **kwargs) -def zeros_like(h): - if isinstance(h, torch.Tensor): - return torch.zeros_like(h) - elif isinstance(h, (list, tuple)): - return [zeros_like(t) for t in h] - elif isinstance(h, dict): - return {k: zeros_like(v) for k, v in h.items()} - else: - raise TypeError("not support type: {}".format(h)) + +class BaseModelWrapper(IModelWrapper): + """ + Overview: + Placeholder class for the model wrapper. This class is used to wrap the model without any extra operations, \ + including a empty ``reset`` method and a ``forward`` method which directly call the wrapped model's forward. + To keep the consistency of the model wrapper interface, we use this class to wrap the model without specific \ + operations in the implementation of DI-engine's policy. + """ + pass class HiddenStateWrapper(IModelWrapper): + """ + Overview: + Maintain the hidden state for RNN-base model. Each sample in a batch has its own state. + Interfaces: + ``__init__``, ``reset``, ``forward``. + """ def __init__( self, @@ -96,9 +120,10 @@ def __init__( Arguments: - model(:obj:`Any`): Wrapped model class, should contain forward method. - state_num (:obj:`int`): Number of states to process. - - save_prev_state (:obj:`bool`): Whether to output the prev state in output['prev_state']. - - init_fn (:obj:`Callable`): The function which is used to init every hidden state when init and reset. \ - Default return None for hidden states. + - save_prev_state (:obj:`bool`): Whether to output the prev state in output. + - init_fn (:obj:`Callable`): The function which is used to init every hidden state when init and reset, \ + default return None for hidden states. + .. note:: 1. This helper must deal with an actual batch with some parts of samples, e.g: 6 samples of state_num 8. 2. This helper must deal with the single sample state reset. @@ -196,11 +221,11 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: """ Arguments: - - input_obs (:obj:`torch.Tensor`): Input observation without sequence shape: (bs, *obs_shape) - - only_last_logit (:obj:`bool`): if True 'logit' only contains the output corresponding to the current - observation (shape: bs, embedding_dim), otherwise logit has shape (seq_len, bs, embedding_dim) - - data_id (:obj:`List`): id of the envs that are currently running. Memory update and logits return has only - effect for those environments. If `None` it is considered that all envs are running. + - input_obs (:obj:`torch.Tensor`): Input observation without sequence shape: ``(bs, *obs_shape)``. + - only_last_logit (:obj:`bool`): if True 'logit' only contains the output corresponding to the current \ + observation (shape: bs, embedding_dim), otherwise logit has shape (seq_len, bs, embedding_dim). + - data_id (:obj:`List`): id of the envs that are currently running. Memory update and logits return has \ + only effect for those environments. If `None` it is considered that all envs are running. Returns: - Dictionary containing the input_sequence 'input_seq' stored in memory and the transformer output 'logit'. """ @@ -387,12 +412,18 @@ def sample_action(logit=None, prob=None): class ArgmaxSampleWrapper(IModelWrapper): - r""" + """ Overview: - Used to help the model to sample argmax action + Used to help the model to sample argmax action. + Interfaces: + ``forward``. """ def forward(self, *args, **kwargs): + """ + Overview: + Employ model forward computation graph, and use the output logit to greedily select max action (argmax). + """ output = self._model.forward(*args, **kwargs) assert isinstance(output, dict), "model output must be dict, but find {}".format(type(output)) logit = output['logit'] @@ -411,11 +442,64 @@ def forward(self, *args, **kwargs): return output +class CombinationArgmaxSampleWrapper(IModelWrapper): + r""" + Overview: + Used to help the model to sample combination argmax action. + Interfaces: + ``forward``. + """ + + def forward(self, shot_number, *args, **kwargs): + output = self._model.forward(*args, **kwargs) + # Generate actions. + act = [] + mask = torch.zeros_like(output['logit']) + for ii in range(shot_number): + masked_logit = output['logit'] + mask + actions = masked_logit.argmax(dim=-1) + act.append(actions) + for jj in range(actions.shape[0]): + mask[jj][actions[jj]] = -1e8 + # `act` is shaped: (B, shot_number) + act = torch.stack(act, dim=1) + output['action'] = act + return output + + +class CombinationMultinomialSampleWrapper(IModelWrapper): + r""" + Overview: + Used to help the model to sample combination multinomial action. + Interfaces: + ``forward``. + """ + + def forward(self, shot_number, *args, **kwargs): + output = self._model.forward(*args, **kwargs) + # Generate actions. + act = [] + mask = torch.zeros_like(output['logit']) + for ii in range(shot_number): + dist = torch.distributions.Categorical(logits=output['logit'] + mask) + actions = dist.sample() + act.append(actions) + for jj in range(actions.shape[0]): + mask[jj][actions[jj]] = -1e8 + + # `act` is shaped: (B, shot_number) + act = torch.stack(act, dim=1) + output['action'] = act + return output + + class HybridArgmaxSampleWrapper(IModelWrapper): r""" Overview: Used to help the model to sample argmax action in hybrid action space, i.e.{'action_type': discrete, 'action_args', continuous} + Interfaces: + ``forward``. """ def forward(self, *args, **kwargs): @@ -440,11 +524,11 @@ def forward(self, *args, **kwargs): class MultinomialSampleWrapper(IModelWrapper): - r""" + """ Overview: - Used to help the model get the corresponding action from the output['logits'] + Used to help the model get the corresponding action from the output['logits']self. Interfaces: - register + ``forward``. """ def forward(self, *args, **kwargs): @@ -482,7 +566,7 @@ class EpsGreedySampleWrapper(IModelWrapper): - float (i.e. python native scalar): for almost normal case - Dict[str, float]: for algorithm NGU Interfaces: - register + ``forward``. """ def forward(self, *args, **kwargs): @@ -536,7 +620,7 @@ class EpsGreedyMultinomialSampleWrapper(IModelWrapper): Epsilon greedy sampler coupled with multinomial sample used in collector_model to help balance exploration and exploitation. Interfaces: - register + ``forward``. """ def forward(self, *args, **kwargs): @@ -583,7 +667,7 @@ class HybridEpsGreedySampleWrapper(IModelWrapper): Epsilon greedy sampler used in collector_model to help balance exploration and exploitation. In hybrid action space, i.e.{'action_type': discrete, 'action_args', continuous} Interfaces: - register, forward + ``forward``. """ def forward(self, *args, **kwargs): @@ -623,7 +707,7 @@ class HybridEpsGreedyMultinomialSampleWrapper(IModelWrapper): to help balance exploration and exploitation. In hybrid action space, i.e.{'action_type': discrete, 'action_args', continuous} Interfaces: - register + ``forward``. """ def forward(self, *args, **kwargs): @@ -712,7 +796,7 @@ def forward(self, *args, **kwargs): return output -class DeterministicSample(IModelWrapper): +class DeterministicSampleWrapper(IModelWrapper): """ Overview: Deterministic sampler (just use mu directly) used in eval_model. @@ -727,7 +811,7 @@ def forward(self, *args, **kwargs): return output -class ReparamSample(IModelWrapper): +class ReparamSampleWrapper(IModelWrapper): """ Overview: Reparameterization gaussian sampler used in collector_model. @@ -749,7 +833,7 @@ class ActionNoiseWrapper(IModelWrapper): Overview: Add noise to collector's action output; Do clips on both generated noise and action after adding noise. Interfaces: - register, __init__, add_noise, reset + ``__init__``, ``forward``. Arguments: - model (:obj:`Any`): Wrapped model class. Should contain ``forward`` method. - noise_type (:obj:`str`): The type of noise that should be generated, support ['gauss', 'ou']. @@ -784,10 +868,14 @@ def forward(self, *args, **kwargs): assert isinstance(output, dict), "model output must be dict, but find {}".format(type(output)) if 'action' in output or 'action_args' in output: key = 'action' if 'action' in output else 'action_args' - action = output[key] + # handle hybrid action space by adding noise to continuous part of model output + action = output[key]['action_args'] if isinstance(output[key], dict) else output[key] assert isinstance(action, torch.Tensor) action = self.add_noise(action) - output[key] = action + if isinstance(output[key], dict): + output[key]['action_args'] = action + else: + output[key] = action return output def add_noise(self, action: torch.Tensor) -> torch.Tensor: @@ -807,13 +895,6 @@ def add_noise(self, action: torch.Tensor) -> torch.Tensor: action = action.clamp(self.action_range['min'], self.action_range['max']) return action - def reset(self) -> None: - r""" - Overview: - Reset noise generator. - """ - pass - class TargetNetworkWrapper(IModelWrapper): r""" @@ -872,17 +953,15 @@ def reset_state(self, target_update_count: int = None) -> None: class TeacherNetworkWrapper(IModelWrapper): - r""" + """ Overview: Set the teacher Network. Set the model's model.teacher_cfg to the input teacher_cfg - - Interfaces: - register """ def __init__(self, model, teacher_cfg): super().__init__(model) self._model._teacher_cfg = teacher_cfg + raise NotImplementedError wrapper_name_map = { @@ -892,8 +971,8 @@ def __init__(self, model, teacher_cfg): 'hybrid_argmax_sample': HybridArgmaxSampleWrapper, 'eps_greedy_sample': EpsGreedySampleWrapper, 'eps_greedy_multinomial_sample': EpsGreedyMultinomialSampleWrapper, - 'deterministic_sample': DeterministicSample, - 'reparam_sample': ReparamSample, + 'deterministic_sample': DeterministicSampleWrapper, + 'reparam_sample': ReparamSampleWrapper, 'hybrid_eps_greedy_sample': HybridEpsGreedySampleWrapper, 'hybrid_eps_greedy_multinomial_sample': HybridEpsGreedyMultinomialSampleWrapper, 'hybrid_reparam_multinomial_sample': HybridReparamMultinomialSampleWrapper, @@ -906,11 +985,24 @@ def __init__(self, model, teacher_cfg): # model wrapper 'target': TargetNetworkWrapper, 'teacher': TeacherNetworkWrapper, + 'combination_argmax_sample': CombinationArgmaxSampleWrapper, + 'combination_multinomial_sample': CombinationMultinomialSampleWrapper, } -def model_wrap(model, wrapper_name: str = None, **kwargs): +def model_wrap(model: Union[nn.Module, IModelWrapper], wrapper_name: str = None, **kwargs): + """ + Overview: + Wrap the model with the specified wrapper and return the wrappered model. + Arguments: + - model (:obj:`Any`): The model to be wrapped. + - wrapper_name (:obj:`str`): The name of the wrapper to be used. + + .. note:: + The arguments of the wrapper should be passed in as kwargs. + """ if wrapper_name in wrapper_name_map: + # TODO test whether to remove this if branch if not isinstance(model, IModelWrapper): model = wrapper_name_map['base'](model) model = wrapper_name_map[wrapper_name](model, **kwargs) @@ -919,13 +1011,15 @@ def model_wrap(model, wrapper_name: str = None, **kwargs): return model -def register_wrapper(name: str, wrapper_type: type): - r""" +def register_wrapper(name: str, wrapper_type: type) -> None: + """ Overview: - Register new wrapper to wrapper_name_map + Register new wrapper to ``wrapper_name_map``. When user implements a new wrapper, they must call this function \ + to complete the registration. Then the wrapper can be called by ``model_wrap``. Arguments: - - name (:obj:`str`): the name of the wrapper - - wrapper_type (subclass of :obj:`IModelWrapper`): the wrapper class added to the plguin_name_map + - name (:obj:`str`): The name of the new wrapper to be registered. + - wrapper_type (:obj:`type`): The wrapper class needs to be added in ``wrapper_name_map``. This argument \ + should be the subclass of ``IModelWrapper``. """ assert isinstance(name, str) assert issubclass(wrapper_type, IModelWrapper) diff --git a/ding/model/wrapper/test_model_wrappers.py b/ding/model/wrapper/test_model_wrappers.py index 274d2801bb..3c0dc73aa2 100644 --- a/ding/model/wrapper/test_model_wrappers.py +++ b/ding/model/wrapper/test_model_wrappers.py @@ -9,7 +9,8 @@ from ding.torch_utils import get_lstm from ding.torch_utils.network.gtrxl import GTrXL -from ding.model import model_wrap, register_wrapper, IModelWrapper, BaseModelWrapper +from ding.model import model_wrap, register_wrapper, IModelWrapper +from ding.model.wrapper.model_wrappers import BaseModelWrapper class TempMLP(torch.nn.Module): @@ -38,7 +39,7 @@ def __init__(self): self.bn1 = nn.BatchNorm1d(4) self.fc2 = nn.Linear(4, 6) self.act = nn.ReLU() - self.out = nn.Softmax() + self.out = nn.Softmax(dim=-1) def forward(self, inputs, tmp=0): x = self.fc1(inputs['obs']) @@ -61,7 +62,7 @@ def __init__(self): self.bn1 = nn.BatchNorm1d(4) self.fc2 = nn.Linear(4, 6) self.act = nn.ReLU() - self.out = nn.Softmax() + self.out = nn.Softmax(dim=-1) self.fc2_cont = nn.Linear(4, 6) self.act_cont = nn.ReLU() @@ -93,7 +94,7 @@ def __init__(self): self.bn1 = nn.BatchNorm1d(4) self.fc2 = nn.Linear(4, 6) self.act = nn.ReLU() - self.out = nn.Softmax() + self.out = nn.Softmax(dim=-1) self.fc2_cont_mu = nn.Linear(4, 6) self.act_cont_mu = nn.ReLU() @@ -131,7 +132,6 @@ def __init__(self): self.bn1 = nn.BatchNorm1d(4) self.fc2 = nn.Linear(4, 6) self.act = nn.ReLU() - self.out = nn.Softmax() self.fc2_cont_mu = nn.Linear(4, 6) self.fc2_cont_sigma = nn.Linear(4, 6) @@ -195,6 +195,20 @@ def forward(self, data): return {'output': output, 'next_state': next_state} +class TempLSTMActor(torch.nn.Module): + + def __init__(self): + super(TempLSTMActor, self).__init__() + self.model = get_lstm(lstm_type='pytorch', input_size=36, hidden_size=32, num_layers=2, norm_type=None) + + def forward(self, data, tmp=0): + output, next_state = self.model(data['f'], data['prev_state'], list_next_state=True) + ret = {'logit': output, 'tmp': tmp, 'action': output + torch.rand_like(output), 'next_state': next_state} + if 'mask' in data: + ret['action_mask'] = data['mask'] + return ret + + @pytest.fixture(scope='function') def setup_model(): return torch.nn.Linear(3, 6) @@ -514,12 +528,23 @@ def test_transformer_input_wrapper(self): model.reset() assert model.obs_memory is None + def test_transformer_segment_wrapper(self): + seq_len, bs, obs_shape = 12, 8, 32 + layer_num, memory_len, emb_dim = 3, 4, 4 + model = GTrXL(input_dim=obs_shape, embedding_dim=emb_dim, memory_len=memory_len, layer_num=layer_num) + model = model_wrap(model, wrapper_name='transformer_segment', seq_len=seq_len) + inputs1 = torch.randn((seq_len, bs, obs_shape)) + out = model.forward(inputs1) + info = model.info('info') + info = model.info('x') + def test_transformer_memory_wrapper(self): seq_len, bs, obs_shape = 12, 8, 32 layer_num, memory_len, emb_dim = 3, 4, 4 model = GTrXL(input_dim=obs_shape, embedding_dim=emb_dim, memory_len=memory_len, layer_num=layer_num) model1 = model_wrap(model, wrapper_name='transformer_memory', batch_size=bs) model2 = model_wrap(model, wrapper_name='transformer_memory', batch_size=bs) + model1.show_memory_occupancy() inputs1 = torch.randn((seq_len, bs, obs_shape)) out = model1.forward(inputs1) new_memory1 = model1.memory @@ -549,3 +574,46 @@ def test_transformer_memory_wrapper(self): assert sum(new_memory2[:, -16:].flatten()) != 0 assert sum(new_memory2[:, :-16].flatten()) == 0 assert torch.all(torch.eq(new_memory1[:, -8:], new_memory2[:, -16:-8])) + + def test_combination_argmax_sample_wrapper(self): + model = model_wrap(ActorMLP(), wrapper_name='combination_argmax_sample') + data = {'obs': torch.randn(4, 3)} + shot_number = 2 + output = model.forward(shot_number=shot_number, inputs=data) + assert output['action'].shape == (4, shot_number) + assert (output['action'] >= 0).all() and (output['action'] < 64).all() + + def test_combination_multinomial_sample_wrapper(self): + model = model_wrap(ActorMLP(), wrapper_name='combination_multinomial_sample') + data = {'obs': torch.randn(4, 3)} + shot_number = 2 + output = model.forward(shot_number=shot_number, inputs=data) + assert output['action'].shape == (4, shot_number) + assert (output['action'] >= 0).all() and (output['action'] < 64).all() + + def test_hidden_state_and_epsilon_greedy_wrapper(self): + model = model_wrap(TempLSTMActor(), wrapper_name='hidden_state', state_num=4, save_prev_state=True) + model = model_wrap(model, wrapper_name='eps_greedy_sample') + model.reset() + # Check that reset properly initializes all states to None + assert all([isinstance(s, type(None)) for s in model._state.values()]) + + data = {'f': torch.randn(2, 4, 36)} + output = model.forward(data, eps=0.8) + assert output['tmp'] == 0 + assert 'logit' in output + assert output['logit'].shape == (2, 4, 32) + assert 'action' in output + assert output['action'].shape == (2, 4) + assert 'prev_state' in output + assert len(output['prev_state']) == 4 + assert output['prev_state'][0]['h'].shape == (2, 1, 32) + assert output['prev_state'][0]['c'].shape == (2, 1, 32) + + assert all([isinstance(s, dict) for s in model._state.values()]) + # Check that reset with specific data_id works + model.reset(data_id=[0, 2]) + assert isinstance(model._state[0], type(None)) + assert isinstance(model._state[2], type(None)) + assert isinstance(model._state[1], dict) + assert isinstance(model._state[3], dict) diff --git a/ding/policy/__init__.py b/ding/policy/__init__.py old mode 100644 new mode 100755 index 3b64eab610..5015823d71 --- a/ding/policy/__init__.py +++ b/ding/policy/__init__.py @@ -1,5 +1,7 @@ from .base_policy import Policy, CommandModePolicy, create_policy, get_policy_cls +from .common_utils import single_env_forward_wrapper, single_env_forward_wrapper_ttorch, default_preprocess_learn from .dqn import DQNSTDIMPolicy, DQNPolicy +from .mdqn import MDQNPolicy from .iqn import IQNPolicy from .fqf import FQFPolicy from .qrdqn import QRDQNPolicy @@ -9,18 +11,22 @@ from .d4pg import D4PGPolicy from .td3 import TD3Policy from .td3_vae import TD3VAEPolicy - from .td3_bc import TD3BCPolicy +from .dt import DTPolicy + +from .pg import PGPolicy from .a2c import A2CPolicy from .ppo import PPOPolicy, PPOPGPolicy, PPOOffPolicy -from .sac import SACPolicy, SACDiscretePolicy, SQILSACPolicy -from .cql import CQLPolicy, CQLDiscretePolicy +from .sac import SACPolicy, DiscreteSACPolicy, SQILSACPolicy +from .cql import CQLPolicy, DiscreteCQLPolicy +from .edac import EDACPolicy from .impala import IMPALAPolicy from .ngu import NGUPolicy from .r2d2 import R2D2Policy from .r2d2_gtrxl import R2D2GTrXLPolicy from .ppg import PPGPolicy, PPGOffPolicy from .sqn import SQNPolicy +from .bdq import BDQPolicy from .qmix import QMIXPolicy from .wqmix import WQMIXPolicy @@ -41,3 +47,14 @@ from .bc import BehaviourCloningPolicy from .ibc import IBCPolicy + +from .pc import ProcedureCloningBFSPolicy + +from .bcq import BCQPolicy +from .qgpo import QGPOPolicy + +# new-type policy +from .ppof import PPOFPolicy +from .prompt_pg import PromptPGPolicy +from .prompt_awr import PromptAWRPolicy +from .happo import HAPPOPolicy diff --git a/ding/policy/a2c.py b/ding/policy/a2c.py index 389440207e..2d2f116afc 100644 --- a/ding/policy/a2c.py +++ b/ding/policy/a2c.py @@ -1,10 +1,12 @@ -from typing import List, Dict, Any, Tuple, Union from collections import namedtuple +from typing import List, Dict, Any, Tuple + import torch -from ding.rl_utils import a2c_data, a2c_error, get_gae_with_default_last_value, get_train_sample -from ding.torch_utils import Adam, to_device from ding.model import model_wrap +from ding.rl_utils import a2c_data, a2c_error, get_gae_with_default_last_value, get_train_sample, \ + a2c_error_continuous +from ding.torch_utils import Adam, to_device from ding.utils import POLICY_REGISTRY, split_data_generator from ding.utils.data import default_collate, default_decollate from .base_policy import Policy @@ -13,68 +15,97 @@ @POLICY_REGISTRY.register('a2c') class A2CPolicy(Policy): - r""" + """ Overview: - Policy class of A2C algorithm. + Policy class of A2C (Advantage Actor-Critic) algorithm, proposed in https://arxiv.org/abs/1602.01783. """ config = dict( - # (string) RL policy register name (refer to function "register_policy"). + # (str) Name of the registered RL policy (refer to the "register_policy" function). type='a2c', - # (bool) Whether to use cuda for network. + # (bool) Flag to enable CUDA for model computation. cuda=False, - # (bool) whether use on-policy training pipeline(behaviour policy and training policy are the same) - on_policy=True, # for a2c strictly on policy algorithm, this line should not be seen by users + # (bool) Flag for using on-policy training (training policy is the same as the behavior policy). + on_policy=True, + # (bool) Flag for enabling priority experience replay. Must be False when priority_IS_weight is False. priority=False, - # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. + # (bool) Flag for using Importance Sampling weights to correct updates. Requires `priority` to be True. priority_IS_weight=False, + # (str) Type of action space used in the policy, with valid options ['discrete', 'continuous']. + action_space='discrete', + # learn_mode configuration learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, - # (int) for a2c, update_per_collect must be 1. - update_per_collect=1, # fixed value, this line should not be modified by users + # (int) Number of updates per data collection. A2C requires this to be set to 1. + update_per_collect=1, + # (int) Batch size for learning. batch_size=64, + # (float) Learning rate for optimizer. learning_rate=0.001, - # (List[float]) + # (Tuple[float, float]) Coefficients used for computing running averages of gradient and its square. betas=(0.9, 0.999), - # (float) + # (float) Term added to the denominator to improve numerical stability in optimizer. eps=1e-8, - # (float) + # (float) Maximum norm for gradients. grad_norm=0.5, - # ============================================================== - # The following configs is algorithm-specific - # ============================================================== - # (float) loss weight of the value network, the weight of policy network is set to 1 + # (float) Scaling factor for value network loss relative to policy network loss. value_weight=0.5, - # (float) loss weight of the entropy regularization, the weight of policy network is set to 1 + # (float) Weight of entropy regularization in the loss function. entropy_weight=0.01, - # (bool) Whether to normalize advantage. Default to False. + # (bool) Flag to enable normalization of advantages. adv_norm=False, + # (bool) If set to True, the 'done' signals that indicate the end of an episode due to environment time + # limits are disregarded. By default, this is set to False. This setting is particularly useful for tasks + # that have a predetermined episode length, such as HalfCheetah and various other MuJoCo environments, + # where the maximum length is capped at 1000 steps. When enabled, any 'done' signal triggered by reaching + # the maximum episode steps will be overridden to 'False'. This ensures the accurate calculation of the + # Temporal Difference (TD) error, using the formula `gamma * (1 - done) * next_v + reward`, + # even when the episode surpasses the predefined step limit. ignore_done=False, ), + # collect_mode configuration collect=dict( - # (int) collect n_sample data, train model n_iteration times - # n_sample=80, + # (int) The length of rollout for data collection. unroll_len=1, - # ============================================================== - # The following configs is algorithm-specific - # ============================================================== - # (float) discount factor for future reward, defaults int [0, 1] + # (float) Discount factor for calculating future rewards, typically in the range [0, 1]. discount_factor=0.9, - # (float) the trade-off factor lambda to balance 1step td and mc + # (float) Trade-off parameter for balancing TD-error and Monte Carlo error in GAE. gae_lambda=0.95, ), + # eval_mode configuration (kept empty for compatibility purposes) eval=dict(), ) def default_model(self) -> Tuple[str, List[str]]: + """ + Overview: + Returns the default model configuration used by the A2C algorithm. ``__init__`` method will \ + automatically call this method to get the default model setting and create model. + + Returns: + - model_info (:obj:`Tuple[str, List[str]]`): \ + Tuple containing the registered model name and model's import_names. + """ return 'vac', ['ding.model.template.vac'] def _init_learn(self) -> None: - r""" + """ Overview: - Learn mode init method. Called by ``self.__init__``. - Init the optimizer, algorithm config, main and target models. + Initialize the learn mode of policy, including related attributes and modules. For A2C, it mainly \ + contains optimizer, algorithm-specific arguments such as value_weight, entropy_weight, adv_norm + and grad_norm, and main model. \ + This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. """ + assert self._cfg.action_space in ["continuous", "discrete"] # Optimizer self._optimizer = Adam( self._model.parameters(), @@ -95,15 +126,32 @@ def _init_learn(self) -> None: self._learn_model = model_wrap(self._model, wrapper_name='base') self._learn_model.reset() - def _forward_learn(self, data: dict) -> Dict[str, Any]: - r""" + def _forward_learn(self, data: List[Dict[str, Any]]) -> Dict[str, Any]: + """ Overview: - Forward and backward function of learn mode. + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data from the replay buffer and then returns the output \ + result, including various training information such as policy_loss, value_loss, entropy_loss. Arguments: - - data (:obj:`dict`): Dict type data, including at least ['obs', 'action', 'reward', 'next_obs','adv'] + - data (:obj:`List[Dict[int, Any]]`): The input data used for policy forward, including a batch of \ + training samples. For each element in the list, the key of the dict is the name of data items and the \ + value is the corresponding data. Usually, the value is torch.Tensor or np.ndarray or there dict/list \ + combinations. In the ``_forward_learn`` method, data often need to first be stacked in the batch \ + dimension by some utility functions such as ``default_preprocess_learn``. \ + For A2C, each element in the list is a dict containing at least the following keys: \ + ['obs', 'action', 'adv', 'value', 'weight']. Returns: - - info_dict (:obj:`Dict[str, Any]`): Including current lr and loss. + - info_dict (:obj:`Dict[str, Any]`): The information dict that indicated training result, which will be \ + recorded in text log and tensorboard, values must be python scalar or a list of scalars. For the \ + detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that is not supported, the main reason is that the corresponding model does not support \ + it. You can implement your own model rather than use the default model. For more information, please \ + raise an issue in GitHub repo, and we will continue to follow up. """ + # Data preprocessing operations, such as stack data, cpu to cuda device data = default_preprocess_learn(data, ignore_done=self._cfg.learn.ignore_done, use_nstep=False) if self._cuda: data = to_device(data, self._device) @@ -121,14 +169,17 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: error_data = a2c_data(output['logit'], batch['action'], output['value'], adv, return_, batch['weight']) # Calculate A2C loss - a2c_loss = a2c_error(error_data) + if self._action_space == 'continuous': + a2c_loss = a2c_error_continuous(error_data) + elif self._action_space == 'discrete': + a2c_loss = a2c_error(error_data) + wv, we = self._value_weight, self._entropy_weight total_loss = a2c_loss.policy_loss + wv * a2c_loss.value_loss - we * a2c_loss.entropy_loss # ==================== # A2C-learning update # ==================== - self._optimizer.zero_grad() total_loss.backward() @@ -153,40 +204,70 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: } def _state_dict_learn(self) -> Dict[str, Any]: + """ + Overview: + Return the state_dict of learn mode, usually including model and optimizer. + Returns: + - state_dict (:obj:`Dict[str, Any]`): The dict of current policy learn state, for saving and restoring. + """ return { 'model': self._learn_model.state_dict(), 'optimizer': self._optimizer.state_dict(), } def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: + """ + Overview: + Load the state_dict variable into policy learn mode. + Arguments: + - state_dict (:obj:`Dict[str, Any]`): The dict of policy learn state saved before. + + .. tip:: + If you want to only load some parts of model, you can simply set the ``strict`` argument in \ + load_state_dict to ``False``, or refer to ``ding.torch_utils.checkpoint_helper`` for more \ + complicated operation. + """ self._learn_model.load_state_dict(state_dict['model']) self._optimizer.load_state_dict(state_dict['optimizer']) def _init_collect(self) -> None: - r""" - Overview: - Collect mode init method. Called by ``self.__init__``. - Init traj and unroll length, collect model. """ + Overview: + Initialize the collect mode of policy, including related attributes and modules. For A2C, it contains the \ + collect_model to balance the exploration and exploitation with ``reparam_sample`` or \ + ``multinomial_sample`` mechanism, and other algorithm-specific arguments such as gamma and gae_lambda. \ + This method will be called in ``__init__`` method if ``collect`` field is in ``enable_field``. + .. note:: + If you want to set some spacial member variables in ``_init_collect`` method, you'd better name them \ + with prefix ``_collect_`` to avoid conflict with other modes, such as ``self._collect_attr1``. + """ + assert self._cfg.action_space in ["continuous", "discrete"] self._unroll_len = self._cfg.collect.unroll_len - self._collect_model = model_wrap(self._model, wrapper_name='multinomial_sample') + + self._action_space = self._cfg.action_space + if self._action_space == 'continuous': + self._collect_model = model_wrap(self._model, wrapper_name='reparam_sample') + elif self._action_space == 'discrete': + self._collect_model = model_wrap(self._model, wrapper_name='multinomial_sample') self._collect_model.reset() # Algorithm self._gamma = self._cfg.collect.discount_factor self._gae_lambda = self._cfg.collect.gae_lambda - def _forward_collect(self, data: dict) -> dict: - r""" + def _forward_collect(self, data: Dict[int, Any]) -> Dict[int, Any]: + """ Overview: - Forward function of collect mode. + Policy forward function of collect mode (collecting training data by interacting with envs). Forward means \ + that the policy gets some necessary data (mainly observation) from the envs and then returns the output \ + data, such as the action to interact with the envs. Arguments: - - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \ - values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer. + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. Returns: - - output (:obj:`Dict[int, Any]`): Dict type data, including at least inferred action according to input obs. - ReturnsKeys - - necessary: ``action`` + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action and \ + other necessary data for learn mode defined in ``self._process_transition`` method. The key of the \ + dict is the same as the input data, i.e. environment id. """ data_id = list(data.keys()) data = default_collate(list(data.values())) @@ -200,66 +281,95 @@ def _forward_collect(self, data: dict) -> dict: output = default_decollate(output) return {i: d for i, d in zip(data_id, output)} - def _process_transition(self, obs: Any, model_output: dict, timestep: namedtuple) -> dict: - r""" + def _process_transition(self, obs: Any, policy_output: Dict[str, torch.Tensor], + timestep: namedtuple) -> Dict[str, torch.Tensor]: + """ Overview: - Generate dict type transition data from inputs. + Process and pack one timestep transition data into a dict, which can be directly used for training and \ + saved in replay buffer. For A2C, it contains obs, next_obs, action, value, reward, done. Arguments: - - obs (:obj:`Any`): Env observation - - model_output (:obj:`dict`): Output of collect model, including at least ['action'] - - timestep (:obj:`namedtuple`): Output after env step, including at least ['obs', 'reward', 'done'] \ - (here 'obs' indicates obs after env step). + - obs (:obj:`torch.Tensor`): The env observation of current timestep, such as stacked 2D image in Atari. + - policy_output (:obj:`Dict[str, torch.Tensor]`): The output of the policy network with the observation \ + as input. For A2C, it contains the action and the value of the state. + - timestep (:obj:`namedtuple`): The execution result namedtuple returned by the environment step method, \ + except all the elements have been transformed into tensor data. Usually, it contains the next obs, \ + reward, done, info, etc. Returns: - - transition (:obj:`dict`): Dict type transition data. + - transition (:obj:`Dict[str, torch.Tensor]`): The processed transition data of the current timestep. """ transition = { 'obs': obs, 'next_obs': timestep.obs, - 'action': model_output['action'], - 'value': model_output['value'], + 'action': policy_output['action'], + 'value': policy_output['value'], 'reward': timestep.reward, 'done': timestep.done, } return transition - def _get_train_sample(self, data: list) -> Union[None, List[Any]]: - r""" + def _get_train_sample(self, transitions: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ Overview: - Get the trajectory and the n step return data, then sample from the n_step return data + For a given trajectory (transitions, a list of transition) data, process it into a list of sample that \ + can be used for training directly. In A2C, a train sample is a processed transition. \ + This method is usually used in collectors to execute necessary \ + RL data preprocessing before training, which can help the learner amortize relevant time consumption. \ + In addition, you can also implement this method as an identity function and do the data processing \ + in ``self._forward_learn`` method. Arguments: - - data (:obj:`list`): The trajectory's buffer list + - transitions (:obj:`List[Dict[str, Any]`): The trajectory data (a list of transition), each element is \ + in the same format as the return value of ``self._process_transition`` method. Returns: - - samples (:obj:`dict`): The training samples generated + - samples (:obj:`List[Dict[str, Any]]`): The processed train samples, each element is similar in format \ + to input transitions, but may contain more data for training, such as advantages. """ - data = get_gae_with_default_last_value( - data, - data[-1]['done'], + transitions = get_gae_with_default_last_value( + transitions, + transitions[-1]['done'], gamma=self._gamma, gae_lambda=self._gae_lambda, cuda=self._cuda, ) - return get_train_sample(data, self._unroll_len) + return get_train_sample(transitions, self._unroll_len) def _init_eval(self) -> None: - r""" + """ Overview: - Evaluate mode init method. Called by ``self.__init__``. - Init eval model with argmax strategy. + Initialize the eval mode of policy, including related attributes and modules. For A2C, it contains the \ + eval model to greedily select action with ``argmax_sample`` mechanism (For discrete action space) and \ + ``deterministic_sample`` mechanism (For continuous action space). \ + This method will be called in ``__init__`` method if ``eval`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_eval`` method, you'd better name them \ + with prefix ``_eval_`` to avoid conflict with other modes, such as ``self._eval_attr1``. """ - self._eval_model = model_wrap(self._model, wrapper_name='argmax_sample') + assert self._cfg.action_space in ["continuous", "discrete"] + self._action_space = self._cfg.action_space + if self._action_space == 'continuous': + self._eval_model = model_wrap(self._model, wrapper_name='deterministic_sample') + elif self._action_space == 'discrete': + self._eval_model = model_wrap(self._model, wrapper_name='argmax_sample') self._eval_model.reset() - def _forward_eval(self, data: dict) -> dict: - r""" + def _forward_eval(self, data: Dict[int, Any]) -> Dict[int, Any]: + """ Overview: - Forward function of eval mode, similar to ``self._forward_collect``. + Policy forward function of eval mode (evaluation policy performance by interacting with envs). Forward \ + means that the policy gets some necessary data (mainly observation) from the envs and then returns the \ + action to interact with the envs. Arguments: - - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \ - values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer. + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. Returns: - - output (:obj:`Dict[int, Any]`): The dict of predicting action for the interaction with env. - ReturnsKeys - - necessary: ``action`` + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action. The \ + key of the dict is the same as the input data, i.e., environment id. + + .. note:: + The input value can be ``torch.Tensor`` or dict/list combinations, current policy supports all of them. \ + For the data type that is not supported, the main reason is that the corresponding model does not \ + support it. You can implement your own model rather than use the default model. For more information, \ + please raise an issue in GitHub repo, and we will continue to follow up. """ data_id = list(data.keys()) data = default_collate(list(data.values())) @@ -274,4 +384,11 @@ def _forward_eval(self, data: dict) -> dict: return {i: d for i, d in zip(data_id, output)} def _monitor_vars_learn(self) -> List[str]: + """ + Overview: + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. + Returns: + - necessary_keys (:obj:`List[str]`): The list of the necessary keys to be logged. + """ return super()._monitor_vars_learn() + ['policy_loss', 'value_loss', 'entropy_loss', 'adv_abs_max', 'grad_norm'] diff --git a/ding/policy/acer.py b/ding/policy/acer.py index 7f490ec606..319b2fe814 100644 --- a/ding/policy/acer.py +++ b/ding/policy/acer.py @@ -47,19 +47,18 @@ class ACERPolicy(Policy): config = dict( type='acer', cuda=False, - # (bool) whether use on-policy training pipeline(behaviour policy and training policy are the same) + # (bool) whether to use on-policy training pipeline (behaviour policy and training policy are the same) # here we follow ppo serial pipeline, the original is False on_policy=False, priority=False, - # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. + # (bool) Whether to use Importance Sampling Weight to correct biased update. If True, priority must be True. priority_IS_weight=False, learn=dict( # (str) the type of gradient clip method grad_clip_type=None, # (float) max value when ACER use gradient clip clip_value=None, - # (bool) Whether to use multi gpu - multi_gpu=False, + # (int) collect n_sample data, train model update_per_collect times # here we follow ppo serial pipeline update_per_collect=4, @@ -296,7 +295,7 @@ def _reshape_data( Update values and rewards with the weight Arguments: - output (:obj:`Dict[int, Any]`): Dict type data, output of learn_model forward. \ - Values are torch.Tensor or np.ndarray or dict/list combinations,keys are value, logit. + Values are torch.Tensor or np.ndarray or dict/list combinations, keys are value, logit. - data (:obj:`Dict[int, Any]`): Dict type data, input of policy._forward_learn \ Values are torch.Tensor or np.ndarray or dict/list combinations. Keys includes at \ least ['logit', 'action', 'reward', 'done',] @@ -379,7 +378,7 @@ def _forward_collect(self, data: Dict[int, Any]) -> Dict[int, Dict[str, Any]]: action, values are torch.Tensor or np.ndarray or dict/list combinations,keys \ are env_id indicated by integer. Returns: - - output (:obj:`Dict[int, Dict[str,Any]]`): Dict of predicting policy_output(logit, action) for each env. + - output (:obj:`Dict[int, Dict[str, Any]]`): Dict of predicting policy_output(logit, action) for each env. ReturnsKeys - necessary: ``logit``, ``action`` """ @@ -480,7 +479,7 @@ def _monitor_vars_learn(self) -> List[str]: Returns: - model_info (:obj:`Tuple[str, List[str]]`): model name and mode import_names .. note:: - The user can define and use customized network model but must obey the same interface definition indicated \ - by import_names path. For IMPALA, ``ding.model.interface.IMPALA`` + The user can define and use a customized network model but must obey the same interface definition \ + indicated by import_names path. For IMPALA, ``ding.model.interface.IMPALA`` """ return ['actor_loss', 'bc_loss', 'policy_loss', 'critic_loss', 'entropy_loss', 'kl_div'] diff --git a/ding/policy/atoc.py b/ding/policy/atoc.py index 011e12f809..8addc32716 100644 --- a/ding/policy/atoc.py +++ b/ding/policy/atoc.py @@ -43,8 +43,6 @@ class ATOCPolicy(Policy): agent_per_group=2, ), learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, # (int) Collect n_sample data, update model n_iteration time update_per_collect=5, # (int) The number of data for a train iteration diff --git a/ding/policy/base_policy.py b/ding/policy/base_policy.py index a30b080015..1c5f32d1db 100644 --- a/ding/policy/base_policy.py +++ b/ding/policy/base_policy.py @@ -1,20 +1,47 @@ +from typing import Optional, List, Dict, Any, Tuple, Union from abc import ABC, abstractmethod from collections import namedtuple -from typing import Optional, List, Dict, Any, Tuple, Union +from easydict import EasyDict -import torch import copy -from easydict import EasyDict +import torch from ding.model import create_model -from ding.utils import import_module, allreduce, broadcast, get_rank, allreduce_async, synchronize, POLICY_REGISTRY +from ding.utils import import_module, allreduce, allreduce_with_indicator, broadcast, get_rank, allreduce_async, \ + synchronize, deep_merge_dicts, POLICY_REGISTRY class Policy(ABC): + """ + Overview: + The basic class of Reinforcement Learning (RL) and Imitation Learning (IL) policy in DI-engine. + Property: + ``cfg``, ``learn_mode``, ``collect_mode``, ``eval_mode`` + """ @classmethod def default_config(cls: type) -> EasyDict: + """ + Overview: + Get the default config of policy. This method is used to create the default config of policy. + Returns: + - cfg (:obj:`EasyDict`): The default config of corresponding policy. For the derived policy class, \ + it will recursively merge the default config of base class and its own default config. + + .. tip:: + This method will deepcopy the ``config`` attribute of the class and return the result. So users don't need \ + to worry about the modification of the returned config. + """ + if cls == Policy: + raise RuntimeError("Basic class Policy doesn't have completed default_config") + + base_cls = cls.__base__ + if base_cls == Policy: + base_policy_cfg = EasyDict(copy.deepcopy(Policy.config)) + else: + base_policy_cfg = copy.deepcopy(base_cls.default_config()) cfg = EasyDict(copy.deepcopy(cls.config)) + cfg = deep_merge_dicts(base_policy_cfg, cfg) cfg.cfg_type = cls.__name__ + 'Dict' return cfg @@ -53,13 +80,52 @@ def default_config(cls: type) -> EasyDict: ] ) total_field = set(['learn', 'collect', 'eval']) + config = dict( + # (bool) Whether the learning policy is the same as the collecting data policy (on-policy). + on_policy=False, + # (bool) Whether to use cuda in policy. + cuda=False, + # (bool) Whether to use data parallel multi-gpu mode in policy. + multi_gpu=False, + # (bool) Whether to synchronize update the model parameters after allreduce the gradients of model parameters. + bp_update_sync=True, + # (bool) Whether to enable infinite trajectory length in data collecting. + traj_len_inf=False, + # neural network model config + model=dict(), + # If resume_training is True, the environment step count (collector.envstep) and training iteration (train_iter) + # will be loaded from the pretrained checkpoint, allowing training to resume seamlessly + # from where the ckpt left off. + learn=dict(resume_training=False), + ) def __init__( self, - cfg: dict, - model: Optional[Union[type, torch.nn.Module]] = None, + cfg: EasyDict, + model: Optional[torch.nn.Module] = None, enable_field: Optional[List[str]] = None ) -> None: + """ + Overview: + Initialize policy instance according to input configures and model. This method will initialize differnent \ + fields in policy, including ``learn``, ``collect``, ``eval``. The ``learn`` field is used to train the \ + policy, the ``collect`` field is used to collect data for training, and the ``eval`` field is used to \ + evaluate the policy. The ``enable_field`` is used to specify which field to initialize, if it is None, \ + then all fields will be initialized. + Arguments: + - cfg (:obj:`EasyDict`): The final merged config used to initialize policy. For the default config, \ + see the ``config`` attribute and its comments of policy class. + - model (:obj:`torch.nn.Module`): The neural network model used to initialize policy. If it \ + is None, then the model will be created according to ``default_model`` method and ``cfg.model`` field. \ + Otherwise, the model will be set to the ``model`` instance created by outside caller. + - enable_field (:obj:`Optional[List[str]]`): The field list to initialize. If it is None, then all fields \ + will be initialized. Otherwise, only the fields in ``enable_field`` will be initialized, which is \ + beneficial to save resources. + + .. note:: + For the derived policy class, it should implement the ``_init_learn``, ``_init_collect``, ``_init_eval`` \ + method to initialize the corresponding field. + """ self._cfg = cfg self._on_policy = self._cfg.on_policy if enable_field is None: @@ -73,18 +139,19 @@ def __init__( self._cuda = cfg.cuda and torch.cuda.is_available() # now only support multi-gpu for only enable learn mode if len(set(self._enable_field).intersection(set(['learn']))) > 0: - self._rank = get_rank() if self._cfg.learn.multi_gpu else 0 + multi_gpu = self._cfg.multi_gpu + self._rank = get_rank() if multi_gpu else 0 if self._cuda: - torch.cuda.set_device(self._rank % torch.cuda.device_count()) + # model.cuda() is an in-place operation. model.cuda() - if self._cfg.learn.multi_gpu: - bp_update_sync = self._cfg.learn.get('bp_update_sync', True) + if multi_gpu: + bp_update_sync = self._cfg.bp_update_sync self._bp_update_sync = bp_update_sync self._init_multi_gpu_setting(model, bp_update_sync) else: self._rank = 0 if self._cuda: - torch.cuda.set_device(self._rank % torch.cuda.device_count()) + # model.cuda() is an in-place operation. model.cuda() self._model = model self._device = 'cuda:{}'.format(self._rank % torch.cuda.device_count()) if self._cuda else 'cpu' @@ -93,13 +160,26 @@ def __init__( self._rank = 0 self._device = 'cpu' + # call the initialization method of different modes, such as ``_init_learn``, ``_init_collect``, ``_init_eval`` for field in self._enable_field: getattr(self, '_init_' + field)() def _init_multi_gpu_setting(self, model: torch.nn.Module, bp_update_sync: bool) -> None: + """ + Overview: + Initialize multi-gpu data parallel training setting, including broadcast model parameters at the beginning \ + of the training, and prepare the hook function to allreduce the gradients of model parameters. + Arguments: + - model (:obj:`torch.nn.Module`): The neural network model to be trained. + - bp_update_sync (:obj:`bool`): Whether to synchronize update the model parameters after allreduce the \ + gradients of model parameters. Async update can be parallel in different network layers like pipeline \ + so that it can save time. + """ for name, param in model.state_dict().items(): assert isinstance(param.data, torch.Tensor), type(param.data) broadcast(param.data, 0) + # here we manually set the gradient to zero tensor at the beginning of the training, which is necessary for + # the case that different GPUs have different computation graph. for name, param in model.named_parameters(): setattr(param, 'grad', torch.zeros_like(param)) if not bp_update_sync: @@ -117,7 +197,23 @@ def hook(*ignore): grad_acc = p_tmp.grad_fn.next_functions[0][0] grad_acc.register_hook(make_hook(name, p)) - def _create_model(self, cfg: dict, model: Optional[torch.nn.Module] = None) -> torch.nn.Module: + def _create_model(self, cfg: EasyDict, model: Optional[torch.nn.Module] = None) -> torch.nn.Module: + """ + Overview: + Create or validate the neural network model according to the input configuration and model. \ + If the input model is None, then the model will be created according to ``default_model`` \ + method and ``cfg.model`` field. Otherwise, the model will be verified as an instance of \ + ``torch.nn.Module`` and set to the ``model`` instance created by outside caller. + Arguments: + - cfg (:obj:`EasyDict`): The final merged config used to initialize policy. + - model (:obj:`torch.nn.Module`): The neural network model used to initialize policy. User can refer to \ + the default model defined in the corresponding policy to customize its own model. + Returns: + - model (:obj:`torch.nn.Module`): The created neural network model. The different modes of policy will \ + add distinct wrappers and plugins to the model, which is used to train, collect and evaluate. + Raises: + - RuntimeError: If the input model is not None and is not an instance of ``torch.nn.Module``. + """ if model is None: model_cfg = cfg.model if 'type' not in model_cfg: @@ -137,18 +233,77 @@ def cfg(self) -> EasyDict: @abstractmethod def _init_learn(self) -> None: + """ + Overview: + Initialize the learn mode of policy, including related attributes and modules. This method will be \ + called in ``__init__`` method if ``learn`` field is in ``enable_field``. Almost different policies have \ + its own learn mode, so this method must be overrided in subclass. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. + """ raise NotImplementedError @abstractmethod def _init_collect(self) -> None: + """ + Overview: + Initialize the collect mode of policy, including related attributes and modules. This method will be \ + called in ``__init__`` method if ``collect`` field is in ``enable_field``. Almost different policies have \ + its own collect mode, so this method must be overrided in subclass. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_collect`` \ + and ``_load_state_dict_collect`` methods. + + .. note:: + If you want to set some spacial member variables in ``_init_collect`` method, you'd better name them \ + with prefix ``_collect_`` to avoid conflict with other modes, such as ``self._collect_attr1``. + """ raise NotImplementedError @abstractmethod def _init_eval(self) -> None: + """ + Overview: + Initialize the eval mode of policy, including related attributes and modules. This method will be \ + called in ``__init__`` method if ``eval`` field is in ``enable_field``. Almost different policies have \ + its own eval mode, so this method must be override in subclass. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_eval`` \ + and ``_load_state_dict_eval`` methods. + + .. note:: + If you want to set some spacial member variables in ``_init_eval`` method, you'd better name them \ + with prefix ``_eval_`` to avoid conflict with other modes, such as ``self._eval_attr1``. + """ raise NotImplementedError @property def learn_mode(self) -> 'Policy.learn_function': # noqa + """ + Overview: + Return the interfaces of learn mode of policy, which is used to train the model. Here we use namedtuple \ + to define immutable interfaces and restrict the usage of policy in different modes. Moreover, derived \ + subclass can override the interfaces to customize its own learn mode. + Returns: + - interfaces (:obj:`Policy.learn_function`): The interfaces of learn mode of policy, it is a namedtuple \ + whose values of distinct fields are different internal methods. + Examples: + >>> policy = Policy(cfg, model) + >>> policy_learn = policy.learn_mode + >>> train_output = policy_learn.forward(data) + >>> state_dict = policy_learn.state_dict() + """ return Policy.learn_function( self._forward_learn, self._reset_learn, @@ -162,6 +317,21 @@ def learn_mode(self) -> 'Policy.learn_function': # noqa @property def collect_mode(self) -> 'Policy.collect_function': # noqa + """ + Overview: + Return the interfaces of collect mode of policy, which is used to train the model. Here we use namedtuple \ + to define immutable interfaces and restrict the usage of policy in different modes. Moreover, derived \ + subclass can override the interfaces to customize its own collect mode. + Returns: + - interfaces (:obj:`Policy.collect_function`): The interfaces of collect mode of policy, it is a \ + namedtuple whose values of distinct fields are different internal methods. + Examples: + >>> policy = Policy(cfg, model) + >>> policy_collect = policy.collect_mode + >>> obs = env_manager.ready_obs + >>> inference_output = policy_collect.forward(obs) + >>> next_obs, rew, done, info = env_manager.step(inference_output.action) + """ return Policy.collect_function( self._forward_collect, self._process_transition, @@ -175,6 +345,21 @@ def collect_mode(self) -> 'Policy.collect_function': # noqa @property def eval_mode(self) -> 'Policy.eval_function': # noqa + """ + Overview: + Return the interfaces of eval mode of policy, which is used to train the model. Here we use namedtuple \ + to define immutable interfaces and restrict the usage of policy in different mode. Moreover, derived \ + subclass can override the interfaces to customize its own eval mode. + Returns: + - interfaces (:obj:`Policy.eval_function`): The interfaces of eval mode of policy, it is a namedtuple \ + whose values of distinct fields are different internal methods. + Examples: + >>> policy = Policy(cfg, model) + >>> policy_eval = policy.eval_mode + >>> obs = env_manager.ready_obs + >>> inference_output = policy_eval.forward(obs) + >>> next_obs, rew, done, info = env_manager.step(inference_output.action) + """ return Policy.eval_function( self._forward_eval, self._reset_eval, @@ -185,9 +370,32 @@ def eval_mode(self) -> 'Policy.eval_function': # noqa ) def _set_attribute(self, name: str, value: Any) -> None: + """ + Overview: + In order to control the access of the policy attributes, we expose different modes to outside rather than \ + directly use the policy instance. And we also provide a method to set the attribute of the policy in \ + different modes. And the new attribute will name as ``_{name}``. + Arguments: + - name (:obj:`str`): The name of the attribute. + - value (:obj:`Any`): The value of the attribute. + """ setattr(self, '_' + name, value) def _get_attribute(self, name: str) -> Any: + """ + Overview: + In order to control the access of the policy attributes, we expose different modes to outside rather than \ + directly use the policy instance. And we also provide a method to get the attribute of the policy in \ + different modes. + Arguments: + - name (:obj:`str`): The name of the attribute. + Returns: + - value (:obj:`Any`): The value of the attribute. + + .. note:: + DI-engine's policy will first try to access `_get_{name}` method, and then try to access `_{name}` \ + attribute. If both of them are not found, it will raise a ``NotImplementedError``. + """ if hasattr(self, '_get_' + name): return getattr(self, '_get_' + name)() elif hasattr(self, '_' + name): @@ -196,117 +404,481 @@ def _get_attribute(self, name: str) -> Any: raise NotImplementedError def __repr__(self) -> str: + """ + Overview: + Get the string representation of the policy. + Returns: + - repr (:obj:`str`): The string representation of the policy. + """ return "DI-engine DRL Policy\n{}".format(repr(self._model)) def sync_gradients(self, model: torch.nn.Module) -> None: + """ + Overview: + Synchronize (allreduce) gradients of model parameters in data-parallel multi-GPU training. + For parameters that did not participate in the forward/backward pass in some GPUs, + assign a zero gradient with an indicator of 0. This ensures that only GPUs which contributed + to the gradient computation are considered when averaging, thereby avoiding an incorrect + division by the total number of GPUs. + Arguments: + - model (:obj:`torch.nn.Module`): The model to synchronize gradients. + + .. note:: + This method is only used in multi-gpu training, and it should be called after the ``backward`` method and \ + before the ``step`` method. The user can also use the ``bp_update_sync`` config to control whether to \ + synchronize gradients allreduce and optimizer updates. + """ if self._bp_update_sync: for name, param in model.named_parameters(): if param.requires_grad: - allreduce(param.grad.data) + # Create an indicator tensor on the same device as the parameter (or its gradient) + if param.grad is not None: + # If the gradient exists, extract its data and set indicator to 1. + grad_tensor = param.grad.data + indicator = torch.tensor(1.0, device=grad_tensor.device) + else: + # If the parameter did not participate in the computation (grad is None), + # create a zero tensor for the gradient and set the indicator to 0. + grad_tensor = torch.zeros_like(param.data) + indicator = torch.tensor(0.0, device=grad_tensor.device) + + # Assign the zero gradient to param.grad to ensure that all GPUs + # participate in the subsequent allreduce call (avoiding deadlock). + param.grad = grad_tensor + + # Use the custom allreduce function to reduce the gradient using the indicator. + allreduce_with_indicator(param.grad, indicator) else: synchronize() # don't need to implement default_model method by force def default_model(self) -> Tuple[str, List[str]]: + """ + Overview: + Return this algorithm default neural network model setting for demonstration. ``__init__`` method will \ + automatically call this method to get the default model setting and create model. + Returns: + - model_info (:obj:`Tuple[str, List[str]]`): The registered model name and model's import_names. + + .. note:: + The user can define and use customized network model but must obey the same inferface definition indicated \ + by import_names path. For example about DQN, its registered name is ``dqn`` and the import_names is \ + ``ding.model.template.q_learning.DQN`` + """ raise NotImplementedError # *************************************** learn function ************************************ @abstractmethod - def _forward_learn(self, data: dict) -> Dict[str, Any]: + def _forward_learn(self, data: List[Dict[str, Any]]) -> Dict[str, Any]: + """ + Overview: + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data from the replay buffer and then returns the output \ + result, including various training information such as loss value, policy entropy, q value, priority, \ + and so on. This method is left to be implemented by the subclass, and more arguments can be added in \ + ``data`` item if necessary. + Arguments: + - data (:obj:`List[Dict[int, Any]]`): The input data used for policy forward, including a batch of \ + training samples. For each element in list, the key of the dict is the name of data items and the \ + value is the corresponding data. Usually, in the ``_forward_learn`` method, data should be stacked in \ + the batch dimension by some utility functions such as ``default_preprocess_learn``. + Returns: + - output (:obj:`Dict[int, Any]`): The training information of policy forward, including some metrics for \ + monitoring training such as loss, priority, q value, policy entropy, and some data for next step \ + training such as priority. Note the output data item should be Python native scalar rather than \ + PyTorch tensor, which is convenient for the outside to use. + """ raise NotImplementedError # don't need to implement _reset_learn method by force def _reset_learn(self, data_id: Optional[List[int]] = None) -> None: + """ + Overview: + Reset some stateful variables for learn mode when necessary, such as the hidden state of RNN or the \ + memory bank of some special algortihms. If ``data_id`` is None, it means to reset all the stateful \ + varaibles. Otherwise, it will reset the stateful variables according to the ``data_id``. For example, \ + different trajectories in ``data_id`` will have different hidden state in RNN. + Arguments: + - data_id (:obj:`Optional[List[int]]`): The id of the data, which is used to reset the stateful variables \ + specified by ``data_id``. + + .. note:: + This method is not mandatory to be implemented. The sub-class can overwrite this method if necessary. + """ pass def _monitor_vars_learn(self) -> List[str]: + """ + Overview: + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. + Returns: + - necessary_keys (:obj:`List[str]`): The list of the necessary keys to be logged. + + .. tip:: + The default implementation is ``['cur_lr', 'total_loss']``. Other derived classes can overwrite this \ + method to add their own keys if necessary. + """ return ['cur_lr', 'total_loss'] def _state_dict_learn(self) -> Dict[str, Any]: - return {'model': self._learn_model.state_dict()} + """ + Overview: + Return the state_dict of learn mode, usually including model and optimizer. + Returns: + - state_dict (:obj:`Dict[str, Any]`): The dict of current policy learn state, for saving and restoring. + """ + return { + 'model': self._learn_model.state_dict(), + 'optimizer': self._optimizer.state_dict(), + } def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: - self._learn_model.load_state_dict(state_dict['model'], strict=True) + """ + Overview: + Load the state_dict variable into policy learn mode. + Arguments: + - state_dict (:obj:`Dict[str, Any]`): The dict of policy learn state saved before. + + .. tip:: + If you want to only load some parts of model, you can simply set the ``strict`` argument in \ + load_state_dict to ``False``, or refer to ``ding.torch_utils.checkpoint_helper`` for more \ + complicated operation. + """ + self._learn_model.load_state_dict(state_dict['model']) + self._optimizer.load_state_dict(state_dict['optimizer']) def _get_batch_size(self) -> Union[int, Dict[str, int]]: - return self._cfg.learn.batch_size + # some specifial algorithms use different batch size for different optimization parts. + if 'batch_size' in self._cfg: + return self._cfg.batch_size + else: # for compatibility + return self._cfg.learn.batch_size # *************************************** collect function ************************************ @abstractmethod - def _forward_collect(self, data: dict, **kwargs) -> dict: + def _forward_collect(self, data: Dict[int, Any], **kwargs) -> Dict[int, Any]: + """ + Overview: + Policy forward function of collect mode (collecting training data by interacting with envs). Forward means \ + that the policy gets some necessary data (mainly observation) from the envs and then returns the output \ + data, such as the action to interact with the envs, or the action logits to calculate the loss in learn \ + mode. This method is left to be implemented by the subclass, and more arguments can be added in ``kwargs`` \ + part if necessary. + Arguments: + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. + Returns: + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action and \ + other necessary data for learn mode defined in ``self._process_transition`` method. The key of the \ + dict is the same as the input data, i.e. environment id. + """ raise NotImplementedError @abstractmethod - def _process_transition(self, obs: Any, model_output: dict, timestep: namedtuple) -> dict: + def _process_transition( + self, obs: Union[torch.Tensor, Dict[str, torch.Tensor]], policy_output: Dict[str, torch.Tensor], + timestep: namedtuple + ) -> Dict[str, torch.Tensor]: + """ + Overview: + Process and pack one timestep transition data into a dict, such as . Some policies \ + need to do some special process and pack its own necessary attributes (e.g. hidden state and logit), \ + so this method is left to be implemented by the subclass. + Arguments: + - obs (:obj:`Union[torch.Tensor, Dict[str, torch.Tensor]]`): The observation of the current timestep. + - policy_output (:obj:`Dict[str, torch.Tensor]`): The output of the policy network with the observation \ + as input. Usually, it contains the action and the logit of the action. + - timestep (:obj:`namedtuple`): The execution result namedtuple returned by the environment step method, \ + except all the elements have been transformed into tensor data. Usually, it contains the next obs, \ + reward, done, info, etc. + Returns: + - transition (:obj:`Dict[str, torch.Tensor]`): The processed transition data of the current timestep. + """ raise NotImplementedError @abstractmethod - def _get_train_sample(self, data: list) -> Union[None, List[Any]]: + def _get_train_sample(self, transitions: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Overview: + For a given trajectory (transitions, a list of transition) data, process it into a list of sample that \ + can be used for training directly. A train sample can be a processed transition (DQN with nstep TD) \ + or some multi-timestep transitions (DRQN). This method is usually used in collectors to execute necessary \ + RL data preprocessing before training, which can help learner amortize revelant time consumption. \ + In addition, you can also implement this method as an identity function and do the data processing \ + in ``self._forward_learn`` method. + Arguments: + - transitions (:obj:`List[Dict[str, Any]`): The trajectory data (a list of transition), each element is \ + the same format as the return value of ``self._process_transition`` method. + Returns: + - samples (:obj:`List[Dict[str, Any]]`): The processed train samples, each element is the similar format \ + as input transitions, but may contain more data for training, such as nstep reward, advantage, etc. + + .. note:: + We will vectorize ``process_transition`` and ``get_train_sample`` method in the following release version. \ + And the user can customize the this data processing procecure by overriding this two methods and collector \ + itself + """ raise NotImplementedError # don't need to implement _reset_collect method by force def _reset_collect(self, data_id: Optional[List[int]] = None) -> None: + """ + Overview: + Reset some stateful variables for collect mode when necessary, such as the hidden state of RNN or the \ + memory bank of some special algortihms. If ``data_id`` is None, it means to reset all the stateful \ + varaibles. Otherwise, it will reset the stateful variables according to the ``data_id``. For example, \ + different environments/episodes in collecting in ``data_id`` will have different hidden state in RNN. + Arguments: + - data_id (:obj:`Optional[List[int]]`): The id of the data, which is used to reset the stateful variables \ + specified by ``data_id``. + + .. note:: + This method is not mandatory to be implemented. The sub-class can overwrite this method if necessary. + """ pass def _state_dict_collect(self) -> Dict[str, Any]: + """ + Overview: + Return the state_dict of collect mode, only including model in usual, which is necessary for distributed \ + training scenarios to auto-recover collectors. + Returns: + - state_dict (:obj:`Dict[str, Any]`): The dict of current policy collect state, for saving and restoring. + + .. tip:: + Not all the scenarios need to auto-recover collectors, sometimes, we can directly shutdown the crashed \ + collector and renew a new one. + """ return {'model': self._collect_model.state_dict()} def _load_state_dict_collect(self, state_dict: Dict[str, Any]) -> None: + """ + Overview: + Load the state_dict variable into policy collect mode, such as load pretrained state_dict, auto-recover \ + checkpoint, or model replica from learner in distributed training scenarios. + Arguments: + - state_dict (:obj:`Dict[str, Any]`): The dict of policy collect state saved before. + + .. tip:: + If you want to only load some parts of model, you can simply set the ``strict`` argument in \ + load_state_dict to ``False``, or refer to ``ding.torch_utils.checkpoint_helper`` for more \ + complicated operation. + """ self._collect_model.load_state_dict(state_dict['model'], strict=True) + def _get_n_sample(self) -> Union[int, None]: + if 'n_sample' in self._cfg: + return self._cfg.n_sample + else: # for compatibility + return self._cfg.collect.get('n_sample', None) # for some adpative collecting data case + + def _get_n_episode(self) -> Union[int, None]: + if 'n_episode' in self._cfg: + return self._cfg.n_episode + else: # for compatibility + return self._cfg.collect.get('n_episode', None) # for some adpative collecting data case + # *************************************** eval function ************************************ @abstractmethod - def _forward_eval(self, data: dict) -> Dict[str, Any]: + def _forward_eval(self, data: Dict[int, Any]) -> Dict[int, Any]: + """ + Overview: + Policy forward function of eval mode (evaluation policy performance, such as interacting with envs or \ + computing metrics on validation dataset). Forward means that the policy gets some necessary data (mainly \ + observation) from the envs and then returns the output data, such as the action to interact with the envs. \ + This method is left to be implemented by the subclass. + Arguments: + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. + Returns: + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action. The \ + key of the dict is the same as the input data, i.e. environment id. + """ raise NotImplementedError # don't need to implement _reset_eval method by force def _reset_eval(self, data_id: Optional[List[int]] = None) -> None: + """ + Overview: + Reset some stateful variables for eval mode when necessary, such as the hidden state of RNN or the \ + memory bank of some special algortihms. If ``data_id`` is None, it means to reset all the stateful \ + varaibles. Otherwise, it will reset the stateful variables according to the ``data_id``. For example, \ + different environments/episodes in evaluation in ``data_id`` will have different hidden state in RNN. + Arguments: + - data_id (:obj:`Optional[List[int]]`): The id of the data, which is used to reset the stateful variables \ + specified by ``data_id``. + + .. note:: + This method is not mandatory to be implemented. The sub-class can overwrite this method if necessary. + """ pass def _state_dict_eval(self) -> Dict[str, Any]: + """ + Overview: + Return the state_dict of eval mode, only including model in usual, which is necessary for distributed \ + training scenarios to auto-recover evaluators. + Returns: + - state_dict (:obj:`Dict[str, Any]`): The dict of current policy eval state, for saving and restoring. + + .. tip:: + Not all the scenarios need to auto-recover evaluators, sometimes, we can directly shutdown the crashed \ + evaluator and renew a new one. + """ return {'model': self._eval_model.state_dict()} def _load_state_dict_eval(self, state_dict: Dict[str, Any]) -> None: + """ + Overview: + Load the state_dict variable into policy eval mode, such as load auto-recover \ + checkpoint, or model replica from learner in distributed training scenarios. + Arguments: + - state_dict (:obj:`Dict[str, Any]`): The dict of policy eval state saved before. + + .. tip:: + If you want to only load some parts of model, you can simply set the ``strict`` argument in \ + load_state_dict to ``False``, or refer to ``ding.torch_utils.checkpoint_helper`` for more \ + complicated operation. + """ self._eval_model.load_state_dict(state_dict['model'], strict=True) class CommandModePolicy(Policy): + """ + Overview: + Policy with command mode, which can be used in old version of DI-engine pipeline: ``serial_pipeline``. \ + ``CommandModePolicy`` uses ``_get_setting_learn``, ``_get_setting_collect``, ``_get_setting_eval`` methods \ + to exchange information between different workers. + + Interface: + ``_init_command``, ``_get_setting_learn``, ``_get_setting_collect``, ``_get_setting_eval`` + Property: + ``command_mode`` + """ command_function = namedtuple('command_function', ['get_setting_learn', 'get_setting_collect', 'get_setting_eval']) total_field = set(['learn', 'collect', 'eval', 'command']) @property def command_mode(self) -> 'Policy.command_function': # noqa + """ + Overview: + Return the interfaces of command mode of policy, which is used to train the model. Here we use namedtuple \ + to define immutable interfaces and restrict the usage of policy in different mode. Moreover, derived \ + subclass can override the interfaces to customize its own command mode. + Returns: + - interfaces (:obj:`Policy.command_function`): The interfaces of command mode, it is a namedtuple \ + whose values of distinct fields are different internal methods. + Examples: + >>> policy = CommandModePolicy(cfg, model) + >>> policy_command = policy.command_mode + >>> settings = policy_command.get_setting_learn(command_info) + """ return CommandModePolicy.command_function( self._get_setting_learn, self._get_setting_collect, self._get_setting_eval ) @abstractmethod def _init_command(self) -> None: + """ + Overview: + Initialize the command mode of policy, including related attributes and modules. This method will be \ + called in ``__init__`` method if ``command`` field is in ``enable_field``. Almost different policies have \ + its own command mode, so this method must be overrided in subclass. + + .. note:: + If you want to set some spacial member variables in ``_init_command`` method, you'd better name them \ + with prefix ``_command_`` to avoid conflict with other modes, such as ``self._command_attr1``. + """ raise NotImplementedError # *************************************** command function ************************************ @abstractmethod - def _get_setting_learn(self, command_info: dict) -> dict: + def _get_setting_learn(self, command_info: Dict[str, Any]) -> Dict[str, Any]: + """ + Overview: + Accoding to ``command_info``, i.e., global training information (e.g. training iteration, collected env \ + step, evaluation results, etc.), return the setting of learn mode, which contains dynamically changed \ + hyperparameters for learn mode, such as ``batch_size``, ``learning_rate``, etc. + Arguments: + - command_info (:obj:`Dict[str, Any]`): The global training information, which is defined in ``commander``. + Returns: + - setting (:obj:`Dict[str, Any]`): The latest setting of learn mode, which is usually used as extra \ + arguments of the ``policy._forward_learn`` method. + """ raise NotImplementedError @abstractmethod - def _get_setting_collect(self, command_info: dict) -> dict: + def _get_setting_collect(self, command_info: Dict[str, Any]) -> Dict[str, Any]: + """ + Overview: + Accoding to ``command_info``, i.e., global training information (e.g. training iteration, collected env \ + step, evaluation results, etc.), return the setting of collect mode, which contains dynamically changed \ + hyperparameters for collect mode, such as ``eps``, ``temperature``, etc. + Arguments: + - command_info (:obj:`Dict[str, Any]`): The global training information, which is defined in ``commander``. + Returns: + - setting (:obj:`Dict[str, Any]`): The latest setting of collect mode, which is usually used as extra \ + arguments of the ``policy._forward_collect`` method. + """ raise NotImplementedError @abstractmethod - def _get_setting_eval(self, command_info: dict) -> dict: + def _get_setting_eval(self, command_info: Dict[str, Any]) -> Dict[str, Any]: + """ + Overview: + Accoding to ``command_info``, i.e., global training information (e.g. training iteration, collected env \ + step, evaluation results, etc.), return the setting of eval mode, which contains dynamically changed \ + hyperparameters for eval mode, such as ``temperature``, etc. + Arguments: + - command_info (:obj:`Dict[str, Any]`): The global training information, which is defined in ``commander``. + Returns: + - setting (:obj:`Dict[str, Any]`): The latest setting of eval mode, which is usually used as extra \ + arguments of the ``policy._forward_eval`` method. + """ raise NotImplementedError -def create_policy(cfg: dict, **kwargs) -> Policy: - cfg = EasyDict(cfg) +def create_policy(cfg: EasyDict, **kwargs) -> Policy: + """ + Overview: + Create a policy instance according to ``cfg`` and other kwargs. + Arguments: + - cfg (:obj:`EasyDict`): Final merged policy config. + ArgumentsKeys: + - type (:obj:`str`): Policy type set in ``POLICY_REGISTRY.register`` method , such as ``dqn`` . + - import_names (:obj:`List[str]`): A list of module names (paths) to import before creating policy, such \ + as ``ding.policy.dqn`` . + Returns: + - policy (:obj:`Policy`): The created policy instance. + + .. tip:: + ``kwargs`` contains other arguments that need to be passed to the policy constructor. You can refer to \ + the ``__init__`` method of the corresponding policy class for details. + + .. note:: + For more details about how to merge config, please refer to the system document of DI-engine \ + (`en link <../03_system/config.html>`_). + """ import_module(cfg.get('import_names', [])) return POLICY_REGISTRY.build(cfg.type, cfg=cfg, **kwargs) def get_policy_cls(cfg: EasyDict) -> type: + """ + Overview: + Get policy class according to ``cfg``, which is used to access related class variables/methods. + Arguments: + - cfg (:obj:`EasyDict`): Final merged policy config. + ArgumentsKeys: + - type (:obj:`str`): Policy type set in ``POLICY_REGISTRY.register`` method , such as ``dqn`` . + - import_names (:obj:`List[str]`): A list of module names (paths) to import before creating policy, such \ + as ``ding.policy.dqn`` . + Returns: + - policy (:obj:`type`): The policy class. + """ import_module(cfg.get('import_names', [])) return POLICY_REGISTRY.get(cfg.type) diff --git a/ding/policy/bc.py b/ding/policy/bc.py index 58a4f52856..0c95b8abec 100644 --- a/ding/policy/bc.py +++ b/ding/policy/bc.py @@ -20,6 +20,11 @@ @POLICY_REGISTRY.register('bc') class BehaviourCloningPolicy(Policy): + """ + Overview: + Behaviour Cloning (BC) policy class, which supports both discrete and continuous action space. \ + The policy is trained by supervised learning, and the data is a offline dataset collected by expert. + """ config = dict( type='bc', @@ -28,7 +33,6 @@ class BehaviourCloningPolicy(Policy): continuous=False, action_shape=19, learn=dict( - multi_gpu=False, update_per_collect=1, batch_size=32, learning_rate=1e-5, @@ -53,18 +57,46 @@ class BehaviourCloningPolicy(Policy): max=0.5, ), ), - eval=dict(), - other=dict(replay_buffer=dict(replay_buffer_size=10000, )), + eval=dict(), # for compatibility ) def default_model(self) -> Tuple[str, List[str]]: + """ + Overview: + Return this algorithm default neural network model setting for demonstration. ``__init__`` method will \ + automatically call this method to get the default model setting and create model. + Returns: + - model_info (:obj:`Tuple[str, List[str]]`): The registered model name and model's import_names. + + .. note:: + The user can define and use customized network model but must obey the same inferface definition indicated \ + by import_names path. For example about discrete BC, its registered name is ``discrete_bc`` and the \ + import_names is ``ding.model.template.bc``. + """ if self._cfg.continuous: return 'continuous_bc', ['ding.model.template.bc'] else: return 'discrete_bc', ['ding.model.template.bc'] - def _init_learn(self): - assert self._cfg.learn.optimizer in ['SGD', 'Adam'] + def _init_learn(self) -> None: + """ + Overview: + Initialize the learn mode of policy, including related attributes and modules. For BC, it mainly contains \ + optimizer, algorithm-specific arguments such as lr_scheduler, loss, etc. \ + This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. + """ + assert self._cfg.learn.optimizer in ['SGD', 'Adam'], self._cfg.learn.optimizer if self._cfg.learn.optimizer == 'SGD': self._optimizer = SGD( self._model.parameters(), @@ -104,20 +136,38 @@ def lr_scheduler_fn(epoch): elif self._cfg.loss_type == 'mse_loss': self._loss = nn.MSELoss() else: - raise KeyError + raise KeyError("not support loss type: {}".format(self._cfg.loss_type)) else: if not self._cfg.learn.ce_label_smooth: self._loss = nn.CrossEntropyLoss() else: self._loss = LabelSmoothCELoss(0.1) - if self._cfg.learn.show_accuracy: - # accuracy statistics for debugging in discrete action space env, e.g. for gfootball - self.total_accuracy_in_dataset = [] - self.action_accuracy_in_dataset = {k: [] for k in range(self._cfg.action_shape)} + def _forward_learn(self, data: List[Dict[str, Any]]) -> Dict[str, Any]: + """ + Overview: + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data from the replay buffer and then returns the output \ + result, including various training information such as loss and time. + Arguments: + - data (:obj:`List[Dict[int, Any]]`): The input data used for policy forward, including a batch of \ + training samples. For each element in list, the key of the dict is the name of data items and the \ + value is the corresponding data. Usually, the value is torch.Tensor or np.ndarray or there dict/list \ + combinations. In the ``_forward_learn`` method, data often need to first be stacked in the batch \ + dimension by some utility functions such as ``default_preprocess_learn``. \ + For BC, each element in list is a dict containing at least the following keys: ``obs``, ``action``. + Returns: + - info_dict (:obj:`Dict[str, Any]`): The information dict that indicated training result, which will be \ + recorded in text log and tensorboard, values must be python scalar or a list of scalars. For the \ + detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. - def _forward_learn(self, data): - if not isinstance(data, dict): + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + """ + if isinstance(data, list): data = default_collate(data) if self._cuda: data = to_device(data, self._device) @@ -126,10 +176,10 @@ def _forward_learn(self, data): obs, action = data['obs'], data['action'].squeeze() if self._cfg.continuous: if self._cfg.learn.tanh_mask: - ''' + """tanh_mask We mask the action out of range of [tanh(-1),tanh(1)], model will learn information and produce action in [-1,1]. So the action won't always converge to -1 or 1. - ''' + """ mu = self._eval_model.forward(data['obs'])['action'] bound = 1 - 2 / (math.exp(2) + 1) # tanh(1): (e-e**(-1))/(e+e**(-1)) mask = mu.ge(-bound) & mu.le(bound) @@ -170,7 +220,7 @@ def _forward_learn(self, data): loss.backward() backward_time = self._timer.value with self._timer: - if self._cfg.learn.multi_gpu: + if self._cfg.multi_gpu: self.sync_gradients(self._learn_model) sync_time = self._timer.value self._optimizer.step() @@ -184,28 +234,57 @@ def _forward_learn(self, data): 'sync_time': sync_time, } - def _monitor_vars_learn(self): + def _monitor_vars_learn(self) -> List[str]: + """ + Overview: + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. + Returns: + - necessary_keys (:obj:`List[str]`): The list of the necessary keys to be logged. + """ return ['cur_lr', 'total_loss', 'forward_time', 'backward_time', 'sync_time'] def _init_eval(self): + """ + Overview: + Initialize the eval mode of policy, including related attributes and modules. For BC, it contains the \ + eval model to greedily select action with argmax q_value mechanism for discrete action space. + This method will be called in ``__init__`` method if ``eval`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_eval`` method, you'd better name them \ + with prefix ``_eval_`` to avoid conflict with other modes, such as ``self._eval_attr1``. + """ if self._cfg.continuous: self._eval_model = model_wrap(self._model, wrapper_name='base') else: self._eval_model = model_wrap(self._model, wrapper_name='argmax_sample') self._eval_model.reset() - def _forward_eval(self, data): - gfootball_flag = False + def _forward_eval(self, data: Dict[int, Any]) -> Dict[int, Any]: + """ + Overview: + Policy forward function of eval mode (evaluation policy performance by interacting with envs). Forward \ + means that the policy gets some necessary data (mainly observation) from the envs and then returns the \ + action to interact with the envs. + Arguments: + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. + Returns: + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action. The \ + key of the dict is the same as the input data, i.e. environment id. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + """ tensor_input = isinstance(data, torch.Tensor) if tensor_input: data = default_collate(list(data)) else: data_id = list(data.keys()) - if data_id == ['processed_obs', 'raw_obs']: - # for gfootball - gfootball_flag = True - data = {0: data} - data_id = list(data.keys()) data = default_collate(list(data.values())) if self._cuda: data = to_device(data, self._device) @@ -214,22 +293,20 @@ def _forward_eval(self, data): output = self._eval_model.forward(data) if self._cuda: output = to_device(output, 'cpu') - if tensor_input or gfootball_flag: + if tensor_input: return output else: output = default_decollate(output) return {i: d for i, d in zip(data_id, output)} def _init_collect(self) -> None: - r""" + """ Overview: - Collect mode init method. Called by ``self.__init__``. - Init traj and unroll length, collect model. - Enable the eps_greedy_sample + BC policy uses offline dataset so it does not need to collect data. However, sometimes we need to use the \ + trained BC policy to collect data for other purposes. """ self._unroll_len = self._cfg.collect.unroll_len if self._cfg.continuous: - # self._collect_model = model_wrap(self._model, wrapper_name='base') self._collect_model = model_wrap( self._model, wrapper_name='action_noise', @@ -245,14 +322,6 @@ def _init_collect(self) -> None: self._collect_model.reset() def _forward_collect(self, data: Dict[int, Any], **kwargs) -> Dict[int, Any]: - r""" - Overview: - Forward function for collect mode with eps_greedy - Arguments: - - data (:obj:`dict`): Dict type data, including at least ['obs']. - Returns: - - data (:obj:`dict`): The collected data - """ data_id = list(data.keys()) data = default_collate(list(data.values())) if self._cuda: @@ -269,43 +338,16 @@ def _forward_collect(self, data: Dict[int, Any], **kwargs) -> Dict[int, Any]: output = default_decollate(output) return {i: d for i, d in zip(data_id, output)} - def _process_transition(self, obs: Any, model_output: dict, timestep: namedtuple) -> dict: - r""" - Overview: - Generate dict type transition data from inputs. - Arguments: - - obs (:obj:`Any`): Env observation - - model_output (:obj:`dict`): Output of collect model, including at least ['action'] - - timestep (:obj:`namedtuple`): Output after env step, including at least ['obs', 'reward', 'done'] \ - (here 'obs' indicates obs after env step). - Returns: - - transition (:obj:`dict`): Dict type transition data. - """ + def _process_transition(self, obs: Any, policy_output: dict, timestep: namedtuple) -> dict: transition = { 'obs': obs, 'next_obs': timestep.obs, - 'action': model_output['action'], + 'action': policy_output['action'], 'reward': timestep.reward, 'done': timestep.done, } return EasyDict(transition) def _get_train_sample(self, data: List[Dict[str, Any]]) -> List[Dict[str, Any]]: - """ - Overview: - For a given trajectory(transitions, a list of transition) data, process it into a list of sample that \ - can be used for training directly. A train sample can be a processed transition(DQN with nstep TD) \ - or some continuous transitions(DRQN). - Arguments: - - data (:obj:`List[Dict[str, Any]`): The trajectory data(a list of transition), each element is the same \ - format as the return value of ``self._process_transition`` method. - Returns: - - samples (:obj:`dict`): The list of training samples. - - .. note:: - We will vectorize ``process_transition`` and ``get_train_sample`` method in the following release version. \ - And the user can customize the this data processing procecure by overriding this two methods and collector \ - itself. - """ data = get_nstep_return_data(data, 1, 1) return get_train_sample(data, self._unroll_len) diff --git a/ding/policy/bcq.py b/ding/policy/bcq.py new file mode 100755 index 0000000000..2f0643e612 --- /dev/null +++ b/ding/policy/bcq.py @@ -0,0 +1,362 @@ +import copy +from collections import namedtuple +from typing import List, Dict, Any, Tuple + +import torch +import torch.nn.functional as F + +from ding.model import model_wrap +from ding.policy import Policy +from ding.rl_utils import v_1step_td_data, v_1step_td_error, get_train_sample, get_nstep_return_data +from ding.torch_utils import Adam, to_device +from ding.utils import POLICY_REGISTRY +from ding.utils.data import default_collate, default_decollate +from .common_utils import default_preprocess_learn + + +@POLICY_REGISTRY.register('bcq') +class BCQPolicy(Policy): + """ + Overview: + Policy class of BCQ (Batch-Constrained deep Q-learning) algorithm, proposed in \ + https://arxiv.org/abs/1812.02900. + """ + + config = dict( + # (str) Name of the registered RL policy (refer to the "register_policy" function). + type='bcq', + # (bool) Indicates if CUDA should be used for network operations. + cuda=False, + # (bool) Determines whether priority sampling is used in the replay buffer. Default is False. + priority=False, + # (bool) If True, Importance Sampling Weight is used to correct updates. Requires 'priority' to be True. + priority_IS_weight=False, + # (int) Number of random samples in replay buffer before training begins. Default is 10000. + random_collect_size=10000, + # (int) The number of steps for calculating target q_value. + nstep=1, + model=dict( + # (List[int]) Sizes of the hidden layers in the actor network. + actor_head_hidden_size=[400, 300], + # (List[int]) Sizes of the hidden layers in the critic network. + critic_head_hidden_size=[400, 300], + # (float) Maximum perturbation for BCQ. Controls exploration in action space. + phi=0.05, + ), + learn=dict( + # (int) Number of policy updates per data collection step. Higher values indicate more off-policy training. + update_per_collect=1, + # (int) Batch size for each gradient descent step. + batch_size=100, + # (float) Learning rate for the Q-network. Set to 1e-3 if `model.value_network` is True. + learning_rate_q=3e-4, + # (float) Learning rate for the policy network. Set to 1e-3 if `model.value_network` is True. + learning_rate_policy=3e-4, + # (float) Learning rate for the VAE network. Initialize if `model.vae_network` is True. + learning_rate_vae=3e-4, + # (bool) If set to True, the 'done' signals that indicate the end of an episode due to environment time + # limits are disregarded. By default, this is set to False. This setting is particularly useful for tasks + # that have a predetermined episode length, such as HalfCheetah and various other MuJoCo environments, + # where the maximum length is capped at 1000 steps. When enabled, any 'done' signal triggered by reaching + # the maximum episode steps will be overridden to 'False'. This ensures the accurate calculation of the + # Temporal Difference (TD) error, using the formula `gamma * (1 - done) * next_v + reward`, + # even when the episode surpasses the predefined step limit. + ignore_done=False, + # (float) Polyak averaging coefficient for the target network update. Typically small. + target_theta=0.005, + # (float) Discount factor for future rewards, often denoted as gamma. + discount_factor=0.99, + # (float) Lambda for TD(lambda) learning. Weighs the trade-off between bias and variance. + lmbda=0.75, + # (float) Range for uniform weight initialization in the output layer. + init_w=3e-3, + ), + collect=dict( + # (int) Length of trajectory segments for unrolling. Set to higher for longer dependencies. + unroll_len=1, + ), + eval=dict(), + other=dict( + replay_buffer=dict( + # (int) Maximum size of the replay buffer. + replay_buffer_size=1000000, + ), + ), + ) + + def default_model(self) -> Tuple[str, List[str]]: + """ + Overview: + Returns the default model configuration used by the BCQ algorithm. ``__init__`` method will \ + automatically call this method to get the default model setting and create model. + + Returns: + - model_info (:obj:`Tuple[str, List[str]]`): \ + Tuple containing the registered model name and model's import_names. + """ + return 'bcq', ['ding.model.template.bcq'] + + def _init_learn(self) -> None: + """ + Overview: + Initialize the learn mode of policy, including related attributes and modules. For BCQ, it mainly \ + contains optimizer, algorithm-specific arguments such as gamma, main and target model. \ + This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. + """ + # Init + self._priority = self._cfg.priority + self._priority_IS_weight = self._cfg.priority_IS_weight + self.lmbda = self._cfg.learn.lmbda + self.latent_dim = self._cfg.model.action_shape * 2 + + # Optimizers + self._optimizer_q = Adam( + self._model.critic.parameters(), + lr=self._cfg.learn.learning_rate_q, + ) + self._optimizer_policy = Adam( + self._model.actor.parameters(), + lr=self._cfg.learn.learning_rate_policy, + ) + self._optimizer_vae = Adam( + self._model.vae.parameters(), + lr=self._cfg.learn.learning_rate_vae, + ) + + # Algorithm config + self._gamma = self._cfg.learn.discount_factor + + # Main and target models + self._target_model = copy.deepcopy(self._model) + self._target_model = model_wrap( + self._target_model, + wrapper_name='target', + update_type='momentum', + update_kwargs={'theta': self._cfg.learn.target_theta} + ) + self._learn_model = model_wrap(self._model, wrapper_name='base') + self._learn_model.reset() + self._target_model.reset() + self._forward_learn_cnt = 0 + + def _forward_learn(self, data: List[Dict[str, Any]]) -> Dict[str, Any]: + """ + Overview: + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data from the replay buffer and then returns the output \ + result, including various training information such as policy_loss, value_loss, entropy_loss. + Arguments: + - data (:obj:`List[Dict[int, Any]]`): The input data used for policy forward, including a batch of \ + training samples. For each element in list, the key of the dict is the name of data items and the \ + value is the corresponding data. Usually, the value is torch.Tensor or np.ndarray or there dict/list \ + combinations. In the ``_forward_learn`` method, data often need to first be stacked in the batch \ + dimension by some utility functions such as ``default_preprocess_learn``. \ + For BCQ, each element in list is a dict containing at least the following keys: \ + ['obs', 'action', 'adv', 'value', 'weight']. + Returns: + - info_dict (:obj:`Dict[str, Any]`): The information dict that indicated training result, which will be \ + recorded in text log and tensorboard, values must be python scalar or a list of scalars. For the \ + detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement your own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + """ + loss_dict = {} + # Data preprocessing operations, such as stack data, cpu to cuda device + data = default_preprocess_learn( + data, + use_priority=self._priority, + use_priority_IS_weight=self._cfg.priority_IS_weight, + ignore_done=self._cfg.learn.ignore_done, + use_nstep=False + ) + if len(data.get('action').shape) == 1: + data['action'] = data['action'].reshape(-1, 1) + + if self._cuda: + data = to_device(data, self._device) + + self._learn_model.train() + self._target_model.train() + obs = data['obs'] + next_obs = data['next_obs'] + reward = data['reward'] + done = data['done'] + batch_size = obs.shape[0] + + # train_vae + vae_out = self._model.forward(data, mode='compute_vae') + recon, mean, log_std = vae_out['recons_action'], vae_out['mu'], vae_out['log_var'] + recons_loss = F.mse_loss(recon, data['action']) + kld_loss = torch.mean(-0.5 * torch.sum(1 + log_std - mean ** 2 - log_std.exp(), dim=1), dim=0) + loss_dict['recons_loss'] = recons_loss + loss_dict['kld_loss'] = kld_loss + vae_loss = recons_loss + 0.5 * kld_loss + loss_dict['vae_loss'] = vae_loss + self._optimizer_vae.zero_grad() + vae_loss.backward() + self._optimizer_vae.step() + + # train_critic + q_value = self._learn_model.forward(data, mode='compute_critic')['q_value'] + + with (torch.no_grad()): + next_obs_rep = torch.repeat_interleave(next_obs, 10, 0) + z = torch.randn((next_obs_rep.shape[0], self.latent_dim)).to(self._device).clamp(-0.5, 0.5) + vae_action = self._model.vae.decode_with_obs(z, next_obs_rep)['reconstruction_action'] + next_action = self._target_model.forward({ + 'obs': next_obs_rep, + 'action': vae_action + }, mode='compute_actor')['action'] + + next_data = {'obs': next_obs_rep, 'action': next_action} + target_q_value = self._target_model.forward(next_data, mode='compute_critic')['q_value'] + # the value of a policy according to the maximum entropy objective + # find min one as target q value + target_q_value = self.lmbda * torch.min(target_q_value[0], target_q_value[1]) \ + + (1 - self.lmbda) * torch.max(target_q_value[0], target_q_value[1]) + target_q_value = target_q_value.reshape(batch_size, -1).max(1)[0].reshape(-1, 1) + + q_data0 = v_1step_td_data(q_value[0], target_q_value, reward, done, data['weight']) + loss_dict['critic_loss'], td_error_per_sample0 = v_1step_td_error(q_data0, self._gamma) + q_data1 = v_1step_td_data(q_value[1], target_q_value, reward, done, data['weight']) + loss_dict['twin_critic_loss'], td_error_per_sample1 = v_1step_td_error(q_data1, self._gamma) + td_error_per_sample = (td_error_per_sample0 + td_error_per_sample1) / 2 + + self._optimizer_q.zero_grad() + (loss_dict['critic_loss'] + loss_dict['twin_critic_loss']).backward() + self._optimizer_q.step() + + # train_policy + z = torch.randn((obs.shape[0], self.latent_dim)).to(self._device).clamp(-0.5, 0.5) + sample_action = self._model.vae.decode_with_obs(z, obs)['reconstruction_action'] + input = {'obs': obs, 'action': sample_action} + perturbed_action = self._model.forward(input, mode='compute_actor')['action'] + q_input = {'obs': obs, 'action': perturbed_action} + q = self._learn_model.forward(q_input, mode='compute_critic')['q_value'][0] + loss_dict['actor_loss'] = -q.mean() + self._optimizer_policy.zero_grad() + loss_dict['actor_loss'].backward() + self._optimizer_policy.step() + self._forward_learn_cnt += 1 + self._target_model.update(self._learn_model.state_dict()) + return { + 'td_error': td_error_per_sample.detach().mean().item(), + 'target_q_value': target_q_value.detach().mean().item(), + **loss_dict + } + + def _monitor_vars_learn(self) -> List[str]: + """ + Overview: + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. + Returns: + - necessary_keys (:obj:`List[str]`): The list of the necessary keys to be logged. + """ + return [ + 'td_error', 'target_q_value', 'critic_loss', 'twin_critic_loss', 'actor_loss', 'recons_loss', 'kld_loss', + 'vae_loss' + ] + + def _state_dict_learn(self) -> Dict[str, Any]: + """ + Overview: + Return the state_dict of learn mode, usually including model and optimizer. + Returns: + - state_dict (:obj:`Dict[str, Any]`): The dict of current policy learn state, for saving and restoring. + """ + ret = { + 'model': self._learn_model.state_dict(), + 'target_model': self._target_model.state_dict(), + 'optimizer_q': self._optimizer_q.state_dict(), + 'optimizer_policy': self._optimizer_policy.state_dict(), + 'optimizer_vae': self._optimizer_vae.state_dict(), + } + return ret + + def _init_eval(self) -> None: + """ + Overview: + Initialize the eval mode of policy, including related attributes and modules. + This method will be called in ``__init__`` method if ``eval`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_eval`` method, you'd better name them \ + with prefix ``_eval_`` to avoid conflict with other modes, such as ``self._eval_attr1``. + """ + self._eval_model = model_wrap(self._model, wrapper_name='base') + self._eval_model.reset() + + def _forward_eval(self, data: Dict[int, Any]) -> Dict[int, Any]: + """ + Overview: + Policy forward function of eval mode (evaluation policy performance by interacting with envs). Forward \ + means that the policy gets some necessary data (mainly observation) from the envs and then returns the \ + action to interact with the envs. + Arguments: + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. + Returns: + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action. The \ + key of the dict is the same as the input data, i.e., environment id. + + .. note:: + The input value can be ``torch.Tensor`` or dict/list combinations, current policy supports all of them. \ + For the data type that is not supported, the main reason is that the corresponding model does not \ + support it. You can implement your own model rather than use the default model. For more information, \ + please raise an issue in GitHub repo, and we will continue to follow up. + """ + data_id = list(data.keys()) + data = default_collate(list(data.values())) + if self._cuda: + data = to_device(data, self._device) + data = {'obs': data} + self._eval_model.eval() + with torch.no_grad(): + output = self._eval_model.forward(data, mode='compute_eval') + if self._cuda: + output = to_device(output, 'cpu') + output = default_decollate(output) + return {i: d for i, d in zip(data_id, output)} + + def _init_collect(self) -> None: + """ + Overview: + Initialize the collect mode of policy, including related attributes and modules. For BCQ, it contains the \ + collect_model to balance the exploration and exploitation with ``eps_greedy_sample`` \ + mechanism, and other algorithm-specific arguments such as gamma and nstep. + This method will be called in ``__init__`` method if ``collect`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_collect`` method, you'd better name them \ + with prefix ``_collect_`` to avoid conflict with other modes, such as ``self._collect_attr1``. + """ + self._unroll_len = self._cfg.collect.unroll_len + self._gamma = self._cfg.discount_factor + self._nstep = self._cfg.nstep + self._collect_model = model_wrap(self._model, wrapper_name='eps_greedy_sample') + self._collect_model.reset() + + def _get_train_sample(self, transitions: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + pass + + def _forward_collect(self, data: dict, **kwargs) -> dict: + pass + + def _process_transition(self, obs: Any, model_output: dict, timestep: namedtuple) -> dict: + pass diff --git a/ding/policy/bdq.py b/ding/policy/bdq.py new file mode 100644 index 0000000000..c994c6cd45 --- /dev/null +++ b/ding/policy/bdq.py @@ -0,0 +1,393 @@ +from typing import List, Dict, Any, Tuple +from collections import namedtuple +import copy +import torch + +from ding.torch_utils import Adam, to_device, ContrastiveLoss +from ding.rl_utils import q_nstep_td_data, bdq_nstep_td_error, get_nstep_return_data, get_train_sample +from ding.model import model_wrap +from ding.utils import POLICY_REGISTRY +from ding.utils.data import default_collate, default_decollate + +from .base_policy import Policy +from .common_utils import default_preprocess_learn + + +@POLICY_REGISTRY.register('bdq') +class BDQPolicy(Policy): + r""" + Overview: + Policy class of BDQ algorithm, extended by PER/multi-step TD. \ + referenced paper Action Branching Architectures for Deep Reinforcement Learning \ + + .. note:: + BDQ algorithm contains a neural architecture featuring a shared decision module \ + followed by several network branches, one for each action dimension. + Config: + == ==================== ======== ============== ======================================== ======================= + ID Symbol Type Default Value Description Other(Shape) + == ==================== ======== ============== ======================================== ======================= + 1 ``type`` str bdq | RL policy register name, refer to | This arg is optional, + | registry ``POLICY_REGISTRY`` | a placeholder + 2 ``cuda`` bool False | Whether to use cuda for network | This arg can be diff- + | erent from modes + 3 ``on_policy`` bool False | Whether the RL algorithm is on-policy + | or off-policy + 4 ``priority`` bool False | Whether use priority(PER) | Priority sample, + | update priority + 5 | ``priority_IS`` bool False | Whether use Importance Sampling Weight + | ``_weight`` | to correct biased update. If True, + | priority must be True. + 6 | ``discount_`` float 0.97, | Reward's future discount factor, aka. | May be 1 when sparse + | ``factor`` [0.95, 0.999] | gamma | reward env + 7 ``nstep`` int 1, | N-step reward discount sum for target + [3, 5] | q_value estimation + 8 | ``learn.update`` int 3 | How many updates(iterations) to train | This args can be vary + | ``per_collect`` | after collector's one collection. Only | from envs. Bigger val + | valid in serial training | means more off-policy + | ``_gpu`` + 10 | ``learn.batch_`` int 64 | The number of samples of an iteration + | ``size`` + 11 | ``learn.learning`` float 0.001 | Gradient step length of an iteration. + | ``_rate`` + 12 | ``learn.target_`` int 100 | Frequence of target network update. | Hard(assign) update + | ``update_freq`` + 13 | ``learn.ignore_`` bool False | Whether ignore done for target value | Enable it for some + | ``done`` | calculation. | fake termination env + 14 ``collect.n_sample`` int [8, 128] | The number of training samples of a | It varies from + | call of collector. | different envs + 15 | ``collect.unroll`` int 1 | unroll length of an iteration | In RNN, unroll_len>1 + | ``_len`` + 16 | ``other.eps.type`` str exp | exploration rate decay type | Support ['exp', + | 'linear']. + 17 | ``other.eps.`` float 0.95 | start value of exploration rate | [0,1] + | ``start`` + 18 | ``other.eps.`` float 0.1 | end value of exploration rate | [0,1] + | ``end`` + 19 | ``other.eps.`` int 10000 | decay length of exploration | greater than 0. set + | ``decay`` | decay=10000 means + | the exploration rate + | decay from start + | value to end value + | during decay length. + == ==================== ======== ============== ======================================== ======================= + """ + + config = dict( + type='bdq', + # (bool) Whether use cuda in policy + cuda=False, + # (bool) Whether learning policy is the same as collecting data policy(on-policy) + on_policy=False, + # (bool) Whether enable priority experience sample + priority=False, + # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. + priority_IS_weight=False, + # (float) Discount factor(gamma) for returns + discount_factor=0.97, + # (int) The number of step for calculating target q_value + nstep=1, + learn=dict( + + # How many updates(iterations) to train after collector's one collection. + # Bigger "update_per_collect" means bigger off-policy. + # collect data -> update policy-> collect data -> ... + update_per_collect=3, + # (int) How many samples in a training batch + batch_size=64, + # (float) The step size of gradient descent + learning_rate=0.001, + # ============================================================== + # The following configs are algorithm-specific + # ============================================================== + # (int) Frequence of target network update. + target_update_freq=100, + # (bool) Whether ignore done(usually for max step termination env) + ignore_done=False, + ), + # collect_mode config + collect=dict( + # (int) Only one of [n_sample, n_episode] shoule be set + # n_sample=8, + # (int) Cut trajectories into pieces with length "unroll_len". + unroll_len=1, + ), + eval=dict(), + # other config + other=dict( + # Epsilon greedy with decay. + eps=dict( + # (str) Decay type. Support ['exp', 'linear']. + type='exp', + # (float) Epsilon start value + start=0.95, + # (float) Epsilon end value + end=0.1, + # (int) Decay length(env step) + decay=10000, + ), + replay_buffer=dict(replay_buffer_size=10000, ), + ), + ) + + def default_model(self) -> Tuple[str, List[str]]: + """ + Overview: + Return this algorithm default model setting for demonstration. + Returns: + - model_info (:obj:`Tuple[str, List[str]]`): model name and mode import_names + + .. note:: + The user can define and use customized network model but must obey the same inferface definition indicated \ + by import_names path. For BDQ, ``ding.model.template.q_learning.BDQ`` + """ + return 'bdq', ['ding.model.template.q_learning'] + + def _init_learn(self) -> None: + """ + Overview: + Learn mode init method. Called by ``self.__init__``, initialize the optimizer, algorithm arguments, main \ + and target model. + """ + self._priority = self._cfg.priority + self._priority_IS_weight = self._cfg.priority_IS_weight + # Optimizer + self._optimizer = Adam(self._model.parameters(), lr=self._cfg.learn.learning_rate) + + self._gamma = self._cfg.discount_factor + self._nstep = self._cfg.nstep + + # use model_wrapper for specialized demands of different modes + self._target_model = copy.deepcopy(self._model) + self._target_model = model_wrap( + self._target_model, + wrapper_name='target', + update_type='assign', + update_kwargs={'freq': self._cfg.learn.target_update_freq} + ) + self._learn_model = model_wrap(self._model, wrapper_name='argmax_sample') + self._learn_model.reset() + self._target_model.reset() + + def _forward_learn(self, data: Dict[str, Any]) -> Dict[str, Any]: + """ + Overview: + Forward computation graph of learn mode(updating policy). + Arguments: + - data (:obj:`Dict[str, Any]`): Dict type data, a batch of data for training, values are torch.Tensor or \ + np.ndarray or dict/list combinations. + Returns: + - info_dict (:obj:`Dict[str, Any]`): Dict type data, a info dict indicated training result, which will be \ + recorded in text log and tensorboard, values are python scalar or a list of scalars. + ArgumentsKeys: + - necessary: ``obs``, ``action``, ``reward``, ``next_obs``, ``done`` + - optional: ``value_gamma``, ``IS`` + ReturnsKeys: + - necessary: ``cur_lr``, ``total_loss``, ``priority`` + - optional: ``action_distribution`` + """ + data = default_preprocess_learn( + data, + use_priority=self._priority, + use_priority_IS_weight=self._cfg.priority_IS_weight, + ignore_done=self._cfg.learn.ignore_done, + use_nstep=True + ) + + if self._cuda: + data = to_device(data, self._device) + # ==================== + # Q-learning forward + # ==================== + self._learn_model.train() + self._target_model.train() + # Current q value (main model) + q_value = self._learn_model.forward(data['obs'])['logit'] + # Target q value + with torch.no_grad(): + target_q_value = self._target_model.forward(data['next_obs'])['logit'] + # Max q value action (main model) + target_q_action = self._learn_model.forward(data['next_obs'])['action'] + if data['action'].shape != target_q_action.shape: + data['action'] = data['action'].unsqueeze(-1) + + data_n = q_nstep_td_data( + q_value, target_q_value, data['action'], target_q_action, data['reward'], data['done'], data['weight'] + ) + value_gamma = data.get('value_gamma') + loss, td_error_per_sample = bdq_nstep_td_error(data_n, self._gamma, nstep=self._nstep, value_gamma=value_gamma) + + # ==================== + # Q-learning update + # ==================== + self._optimizer.zero_grad() + loss.backward() + if self._cfg.multi_gpu: + self.sync_gradients(self._learn_model) + self._optimizer.step() + + # ============= + # after update + # ============= + self._target_model.update(self._learn_model.state_dict()) + update_info = { + 'cur_lr': self._optimizer.defaults['lr'], + 'total_loss': loss.item(), + 'q_value': q_value.mean().item(), + 'target_q_value': target_q_value.mean().item(), + 'priority': td_error_per_sample.abs().tolist(), + # Only discrete action satisfying len(data['action'])==1 can return this and draw histogram on tensorboard. + # '[histogram]action_distribution': data['action'], + } + q_value_per_branch = torch.mean(q_value, 2, keepdim=False) + for i in range(self._model.num_branches): + update_info['q_value_b_' + str(i)] = q_value_per_branch[:, i].mean().item() + return update_info + + def _monitor_vars_learn(self) -> List[str]: + return ['cur_lr', 'total_loss', 'q_value'] + ['q_value_b_' + str(i) for i in range(self._model.num_branches)] + + def _state_dict_learn(self) -> Dict[str, Any]: + """ + Overview: + Return the state_dict of learn mode, usually including model and optimizer. + Returns: + - state_dict (:obj:`Dict[str, Any]`): the dict of current policy learn state, for saving and restoring. + """ + return { + 'model': self._learn_model.state_dict(), + 'target_model': self._target_model.state_dict(), + 'optimizer': self._optimizer.state_dict(), + } + + def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: + """ + Overview: + Load the state_dict variable into policy learn mode. + Arguments: + - state_dict (:obj:`Dict[str, Any]`): the dict of policy learn state saved before. + + .. tip:: + If you want to only load some parts of model, you can simply set the ``strict`` argument in \ + load_state_dict to ``False``, or refer to ``ding.torch_utils.checkpoint_helper`` for more \ + complicated operation. + """ + self._learn_model.load_state_dict(state_dict['model']) + self._target_model.load_state_dict(state_dict['target_model']) + self._optimizer.load_state_dict(state_dict['optimizer']) + + def _init_collect(self) -> None: + """ + Overview: + Collect mode init method. Called by ``self.__init__``, initialize algorithm arguments and collect_model, \ + enable the eps_greedy_sample for exploration. + """ + self._unroll_len = self._cfg.collect.unroll_len + self._gamma = self._cfg.discount_factor # necessary for parallel + self._nstep = self._cfg.nstep # necessary for parallel + self._collect_model = model_wrap(self._model, wrapper_name='eps_greedy_sample') + self._collect_model.reset() + + def _forward_collect(self, data: Dict[int, Any], eps: float) -> Dict[int, Any]: + """ + Overview: + Forward computation graph of collect mode(collect training data), with eps_greedy for exploration. + Arguments: + - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \ + values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer. + - eps (:obj:`float`): epsilon value for exploration, which is decayed by collected env step. + Returns: + - output (:obj:`Dict[int, Any]`): The dict of predicting policy_output(action) for the interaction with \ + env and the constructing of transition. + ArgumentsKeys: + - necessary: ``obs`` + ReturnsKeys + - necessary: ``logit``, ``action`` + """ + data_id = list(data.keys()) + data = default_collate(list(data.values())) + if self._cuda: + data = to_device(data, self._device) + self._collect_model.eval() + with torch.no_grad(): + output = self._collect_model.forward(data, eps=eps) + if self._cuda: + output = to_device(output, 'cpu') + output = default_decollate(output) + return {i: d for i, d in zip(data_id, output)} + + def _get_train_sample(self, data: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Overview: + For a given trajectory(transitions, a list of transition) data, process it into a list of sample that \ + can be used for training directly. A train sample can be a processed transition(BDQ with nstep TD). + Arguments: + - data (:obj:`List[Dict[str, Any]`): The trajectory data(a list of transition), each element is the same \ + format as the return value of ``self._process_transition`` method. + Returns: + - samples (:obj:`dict`): The list of training samples. + + .. note:: + We will vectorize ``process_transition`` and ``get_train_sample`` method in the following release version. \ + And the user can customize the this data processing procecure by overriding this two methods and collector \ + itself. + """ + data = get_nstep_return_data(data, self._nstep, gamma=self._gamma) + return get_train_sample(data, self._unroll_len) + + def _process_transition(self, obs: Any, policy_output: Dict[str, Any], timestep: namedtuple) -> Dict[str, Any]: + """ + Overview: + Generate a transition(e.g.: ) for this algorithm training. + Arguments: + - obs (:obj:`Any`): Env observation. + - policy_output (:obj:`Dict[str, Any]`): The output of policy collect mode(``self._forward_collect``),\ + including at least ``action``. + - timestep (:obj:`namedtuple`): The output after env step(execute policy output action), including at \ + least ``obs``, ``reward``, ``done``, (here obs indicates obs after env step). + Returns: + - transition (:obj:`dict`): Dict type transition data. + """ + transition = { + 'obs': obs, + 'next_obs': timestep.obs, + 'action': policy_output['action'], + 'reward': timestep.reward, + 'done': timestep.done, + } + return transition + + def _init_eval(self) -> None: + r""" + Overview: + Evaluate mode init method. Called by ``self.__init__``, initialize eval_model. + """ + self._eval_model = model_wrap(self._model, wrapper_name='argmax_sample') + self._eval_model.reset() + + def _forward_eval(self, data: Dict[int, Any]) -> Dict[int, Any]: + """ + Overview: + Forward computation graph of eval mode(evaluate policy performance), at most cases, it is similar to \ + ``self._forward_collect``. + Arguments: + - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \ + values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer. + Returns: + - output (:obj:`Dict[int, Any]`): The dict of predicting action for the interaction with env. + ArgumentsKeys: + - necessary: ``obs`` + ReturnsKeys + - necessary: ``action`` + """ + data_id = list(data.keys()) + data = default_collate(list(data.values())) + if self._cuda: + data = to_device(data, self._device) + self._eval_model.eval() + with torch.no_grad(): + output = self._eval_model.forward(data) + if self._cuda: + output = to_device(output, 'cpu') + output = default_decollate(output) + return {i: d for i, d in zip(data_id, output)} diff --git a/ding/policy/c51.py b/ding/policy/c51.py index 441c1ee3de..0b6f36d68e 100644 --- a/ding/policy/c51.py +++ b/ding/policy/c51.py @@ -68,8 +68,7 @@ class C51Policy(DQNPolicy): n_atom=51, ), learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, + # How many updates(iterations) to train after collector's one collection. # Bigger "update_per_collect" means bigger off-policy. # collect data -> update policy-> collect data -> ... @@ -158,15 +157,20 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: self._learn_model.train() self._target_model.train() # Current q value (main model) - q_value = self._learn_model.forward(data['obs'])['distribution'] + output = self._learn_model.forward(data['obs']) + q_value = output['logit'] + q_value_dist = output['distribution'] # Target q value with torch.no_grad(): - target_q_value = self._target_model.forward(data['next_obs'])['distribution'] + target_output = self._target_model.forward(data['next_obs']) + target_q_value_dist = target_output['distribution'] + target_q_value = target_output['logit'] # Max q value action (main model) target_q_action = self._learn_model.forward(data['next_obs'])['action'] data_n = dist_nstep_td_data( - q_value, target_q_value, data['action'], target_q_action, data['reward'], data['done'], data['weight'] + q_value_dist, target_q_value_dist, data['action'], target_q_action, data['reward'], data['done'], + data['weight'] ) value_gamma = data.get('value_gamma') loss, td_error_per_sample = dist_nstep_td_error( @@ -178,7 +182,7 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: # ==================== self._optimizer.zero_grad() loss.backward() - if self._cfg.learn.multi_gpu: + if self._cfg.multi_gpu: self.sync_gradients(self._learn_model) self._optimizer.step() @@ -189,11 +193,16 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: return { 'cur_lr': self._optimizer.defaults['lr'], 'total_loss': loss.item(), + 'q_value': q_value.mean().item(), + 'target_q_value': target_q_value.mean().item(), 'priority': td_error_per_sample.abs().tolist(), # Only discrete action satisfying len(data['action'])==1 can return this and draw histogram on tensorboard. # '[histogram]action_distribution': data['action'], } + def _monitor_vars_learn(self) -> List[str]: + return ['cur_lr', 'total_loss', 'q_value', 'target_q_value'] + def _state_dict_learn(self) -> Dict[str, Any]: return { 'model': self._learn_model.state_dict(), diff --git a/ding/policy/collaq.py b/ding/policy/collaq.py index ea738f61aa..961d4fed80 100644 --- a/ding/policy/collaq.py +++ b/ding/policy/collaq.py @@ -58,8 +58,7 @@ class CollaQPolicy(Policy): # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. priority_IS_weight=False, learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, + # (int) Collect n_episode data, update_model n_iteration times update_per_collect=20, # (int) The number of data for a train iteration diff --git a/ding/policy/coma.py b/ding/policy/coma.py index f9a5f5ffbe..5940a25f0e 100644 --- a/ding/policy/coma.py +++ b/ding/policy/coma.py @@ -61,8 +61,6 @@ class COMAPolicy(Policy): # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. priority_IS_weight=False, learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, update_per_collect=20, batch_size=32, learning_rate=0.0005, diff --git a/ding/policy/command_mode_policy_instance.py b/ding/policy/command_mode_policy_instance.py index 92f6ec1d90..acf224185b 100644 --- a/ding/policy/command_mode_policy_instance.py +++ b/ding/policy/command_mode_policy_instance.py @@ -3,6 +3,7 @@ from .base_policy import CommandModePolicy from .dqn import DQNPolicy, DQNSTDIMPolicy +from .mdqn import MDQNPolicy from .c51 import C51Policy from .qrdqn import QRDQNPolicy from .iqn import IQNPolicy @@ -15,6 +16,7 @@ from .ppo import PPOPolicy, PPOOffPolicy, PPOPGPolicy, PPOSTDIMPolicy from .offppo_collect_traj import OffPPOCollectTrajPolicy from .ppg import PPGPolicy, PPGOffPolicy +from .pg import PGPolicy from .a2c import A2CPolicy from .impala import IMPALAPolicy from .ngu import NGUPolicy @@ -22,8 +24,9 @@ from .td3 import TD3Policy from .td3_vae import TD3VAEPolicy from .td3_bc import TD3BCPolicy -from .sac import SACPolicy, SACDiscretePolicy +from .sac import SACPolicy, DiscreteSACPolicy, SQILSACPolicy from .mbpolicy.mbsac import MBSACPolicy, STEVESACPolicy +from .mbpolicy.dreamer import DREAMERPolicy from .qmix import QMIXPolicy from .wqmix import WQMIXPolicy from .collaq import CollaQPolicy @@ -39,10 +42,18 @@ from .r2d3 import R2D3Policy from .d4pg import D4PGPolicy -from .cql import CQLPolicy, CQLDiscretePolicy -from .decision_transformer import DTPolicy +from .cql import CQLPolicy, DiscreteCQLPolicy +from .iql import IQLPolicy +from .dt import DTPolicy from .pdqn import PDQNPolicy -from .sac import SQILSACPolicy +from .madqn import MADQNPolicy +from .bdq import BDQPolicy +from .bcq import BCQPolicy +from .edac import EDACPolicy +from .prompt_pg import PromptPGPolicy +from .plan_diffuser import PDPolicy +from .happo import HAPPOPolicy +from .prompt_awr import PromptAWRPolicy class EpsCommandModePolicy(CommandModePolicy): @@ -93,6 +104,16 @@ def _get_setting_eval(self, command_info: dict) -> dict: return {} +@POLICY_REGISTRY.register('bdq_command') +class BDQCommandModePolicy(BDQPolicy, EpsCommandModePolicy): + pass + + +@POLICY_REGISTRY.register('mdqn_command') +class MDQNCommandModePolicy(MDQNPolicy, EpsCommandModePolicy): + pass + + @POLICY_REGISTRY.register('dqn_command') class DQNCommandModePolicy(DQNPolicy, EpsCommandModePolicy): pass @@ -168,6 +189,11 @@ class PPOCommandModePolicy(PPOPolicy, DummyCommandModePolicy): pass +@POLICY_REGISTRY.register('happo_command') +class HAPPOCommandModePolicy(HAPPOPolicy, DummyCommandModePolicy): + pass + + @POLICY_REGISTRY.register('ppo_stdim_command') class PPOSTDIMCommandModePolicy(PPOSTDIMPolicy, DummyCommandModePolicy): pass @@ -188,6 +214,11 @@ class PPOOffCollectTrajCommandModePolicy(OffPPOCollectTrajPolicy, DummyCommandMo pass +@POLICY_REGISTRY.register('pg_command') +class PGCommandModePolicy(PGPolicy, DummyCommandModePolicy): + pass + + @POLICY_REGISTRY.register('a2c_command') class A2CCommandModePolicy(A2CPolicy, DummyCommandModePolicy): pass @@ -208,6 +239,11 @@ class PPGCommandModePolicy(PPGPolicy, DummyCommandModePolicy): pass +@POLICY_REGISTRY.register('madqn_command') +class MADQNCommandModePolicy(MADQNPolicy, EpsCommandModePolicy): + pass + + @POLICY_REGISTRY.register('ddpg_command') class DDPGCommandModePolicy(DDPGPolicy, CommandModePolicy): @@ -277,13 +313,23 @@ class STEVESACCommandModePolicy(STEVESACPolicy, DummyCommandModePolicy): pass +@POLICY_REGISTRY.register('dreamer_command') +class DREAMERCommandModePolicy(DREAMERPolicy, DummyCommandModePolicy): + pass + + @POLICY_REGISTRY.register('cql_command') class CQLCommandModePolicy(CQLPolicy, DummyCommandModePolicy): pass -@POLICY_REGISTRY.register('cql_discrete_command') -class CQLDiscreteCommandModePolicy(CQLDiscretePolicy, EpsCommandModePolicy): +@POLICY_REGISTRY.register('iql_command') +class IQLCommandModePolicy(IQLPolicy, DummyCommandModePolicy): + pass + + +@POLICY_REGISTRY.register('discrete_cql_command') +class DiscreteCQLCommandModePolicy(DiscreteCQLPolicy, EpsCommandModePolicy): pass @@ -342,8 +388,8 @@ class PDQNCommandModePolicy(PDQNPolicy, EpsCommandModePolicy): pass -@POLICY_REGISTRY.register('sac_discrete_command') -class SACDiscreteCommandModePolicy(SACDiscretePolicy, EpsCommandModePolicy): +@POLICY_REGISTRY.register('discrete_sac_command') +class DiscreteSACCommandModePolicy(DiscreteSACPolicy, EpsCommandModePolicy): pass @@ -357,6 +403,21 @@ class IBCCommandModePolicy(IBCPolicy, DummyCommandModePolicy): pass +@POLICY_REGISTRY.register('bcq_command') +class BCQCommandModelPolicy(BCQPolicy, DummyCommandModePolicy): + pass + + +@POLICY_REGISTRY.register('edac_command') +class EDACCommandModelPolicy(EDACPolicy, DummyCommandModePolicy): + pass + + +@POLICY_REGISTRY.register('pd_command') +class PDCommandModelPolicy(PDPolicy, DummyCommandModePolicy): + pass + + @POLICY_REGISTRY.register('bc_command') class BCCommandModePolicy(BehaviourCloningPolicy, DummyCommandModePolicy): @@ -396,3 +457,13 @@ def _get_setting_learn(self, command_info: dict) -> dict: def _get_setting_eval(self, command_info: dict) -> dict: return {} + + +@POLICY_REGISTRY.register('prompt_pg_command') +class PromptPGCommandModePolicy(PromptPGPolicy, DummyCommandModePolicy): + pass + + +@POLICY_REGISTRY.register('prompt_awr_command') +class PromptAWRCommandModePolicy(PromptAWRPolicy, DummyCommandModePolicy): + pass diff --git a/ding/policy/common_utils.py b/ding/policy/common_utils.py index 6a20f1aaa8..f5b733fa35 100644 --- a/ding/policy/common_utils.py +++ b/ding/policy/common_utils.py @@ -1,6 +1,28 @@ -from typing import List, Any +from typing import List, Any, Dict, Callable import torch +import torch.nn as nn +import numpy as np +import treetensor.torch as ttorch from ding.utils.data import default_collate +from ding.torch_utils import to_tensor, to_ndarray, unsqueeze, squeeze +from ding.torch_utils import NoiseLinearLayer + + +def set_noise_mode(module: nn.Module, noise_enabled: bool): + """ + Overview: + Recursively set the 'enable_noise' attribute for all NoiseLinearLayer modules within the given module. + This function is typically used in algorithms such as NoisyNet and Rainbow. + During training, 'enable_noise' should be set to True to enable noise for exploration. + During inference or evaluation, it should be set to False to disable noise for deterministic behavior. + + Arguments: + - module (:obj:`nn.Module`): The root module to search for NoiseLinearLayer instances. + - noise_enabled (:obj:`bool`): Whether to enable or disable noise. + """ + for m in module.modules(): + if isinstance(m, NoiseLinearLayer): + m.enable_noise = noise_enabled def default_preprocess_learn( @@ -9,13 +31,41 @@ def default_preprocess_learn( use_priority: bool = False, use_nstep: bool = False, ignore_done: bool = False, -) -> dict: +) -> Dict[str, torch.Tensor]: + """ + Overview: + Default data pre-processing in policy's ``_forward_learn`` method, including stacking batch data, preprocess \ + ignore done, nstep and priority IS weight. + Arguments: + - data (:obj:`List[Any]`): The list of a training batch samples, each sample is a dict of PyTorch Tensor. + - use_priority_IS_weight (:obj:`bool`): Whether to use priority IS weight correction, if True, this function \ + will set the weight of each sample to the priority IS weight. + - use_priority (:obj:`bool`): Whether to use priority, if True, this function will set the priority IS weight. + - use_nstep (:obj:`bool`): Whether to use nstep TD error, if True, this function will reshape the reward. + - ignore_done (:obj:`bool`): Whether to ignore done, if True, this function will set the done to 0. + Returns: + - data (:obj:`Dict[str, torch.Tensor]`): The preprocessed dict data whose values can be directly used for \ + the following model forward and loss computation. + """ # data preprocess - data = default_collate(data) + elem = data[0] + if isinstance(elem['action'], (np.ndarray, torch.Tensor)) and elem['action'].dtype in [np.int64, torch.int64]: + data = default_collate(data, cat_1dim=True) # for discrete action + else: + data = default_collate(data, cat_1dim=False) # for continuous action + if 'value' in data and data['value'].dim() == 2 and data['value'].shape[1] == 1: + data['value'] = data['value'].squeeze(-1) + if 'adv' in data and data['adv'].dim() == 2 and data['adv'].shape[1] == 1: + data['adv'] = data['adv'].squeeze(-1) + if ignore_done: data['done'] = torch.zeros_like(data['done']).float() else: data['done'] = data['done'].float() + + if data['done'].dim() == 2 and data['done'].shape[1] == 1: + data['done'] = data['done'].squeeze(-1) + if use_priority_IS_weight: assert use_priority, "Use IS Weight correction, but Priority is not used." if use_priority and use_priority_IS_weight: @@ -26,11 +76,82 @@ def default_preprocess_learn( else: data['weight'] = data.get('weight', None) if use_nstep: - # Reward reshaping for n-step + # reward reshaping for n-step reward = data['reward'] if len(reward.shape) == 1: reward = reward.unsqueeze(1) - # reward: (batch_size, nstep) -> (nstep, batch_size) - data['reward'] = reward.permute(1, 0).contiguous() + # single agent reward: (batch_size, nstep) -> (nstep, batch_size) + # multi-agent reward: (batch_size, agent_dim, nstep) -> (nstep, batch_size, agent_dim) + # Assuming 'reward' is a PyTorch tensor with shape (batch_size, nstep) or (batch_size, agent_dim, nstep) + if reward.ndim == 2: + # For a 2D tensor, simply transpose it to get (nstep, batch_size) + data['reward'] = reward.transpose(0, 1).contiguous() + elif reward.ndim == 3: + # For a 3D tensor, move the last dimension to the front to get (nstep, batch_size, agent_dim) + data['reward'] = reward.permute(2, 0, 1).contiguous() + else: + raise ValueError("The 'reward' tensor must be either 2D or 3D. Got shape: {}".format(reward.shape)) + else: + if data['reward'].dim() == 2 and data['reward'].shape[1] == 1: + data['reward'] = data['reward'].squeeze(-1) return data + + +def single_env_forward_wrapper(forward_fn: Callable) -> Callable: + """ + Overview: + Wrap policy to support gym-style interaction between policy and single environment. + Arguments: + - forward_fn (:obj:`Callable`): The original forward function of policy. + Returns: + - wrapped_forward_fn (:obj:`Callable`): The wrapped forward function of policy. + Examples: + >>> env = gym.make('CartPole-v0') + >>> policy = DQNPolicy(...) + >>> forward_fn = single_env_forward_wrapper(policy.eval_mode.forward) + >>> obs = env.reset() + >>> action = forward_fn(obs) + >>> next_obs, rew, done, info = env.step(action) + + """ + + def _forward(obs): + obs = {0: unsqueeze(to_tensor(obs))} + action = forward_fn(obs)[0]['action'] + action = to_ndarray(squeeze(action)) + return action + + return _forward + + +def single_env_forward_wrapper_ttorch(forward_fn: Callable, cuda: bool = True) -> Callable: + """ + Overview: + Wrap policy to support gym-style interaction between policy and single environment for treetensor (ttorch) data. + Arguments: + - forward_fn (:obj:`Callable`): The original forward function of policy. + - cuda (:obj:`bool`): Whether to use cuda in policy, if True, this function will move the input data to cuda. + Returns: + - wrapped_forward_fn (:obj:`Callable`): The wrapped forward function of policy. + + Examples: + >>> env = gym.make('CartPole-v0') + >>> policy = PPOFPolicy(...) + >>> forward_fn = single_env_forward_wrapper_ttorch(policy.eval) + >>> obs = env.reset() + >>> action = forward_fn(obs) + >>> next_obs, rew, done, info = env.step(action) + """ + + def _forward(obs): + # unsqueeze means add batch dim, i.e. (O, ) -> (1, O) + obs = ttorch.as_tensor(obs).unsqueeze(0) + if cuda and torch.cuda.is_available(): + obs = obs.cuda() + action = forward_fn(obs).action + # squeeze means delete batch dim, i.e. (1, A) -> (A, ) + action = action.squeeze(0).cpu().numpy() + return action + + return _forward diff --git a/ding/policy/cql.py b/ding/policy/cql.py index 786fe92b13..b82ffd65df 100644 --- a/ding/policy/cql.py +++ b/ding/policy/cql.py @@ -12,148 +12,118 @@ from ding.utils import POLICY_REGISTRY from ding.utils.data import default_collate, default_decollate from .sac import SACPolicy -from .dqn import DQNPolicy +from .qrdqn import QRDQNPolicy from .common_utils import default_preprocess_learn @POLICY_REGISTRY.register('cql') class CQLPolicy(SACPolicy): """ - Overview: - Policy class of CQL algorithm. - - Config: - == ==================== ======== ============= ================================= ======================= - ID Symbol Type Default Value Description Other(Shape) - == ==================== ======== ============= ================================= ======================= - 1 ``type`` str td3 | RL policy register name, refer | this arg is optional, - | to registry ``POLICY_REGISTRY`` | a placeholder - 2 ``cuda`` bool True | Whether to use cuda for network | - 3 | ``random_`` int 10000 | Number of randomly collected | Default to 10000 for - | ``collect_size`` | training samples in replay | SAC, 25000 for DDPG/ - | | buffer when training starts. | TD3. - 4 | ``model.policy_`` int 256 | Linear layer size for policy | - | ``embedding_size`` | network. | - 5 | ``model.soft_q_`` int 256 | Linear layer size for soft q | - | ``embedding_size`` | network. | - 6 | ``model.value_`` int 256 | Linear layer size for value | Defalut to None when - | ``embedding_size`` | network. | model.value_network - | | | is False. - 7 | ``learn.learning`` float 3e-4 | Learning rate for soft q | Defalut to 1e-3, when - | ``_rate_q`` | network. | model.value_network - | | | is True. - 8 | ``learn.learning`` float 3e-4 | Learning rate for policy | Defalut to 1e-3, when - | ``_rate_policy`` | network. | model.value_network - | | | is True. - 9 | ``learn.learning`` float 3e-4 | Learning rate for policy | Defalut to None when - | ``_rate_value`` | network. | model.value_network - | | | is False. - 10 | ``learn.alpha`` float 0.2 | Entropy regularization | alpha is initiali- - | | coefficient. | zation for auto - | | | `alpha`, when - | | | auto_alpha is True - 11 | ``learn.repara_`` bool True | Determine whether to use | - | ``meterization`` | reparameterization trick. | - 12 | ``learn.`` bool False | Determine whether to use | Temperature parameter - | ``auto_alpha`` | auto temperature parameter | determines the - | | `alpha`. | relative importance - | | | of the entropy term - | | | against the reward. - 13 | ``learn.-`` bool False | Determine whether to ignore | Use ignore_done only - | ``ignore_done`` | done flag. | in halfcheetah env. - 14 | ``learn.-`` float 0.005 | Used for soft update of the | aka. Interpolation - | ``target_theta`` | target network. | factor in polyak aver - | | | aging for target - | | | networks. - == ==================== ======== ============= ================================= ======================= - """ + Overview: + Policy class of CQL algorithm for continuous control. Paper link: https://arxiv.org/abs/2006.04779. + + Config: + == ==================== ======== ============= ================================= ======================= + ID Symbol Type Default Value Description Other(Shape) + == ==================== ======== ============= ================================= ======================= + 1 ``type`` str cql | RL policy register name, refer | this arg is optional, + | to registry ``POLICY_REGISTRY`` | a placeholder + 2 ``cuda`` bool True | Whether to use cuda for network | + 3 | ``random_`` int 10000 | Number of randomly collected | Default to 10000 for + | ``collect_size`` | training samples in replay | SAC, 25000 for DDPG/ + | | buffer when training starts. | TD3. + 4 | ``model.policy_`` int 256 | Linear layer size for policy | + | ``embedding_size`` | network. | + 5 | ``model.soft_q_`` int 256 | Linear layer size for soft q | + | ``embedding_size`` | network. | + 6 | ``model.value_`` int 256 | Linear layer size for value | Defalut to None when + | ``embedding_size`` | network. | model.value_network + | | | is False. + 7 | ``learn.learning`` float 3e-4 | Learning rate for soft q | Defalut to 1e-3, when + | ``_rate_q`` | network. | model.value_network + | | | is True. + 8 | ``learn.learning`` float 3e-4 | Learning rate for policy | Defalut to 1e-3, when + | ``_rate_policy`` | network. | model.value_network + | | | is True. + 9 | ``learn.learning`` float 3e-4 | Learning rate for policy | Defalut to None when + | ``_rate_value`` | network. | model.value_network + | | | is False. + 10 | ``learn.alpha`` float 0.2 | Entropy regularization | alpha is initiali- + | | coefficient. | zation for auto + | | | `alpha`, when + | | | auto_alpha is True + 11 | ``learn.repara_`` bool True | Determine whether to use | + | ``meterization`` | reparameterization trick. | + 12 | ``learn.`` bool False | Determine whether to use | Temperature parameter + | ``auto_alpha`` | auto temperature parameter | determines the + | | `alpha`. | relative importance + | | | of the entropy term + | | | against the reward. + 13 | ``learn.-`` bool False | Determine whether to ignore | Use ignore_done only + | ``ignore_done`` | done flag. | in halfcheetah env. + 14 | ``learn.-`` float 0.005 | Used for soft update of the | aka. Interpolation + | ``target_theta`` | target network. | factor in polyak aver + | | | aging for target + | | | networks. + == ==================== ======== ============= ================================= ======================= + """ config = dict( # (str) RL policy register name (refer to function "POLICY_REGISTRY"). - type='sac', - # (bool) Whether to use cuda for network. + type='cql', + # (bool) Whether to use cuda for policy. cuda=False, - # (bool type) on_policy: Determine whether on-policy or off-policy. + # (bool) on_policy: Determine whether on-policy or off-policy. # on-policy setting influences the behaviour of buffer. - # Default False in SAC. on_policy=False, - multi_agent=False, - # (bool type) priority: Determine whether to use priority in buffer sample. - # Default False in SAC. + # (bool) priority: Determine whether to use priority in buffer sample. priority=False, # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. priority_IS_weight=False, # (int) Number of training samples(randomly collected) in replay buffer when training starts. - # Default 10000 in SAC. random_collect_size=10000, model=dict( # (bool type) twin_critic: Determine whether to use double-soft-q-net for target q computation. # Please refer to TD3 about Clipped Double-Q Learning trick, which learns two Q-functions instead of one . # Default to True. twin_critic=True, - - # (bool type) value_network: Determine whether to use value network as the - # original SAC paper (arXiv 1801.01290). - # using value_network needs to set learning_rate_value, learning_rate_q, - # and learning_rate_policy in `cfg.policy.learn`. - # Default to False. - # value_network=False, - # (str type) action_space: Use reparameterization trick for continous action action_space='reparameterization', - # (int) Hidden size for actor network head. actor_head_hidden_size=256, - # (int) Hidden size for critic network head. critic_head_hidden_size=256, ), + # learn_mode config learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, - # How many updates(iterations) to train after collector's one collection. + # (int) How many updates (iterations) to train after collector's one collection. # Bigger "update_per_collect" means bigger off-policy. - # collect data -> update policy-> collect data -> ... update_per_collect=1, # (int) Minibatch size for gradient descent. batch_size=256, - - # (float type) learning_rate_q: Learning rate for soft q network. - # Default to 3e-4. - # Please set to 1e-3, when model.value_network is True. + # (float) learning_rate_q: Learning rate for soft q network. learning_rate_q=3e-4, - # (float type) learning_rate_policy: Learning rate for policy network. - # Default to 3e-4. - # Please set to 1e-3, when model.value_network is True. + # (float) learning_rate_policy: Learning rate for policy network. learning_rate_policy=3e-4, - # (float type) learning_rate_value: Learning rate for value network. - # `learning_rate_value` should be initialized, when model.value_network is True. - # Please set to 3e-4, when model.value_network is True. - learning_rate_value=3e-4, - - # (float type) learning_rate_alpha: Learning rate for auto temperature parameter `\alpha`. - # Default to 3e-4. + # (float) learning_rate_alpha: Learning rate for auto temperature parameter ``alpha``. learning_rate_alpha=3e-4, - # (float type) target_theta: Used for soft update of the target network, + # (float) target_theta: Used for soft update of the target network, # aka. Interpolation factor in polyak averaging for target networks. - # Default to 0.005. target_theta=0.005, # (float) discount factor for the discounted sum of rewards, aka. gamma. discount_factor=0.99, - - # (float type) alpha: Entropy regularization coefficient. + # (float) alpha: Entropy regularization coefficient. # Please check out the original SAC paper (arXiv 1801.01290): Eq 1 for more details. # If auto_alpha is set to `True`, alpha is initialization for auto `\alpha`. # Default to 0.2. alpha=0.2, - - # (bool type) auto_alpha: Determine whether to use auto temperature parameter `\alpha` . + # (bool) auto_alpha: Determine whether to use auto temperature parameter `\alpha` . # Temperature parameter determines the relative importance of the entropy term against the reward. # Please check out the original SAC paper (arXiv 1801.01290): Eq 1 for more details. # Default to False. # Note that: Using auto alpha needs to set learning_rate_alpha in `cfg.policy.learn`. auto_alpha=True, - # (bool type) log_space: Determine whether to use auto `\alpha` in log space. + # (bool) log_space: Determine whether to use auto `\alpha` in log space. log_space=True, # (bool) Whether ignore done(usually for max step termination env. e.g. pendulum) # Note: Gym wraps the MuJoCo envs by default with TimeLimit environment wrappers. @@ -163,50 +133,44 @@ class CQLPolicy(SACPolicy): # TD-error accurate computation(``gamma * (1 - done) * next_v + reward``), # when the episode step is greater than max episode step. ignore_done=False, - # (float) Weight uniform initialization range in the last output layer + # (float) Weight uniform initialization range in the last output layer. init_w=3e-3, - # (int) The numbers of action sample each at every state s from a uniform-at-random + # (int) The numbers of action sample each at every state s from a uniform-at-random. num_actions=10, # (bool) Whether use lagrange multiplier in q value loss. with_lagrange=False, - # (float) The threshold for difference in Q-values + # (float) The threshold for difference in Q-values. lagrange_thresh=-1, # (float) Loss weight for conservative item. min_q_weight=1.0, # (bool) Whether to use entropy in target q. with_q_entropy=False, ), - collect=dict( - # (int) Cut trajectories into pieces with length "unroll_len". - unroll_len=1, - ), - eval=dict(), - other=dict( - replay_buffer=dict( - # (int type) replay_buffer_size: Max size of replay buffer. - replay_buffer_size=1000000, - # (int type) max_use: Max use times of one data in the buffer. - # Data will be removed once used for too many times. - # Default to infinite. - # max_use=256, - ), - ), + eval=dict(), # for compatibility ) - r""" - Overview: - Policy class of SAC algorithm. - """ def _init_learn(self) -> None: - r""" + """ Overview: - Learn mode init method. Called by ``self.__init__``. - Init q, value and policy's optimizers, algorithm config, main and target models. + Initialize the learn mode of policy, including related attributes and modules. For SAC, it mainly \ + contains three optimizers, algorithm-specific arguments such as gamma, min_q_weight, with_lagrange and \ + with_q_entropy, main and target model. Especially, the ``auto_alpha`` mechanism for balancing max entropy \ + target is also initialized here. + This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. """ - # Init self._priority = self._cfg.priority self._priority_IS_weight = self._cfg.priority_IS_weight - self._value_network = False self._twin_critic = self._cfg.model.twin_critic self._num_actions = self._cfg.learn.num_actions @@ -226,25 +190,20 @@ def _init_learn(self) -> None: # Weight Init init_w = self._cfg.learn.init_w - self._model.actor[2].mu.weight.data.uniform_(-init_w, init_w) - self._model.actor[2].mu.bias.data.uniform_(-init_w, init_w) - self._model.actor[2].log_sigma_layer.weight.data.uniform_(-init_w, init_w) - self._model.actor[2].log_sigma_layer.bias.data.uniform_(-init_w, init_w) + self._model.actor_head[-1].mu.weight.data.uniform_(-init_w, init_w) + self._model.actor_head[-1].mu.bias.data.uniform_(-init_w, init_w) + self._model.actor_head[-1].log_sigma_layer.weight.data.uniform_(-init_w, init_w) + self._model.actor_head[-1].log_sigma_layer.bias.data.uniform_(-init_w, init_w) if self._twin_critic: - self._model.critic[0][2].last.weight.data.uniform_(-init_w, init_w) - self._model.critic[0][2].last.bias.data.uniform_(-init_w, init_w) - self._model.critic[1][2].last.weight.data.uniform_(-init_w, init_w) - self._model.critic[1][2].last.bias.data.uniform_(-init_w, init_w) + self._model.critic_head[0][-1].last.weight.data.uniform_(-init_w, init_w) + self._model.critic_head[0][-1].last.bias.data.uniform_(-init_w, init_w) + self._model.critic_head[1][-1].last.weight.data.uniform_(-init_w, init_w) + self._model.critic_head[1][-1].last.bias.data.uniform_(-init_w, init_w) else: - self._model.critic[2].last.weight.data.uniform_(-init_w, init_w) - self._model.critic[2].last.bias.data.uniform_(-init_w, init_w) + self._model.critic_head[2].last.weight.data.uniform_(-init_w, init_w) + self._model.critic_head[-1].last.bias.data.uniform_(-init_w, init_w) # Optimizers - if self._value_network: - self._optimizer_value = Adam( - self._model.value_critic.parameters(), - lr=self._cfg.learn.learning_rate_value, - ) self._optimizer_q = Adam( self._model.critic.parameters(), lr=self._cfg.learn.learning_rate_q, @@ -258,7 +217,11 @@ def _init_learn(self) -> None: self._gamma = self._cfg.learn.discount_factor # Init auto alpha if self._cfg.learn.auto_alpha: - self._target_entropy = self._cfg.learn.get('target_entropy', -np.prod(self._cfg.model.action_shape)) + if self._cfg.learn.target_entropy is None: + assert 'action_shape' in self._cfg.model, "CQL need network model with action_shape variable" + self._target_entropy = -np.prod(self._cfg.model.action_shape) + else: + self._target_entropy = self._cfg.learn.target_entropy if self._cfg.learn.log_space: self._log_alpha = torch.log(torch.FloatTensor([self._cfg.learn.alpha])) self._log_alpha = self._log_alpha.to(self._device).requires_grad_() @@ -292,14 +255,30 @@ def _init_learn(self) -> None: self._forward_learn_cnt = 0 - def _forward_learn(self, data: dict) -> Dict[str, Any]: - r""" + def _forward_learn(self, data: List[Dict[str, Any]]) -> Dict[str, Any]: + """ Overview: - Forward and backward function of learn mode. + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data from the offline dataset and then returns the output \ + result, including various training information such as loss, action, priority. Arguments: - - data (:obj:`dict`): Dict type data, including at least ['obs', 'action', 'reward', 'next_obs'] + - data (:obj:`List[Dict[int, Any]]`): The input data used for policy forward, including a batch of \ + training samples. For each element in list, the key of the dict is the name of data items and the \ + value is the corresponding data. Usually, the value is torch.Tensor or np.ndarray or there dict/list \ + combinations. In the ``_forward_learn`` method, data often need to first be stacked in the batch \ + dimension by some utility functions such as ``default_preprocess_learn``. \ + For CQL, each element in list is a dict containing at least the following keys: ``obs``, ``action``, \ + ``reward``, ``next_obs``, ``done``. Sometimes, it also contains other keys such as ``weight``. Returns: - - info_dict (:obj:`Dict[str, Any]`): Including current lr and loss. + - info_dict (:obj:`Dict[str, Any]`): The information dict that indicated training result, which will be \ + recorded in text log and tensorboard, values must be python scalar or a list of scalars. For the \ + detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. """ loss_dict = {} data = default_preprocess_learn( @@ -326,37 +305,29 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: q_value = self._learn_model.forward(data, mode='compute_critic')['q_value'] # 2. predict target value - if self._value_network: - # predict v value - v_value = self._learn_model.forward(obs, mode='compute_value_critic')['v_value'] - with torch.no_grad(): - next_v_value = self._target_model.forward(next_obs, mode='compute_value_critic')['v_value'] - target_q_value = next_v_value - else: - # target q value. - with torch.no_grad(): - (mu, sigma) = self._learn_model.forward(next_obs, mode='compute_actor')['logit'] - - dist = Independent(Normal(mu, sigma), 1) - pred = dist.rsample() - next_action = torch.tanh(pred) - y = 1 - next_action.pow(2) + 1e-6 - next_log_prob = dist.log_prob(pred).unsqueeze(-1) - next_log_prob = next_log_prob - torch.log(y).sum(-1, keepdim=True) - - next_data = {'obs': next_obs, 'action': next_action} - target_q_value = self._target_model.forward(next_data, mode='compute_critic')['q_value'] - # the value of a policy according to the maximum entropy objective - if self._twin_critic: - # find min one as target q value - if self._with_q_entropy: - target_q_value = torch.min(target_q_value[0], - target_q_value[1]) - self._alpha * next_log_prob.squeeze(-1) - else: - target_q_value = torch.min(target_q_value[0], target_q_value[1]) + with torch.no_grad(): + (mu, sigma) = self._learn_model.forward(next_obs, mode='compute_actor')['logit'] + + dist = Independent(Normal(mu, sigma), 1) + pred = dist.rsample() + next_action = torch.tanh(pred) + y = 1 - next_action.pow(2) + 1e-6 + next_log_prob = dist.log_prob(pred).unsqueeze(-1) + next_log_prob = next_log_prob - torch.log(y).sum(-1, keepdim=True) + + next_data = {'obs': next_obs, 'action': next_action} + target_q_value = self._target_model.forward(next_data, mode='compute_critic')['q_value'] + # the value of a policy according to the maximum entropy objective + if self._twin_critic: + # find min one as target q value + if self._with_q_entropy: + target_q_value = torch.min(target_q_value[0], + target_q_value[1]) - self._alpha * next_log_prob.squeeze(-1) else: - if self._with_q_entropy: - target_q_value = target_q_value - self._alpha * next_log_prob.squeeze(-1) + target_q_value = torch.min(target_q_value[0], target_q_value[1]) + else: + if self._with_q_entropy: + target_q_value = target_q_value - self._alpha * next_log_prob.squeeze(-1) # 3. compute q loss if self._twin_critic: @@ -451,20 +422,6 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: if self._twin_critic: new_q_value = torch.min(new_q_value[0], new_q_value[1]) - # 7. (optional)compute value loss - if self._value_network: - # new_q_value: (bs, ), log_prob: (bs, act_shape) -> target_v_value: (bs, ) - if self._with_q_entropy: - target_v_value = (new_q_value.unsqueeze(-1) - self._alpha * log_prob).mean(dim=-1) - else: - target_v_value = new_q_value.unsqueeze(-1).mean(dim=-1) - loss_dict['value_loss'] = F.mse_loss(v_value, target_v_value.detach()) - - # update value network - self._optimizer_value.zero_grad() - loss_dict['value_loss'].backward() - self._optimizer_value.step() - # 8. compute policy loss policy_loss = (self._alpha * log_prob - new_q_value.unsqueeze(-1)).mean() @@ -512,8 +469,7 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: **loss_dict } - def _get_policy_actions(self, data: Dict, num_actions=10, epsilon: float = 1e-6) -> List: - + def _get_policy_actions(self, data: Dict, num_actions: int = 10, epsilon: float = 1e-6) -> List: # evaluate to get action distribution obs = data['obs'] obs = obs.unsqueeze(1).repeat(1, num_actions, 1).view(obs.shape[0] * num_actions, obs.shape[1]) @@ -529,7 +485,7 @@ def _get_policy_actions(self, data: Dict, num_actions=10, epsilon: float = 1e-6) return action, log_prob.view(-1, num_actions, 1) - def _get_q_value(self, data: Dict, keep=True) -> torch.Tensor: + def _get_q_value(self, data: Dict, keep: bool = True) -> torch.Tensor: new_q_value = self._learn_model.forward(data, mode='compute_critic')['q_value'] if self._twin_critic: new_q_value = [value.view(-1, self._num_actions, 1) for value in new_q_value] @@ -540,17 +496,18 @@ def _get_q_value(self, data: Dict, keep=True) -> torch.Tensor: return new_q_value -@POLICY_REGISTRY.register('cql_discrete') -class CQLDiscretePolicy(DQNPolicy): +@POLICY_REGISTRY.register('discrete_cql') +class DiscreteCQLPolicy(QRDQNPolicy): """ - Overview: - Policy class of CQL algorithm in discrete environments. + Overview: + Policy class of discrete CQL algorithm in discrete action space environments. + Paper link: https://arxiv.org/abs/2006.04779. """ config = dict( # (str) RL policy register name (refer to function "POLICY_REGISTRY"). - type='cql_discrete', - # (bool) Whether to use cuda for network. + type='discrete_cql', + # (bool) Whether to use cuda for policy. cuda=False, # (bool) Whether the RL algorithm is on-policy or off-policy. on_policy=False, @@ -560,55 +517,43 @@ class CQLDiscretePolicy(DQNPolicy): discount_factor=0.97, # (int) N-step reward for target q_value estimation nstep=1, + # learn_mode config learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, - # How many updates(iterations) to train after collector's one collection. + # (int) How many updates (iterations) to train after collector's one collection. # Bigger "update_per_collect" means bigger off-policy. - # collect data -> update policy-> collect data -> ... update_per_collect=1, + # (int) Minibatch size for one gradient descent. batch_size=64, + # (float) Learning rate for soft q network. learning_rate=0.001, - # ============================================================== - # The following configs are algorithm-specific - # ============================================================== # (int) Frequence of target network update. target_update_freq=100, - # (bool) Whether ignore done(usually for max step termination env) + # (bool) Whether ignore done(usually for max step termination env). ignore_done=False, # (float) Loss weight for conservative item. min_q_weight=1.0, ), - # collect_mode config - collect=dict( - # (int) Cut trajectories into pieces with length "unroll_len". - unroll_len=1, - ), - eval=dict(), - # other config - other=dict( - # Epsilon greedy with decay. - eps=dict( - # (str) Decay type. Support ['exp', 'linear']. - type='exp', - start=0.95, - end=0.1, - # (int) Decay length(env step) - decay=10000, - ), - replay_buffer=dict(replay_buffer_size=10000, ) - ), + eval=dict(), # for compatibility ) - def default_model(self) -> Tuple[str, List[str]]: - return 'qrdqn', ['ding.model.template.q_learning'] - def _init_learn(self) -> None: - r""" - Overview: - Learn mode init method. Called by ``self.__init__``. - Init the optimizer, algorithm config, main and target models. - """ + """ + Overview: + Initialize the learn mode of policy, including related attributes and modules. For DiscreteCQL, it mainly \ + contains the optimizer, algorithm-specific arguments such as gamma, nstep and min_q_weight, main and \ + target model. This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. + """ self._min_q_weight = self._cfg.learn.min_q_weight self._priority = self._cfg.priority # Optimizer @@ -629,20 +574,39 @@ def _init_learn(self) -> None: self._learn_model.reset() self._target_model.reset() - def _forward_learn(self, data: dict) -> Dict[str, Any]: - r""" - Overview: - Forward and backward function of learn mode. - Arguments: - - data (:obj:`dict`): Dict type data, including at least ['obs', 'action', 'reward', 'next_obs'] - Returns: - - info_dict (:obj:`Dict[str, Any]`): Including current lr and loss. - """ + def _forward_learn(self, data: List[Dict[str, Any]]) -> Dict[str, Any]: + """ + Overview: + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data from the offline dataset and then returns the output \ + result, including various training information such as loss, action, priority. + Arguments: + - data (:obj:`List[Dict[int, Any]]`): The input data used for policy forward, including a batch of \ + training samples. For each element in list, the key of the dict is the name of data items and the \ + value is the corresponding data. Usually, the value is torch.Tensor or np.ndarray or there dict/list \ + combinations. In the ``_forward_learn`` method, data often need to first be stacked in the batch \ + dimension by some utility functions such as ``default_preprocess_learn``. \ + For DiscreteCQL, each element in list is a dict containing at least the following keys: ``obs``, \ + ``action``, ``reward``, ``next_obs``, ``done``. Sometimes, it also contains other keys like ``weight`` \ + and ``value_gamma`` for nstep return computation. + Returns: + - info_dict (:obj:`Dict[str, Any]`): The information dict that indicated training result, which will be \ + recorded in text log and tensorboard, values must be python scalar or a list of scalars. For the \ + detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + """ data = default_preprocess_learn( data, use_priority=self._priority, ignore_done=self._cfg.learn.ignore_done, use_nstep=True ) if self._cuda: data = to_device(data, self._device) + if data['action'].dim() == 2 and data['action'].shape[-1] == 1: + data['action'] = data['action'].squeeze(-1) # ==================== # Q-learning forward # ==================== @@ -684,7 +648,7 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: # ==================== self._optimizer.zero_grad() loss.backward() - if self._cfg.learn.multi_gpu: + if self._cfg.multi_gpu: self.sync_gradients(self._learn_model) self._optimizer.step() @@ -702,70 +666,12 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: # '[histogram]action_distribution': data['action'], } - def _state_dict_learn(self) -> Dict[str, Any]: - return { - 'model': self._learn_model.state_dict(), - 'target_model': self._target_model.state_dict(), - 'optimizer': self._optimizer.state_dict(), - } - - def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: - self._learn_model.load_state_dict(state_dict['model']) - self._target_model.load_state_dict(state_dict['target_model']) - self._optimizer.load_state_dict(state_dict['optimizer']) - - def _init_collect(self) -> None: - r""" - Overview: - Collect mode init method. Called by ``self.__init__``. - Init traj and unroll length, collect model. - Enable the eps_greedy_sample - """ - self._unroll_len = self._cfg.collect.unroll_len - self._gamma = self._cfg.discount_factor # necessary for parallel - self._nstep = self._cfg.nstep # necessary for parallel - self._collect_model = model_wrap(self._model, wrapper_name='eps_greedy_sample') - self._collect_model.reset() - - def _forward_collect(self, data: Dict[int, Any], eps: float) -> Dict[int, Any]: + def _monitor_vars_learn(self) -> List[str]: """ Overview: - Forward computation graph of collect mode(collect training data), with eps_greedy for exploration. - Arguments: - - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \ - values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer. - - eps (:obj:`float`): epsilon value for exploration, which is decayed by collected env step. + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. Returns: - - output (:obj:`Dict[int, Any]`): The dict of predicting policy_output(action) for the interaction with \ - env and the constructing of transition. - ArgumentsKeys: - - necessary: ``obs`` - ReturnsKeys - - necessary: ``logit``, ``action`` + - necessary_keys (:obj:`List[str]`): The list of the necessary keys to be logged. """ - data_id = list(data.keys()) - data = default_collate(list(data.values())) - if self._cuda: - data = to_device(data, self._device) - self._collect_model.eval() - with torch.no_grad(): - output = self._collect_model.forward(data, eps=eps) - if self._cuda: - output = to_device(output, 'cpu') - output = default_decollate(output) - return {i: d for i, d in zip(data_id, output)} - - def _get_train_sample(self, data: list) -> Union[None, List[Any]]: - r""" - Overview: - Get the trajectory and the n step return data, then sample from the n_step return data - Arguments: - - data (:obj:`list`): The trajectory's cache - Returns: - - samples (:obj:`dict`): The training samples generated - """ - data = get_nstep_return_data(data, self._nstep, gamma=self._gamma) - return get_train_sample(data, self._unroll_len) - - def _monitor_vars_learn(self) -> List[str]: return ['cur_lr', 'total_loss', 'q_target', 'q_value'] diff --git a/ding/policy/d4pg.py b/ding/policy/d4pg.py index fef0f412ca..c7f4f4ebd7 100644 --- a/ding/policy/d4pg.py +++ b/ding/policy/d4pg.py @@ -14,9 +14,12 @@ @POLICY_REGISTRY.register('d4pg') class D4PGPolicy(DDPGPolicy): - r""" + """ Overview: - Policy class of D4PG algorithm. + Policy class of D4PG algorithm. D4PG is a variant of DDPG, which uses distributional critic. \ + The distributional critic is implemented by using quantile regression. \ + Paper link: https://arxiv.org/abs/1804.08617. + Property: learn_mode, collect_mode, eval_mode Config: @@ -100,8 +103,7 @@ class D4PGPolicy(DDPGPolicy): n_atom=51 ), learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, + # How many updates(iterations) to train after collector's one collection. # Bigger "update_per_collect" means bigger off-policy. # collect data -> update policy-> collect data -> ... @@ -150,13 +152,38 @@ class D4PGPolicy(DDPGPolicy): ) def default_model(self) -> Tuple[str, List[str]]: + """ + Overview: + Return the default neural network model class for D4PGPolicy. ``__init__`` method will \ + automatically call this method to get the default model setting and create model. + + Returns: + - model_info (:obj:`Tuple[str, List[str]]`): The registered model name and model's import_names. + """ return 'qac_dist', ['ding.model.template.qac_dist'] def _init_learn(self) -> None: - r""" + """ Overview: - Learn mode init method. Called by ``self.__init__``. - Init actor and critic optimizers, algorithm config, main and target models. + Initialize the D4PG policy's learning mode, which involves setting up key components \ + specific to the D4PG algorithm. This includes creating separate optimizers for the actor \ + and critic networks, a distinctive trait of D4PG's actor-critic approach, and configuring \ + algorithm-specific parameters such as v_min, v_max, and n_atom for the distributional aspect \ + of the critic. Additionally, the method sets up the target model with momentum-based updates, \ + crucial for stabilizing learning, and optionally integrates noise into the target model for \ + effective exploration. This method is invoked during the '__init__' if 'learn' is specified \ + in 'enable_field'. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. """ self._priority = self._cfg.priority self._priority_IS_weight = self._cfg.priority_IS_weight @@ -204,14 +231,36 @@ def _init_learn(self) -> None: self._forward_learn_cnt = 0 # count iterations - def _forward_learn(self, data: dict) -> Dict[str, Any]: - r""" + def _forward_learn(self, data: List[Dict[str, Any]]) -> Dict[str, Any]: + """ Overview: - Forward and backward function of learn mode. + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data from the replay buffer and then returns the output \ + result, including various training information such as different loss, actor and critic lr. Arguments: - - data (:obj:`dict`): Dict type data, including at least ['obs', 'action', 'reward', 'next_obs'] + - data (:obj:`dict`): Input data used for policy forward, including the \ + collected training samples from replay buffer. For each element in dict, the key of the \ + dict is the name of data items and the value is the corresponding data. Usually, the value is \ + torch.Tensor or np.ndarray or there dict/list combinations. In the ``_forward_learn`` method, data \ + often need to first be stacked in the batch dimension by some utility functions such as \ + ``default_preprocess_learn``. \ + For D4PG, each element in list is a dict containing at least the following keys: ``obs``, \ + ``action``, ``reward``, ``next_obs``. Sometimes, it also contains other keys such as ``weight``. + Returns: - - info_dict (:obj:`Dict[str, Any]`): Including at least actor and critic lr, different losses. + - info_dict (:obj:`Dict[str, Any]`): The output result dict of forward learn, containing at \ + least the "cur_lr_actor", "cur_lr_critic", "different losses", "q_value", "action", "priority", \ + keys. Additionally, loss_dict also contains other keys, which are mainly used for monitoring and \ + debugging. "q_value_dict" is used to record the q_value statistics. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for D4PGPolicy: ``ding.policy.tests.test_d4pg``. """ loss_dict = {} data = default_preprocess_learn( @@ -292,23 +341,36 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: } def _get_train_sample(self, traj: list) -> Union[None, List[Any]]: - r""" - Overview: - Get the trajectory and the n step return data, then sample from the n_step return data - Arguments: - - traj (:obj:`list`): The trajectory's buffer list - Returns: - - samples (:obj:`dict`): The training samples generated + """ + Overview: + Process the data of a given trajectory (transitions, a list of transition) into a list of sample that \ + can be used for training directly. The sample is generated by the following steps: \ + 1. Calculate the nstep return data. \ + 2. Sample the data from the nstep return data. \ + 3. Stack the data in the batch dimension. \ + 4. Return the sample data. \ + For D4PG, the nstep return data is generated by ``get_nstep_return_data`` and the sample data is \ + generated by ``get_train_sample``. + + Arguments: + - traj (:obj:`list`): The trajectory data (a list of transition), each element is \ + the same format as the return value of ``self._process_transition`` method. + + Returns: + - samples (:obj:`dict`): The training samples generated, including at least the following keys: \ + ``'obs'``, ``'next_obs'``, ``'action'``, ``'reward'``, ``'done'``, ``'weight'``, ``'value_gamma'``. \ + For more information, please refer to the ``get_train_sample`` method. """ data = get_nstep_return_data(traj, self._nstep, gamma=self._gamma) return get_train_sample(data, self._unroll_len) def _monitor_vars_learn(self) -> List[str]: - r""" + """ Overview: - Return variables' name if variables are to used in monitor. + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. Returns: - - vars (:obj:`List[str]`): Variables' name list. + - necessary_keys (:obj:`List[str]`): The list of the necessary keys to be logged. """ ret = ['cur_lr_actor', 'cur_lr_critic', 'critic_loss', 'actor_loss', 'total_loss', 'q_value', 'action'] return ret diff --git a/ding/policy/ddpg.py b/ding/policy/ddpg.py index 6bc4955b87..6f62c59795 100644 --- a/ding/policy/ddpg.py +++ b/ding/policy/ddpg.py @@ -14,17 +14,11 @@ @POLICY_REGISTRY.register('ddpg') class DDPGPolicy(Policy): - r""" + """ Overview: - Policy class of DDPG algorithm. - - https://arxiv.org/pdf/1509.02971.pdf - - Property: - learn_mode, collect_mode, eval_mode + Policy class of DDPG algorithm. Paper link: https://arxiv.org/abs/1509.02971. Config: - == ==================== ======== ============= ================================= ======================= ID Symbol Type Default Value Description Other(Shape) == ==================== ======== ============= ================================= ======================= @@ -68,36 +62,28 @@ class DDPGPolicy(Policy): config = dict( # (str) RL policy register name (refer to function "POLICY_REGISTRY"). type='ddpg', - # (bool) Whether to use cuda for network. + # (bool) Whether to use cuda in policy. cuda=False, - # (bool type) on_policy: Determine whether on-policy or off-policy. - # on-policy setting influences the behaviour of buffer. - # Default False in DDPG. + # (bool) Whether learning policy is the same as collecting data policy(on-policy). Default False in DDPG. on_policy=False, - # (bool) Whether use priority(priority sample, IS weight, update priority) - # Default False in DDPG. + # (bool) Whether to enable priority experience sample. priority=False, - # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. + # (bool) Whether to use Importance Sampling Weight to correct biased update. If True, priority must be True. priority_IS_weight=False, # (int) Number of training samples(randomly collected) in replay buffer when training starts. # Default 25000 in DDPG/TD3. random_collect_size=25000, - # (bool) Whether to need policy data in process transition + # (bool) Whether to need policy data in process transition. transition_with_policy_data=False, - # (str) Action space type - action_space='continuous', # ['continuous', 'hybrid'] - # (bool) Whether use batch normalization for reward + # (str) Action space type, including ['continuous', 'hybrid']. + action_space='continuous', + # (bool) Whether use batch normalization for reward. reward_batch_norm=False, - model=dict( - # (bool) Whether to use two critic networks or only one. - # Clipped Double Q-Learning for Actor-Critic in original TD3 paper(https://arxiv.org/pdf/1802.09477.pdf). - # Default True for TD3, False for DDPG. - twin_critic=False, - ), + # (bool) Whether to enable multi-agent training setting. + multi_agent=False, + # learn_mode config learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, - # How many updates(iterations) to train after collector's one collection. + # (int) How many updates(iterations) to train after collector's one collection. # Bigger "update_per_collect" means bigger off-policy. # collect data -> update policy-> collect data -> ... update_per_collect=1, @@ -115,7 +101,7 @@ class DDPGPolicy(Policy): # TD-error accurate computation(``gamma * (1 - done) * next_v + reward``), # when the episode step is greater than max episode step. ignore_done=False, - # (float type) target_theta: Used for soft update of the target network, + # (float) target_theta: Used for soft update of the target network, # aka. Interpolation factor in polyak averaging for target networks. # Default to 0.005. target_theta=0.005, @@ -130,36 +116,55 @@ class DDPGPolicy(Policy): # Default True for TD3, False for DDPG. noise=False, ), + # collect_mode config collect=dict( - # (int) Only one of [n_sample, n_episode] shoule be set + # (int) How many training samples collected in one collection procedure. + # Only one of [n_sample, n_episode] shoule be set. # n_sample=1, - # (int) Cut trajectories into pieces with length "unroll_len". + # (int) Split episodes or trajectories into pieces with length `unroll_len`. unroll_len=1, # (float) It is a must to add noise during collection. So here omits "noise" and only set "noise_sigma". noise_sigma=0.1, ), - eval=dict( - evaluator=dict( - # (int) Evaluate every "eval_freq" training iterations. - eval_freq=5000, - ), - ), + eval=dict(), # for compability other=dict( replay_buffer=dict( - # (int) Maximum size of replay buffer. + # (int) Maximum size of replay buffer. Usually, larger buffer size is better. replay_buffer_size=100000, ), ), ) def default_model(self) -> Tuple[str, List[str]]: - return 'qac', ['ding.model.template.qac'] + """ + Overview: + Return this algorithm default neural network model setting for demonstration. ``__init__`` method will \ + automatically call this method to get the default model setting and create model. + Returns: + - model_info (:obj:`Tuple[str, List[str]]`): The registered model name and model's import_names. + """ + if self._cfg.multi_agent: + return 'continuous_maqac', ['ding.model.template.maqac'] + else: + return 'continuous_qac', ['ding.model.template.qac'] def _init_learn(self) -> None: - r""" + """ Overview: - Learn mode init method. Called by ``self.__init__``. - Init actor and critic optimizers, algorithm config, main and target models. + Initialize the learn mode of policy, including related attributes and modules. For DDPG, it mainly \ + contains two optimizers, algorithm-specific arguments such as gamma and twin_critic, main and target model. + This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. """ self._priority = self._cfg.priority self._priority_IS_weight = self._cfg.priority_IS_weight @@ -180,7 +185,9 @@ def _init_learn(self) -> None: # main and target models self._target_model = copy.deepcopy(self._model) + self._learn_model = model_wrap(self._model, wrapper_name='base') if self._cfg.action_space == 'hybrid': + self._learn_model = model_wrap(self._learn_model, wrapper_name='hybrid_argmax_sample') self._target_model = model_wrap(self._target_model, wrapper_name='hybrid_argmax_sample') self._target_model = model_wrap( self._target_model, @@ -199,22 +206,39 @@ def _init_learn(self) -> None: }, noise_range=self._cfg.learn.noise_range ) - self._learn_model = model_wrap(self._model, wrapper_name='base') - if self._cfg.action_space == 'hybrid': - self._learn_model = model_wrap(self._learn_model, wrapper_name='hybrid_argmax_sample') self._learn_model.reset() self._target_model.reset() self._forward_learn_cnt = 0 # count iterations - def _forward_learn(self, data: dict) -> Dict[str, Any]: - r""" + def _forward_learn(self, data: List[Dict[str, Any]]) -> Dict[str, Any]: + """ Overview: - Forward and backward function of learn mode. + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data from the replay buffer and then returns the output \ + result, including various training information such as loss, action, priority. Arguments: - - data (:obj:`dict`): Dict type data, including at least ['obs', 'action', 'reward', 'next_obs'] + - data (:obj:`List[Dict[int, Any]]`): The input data used for policy forward, including a batch of \ + training samples. For each element in list, the key of the dict is the name of data items and the \ + value is the corresponding data. Usually, the value is torch.Tensor or np.ndarray or there dict/list \ + combinations. In the ``_forward_learn`` method, data often need to first be stacked in the batch \ + dimension by some utility functions such as ``default_preprocess_learn``. \ + For DDPG, each element in list is a dict containing at least the following keys: ``obs``, ``action``, \ + ``reward``, ``next_obs``, ``done``. Sometimes, it also contains other keys such as ``weight`` \ + and ``logit`` which is used for hybrid action space. Returns: - - info_dict (:obj:`Dict[str, Any]`): Including at least actor and critic lr, different losses. + - info_dict (:obj:`Dict[str, Any]`): The information dict that indicated training result, which will be \ + recorded in text log and tensorboard, values must be python scalar or a list of scalars. For the \ + detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for DDPGPolicy: ``ding.policy.tests.test_ddpg``. """ loss_dict = {} data = default_preprocess_learn( @@ -237,20 +261,22 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: reward = (reward - reward.mean()) / (reward.std() + 1e-8) # current q value q_value = self._learn_model.forward(data, mode='compute_critic')['q_value'] - q_value_dict = {} - if self._twin_critic: - q_value_dict['q_value'] = q_value[0].mean() - q_value_dict['q_value_twin'] = q_value[1].mean() - else: - q_value_dict['q_value'] = q_value.mean() + # target q value. with torch.no_grad(): next_actor_data = self._target_model.forward(next_obs, mode='compute_actor') next_actor_data['obs'] = next_obs target_q_value = self._target_model.forward(next_actor_data, mode='compute_critic')['q_value'] + + q_value_dict = {} + target_q_value_dict = {} + if self._twin_critic: # TD3: two critic networks target_q_value = torch.min(target_q_value[0], target_q_value[1]) # find min one as target q value + q_value_dict['q_value'] = q_value[0].mean().data.item() + q_value_dict['q_value_twin'] = q_value[1].mean().data.item() + target_q_value_dict['target q_value'] = target_q_value.mean().data.item() # critic network1 td_data = v_1step_td_data(q_value[0], target_q_value, reward, data['done'], data['weight']) critic_loss, td_error_per_sample1 = v_1step_td_error(td_data, self._gamma) @@ -262,6 +288,8 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: td_error_per_sample = (td_error_per_sample1 + td_error_per_sample2) / 2 else: # DDPG: single critic network + q_value_dict['q_value'] = q_value.mean().data.item() + target_q_value_dict['target q_value'] = target_q_value.mean().data.item() td_data = v_1step_td_data(q_value, target_q_value, reward, data['done'], data['weight']) critic_loss, td_error_per_sample = v_1step_td_error(td_data, self._gamma) loss_dict['critic_loss'] = critic_loss @@ -309,9 +337,16 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: 'td_error': td_error_per_sample.abs().mean(), **loss_dict, **q_value_dict, + **target_q_value_dict, } def _state_dict_learn(self) -> Dict[str, Any]: + """ + Overview: + Return the state_dict of learn mode, usually including model, target_model and optimizers. + Returns: + - state_dict (:obj:`Dict[str, Any]`): The dict of current policy learn state, for saving and restoring. + """ return { 'model': self._learn_model.state_dict(), 'target_model': self._target_model.state_dict(), @@ -320,16 +355,33 @@ def _state_dict_learn(self) -> Dict[str, Any]: } def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: + """ + Overview: + Load the state_dict variable into policy learn mode. + Arguments: + - state_dict (:obj:`Dict[str, Any]`): The dict of policy learn state saved before. + + .. tip:: + If you want to only load some parts of model, you can simply set the ``strict`` argument in \ + load_state_dict to ``False``, or refer to ``ding.torch_utils.checkpoint_helper`` for more \ + complicated operation. + """ self._learn_model.load_state_dict(state_dict['model']) self._target_model.load_state_dict(state_dict['target_model']) self._optimizer_actor.load_state_dict(state_dict['optimizer_actor']) self._optimizer_critic.load_state_dict(state_dict['optimizer_critic']) def _init_collect(self) -> None: - r""" + """ Overview: - Collect mode init method. Called by ``self.__init__``. - Init traj and unroll length, collect model. + Initialize the collect mode of policy, including related attributes and modules. For DDPG, it contains the \ + collect_model to balance the exploration and exploitation with the perturbed noise mechanism, and other \ + algorithm-specific arguments such as unroll_len. \ + This method will be called in ``__init__`` method if ``collect`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_collect`` method, you'd better name them \ + with prefix ``_collect_`` to avoid conflict with other modes, such as ``self._collect_attr1``. """ self._unroll_len = self._cfg.collect.unroll_len # collect model @@ -347,18 +399,28 @@ def _init_collect(self) -> None: self._collect_model = model_wrap(self._collect_model, wrapper_name='hybrid_eps_greedy_multinomial_sample') self._collect_model.reset() - def _forward_collect(self, data: dict, **kwargs) -> dict: - r""" + def _forward_collect(self, data: Dict[int, Any], **kwargs) -> Dict[int, Any]: + """ Overview: - Forward function of collect mode. + Policy forward function of collect mode (collecting training data by interacting with envs). Forward means \ + that the policy gets some necessary data (mainly observation) from the envs and then returns the output \ + data, such as the action to interact with the envs. Arguments: - - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \ - values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer. + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. Returns: - - output (:obj:`Dict[int, Any]`): Dict type data, including at least inferred action according to input obs. - ReturnsKeys - - necessary: ``action`` - - optional: ``logit`` + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action and \ + other necessary data for learn mode defined in ``self._process_transition`` method. The key of the \ + dict is the same as the input data, i.e., environment id. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for DDPGPolicy: ``ding.policy.tests.test_ddpg``. """ data_id = list(data.keys()) data = default_collate(list(data.values())) @@ -372,55 +434,84 @@ def _forward_collect(self, data: dict, **kwargs) -> dict: output = default_decollate(output) return {i: d for i, d in zip(data_id, output)} - def _process_transition(self, obs: Any, model_output: dict, timestep: namedtuple) -> Dict[str, Any]: - r""" + def _process_transition(self, obs: torch.Tensor, policy_output: Dict[str, torch.Tensor], + timestep: namedtuple) -> Dict[str, torch.Tensor]: + """ Overview: - Generate dict type transition data from inputs. + Process and pack one timestep transition data into a dict, which can be directly used for training and \ + saved in replay buffer. For DDPG, it contains obs, next_obs, action, reward, done. Arguments: - - obs (:obj:`Any`): Env observation - - model_output (:obj:`dict`): Output of collect model, including at least ['action'] - - timestep (:obj:`namedtuple`): Output after env step, including at least ['obs', 'reward', 'done'] \ - (here 'obs' indicates obs after env step, i.e. next_obs). - Return: - - transition (:obj:`Dict[str, Any]`): Dict type transition data. + - obs (:obj:`torch.Tensor`): The env observation of current timestep, such as stacked 2D image in Atari. + - policy_output (:obj:`Dict[str, torch.Tensor]`): The output of the policy network with the observation \ + as input. For DDPG, it contains the action and the logit of the action (in hybrid action space). + - timestep (:obj:`namedtuple`): The execution result namedtuple returned by the environment step method, \ + except all the elements have been transformed into tensor data. Usually, it contains the next obs, \ + reward, done, info, etc. + Returns: + - transition (:obj:`Dict[str, torch.Tensor]`): The processed transition data of the current timestep. """ transition = { 'obs': obs, 'next_obs': timestep.obs, - 'action': model_output['action'], + 'action': policy_output['action'], 'reward': timestep.reward, 'done': timestep.done, } if self._cfg.action_space == 'hybrid': - transition['logit'] = model_output['logit'] + transition['logit'] = policy_output['logit'] return transition - def _get_train_sample(self, data: list) -> Union[None, List[Any]]: - return get_train_sample(data, self._unroll_len) + def _get_train_sample(self, transitions: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Overview: + For a given trajectory (transitions, a list of transition) data, process it into a list of sample that \ + can be used for training directly. In DDPG, a train sample is a processed transition (unroll_len=1). + Arguments: + - transitions (:obj:`List[Dict[str, Any]`): The trajectory data (a list of transition), each element is \ + the same format as the return value of ``self._process_transition`` method. + Returns: + - samples (:obj:`List[Dict[str, Any]]`): The processed train samples, each element is the similar format \ + as input transitions, but may contain more data for training. + """ + return get_train_sample(transitions, self._unroll_len) def _init_eval(self) -> None: - r""" + """ Overview: - Evaluate mode init method. Called by ``self.__init__``. - Init eval model. Unlike learn and collect model, eval model does not need noise. + Initialize the eval mode of policy, including related attributes and modules. For DDPG, it contains the \ + eval model to greedily select action type with argmax q_value mechanism for hybrid action space. \ + This method will be called in ``__init__`` method if ``eval`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_eval`` method, you'd better name them \ + with prefix ``_eval_`` to avoid conflict with other modes, such as ``self._eval_attr1``. """ self._eval_model = model_wrap(self._model, wrapper_name='base') if self._cfg.action_space == 'hybrid': self._eval_model = model_wrap(self._eval_model, wrapper_name='hybrid_argmax_sample') self._eval_model.reset() - def _forward_eval(self, data: dict) -> dict: - r""" + def _forward_eval(self, data: Dict[int, Any]) -> Dict[int, Any]: + """ Overview: - Forward function of eval mode, similar to ``self._forward_collect``. + Policy forward function of eval mode (evaluation policy performance by interacting with envs). Forward \ + means that the policy gets some necessary data (mainly observation) from the envs and then returns the \ + action to interact with the envs. Arguments: - - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \ - values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer. + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. Returns: - - output (:obj:`Dict[int, Any]`): The dict of predicting action for the interaction with env. - ReturnsKeys - - necessary: ``action`` - - optional: ``logit`` + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action. The \ + key of the dict is the same as the input data, i.e. environment id. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for DDPGPolicy: ``ding.policy.tests.test_ddpg``. """ data_id = list(data.keys()) data = default_collate(list(data.values())) @@ -435,11 +526,12 @@ def _forward_eval(self, data: dict) -> dict: return {i: d for i, d in zip(data_id, output)} def _monitor_vars_learn(self) -> List[str]: - r""" + """ Overview: - Return variables' names if variables are to used in monitor. + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. Returns: - - vars (:obj:`List[str]`): Variables' name list. + - necessary_keys (:obj:`List[str]`): The list of the necessary keys to be logged. """ ret = [ 'cur_lr_actor', 'cur_lr_critic', 'critic_loss', 'actor_loss', 'total_loss', 'q_value', 'q_value_twin', diff --git a/ding/policy/decision_transformer.py b/ding/policy/decision_transformer.py deleted file mode 100644 index 8a59fe35d1..0000000000 --- a/ding/policy/decision_transformer.py +++ /dev/null @@ -1,379 +0,0 @@ -"""The code is adapted from https://github.com/nikhilbarhate99/min-decision-transformer -""" - -from typing import List, Dict, Any, Tuple, Union -from collections import namedtuple -from torch.distributions import Normal, Independent -from ding.torch_utils import Adam, to_device -from ditk import logging -from ding.rl_utils import v_1step_td_data, v_1step_td_error, get_train_sample, \ - qrdqn_nstep_td_data, qrdqn_nstep_td_error, get_nstep_return_data -from ding.model import model_wrap -from ding.utils.data.dataset import D4RLTrajectoryDataset -from ding.utils import POLICY_REGISTRY -from ding.utils.data import default_collate, default_decollate -from datetime import datetime -from ding.torch_utils import one_hot -import numpy as np -import torch.nn.functional as F -import torch -import gym -import copy -import os -import csv -from .dqn import DQNPolicy - - -@POLICY_REGISTRY.register('dt') -class DTPolicy(DQNPolicy): - r""" - Overview: - Policy class of DT algorithm in discrete environments. - """ - config = dict( - # (str) RL policy register name (refer to function "POLICY_REGISTRY"). - type='dt', - # (bool) Whether to use cuda for network. - cuda=False, - # (bool) Whether the RL algorithm is on-policy or off-policy. - on_policy=False, - # (bool) Whether use priority(priority sample, IS weight, update priority) - priority=False, - # (float) Reward's future discount factor, aka. gamma. - discount_factor=0.97, - # (int) N-step reward for target q_value estimation - nstep=1, - obs_shape=4, - action_shape=2, - # encoder_hidden_size_list=[128, 128, 64], - dataset='medium', # medium / medium-replay / medium-expert - rtg_scale=1000, # normalize returns to go - max_eval_ep_len=1000, # max len of one episode - num_eval_ep=10, # num of evaluation episodes - batch_size=64, # training batch size - wt_decay=1e-4, - warmup_steps=10000, - max_train_iters=200, - context_len=20, - n_blocks=3, - embed_dim=128, - dropout_p=0.1, - learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, - # batch_size=64, - learning_rate=1e-4, - # ============================================================== - # The following configs are algorithm-specific - # ============================================================== - ), - # collect_mode config - collect=dict(), - eval=dict(), - # other config - other=dict(), - ) - - def default_model(self) -> Tuple[str, List[str]]: - return 'dt', ['ding.model.template.decision_transformer'] - - def _init_learn(self) -> None: - r""" - Overview: - Learn mode init method. Called by ``self.__init__``. - Init the optimizer, algorithm config, main and target models. - """ - - self.stop_value = self._cfg.stop_value - self.env_name = self._cfg.env_name - dataset = self._cfg.dataset # medium / medium-replay / medium-expert - # rtg_scale: scale of `return to go` - # rtg_target: max target of `return to go` - # Our goal is normalize `return to go` to (0, 1), which will favour the covergence. - # As a result, we usually set rtg_scale == rtg_target. - self.rtg_scale = self._cfg.rtg_target # normalize returns to go - self.rtg_target = self._cfg.rtg_target # max target reward_to_go - self.max_eval_ep_len = self._cfg.max_eval_ep_len # max len of one episode - self.num_eval_ep = self._cfg.num_eval_ep # num of evaluation episodes - - lr = self._cfg.learn.learning_rate # learning rate - wt_decay = self._cfg.wt_decay # weight decay - warmup_steps = self._cfg.warmup_steps # warmup steps for lr scheduler - - max_train_iters = self._cfg.max_train_iters - - self.context_len = self._cfg.context_len # K in decision transformer - n_blocks = self._cfg.n_blocks # num of transformer blocks - embed_dim = self._cfg.embed_dim # embedding (hidden) dim of transformer - dropout_p = self._cfg.dropout_p # dropout probability - - # # load data from this file - # dataset_path = f'{self._cfg.dataset_dir}/{env_d4rl_name}.pkl' - - # saves model and csv in this directory - self.log_dir = self._cfg.log_dir - if not os.path.exists(self.log_dir): - os.makedirs(self.log_dir) - - # training and evaluation device - self.device = torch.device(self._device) - - self.start_time = datetime.now().replace(microsecond=0) - self.start_time_str = self.start_time.strftime("%y-%m-%d-%H-%M-%S") - - # prefix = "dt_" + env_d4rl_name - self.prefix = "dt_" + self.env_name - - save_model_name = self.prefix + "_model_" + self.start_time_str + ".pt" - self.save_model_path = os.path.join(self.log_dir, save_model_name) - self.save_best_model_path = self.save_model_path[:-3] + "_best.pt" - - log_csv_name = self.prefix + "_log_" + self.start_time_str + ".csv" - log_csv_path = os.path.join(self.log_dir, log_csv_name) - - self.csv_writer = csv.writer(open(log_csv_path, 'a', 1)) - csv_header = (["duration", "num_updates", "eval_avg_reward", "eval_avg_ep_len", "eval_d4rl_score"]) - - self.csv_writer.writerow(csv_header) - - dataset_path = self._cfg.learn.dataset_path - logging.info("=" * 60) - logging.info("start time: " + self.start_time_str) - logging.info("=" * 60) - - logging.info("device set to: " + str(self.device)) - logging.info("dataset path: " + dataset_path) - logging.info("model save path: " + self.save_model_path) - logging.info("log csv save path: " + log_csv_path) - - self._env = gym.make(self.env_name) - - self.state_dim = self._cfg.model.state_dim - self.act_dim = self._cfg.model.act_dim - - self._learn_model = self._model - self._optimizer = torch.optim.AdamW(self._learn_model.parameters(), lr=lr, weight_decay=wt_decay) - - self._scheduler = torch.optim.lr_scheduler.LambdaLR( - self._optimizer, lambda steps: min((steps + 1) / warmup_steps, 1) - ) - - self.max_env_score = -1.0 - - def _forward_learn(self, data: list) -> Dict[str, Any]: - r""" - Overview: - Forward and backward function of learn mode. - Arguments: - - data (:obj:`dict`): Dict type data, including at least ['obs', 'action', 'reward', 'next_obs'] - Returns: - - info_dict (:obj:`Dict[str, Any]`): Including current lr and loss. - """ - - self._learn_model.train() - - timesteps, states, actions, returns_to_go, traj_mask = data - - timesteps = timesteps.to(self.device) # B x T - states = states.to(self.device) # B x T x state_dim - actions = actions.to(self.device) # B x T x act_dim - returns_to_go = returns_to_go.to(self.device) # B x T x 1 - traj_mask = traj_mask.to(self.device) # B x T - action_target = torch.clone(actions).detach().to(self.device) - - # The shape of `returns_to_go` may differ with different dataset (B x T or B x T x 1), - # and we need a 3-dim tensor - if len(returns_to_go.shape) == 2: - returns_to_go = returns_to_go.unsqueeze(-1) - - # if discrete - if not self._cfg.model.continuous: - actions = one_hot(actions.squeeze(-1), num=self.act_dim) - - state_preds, action_preds, return_preds = self._learn_model.forward( - timesteps=timesteps, states=states, actions=actions, returns_to_go=returns_to_go - ) - - traj_mask = traj_mask.view(-1, ) - - # only consider non padded elements - action_preds = action_preds.view(-1, self.act_dim)[traj_mask > 0] - - if self._cfg.model.continuous: - action_target = action_target.view(-1, self.act_dim)[traj_mask > 0] - else: - action_target = action_target.view(-1)[traj_mask > 0] - - if self._cfg.model.continuous: - action_loss = F.mse_loss(action_preds, action_target) - else: - action_loss = F.cross_entropy(action_preds, action_target) - - self._optimizer.zero_grad() - action_loss.backward() - torch.nn.utils.clip_grad_norm_(self._learn_model.parameters(), 0.25) - self._optimizer.step() - self._scheduler.step() - - return { - 'cur_lr': self._optimizer.state_dict()['param_groups'][0]['lr'], - 'action_loss': action_loss.detach().cpu().item(), - } - - def evaluate_on_env(self, state_mean=None, state_std=None, render=False): - - eval_batch_size = 1 # required for forward pass - - results = {} - total_reward = 0 - total_timesteps = 0 - - # state_dim = env.observation_space.shape[0] - # act_dim = env.action_space.shape[0] - - if state_mean is None: - self.state_mean = torch.zeros((self.state_dim, )).to(self.device) - else: - self.state_mean = torch.from_numpy(state_mean).to(self.device) - - if state_std is None: - self.state_std = torch.ones((self.state_dim, )).to(self.device) - else: - self.state_std = torch.from_numpy(state_std).to(self.device) - - # same as timesteps used for training the transformer - # also, crashes if device is passed to arange() - timesteps = torch.arange(start=0, end=self.max_eval_ep_len, step=1) - timesteps = timesteps.repeat(eval_batch_size, 1).to(self.device) - - self._learn_model.eval() - - with torch.no_grad(): - - for _ in range(self.num_eval_ep): - - # zeros place holders - # continuous action - actions = torch.zeros( - (eval_batch_size, self.max_eval_ep_len, self.act_dim), dtype=torch.float32, device=self.device - ) - - # discrete action # TODO - # actions = torch.randint(0,self.act_dim,[eval_batch_size, self.max_eval_ep_len, 1], - # dtype=torch.long, device=self.device) - - states = torch.zeros( - (eval_batch_size, self.max_eval_ep_len, self.state_dim), dtype=torch.float32, device=self.device - ) - rewards_to_go = torch.zeros( - (eval_batch_size, self.max_eval_ep_len, 1), dtype=torch.float32, device=self.device - ) - - # init episode - running_state = self._env.reset() - running_reward = 0 - running_rtg = self.rtg_target / self.rtg_scale - - for t in range(self.max_eval_ep_len): - - total_timesteps += 1 - - # add state in placeholder and normalize - states[0, t] = torch.from_numpy(running_state).to(self.device) - # states[0, t] = (states[0, t].cpu() - self.state_mean.cpu().numpy()) / self.state_std.cpu().numpy() - states[0, t] = (states[0, t] - self.state_mean) / self.state_std - - # calcualate running rtg and add it in placeholder - running_rtg = running_rtg - (running_reward / self.rtg_scale) - rewards_to_go[0, t] = running_rtg - - if t < self.context_len: - _, act_preds, _ = self._learn_model.forward( - timesteps[:, :self.context_len], states[:, :self.context_len], - actions[:, :self.context_len], rewards_to_go[:, :self.context_len] - ) - act = act_preds[0, t].detach() - else: - _, act_preds, _ = self._learn_model.forward( - timesteps[:, t - self.context_len + 1:t + 1], states[:, t - self.context_len + 1:t + 1], - actions[:, t - self.context_len + 1:t + 1], rewards_to_go[:, t - self.context_len + 1:t + 1] - ) - act = act_preds[0, -1].detach() - - # if discrete - if not self._cfg.model.continuous: - act = torch.argmax(act) - running_state, running_reward, done, _ = self._env.step(act.cpu().numpy()) - - # add action in placeholder - actions[0, t] = act - - total_reward += running_reward - - if render: - self._env.render() - if done: - break - - results['eval/avg_reward'] = total_reward / self.num_eval_ep - results['eval/avg_ep_len'] = total_timesteps / self.num_eval_ep - - return results - - def evaluate(self, total_update_times, state_mean=None, state_std=None, render=False): - results = self.evaluate_on_env(state_mean, state_std, render) - - eval_avg_reward = results['eval/avg_reward'] - eval_avg_ep_len = results['eval/avg_ep_len'] - eval_d4rl_score = self.get_d4rl_normalized_score(results['eval/avg_reward'], self.env_name) * 100 - - time_elapsed = str(datetime.now().replace(microsecond=0) - self.start_time) - - log_str = ( - "=" * 60 + '\n' + "time elapsed: " + time_elapsed + '\n' + "num of updates: " + str(total_update_times) + - '\n' + '\n' + "eval avg reward: " + format(eval_avg_reward, ".5f") + '\n' + "eval avg ep len: " + - format(eval_avg_ep_len, ".5f") + '\n' + "eval d4rl score: " + format(eval_d4rl_score, ".5f") - ) - - logging.info(log_str) - - log_data = [time_elapsed, total_update_times, eval_avg_reward, eval_avg_ep_len, eval_d4rl_score] - log_csv_name = self.prefix + "_log_" + self.start_time_str + ".csv" - log_csv_path = os.path.join(self.log_dir, log_csv_name) - - self.csv_writer.writerow(log_data) - - # save model - logging.info("eval_avg_reward: " + format(eval_avg_reward, ".5f")) - eval_env_score = eval_avg_reward - if eval_env_score >= self.max_env_score: - logging.info("saving max env score model at: " + self.save_best_model_path) - torch.save(self._learn_model.state_dict(), self.save_best_model_path) - self.max_env_score = eval_env_score - - logging.info("saving current model at: " + self.save_model_path) - torch.save(self._learn_model.state_dict(), self.save_model_path) - - return self.max_env_score >= self.stop_value - - def get_d4rl_normalized_score(self, score, env_name): - env_key = env_name.split('-')[0].lower() - assert env_key in D4RLTrajectoryDataset.REF_MAX_SCORE, \ - f'no reference score for {env_key} env to calculate d4rl score' - d4rl_max_score, d4rl_min_score = D4RLTrajectoryDataset.REF_MAX_SCORE, D4RLTrajectoryDataset.REF_MIN_SCORE - return (score - d4rl_min_score[env_key]) / (d4rl_max_score[env_key] - d4rl_min_score[env_key]) - - def _state_dict_learn(self) -> Dict[str, Any]: - return { - 'model': self._learn_model.state_dict(), - # 'target_model': self._target_model.state_dict(), - 'optimizer': self._optimizer.state_dict(), - } - - def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: - self._learn_model.load_state_dict(state_dict['model']) - # self._target_model.load_state_dict(state_dict['target_model']) - self._optimizer.load_state_dict(state_dict['optimizer']) - - def _monitor_vars_learn(self) -> List[str]: - return ['cur_lr', 'action_loss'] diff --git a/ding/policy/dqfd.py b/ding/policy/dqfd.py index 7fcb6f77ff..9e9ecab853 100644 --- a/ding/policy/dqfd.py +++ b/ding/policy/dqfd.py @@ -83,8 +83,7 @@ class DQFDPolicy(DQNPolicy): margin_function=0.8, # number of pertraining iterations per_train_iter_k=10, - # (bool) Whether to use multi gpu - multi_gpu=False, + # How many updates(iterations) to train after collector's one collection. # Bigger "update_per_collect" means bigger off-policy. # collect data -> update policy-> collect data -> ... @@ -231,7 +230,7 @@ def _forward_learn(self, data: Dict[str, Any]) -> Dict[str, Any]: # ==================== self._optimizer.zero_grad() loss.backward() - if self._cfg.learn.multi_gpu: + if self._cfg.multi_gpu: self.sync_gradients(self._learn_model) self._optimizer.step() diff --git a/ding/policy/dqn.py b/ding/policy/dqn.py index e758079095..116ba93cf6 100644 --- a/ding/policy/dqn.py +++ b/ding/policy/dqn.py @@ -10,142 +10,186 @@ from ding.utils.data import default_collate, default_decollate from .base_policy import Policy -from .common_utils import default_preprocess_learn +from .common_utils import default_preprocess_learn, set_noise_mode @POLICY_REGISTRY.register('dqn') class DQNPolicy(Policy): - r""" + """ Overview: Policy class of DQN algorithm, extended by Double DQN/Dueling DQN/PER/multi-step TD. Config: - == ==================== ======== ============== ======================================== ======================= - ID Symbol Type Default Value Description Other(Shape) - == ==================== ======== ============== ======================================== ======================= - 1 ``type`` str dqn | RL policy register name, refer to | This arg is optional, - | registry ``POLICY_REGISTRY`` | a placeholder - 2 ``cuda`` bool False | Whether to use cuda for network | This arg can be diff- + == ===================== ======== ============== ======================================= ======================= + ID Symbol Type Default Value Description Other(Shape) + == ===================== ======== ============== ======================================= ======================= + 1 ``type`` str dqn | RL policy register name, refer to | This arg is optional, + | registry ``POLICY_REGISTRY`` | a placeholder + 2 ``cuda`` bool False | Whether to use cuda for network | This arg can be diff- | erent from modes - 3 ``on_policy`` bool False | Whether the RL algorithm is on-policy - | or off-policy - 4 ``priority`` bool False | Whether use priority(PER) | Priority sample, + 3 ``on_policy`` bool False | Whether the RL algorithm is on-policy + | or off-policy + 4 ``priority`` bool False | Whether use priority(PER) | Priority sample, | update priority - 5 | ``priority_IS`` bool False | Whether use Importance Sampling Weight - | ``_weight`` | to correct biased update. If True, - | priority must be True. - 6 | ``discount_`` float 0.97, | Reward's future discount factor, aka. | May be 1 when sparse - | ``factor`` [0.95, 0.999] | gamma | reward env - 7 ``nstep`` int 1, | N-step reward discount sum for target - [3, 5] | q_value estimation - 8 | ``learn.update`` int 3 | How many updates(iterations) to train | This args can be vary - | ``per_collect`` | after collector's one collection. Only | from envs. Bigger val - | valid in serial training | means more off-policy - 9 | ``learn.multi`` bool False | whether to use multi gpu during - | ``_gpu`` - 10 | ``learn.batch_`` int 64 | The number of samples of an iteration + 5 | ``priority_IS`` bool False | Whether use Importance Sampling + | ``_weight`` | Weight to correct biased update. If + | True, priority must be True. + 6 | ``discount_`` float 0.97, | Reward's future discount factor, aka. | May be 1 when sparse + | ``factor`` [0.95, 0.999] | gamma | reward env + 7 ``nstep`` int 1, | N-step reward discount sum for target + [3, 5] | q_value estimation + 8 | ``model.dueling`` bool True | dueling head architecture + 9 | ``model.encoder`` list [32, 64, | Sequence of ``hidden_size`` of | default kernel_size + | ``_hidden`` (int) 64, 128] | subsequent conv layers and the | is [8, 4, 3] + | ``_size_list`` | final dense layer. | default stride is + | [4, 2 ,1] + 10 | ``model.dropout`` float None | Dropout rate for dropout layers. | [0,1] + | If set to ``None`` + | means no dropout + 11 | ``learn.update`` int 3 | How many updates(iterations) to train | This args can be vary + | ``per_collect`` | after collector's one collection. | from envs. Bigger val + | Only valid in serial training | means more off-policy + 12 | ``learn.batch_`` int 64 | The number of samples of an iteration | ``size`` - 11 | ``learn.learning`` float 0.001 | Gradient step length of an iteration. + 13 | ``learn.learning`` float 0.001 | Gradient step length of an iteration. | ``_rate`` - 12 | ``learn.target_`` int 100 | Frequence of target network update. | Hard(assign) update + 14 | ``learn.target_`` int 100 | Frequence of target network update. | Hard(assign) update | ``update_freq`` - 13 | ``learn.ignore_`` bool False | Whether ignore done for target value | Enable it for some - | ``done`` | calculation. | fake termination env - 14 ``collect.n_sample`` int [8, 128] | The number of training samples of a | It varies from - | call of collector. | different envs - 15 | ``collect.unroll`` int 1 | unroll length of an iteration | In RNN, unroll_len>1 + 15 | ``learn.target_`` float 0.005 | Frequence of target network update. | Soft(assign) update + | ``theta`` | Only one of [target_update_freq, + | | target_theta] should be set + 16 | ``learn.ignore_`` bool False | Whether ignore done for target value | Enable it for some + | ``done`` | calculation. | fake termination env + 17 ``collect.n_sample`` int [8, 128] | The number of training samples of a | It varies from + | call of collector. | different envs + 18 ``collect.n_episode`` int 8 | The number of training episodes of a | only one of [n_sample + | call of collector | ,n_episode] should + | | be set + 19 | ``collect.unroll`` int 1 | unroll length of an iteration | In RNN, unroll_len>1 | ``_len`` - 16 | ``other.eps.type`` str exp | exploration rate decay type | Support ['exp', + 20 | ``other.eps.type`` str exp | exploration rate decay type | Support ['exp', | 'linear']. - 17 | ``other.eps.`` float 0.95 | start value of exploration rate | [0,1] + 21 | ``other.eps.`` float 0.95 | start value of exploration rate | [0,1] | ``start`` - 18 | ``other.eps.`` float 0.1 | end value of exploration rate | [0,1] + 22 | ``other.eps.`` float 0.1 | end value of exploration rate | [0,1] | ``end`` - 19 | ``other.eps.`` int 10000 | decay length of exploration | greater than 0. set + 23 | ``other.eps.`` int 10000 | decay length of exploration | greater than 0. set | ``decay`` | decay=10000 means | the exploration rate | decay from start | value to end value | during decay length. - == ==================== ======== ============== ======================================== ======================= + == ===================== ======== ============== ======================================= ======================= """ config = dict( + # (str) RL policy register name (refer to function "POLICY_REGISTRY"). type='dqn', - # (bool) Whether use cuda in policy + # (bool) Whether to use cuda in policy. cuda=False, - # (bool) Whether learning policy is the same as collecting data policy(on-policy) + # (bool) Whether learning policy is the same as collecting data policy(on-policy). on_policy=False, - # (bool) Whether enable priority experience sample + # (bool) Whether to enable priority experience sample. priority=False, - # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. + # (bool) Whether to use Importance Sampling Weight to correct biased update. If True, priority must be True. priority_IS_weight=False, - # (float) Discount factor(gamma) for returns + # (float) Discount factor(gamma) for returns. discount_factor=0.97, - # (int) The number of step for calculating target q_value + # (int) The number of steps for calculating target q_value. nstep=1, + # (bool) Whether to use NoisyNet for exploration in both learning and collecting. Default is False. + noisy_net=False, + model=dict( + # (list(int)) Sequence of ``hidden_size`` of subsequent conv layers and the final dense layer. + encoder_hidden_size_list=[128, 128, 64], + ), + # learn_mode config learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, - # How many updates(iterations) to train after collector's one collection. + # (int) How many updates(iterations) to train after collector's one collection. # Bigger "update_per_collect" means bigger off-policy. # collect data -> update policy-> collect data -> ... update_per_collect=3, - # (int) How many samples in a training batch + # (int) How many samples in a training batch. batch_size=64, - # (float) The step size of gradient descent + # (float) The step size of gradient descent. learning_rate=0.001, - # ============================================================== - # The following configs are algorithm-specific - # ============================================================== - # (int) Frequence of target network update. + # (int) Frequency of target network update. + # Only one of [target_update_freq, target_theta] should be set. target_update_freq=100, - # (bool) Whether ignore done(usually for max step termination env) + # (float) Used for soft update of the target network. + # aka. Interpolation factor in EMA update for target network. + # Only one of [target_update_freq, target_theta] should be set. + target_theta=0.005, + # (bool) If set to True, the 'done' signals that indicate the end of an episode due to environment time + # limits are disregarded. By default, this is set to False. This setting is particularly useful for tasks + # that have a predetermined episode length, such as HalfCheetah and various other MuJoCo environments, + # where the maximum length is capped at 1000 steps. When enabled, any 'done' signal triggered by reaching + # the maximum episode steps will be overridden to 'False'. This ensures the accurate calculation of the + # Temporal Difference (TD) error, using the formula `gamma * (1 - done) * next_v + reward`, + # even when the episode surpasses the predefined step limit. ignore_done=False, ), # collect_mode config collect=dict( - # (int) Only one of [n_sample, n_episode] shoule be set - # n_sample=8, - # (int) Cut trajectories into pieces with length "unroll_len". + # (int) How many training samples collected in one collection procedure. + # Only one of [n_sample, n_episode] should be set. + n_sample=8, + # (int) Split episodes or trajectories into pieces with length `unroll_len`. unroll_len=1, ), - eval=dict(), + eval=dict(), # for compatibility # other config other=dict( # Epsilon greedy with decay. eps=dict( # (str) Decay type. Support ['exp', 'linear']. type='exp', - # (float) Epsilon start value + # (float) Epsilon start value. start=0.95, - # (float) Epsilon end value + # (float) Epsilon end value. end=0.1, - # (int) Decay length(env step) + # (int) Decay length(env step). decay=10000, ), - replay_buffer=dict(replay_buffer_size=10000, ), + replay_buffer=dict( + # (int) Maximum size of replay buffer. Usually, larger buffer size is better. + replay_buffer_size=10000, + ), ), ) def default_model(self) -> Tuple[str, List[str]]: """ Overview: - Return this algorithm default model setting for demonstration. + Return this algorithm default neural network model setting for demonstration. ``__init__`` method will \ + automatically call this method to get the default model setting and create model. Returns: - - model_info (:obj:`Tuple[str, List[str]]`): model name and mode import_names + - model_info (:obj:`Tuple[str, List[str]]`): The registered model name and model's import_names. .. note:: - The user can define and use customized network model but must obey the same inferface definition indicated \ - by import_names path. For DQN, ``ding.model.template.q_learning.DQN`` + The user can define and use customized network model but must obey the same interface definition indicated \ + by import_names path. For example about DQN, its registered name is ``dqn`` and the import_names is \ + ``ding.model.template.q_learning``. """ return 'dqn', ['ding.model.template.q_learning'] def _init_learn(self) -> None: """ Overview: - Learn mode init method. Called by ``self.__init__``, initialize the optimizer, algorithm arguments, main \ - and target model. + Initialize the learn mode of policy, including related attributes and modules. For DQN, it mainly contains \ + optimizer, algorithm-specific arguments such as nstep and gamma, main and target model. + This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. """ self._priority = self._cfg.priority self._priority_IS_weight = self._cfg.priority_IS_weight @@ -157,33 +201,71 @@ def _init_learn(self) -> None: # use model_wrapper for specialized demands of different modes self._target_model = copy.deepcopy(self._model) - self._target_model = model_wrap( - self._target_model, - wrapper_name='target', - update_type='assign', - update_kwargs={'freq': self._cfg.learn.target_update_freq} - ) + if 'target_update_freq' in self._cfg.learn: + self._target_model = model_wrap( + self._target_model, + wrapper_name='target', + update_type='assign', + update_kwargs={'freq': self._cfg.learn.target_update_freq} + ) + elif 'target_theta' in self._cfg.learn: + self._target_model = model_wrap( + self._target_model, + wrapper_name='target', + update_type='momentum', + update_kwargs={'theta': self._cfg.learn.target_theta} + ) + else: + raise RuntimeError("DQN needs target network, please either indicate target_update_freq or target_theta") self._learn_model = model_wrap(self._model, wrapper_name='argmax_sample') self._learn_model.reset() self._target_model.reset() - def _forward_learn(self, data: Dict[str, Any]) -> Dict[str, Any]: + def _forward_learn(self, data: List[Dict[str, Any]]) -> Dict[str, Any]: """ Overview: - Forward computation graph of learn mode(updating policy). + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data from the replay buffer and then returns the output \ + result, including various training information such as loss, q value, priority. Arguments: - - data (:obj:`Dict[str, Any]`): Dict type data, a batch of data for training, values are torch.Tensor or \ - np.ndarray or dict/list combinations. + - data (:obj:`List[Dict[int, Any]]`): The input data used for policy forward, including a batch of \ + training samples. For each element in list, the key of the dict is the name of data items and the \ + value is the corresponding data. Usually, the value is torch.Tensor or np.ndarray or there dict/list \ + combinations. In the ``_forward_learn`` method, data often need to first be stacked in the batch \ + dimension by some utility functions such as ``default_preprocess_learn``. \ + For DQN, each element in list is a dict containing at least the following keys: ``obs``, ``action``, \ + ``reward``, ``next_obs``, ``done``. Sometimes, it also contains other keys such as ``weight`` \ + and ``value_gamma``. Returns: - - info_dict (:obj:`Dict[str, Any]`): Dict type data, a info dict indicated training result, which will be \ - recorded in text log and tensorboard, values are python scalar or a list of scalars. - ArgumentsKeys: - - necessary: ``obs``, ``action``, ``reward``, ``next_obs``, ``done`` - - optional: ``value_gamma``, ``IS`` - ReturnsKeys: - - necessary: ``cur_lr``, ``total_loss``, ``priority`` - - optional: ``action_distribution`` - """ + - info_dict (:obj:`Dict[str, Any]`): The information dict that indicated training result, which will be \ + recorded in text log and tensorboard, values must be python scalar or a list of scalars. For the \ + detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement your own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for DQNPolicy: ``ding.policy.tests.test_dqn``. + """ + # Set noise mode for NoisyNet for exploration in learning if enabled in config + # We need to reset set_noise_mode every _forward_xxx because the model is reused across different + # phases (learn/collect/eval). + if self._cfg.noisy_net: + set_noise_mode(self._learn_model, True) + set_noise_mode(self._target_model, True) + + # A noisy network agent samples a new set of parameters after every step of optimisation. + # Between optimisation steps, the agent acts according to a fixed set of parameters (weights and biases). + # This ensures that the agent always acts according to parameters that are drawn from + # the current noise distribution. + if self._cfg.noisy_net: + self._reset_noise(self._learn_model) + self._reset_noise(self._target_model) + + # Data preprocessing operations, such as stack data, cpu to cuda device data = default_preprocess_learn( data, use_priority=self._priority, @@ -193,9 +275,7 @@ def _forward_learn(self, data: Dict[str, Any]) -> Dict[str, Any]: ) if self._cuda: data = to_device(data, self._device) - # ==================== # Q-learning forward - # ==================== self._learn_model.train() self._target_model.train() # Current q value (main model) @@ -203,7 +283,7 @@ def _forward_learn(self, data: Dict[str, Any]) -> Dict[str, Any]: # Target q value with torch.no_grad(): target_q_value = self._target_model.forward(data['next_obs'])['logit'] - # Max q value action (main model) + # Max q value action (main model), i.e. Double DQN target_q_action = self._learn_model.forward(data['next_obs'])['action'] data_n = q_nstep_td_data( @@ -212,37 +292,41 @@ def _forward_learn(self, data: Dict[str, Any]) -> Dict[str, Any]: value_gamma = data.get('value_gamma') loss, td_error_per_sample = q_nstep_td_error(data_n, self._gamma, nstep=self._nstep, value_gamma=value_gamma) - # ==================== - # Q-learning update - # ==================== + # Update network parameters self._optimizer.zero_grad() loss.backward() - if self._cfg.learn.multi_gpu: + if self._cfg.multi_gpu: self.sync_gradients(self._learn_model) self._optimizer.step() - # ============= - # after update - # ============= + # Postprocessing operations, such as updating target model, return logged values and priority. self._target_model.update(self._learn_model.state_dict()) return { 'cur_lr': self._optimizer.defaults['lr'], 'total_loss': loss.item(), 'q_value': q_value.mean().item(), + 'target_q_value': target_q_value.mean().item(), 'priority': td_error_per_sample.abs().tolist(), # Only discrete action satisfying len(data['action'])==1 can return this and draw histogram on tensorboard. # '[histogram]action_distribution': data['action'], } def _monitor_vars_learn(self) -> List[str]: - return ['cur_lr', 'total_loss', 'q_value'] + """ + Overview: + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. + Returns: + - necessary_keys (:obj:`List[str]`): The list of the necessary keys to be logged. + """ + return ['cur_lr', 'total_loss', 'q_value', 'target_q_value'] def _state_dict_learn(self) -> Dict[str, Any]: """ Overview: - Return the state_dict of learn mode, usually including model and optimizer. + Return the state_dict of learn mode, usually including model, target_model and optimizer. Returns: - - state_dict (:obj:`Dict[str, Any]`): the dict of current policy learn state, for saving and restoring. + - state_dict (:obj:`Dict[str, Any]`): The dict of current policy learn state, for saving and restoring. """ return { 'model': self._learn_model.state_dict(), @@ -255,7 +339,7 @@ def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: Overview: Load the state_dict variable into policy learn mode. Arguments: - - state_dict (:obj:`Dict[str, Any]`): the dict of policy learn state saved before. + - state_dict (:obj:`Dict[str, Any]`): The dict of policy learn state saved before. .. tip:: If you want to only load some parts of model, you can simply set the ``strict`` argument in \ @@ -269,8 +353,18 @@ def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: def _init_collect(self) -> None: """ Overview: - Collect mode init method. Called by ``self.__init__``, initialize algorithm arguments and collect_model, \ - enable the eps_greedy_sample for exploration. + Initialize the collect mode of policy, including related attributes and modules. For DQN, it contains the \ + collect_model to balance the exploration and exploitation with epsilon-greedy sample mechanism, and other \ + algorithm-specific arguments such as unroll_len and nstep. + This method will be called in ``__init__`` method if ``collect`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_collect`` method, you'd better name them \ + with prefix ``_collect_`` to avoid conflict with other modes, such as ``self._collect_attr1``. + + .. tip:: + Some variables need to initialize independently in different modes, such as gamma and nstep in DQN. This \ + design is for the convenience of parallel execution of different policy modes. """ self._unroll_len = self._cfg.collect.unroll_len self._gamma = self._cfg.discount_factor # necessary for parallel @@ -281,23 +375,39 @@ def _init_collect(self) -> None: def _forward_collect(self, data: Dict[int, Any], eps: float) -> Dict[int, Any]: """ Overview: - Forward computation graph of collect mode(collect training data), with eps_greedy for exploration. + Policy forward function of collect mode (collecting training data by interacting with envs). Forward means \ + that the policy gets some necessary data (mainly observation) from the envs and then returns the output \ + data, such as the action to interact with the envs. Besides, this policy also needs ``eps`` argument for \ + exploration, i.e., classic epsilon-greedy exploration strategy. Arguments: - - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \ - values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer. - - eps (:obj:`float`): epsilon value for exploration, which is decayed by collected env step. + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. + - eps (:obj:`float`): The epsilon value for exploration. Returns: - - output (:obj:`Dict[int, Any]`): The dict of predicting policy_output(action) for the interaction with \ - env and the constructing of transition. - ArgumentsKeys: - - necessary: ``obs`` - ReturnsKeys - - necessary: ``logit``, ``action`` + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action and \ + other necessary data for learn mode defined in ``self._process_transition`` method. The key of the \ + dict is the same as the input data, i.e. environment id. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for DQNPolicy: ``ding.policy.tests.test_dqn``. """ + # Set noise mode for NoisyNet for exploration in collecting if enabled in config. + # We need to reset set_noise_mode every _forward_xxx because the model is reused across different + # phases (learn/collect/eval). + if self._cfg.noisy_net: + set_noise_mode(self._collect_model, True) + data_id = list(data.keys()) data = default_collate(list(data.values())) if self._cuda: data = to_device(data, self._device) + self._collect_model.eval() with torch.no_grad(): output = self._collect_model.forward(data, eps=eps) @@ -306,38 +416,40 @@ def _forward_collect(self, data: Dict[int, Any], eps: float) -> Dict[int, Any]: output = default_decollate(output) return {i: d for i, d in zip(data_id, output)} - def _get_train_sample(self, data: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + def _get_train_sample(self, transitions: List[Dict[str, Any]]) -> List[Dict[str, Any]]: """ Overview: - For a given trajectory(transitions, a list of transition) data, process it into a list of sample that \ - can be used for training directly. A train sample can be a processed transition(DQN with nstep TD) \ - or some continuous transitions(DRQN). + For a given trajectory (transitions, a list of transition) data, process it into a list of sample that \ + can be used for training directly. In DQN with nstep TD, a train sample is a processed transition. \ + This method is usually used in collectors to execute necessary \ + RL data preprocessing before training, which can help learner amortize relevant time consumption. \ + In addition, you can also implement this method as an identity function and do the data processing \ + in ``self._forward_learn`` method. Arguments: - - data (:obj:`List[Dict[str, Any]`): The trajectory data(a list of transition), each element is the same \ - format as the return value of ``self._process_transition`` method. + - transitions (:obj:`List[Dict[str, Any]`): The trajectory data (a list of transition), each element is \ + in the same format as the return value of ``self._process_transition`` method. Returns: - - samples (:obj:`dict`): The list of training samples. - - .. note:: - We will vectorize ``process_transition`` and ``get_train_sample`` method in the following release version. \ - And the user can customize the this data processing procecure by overriding this two methods and collector \ - itself. + - samples (:obj:`List[Dict[str, Any]]`): The processed train samples, each element is similar in format \ + to input transitions, but may contain more data for training, such as nstep reward and target obs. """ - data = get_nstep_return_data(data, self._nstep, gamma=self._gamma) - return get_train_sample(data, self._unroll_len) + transitions = get_nstep_return_data(transitions, self._nstep, gamma=self._gamma) + return get_train_sample(transitions, self._unroll_len) - def _process_transition(self, obs: Any, policy_output: Dict[str, Any], timestep: namedtuple) -> Dict[str, Any]: + def _process_transition(self, obs: torch.Tensor, policy_output: Dict[str, torch.Tensor], + timestep: namedtuple) -> Dict[str, torch.Tensor]: """ Overview: - Generate a transition(e.g.: ) for this algorithm training. + Process and pack one timestep transition data into a dict, which can be directly used for training and \ + saved in replay buffer. For DQN, it contains obs, next_obs, action, reward, done. Arguments: - - obs (:obj:`Any`): Env observation. - - policy_output (:obj:`Dict[str, Any]`): The output of policy collect mode(``self._forward_collect``),\ - including at least ``action``. - - timestep (:obj:`namedtuple`): The output after env step(execute policy output action), including at \ - least ``obs``, ``reward``, ``done``, (here obs indicates obs after env step). + - obs (:obj:`torch.Tensor`): The env observation of current timestep, such as stacked 2D image in Atari. + - policy_output (:obj:`Dict[str, torch.Tensor]`): The output of the policy network with the observation \ + as input. For DQN, it contains the action and the logit (q_value) of the action. + - timestep (:obj:`namedtuple`): The execution result namedtuple returned by the environment step method, \ + except all the elements have been transformed into tensor data. Usually, it contains the next obs, \ + reward, done, info, etc. Returns: - - transition (:obj:`dict`): Dict type transition data. + - transition (:obj:`Dict[str, torch.Tensor]`): The processed transition data of the current timestep. """ transition = { 'obs': obs, @@ -349,9 +461,15 @@ def _process_transition(self, obs: Any, policy_output: Dict[str, Any], timestep: return transition def _init_eval(self) -> None: - r""" + """ Overview: - Evaluate mode init method. Called by ``self.__init__``, initialize eval_model. + Initialize the eval mode of policy, including related attributes and modules. For DQN, it contains the \ + eval model to greedily select action with argmax q_value mechanism. + This method will be called in ``__init__`` method if ``eval`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_eval`` method, you'd better name them \ + with prefix ``_eval_`` to avoid conflict with other modes, such as ``self._eval_attr1``. """ self._eval_model = model_wrap(self._model, wrapper_name='argmax_sample') self._eval_model.reset() @@ -359,22 +477,35 @@ def _init_eval(self) -> None: def _forward_eval(self, data: Dict[int, Any]) -> Dict[int, Any]: """ Overview: - Forward computation graph of eval mode(evaluate policy performance), at most cases, it is similar to \ - ``self._forward_collect``. + Policy forward function of eval mode (evaluation policy performance by interacting with envs). Forward \ + means that the policy gets some necessary data (mainly observation) from the envs and then returns the \ + action to interact with the envs. Arguments: - - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \ - values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer. + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. Returns: - - output (:obj:`Dict[int, Any]`): The dict of predicting action for the interaction with env. - ArgumentsKeys: - - necessary: ``obs`` - ReturnsKeys - - necessary: ``action`` + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action. The \ + key of the dict is the same as the input data, i.e. environment id. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for DQNPolicy: ``ding.policy.tests.test_dqn``. """ + # We need to reset set_noise_mode every _forward_xxx because the model is reused across different + # phases (learn/collect/eval). + # Ensure that in evaluation mode noise is disabled. + set_noise_mode(self._eval_model, False) + data_id = list(data.keys()) data = default_collate(list(data.values())) if self._cuda: data = to_device(data, self._device) + self._eval_model.eval() with torch.no_grad(): output = self._eval_model.forward(data) @@ -383,12 +514,74 @@ def _forward_eval(self, data: Dict[int, Any]) -> Dict[int, Any]: output = default_decollate(output) return {i: d for i, d in zip(data_id, output)} + def calculate_priority(self, data: Dict[int, Any], update_target_model: bool = False) -> Dict[str, Any]: + """ + Overview: + Calculate priority for replay buffer. + Arguments: + - data (:obj:`Dict[str, Any]`): Dict type data, a batch of data for training. + - update_target_model (:obj:`bool`): Whether to update target model. + Returns: + - priority (:obj:`Dict[str, Any]`): Dict type priority data, values are python scalar or a list of scalars. + ArgumentsKeys: + - necessary: ``obs``, ``action``, ``reward``, ``next_obs``, ``done`` + - optional: ``value_gamma`` + ReturnsKeys: + - necessary: ``priority`` + """ + + if update_target_model: + self._target_model.load_state_dict(self._learn_model.state_dict()) + + data = default_preprocess_learn( + data, + use_priority=False, + use_priority_IS_weight=False, + ignore_done=self._cfg.learn.ignore_done, + use_nstep=True + ) + if self._cuda: + data = to_device(data, self._device) + # ==================== + # Q-learning forward + # ==================== + self._learn_model.eval() + self._target_model.eval() + with torch.no_grad(): + # Current q value (main model) + q_value = self._learn_model.forward(data['obs'])['logit'] + # Target q value + target_q_value = self._target_model.forward(data['next_obs'])['logit'] + # Max q value action (main model), i.e. Double DQN + target_q_action = self._learn_model.forward(data['next_obs'])['action'] + data_n = q_nstep_td_data( + q_value, target_q_value, data['action'], target_q_action, data['reward'], data['done'], data['weight'] + ) + value_gamma = data.get('value_gamma') + loss, td_error_per_sample = q_nstep_td_error( + data_n, self._gamma, nstep=self._nstep, value_gamma=value_gamma + ) + return {'priority': td_error_per_sample.abs().tolist()} + + def _reset_noise(self, model: torch.nn.Module): + r""" + Overview: + Reset the noise of model. + + Arguments: + - model (:obj:`torch.nn.Module`): the model to reset, must contain reset_noise method + """ + for m in model.modules(): + if hasattr(m, 'reset_noise'): + m.reset_noise() + @POLICY_REGISTRY.register('dqn_stdim') class DQNSTDIMPolicy(DQNPolicy): """ Overview: - Policy class of DQN algorithm, extended by auxiliary objectives. + Policy class of DQN algorithm, extended by ST-DIM auxiliary objectives. + ST-DIM paper link: https://arxiv.org/abs/1906.08226. Config: == ==================== ======== ============== ======================================== ======================= ID Symbol Type Default Value Description Other(Shape) @@ -411,7 +604,6 @@ class DQNSTDIMPolicy(DQNPolicy): 8 | ``learn.update`` int 3 | How many updates(iterations) to train | This args can be vary | ``per_collect`` | after collector's one collection. Only | from envs. Bigger val | valid in serial training | means more off-policy - 9 | ``learn.multi`` bool False | whether to use multi gpu during | ``_gpu`` 10 | ``learn.batch_`` int 64 | The number of samples of an iteration | ``size`` @@ -444,69 +636,83 @@ class DQNSTDIMPolicy(DQNPolicy): """ config = dict( + # (str) RL policy register name (refer to function "POLICY_REGISTRY"). type='dqn_stdim', - # (bool) Whether use cuda in policy + # (bool) Whether to use cuda in policy. cuda=False, - # (bool) Whether learning policy is the same as collecting data policy(on-policy) + # (bool) Whether to learning policy is the same as collecting data policy (on-policy). on_policy=False, - # (bool) Whether enable priority experience sample + # (bool) Whether to enable priority experience sample. priority=False, - # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. + # (bool) Whether to use Importance Sampling Weight to correct biased update. If True, priority must be True. priority_IS_weight=False, - # (float) Discount factor(gamma) for returns + # (float) Discount factor(gamma) for returns. discount_factor=0.97, - # (int) The number of step for calculating target q_value + # (int) The number of step for calculating target q_value. nstep=1, + # (float) The weight of auxiliary loss to main loss. + aux_loss_weight=0.001, + # learn_mode config learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, # How many updates(iterations) to train after collector's one collection. # Bigger "update_per_collect" means bigger off-policy. # collect data -> update policy-> collect data -> ... update_per_collect=3, - # (int) How many samples in a training batch + # (int) How many samples in a training batch. batch_size=64, - # (float) The step size of gradient descent + # (float) The step size of gradient descent. learning_rate=0.001, - # ============================================================== - # The following configs are algorithm-specific - # ============================================================== # (int) Frequence of target network update. target_update_freq=100, - # (bool) Whether ignore done(usually for max step termination env) + # (bool) Whether ignore done(usually for max step termination env). ignore_done=False, ), # collect_mode config collect=dict( - # (int) Only one of [n_sample, n_episode] shoule be set + # (int) How many training samples collected in one collection procedure. + # Only one of [n_sample, n_episode] shoule be set. # n_sample=8, # (int) Cut trajectories into pieces with length "unroll_len". unroll_len=1, ), - eval=dict(), + eval=dict(), # for compability # other config other=dict( # Epsilon greedy with decay. eps=dict( # (str) Decay type. Support ['exp', 'linear']. type='exp', - # (float) Epsilon start value + # (float) Epsilon start value. start=0.95, - # (float) Epsilon end value + # (float) Epsilon end value. end=0.1, - # (int) Decay length(env step) + # (int) Decay length (env step). decay=10000, ), - replay_buffer=dict(replay_buffer_size=10000, ), + replay_buffer=dict( + # (int) Maximum size of replay buffer. Usually, larger buffer size is better. + replay_buffer_size=10000, + ), ), - aux_loss_weight=0.001, ) def _init_learn(self) -> None: """ Overview: - Learn mode init method. Called by ``self.__init__``. - Init the auxiliary model, its optimizer, and the axuliary loss weight to the main loss. + Initialize the learn mode of policy, including related attributes and modules. For DQNSTDIM, it first \ + call super class's ``_init_learn`` method, then initialize extra auxiliary model, its optimizer, and the \ + loss weight. This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. """ super()._init_learn() x_size, y_size = self._get_encoding_size() @@ -516,12 +722,12 @@ def _init_learn(self) -> None: self._aux_optimizer = Adam(self._aux_model.parameters(), lr=self._cfg.learn.learning_rate) self._aux_loss_weight = self._cfg.aux_loss_weight - def _get_encoding_size(self): + def _get_encoding_size(self) -> Tuple[Tuple[int], Tuple[int]]: """ Overview: Get the input encoding size of the ST-DIM axuiliary model. Returns: - - info_dict (:obj:`[Tuple, Tuple]`): The encoding size without the first (Batch) dimension. + - info_dict (:obj:`Tuple[Tuple[int], Tuple[int]]`): The encoding size without the first (Batch) dimension. """ obs = self._cfg.model.obs_shape if isinstance(obs, int): @@ -536,16 +742,15 @@ def _get_encoding_size(self): x, y = self._model_encode(test_data) return x.size()[1:], y.size()[1:] - def _model_encode(self, data): + def _model_encode(self, data: dict) -> Tuple[torch.Tensor]: """ Overview: Get the encoding of the main model as input for the auxiliary model. Arguments: - data (:obj:`dict`): Dict type data, same as the _forward_learn input. Returns: - - (:obj:`Tuple[Tensor]`): the tuple of two tensors to apply contrastive embedding learning. - In ST-DIM algorithm, these two variables are the dqn encoding of `obs` and `next_obs`\ - respectively. + - (:obj:`Tuple[torch.Tensor]`): the tuple of two tensors to apply contrastive embedding learning. \ + In ST-DIM algorithm, these two variables are the dqn encoding of `obs` and `next_obs` respectively. """ assert hasattr(self._model, "encoder") x = self._model.encoder(data["obs"]) @@ -555,19 +760,28 @@ def _model_encode(self, data): def _forward_learn(self, data: Dict[str, Any]) -> Dict[str, Any]: """ Overview: - Forward computation graph of learn mode(updating policy). + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data from the replay buffer and then returns the output \ + result, including various training information such as loss, q value, priority, aux_loss. Arguments: - - data (:obj:`Dict[str, Any]`): Dict type data, a batch of data for training, values are torch.Tensor or \ - np.ndarray or dict/list combinations. + - data (:obj:`List[Dict[int, Any]]`): The input data used for policy forward, including a batch of \ + training samples. For each element in list, the key of the dict is the name of data items and the \ + value is the corresponding data. Usually, the value is torch.Tensor or np.ndarray or there dict/list \ + combinations. In the ``_forward_learn`` method, data often need to first be stacked in the batch \ + dimension by some utility functions such as ``default_preprocess_learn``. \ + For DQNSTDIM, each element in list is a dict containing at least the following keys: ``obs``, \ + ``action``, ``reward``, ``next_obs``, ``done``. Sometimes, it also contains other keys such as \ + ``weight`` and ``value_gamma``. Returns: - - info_dict (:obj:`Dict[str, Any]`): Dict type data, a info dict indicated training result, which will be \ - recorded in text log and tensorboard, values are python scalar or a list of scalars. - ArgumentsKeys: - - necessary: ``obs``, ``action``, ``reward``, ``next_obs``, ``done`` - - optional: ``value_gamma``, ``IS`` - ReturnsKeys: - - necessary: ``cur_lr``, ``total_loss``, ``priority`` - - optional: ``action_distribution`` + - info_dict (:obj:`Dict[str, Any]`): The information dict that indicated training result, which will be \ + recorded in text log and tensorboard, values must be python scalar or a list of scalars. For the \ + detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. """ data = default_preprocess_learn( data, @@ -592,7 +806,7 @@ def _forward_learn(self, data: Dict[str, Any]) -> Dict[str, Any]: # the BP process of the auxiliary network self._aux_optimizer.zero_grad() aux_loss_learn.backward() - if self._cfg.learn.multi_gpu: + if self._cfg.multi_gpu: self.sync_gradients(self._aux_model) self._aux_optimizer.step() @@ -630,7 +844,7 @@ def _forward_learn(self, data: Dict[str, Any]) -> Dict[str, Any]: # ==================== self._optimizer.zero_grad() loss.backward() - if self._cfg.learn.multi_gpu: + if self._cfg.multi_gpu: self.sync_gradients(self._learn_model) self._optimizer.step() @@ -651,6 +865,13 @@ def _forward_learn(self, data: Dict[str, Any]) -> Dict[str, Any]: } def _monitor_vars_learn(self) -> List[str]: + """ + Overview: + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. + Returns: + - necessary_keys (:obj:`List[str]`): The list of the necessary keys to be logged. + """ return ['cur_lr', 'bellman_loss', 'aux_loss_learn', 'aux_loss_eval', 'total_loss', 'q_value'] def _state_dict_learn(self) -> Dict[str, Any]: diff --git a/ding/policy/dt.py b/ding/policy/dt.py new file mode 100644 index 0000000000..005a624644 --- /dev/null +++ b/ding/policy/dt.py @@ -0,0 +1,433 @@ +from typing import List, Dict, Any, Tuple, Optional +from collections import namedtuple +import torch.nn.functional as F +import torch +import numpy as np +from ding.torch_utils import to_device +from ding.utils import POLICY_REGISTRY +from ding.utils.data import default_decollate +from .base_policy import Policy + + +@POLICY_REGISTRY.register('dt') +class DTPolicy(Policy): + """ + Overview: + Policy class of Decision Transformer algorithm in discrete environments. + Paper link: https://arxiv.org/abs/2106.01345. + """ + config = dict( + # (str) RL policy register name (refer to function "POLICY_REGISTRY"). + type='dt', + # (bool) Whether to use cuda for network. + cuda=False, + # (bool) Whether the RL algorithm is on-policy or off-policy. + on_policy=False, + # (bool) Whether use priority(priority sample, IS weight, update priority) + priority=False, + # (int) N-step reward for target q_value estimation + obs_shape=4, + action_shape=2, + rtg_scale=1000, # normalize returns to go + max_eval_ep_len=1000, # max len of one episode + batch_size=64, # training batch size + wt_decay=1e-4, # decay weight in optimizer + warmup_steps=10000, # steps for learning rate warmup + context_len=20, # length of transformer input + learning_rate=1e-4, + ) + + def default_model(self) -> Tuple[str, List[str]]: + """ + Overview: + Return this algorithm default neural network model setting for demonstration. ``__init__`` method will \ + automatically call this method to get the default model setting and create model. + Returns: + - model_info (:obj:`Tuple[str, List[str]]`): The registered model name and model's import_names. + + .. note:: + The user can define and use customized network model but must obey the same inferface definition indicated \ + by import_names path. For example about DQN, its registered name is ``dqn`` and the import_names is \ + ``ding.model.template.q_learning``. + """ + return 'dt', ['ding.model.template.dt'] + + def _init_learn(self) -> None: + """ + Overview: + Initialize the learn mode of policy, including related attributes and modules. For Decision Transformer, \ + it mainly contains the optimizer, algorithm-specific arguments such as rtg_scale and lr scheduler. + This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. + """ + # rtg_scale: scale of `return to go` + # rtg_target: max target of `return to go` + # Our goal is normalize `return to go` to (0, 1), which will favour the covergence. + # As a result, we usually set rtg_scale == rtg_target. + self.rtg_scale = self._cfg.rtg_scale # normalize returns to go + self.rtg_target = self._cfg.rtg_target # max target reward_to_go + self.max_eval_ep_len = self._cfg.max_eval_ep_len # max len of one episode + + lr = self._cfg.learning_rate # learning rate + wt_decay = self._cfg.wt_decay # weight decay + warmup_steps = self._cfg.warmup_steps # warmup steps for lr scheduler + + self.clip_grad_norm_p = self._cfg.clip_grad_norm_p + self.context_len = self._cfg.model.context_len # K in decision transformer + + self.state_dim = self._cfg.model.state_dim + self.act_dim = self._cfg.model.act_dim + + self._learn_model = self._model + self._atari_env = 'state_mean' not in self._cfg + self._basic_discrete_env = not self._cfg.model.continuous and 'state_mean' in self._cfg + + if self._atari_env: + self._optimizer = self._learn_model.configure_optimizers(wt_decay, lr) + else: + self._optimizer = torch.optim.AdamW(self._learn_model.parameters(), lr=lr, weight_decay=wt_decay) + + self._scheduler = torch.optim.lr_scheduler.LambdaLR( + self._optimizer, lambda steps: min((steps + 1) / warmup_steps, 1) + ) + + self.max_env_score = -1.0 + + def _forward_learn(self, data: List[torch.Tensor]) -> Dict[str, Any]: + """ + Overview: + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data from the offline dataset and then returns the output \ + result, including various training information such as loss, current learning rate. + Arguments: + - data (:obj:`List[torch.Tensor]`): The input data used for policy forward, including a series of \ + processed torch.Tensor data, i.e., timesteps, states, actions, returns_to_go, traj_mask. + Returns: + - info_dict (:obj:`Dict[str, Any]`): The information dict that indicated training result, which will be \ + recorded in text log and tensorboard, values must be python scalar or a list of scalars. For the \ + detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + """ + self._learn_model.train() + + timesteps, states, actions, returns_to_go, traj_mask = data + + # The shape of `returns_to_go` may differ with different dataset (B x T or B x T x 1), + # and we need a 3-dim tensor + if len(returns_to_go.shape) == 2: + returns_to_go = returns_to_go.unsqueeze(-1) + + if self._basic_discrete_env: + actions = actions.to(torch.long) + actions = actions.squeeze(-1) + action_target = torch.clone(actions).detach().to(self._device) + + if self._atari_env: + state_preds, action_preds, return_preds = self._learn_model.forward( + timesteps=timesteps, states=states, actions=actions, returns_to_go=returns_to_go, tar=1 + ) + else: + state_preds, action_preds, return_preds = self._learn_model.forward( + timesteps=timesteps, states=states, actions=actions, returns_to_go=returns_to_go + ) + + if self._atari_env: + action_loss = F.cross_entropy(action_preds.reshape(-1, action_preds.size(-1)), action_target.reshape(-1)) + else: + traj_mask = traj_mask.view(-1, ) + + # only consider non padded elements + action_preds = action_preds.view(-1, self.act_dim)[traj_mask > 0] + + if self._cfg.model.continuous: + action_target = action_target.view(-1, self.act_dim)[traj_mask > 0] + action_loss = F.mse_loss(action_preds, action_target) + else: + action_target = action_target.view(-1)[traj_mask > 0] + action_loss = F.cross_entropy(action_preds, action_target) + + self._optimizer.zero_grad() + action_loss.backward() + if self._cfg.multi_gpu: + self.sync_gradients(self._learn_model) + torch.nn.utils.clip_grad_norm_(self._learn_model.parameters(), self.clip_grad_norm_p) + self._optimizer.step() + self._scheduler.step() + + return { + 'cur_lr': self._optimizer.state_dict()['param_groups'][0]['lr'], + 'action_loss': action_loss.detach().cpu().item(), + 'total_loss': action_loss.detach().cpu().item(), + } + + def _init_eval(self) -> None: + """ + Overview: + Initialize the eval mode of policy, including related attributes and modules. For DQN, it contains the \ + eval model, some algorithm-specific parameters such as context_len, max_eval_ep_len, etc. + This method will be called in ``__init__`` method if ``eval`` field is in ``enable_field``. + + .. tip:: + For the evaluation of complete episodes, we need to maintain some historical information for transformer \ + inference. These variables need to be initialized in ``_init_eval`` and reset in ``_reset_eval`` when \ + necessary. + + .. note:: + If you want to set some spacial member variables in ``_init_eval`` method, you'd better name them \ + with prefix ``_eval_`` to avoid conflict with other modes, such as ``self._eval_attr1``. + """ + self._eval_model = self._model + # init data + self._device = torch.device(self._device) + self.rtg_scale = self._cfg.rtg_scale # normalize returns to go + self.rtg_target = self._cfg.rtg_target # max target reward_to_go + self.state_dim = self._cfg.model.state_dim + self.act_dim = self._cfg.model.act_dim + self.eval_batch_size = self._cfg.evaluator_env_num + self.max_eval_ep_len = self._cfg.max_eval_ep_len + self.context_len = self._cfg.model.context_len # K in decision transformer + + self.t = [0 for _ in range(self.eval_batch_size)] + if self._cfg.model.continuous: + self.actions = torch.zeros( + (self.eval_batch_size, self.max_eval_ep_len, self.act_dim), dtype=torch.float32, device=self._device + ) + else: + self.actions = torch.zeros( + (self.eval_batch_size, self.max_eval_ep_len, 1), dtype=torch.long, device=self._device + ) + self._atari_env = 'state_mean' not in self._cfg + self._basic_discrete_env = not self._cfg.model.continuous and 'state_mean' in self._cfg + if self._atari_env: + self.states = torch.zeros( + ( + self.eval_batch_size, + self.max_eval_ep_len, + ) + tuple(self.state_dim), + dtype=torch.float32, + device=self._device + ) + self.running_rtg = [self.rtg_target for _ in range(self.eval_batch_size)] + else: + self.running_rtg = [self.rtg_target / self.rtg_scale for _ in range(self.eval_batch_size)] + self.states = torch.zeros( + (self.eval_batch_size, self.max_eval_ep_len, self.state_dim), dtype=torch.float32, device=self._device + ) + self.state_mean = torch.from_numpy(np.array(self._cfg.state_mean)).to(self._device) + self.state_std = torch.from_numpy(np.array(self._cfg.state_std)).to(self._device) + self.timesteps = torch.arange( + start=0, end=self.max_eval_ep_len, step=1 + ).repeat(self.eval_batch_size, 1).to(self._device) + self.rewards_to_go = torch.zeros( + (self.eval_batch_size, self.max_eval_ep_len, 1), dtype=torch.float32, device=self._device + ) + + def _forward_eval(self, data: Dict[int, Any]) -> Dict[int, Any]: + """ + Overview: + Policy forward function of eval mode (evaluation policy performance, such as interacting with envs. \ + Forward means that the policy gets some input data (current obs/return-to-go and historical information) \ + from the envs and then returns the output data, such as the action to interact with the envs. \ + Arguments: + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs and \ + reward to calculate running return-to-go. The key of the dict is environment id and the value is the \ + corresponding data of the env. + Returns: + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action. The \ + key of the dict is the same as the input data, i.e. environment id. + + .. note:: + Decision Transformer will do different operations for different types of envs in evaluation. + """ + # save and forward + data_id = list(data.keys()) + + self._eval_model.eval() + with torch.no_grad(): + if self._atari_env: + states = torch.zeros( + ( + self.eval_batch_size, + self.context_len, + ) + tuple(self.state_dim), + dtype=torch.float32, + device=self._device + ) + timesteps = torch.zeros((self.eval_batch_size, 1, 1), dtype=torch.long, device=self._device) + else: + states = torch.zeros( + (self.eval_batch_size, self.context_len, self.state_dim), dtype=torch.float32, device=self._device + ) + timesteps = torch.zeros((self.eval_batch_size, self.context_len), dtype=torch.long, device=self._device) + if not self._cfg.model.continuous: + actions = torch.zeros( + (self.eval_batch_size, self.context_len, 1), dtype=torch.long, device=self._device + ) + else: + actions = torch.zeros( + (self.eval_batch_size, self.context_len, self.act_dim), dtype=torch.float32, device=self._device + ) + rewards_to_go = torch.zeros( + (self.eval_batch_size, self.context_len, 1), dtype=torch.float32, device=self._device + ) + for i in data_id: + if self._atari_env: + self.states[i, self.t[i]] = data[i]['obs'].to(self._device) + else: + self.states[i, self.t[i]] = (data[i]['obs'].to(self._device) - self.state_mean) / self.state_std + self.running_rtg[i] = self.running_rtg[i] - (data[i]['reward'] / self.rtg_scale).to(self._device) + self.rewards_to_go[i, self.t[i]] = self.running_rtg[i] + + if self.t[i] <= self.context_len: + if self._atari_env: + timesteps[i] = min(self.t[i], self._cfg.model.max_timestep) * torch.ones( + (1, 1), dtype=torch.int64 + ).to(self._device) + else: + timesteps[i] = self.timesteps[i, :self.context_len] + states[i] = self.states[i, :self.context_len] + actions[i] = self.actions[i, :self.context_len] + rewards_to_go[i] = self.rewards_to_go[i, :self.context_len] + else: + if self._atari_env: + timesteps[i] = min(self.t[i], self._cfg.model.max_timestep) * torch.ones( + (1, 1), dtype=torch.int64 + ).to(self._device) + else: + timesteps[i] = self.timesteps[i, self.t[i] - self.context_len + 1:self.t[i] + 1] + states[i] = self.states[i, self.t[i] - self.context_len + 1:self.t[i] + 1] + actions[i] = self.actions[i, self.t[i] - self.context_len + 1:self.t[i] + 1] + rewards_to_go[i] = self.rewards_to_go[i, self.t[i] - self.context_len + 1:self.t[i] + 1] + if self._basic_discrete_env: + actions = actions.squeeze(-1) + _, act_preds, _ = self._eval_model.forward(timesteps, states, actions, rewards_to_go) + del timesteps, states, actions, rewards_to_go + + logits = act_preds[:, -1, :] + if not self._cfg.model.continuous: + if self._atari_env: + probs = F.softmax(logits, dim=-1) + act = torch.zeros((self.eval_batch_size, 1), dtype=torch.long, device=self._device) + for i in data_id: + act[i] = torch.multinomial(probs[i], num_samples=1) + else: + act = torch.argmax(logits, axis=1).unsqueeze(1) + else: + act = logits + for i in data_id: + self.actions[i, self.t[i]] = act[i] # TODO: self.actions[i] should be a queue when exceed max_t + self.t[i] += 1 + + if self._cuda: + act = to_device(act, 'cpu') + output = {'action': act} + output = default_decollate(output) + return {i: d for i, d in zip(data_id, output)} + + def _reset_eval(self, data_id: Optional[List[int]] = None) -> None: + """ + Overview: + Reset some stateful variables for eval mode when necessary, such as the historical info of transformer \ + for decision transformer. If ``data_id`` is None, it means to reset all the stateful \ + varaibles. Otherwise, it will reset the stateful variables according to the ``data_id``. For example, \ + different environments/episodes in evaluation in ``data_id`` will have different history. + Arguments: + - data_id (:obj:`Optional[List[int]]`): The id of the data, which is used to reset the stateful variables \ + specified by ``data_id``. + """ + # clean data + if data_id is None: + self.t = [0 for _ in range(self.eval_batch_size)] + self.timesteps = torch.arange( + start=0, end=self.max_eval_ep_len, step=1 + ).repeat(self.eval_batch_size, 1).to(self._device) + if not self._cfg.model.continuous: + self.actions = torch.zeros( + (self.eval_batch_size, self.max_eval_ep_len, 1), dtype=torch.long, device=self._device + ) + else: + self.actions = torch.zeros( + (self.eval_batch_size, self.max_eval_ep_len, self.act_dim), + dtype=torch.float32, + device=self._device + ) + if self._atari_env: + self.states = torch.zeros( + ( + self.eval_batch_size, + self.max_eval_ep_len, + ) + tuple(self.state_dim), + dtype=torch.float32, + device=self._device + ) + self.running_rtg = [self.rtg_target for _ in range(self.eval_batch_size)] + else: + self.states = torch.zeros( + (self.eval_batch_size, self.max_eval_ep_len, self.state_dim), + dtype=torch.float32, + device=self._device + ) + self.running_rtg = [self.rtg_target / self.rtg_scale for _ in range(self.eval_batch_size)] + + self.rewards_to_go = torch.zeros( + (self.eval_batch_size, self.max_eval_ep_len, 1), dtype=torch.float32, device=self._device + ) + else: + for i in data_id: + self.t[i] = 0 + if not self._cfg.model.continuous: + self.actions[i] = torch.zeros((self.max_eval_ep_len, 1), dtype=torch.long, device=self._device) + else: + self.actions[i] = torch.zeros( + (self.max_eval_ep_len, self.act_dim), dtype=torch.float32, device=self._device + ) + if self._atari_env: + self.states[i] = torch.zeros( + (self.max_eval_ep_len, ) + tuple(self.state_dim), dtype=torch.float32, device=self._device + ) + self.running_rtg[i] = self.rtg_target + else: + self.states[i] = torch.zeros( + (self.max_eval_ep_len, self.state_dim), dtype=torch.float32, device=self._device + ) + self.running_rtg[i] = self.rtg_target / self.rtg_scale + self.timesteps[i] = torch.arange(start=0, end=self.max_eval_ep_len, step=1).to(self._device) + self.rewards_to_go[i] = torch.zeros((self.max_eval_ep_len, 1), dtype=torch.float32, device=self._device) + + def _monitor_vars_learn(self) -> List[str]: + """ + Overview: + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. + Returns: + - necessary_keys (:obj:`List[str]`): The list of the necessary keys to be logged. + """ + return ['cur_lr', 'action_loss'] + + def _init_collect(self) -> None: + pass + + def _forward_collect(self, data: Dict[int, Any], eps: float) -> Dict[int, Any]: + pass + + def _get_train_sample(self, data: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + pass + + def _process_transition(self, obs: Any, policy_output: Dict[str, Any], timestep: namedtuple) -> Dict[str, Any]: + pass diff --git a/ding/policy/edac.py b/ding/policy/edac.py new file mode 100755 index 0000000000..a8665b494d --- /dev/null +++ b/ding/policy/edac.py @@ -0,0 +1,340 @@ +from typing import List, Dict, Any, Tuple, Union +import copy +import numpy as np +import torch +import torch.nn as nn +from torch.distributions import Normal, Independent + +from ding.torch_utils import Adam, to_device +from ding.rl_utils import v_1step_td_data, v_1step_td_error, get_train_sample, \ + qrdqn_nstep_td_data, qrdqn_nstep_td_error, get_nstep_return_data +from ding.model import model_wrap +from ding.utils import POLICY_REGISTRY +from ding.utils.data import default_collate, default_decollate +from .sac import SACPolicy +from .dqn import DQNPolicy +from .common_utils import default_preprocess_learn + + +@POLICY_REGISTRY.register('edac') +class EDACPolicy(SACPolicy): + """ + Overview: + Policy class of EDAC algorithm. Paper link: https://arxiv.org/pdf/2110.01548.pdf + + Config: + == ==================== ======== ============= ================================= ======================= + ID Symbol Type Default Value Description Other(Shape) + == ==================== ======== ============= ================================= ======================= + 1 ``type`` str td3 | RL policy register name, refer | this arg is optional, + | to registry ``POLICY_REGISTRY`` | a placeholder + 2 ``cuda`` bool True | Whether to use cuda for network | + 3 | ``random_`` int 10000 | Number of randomly collected | Default to 10000 for + | ``collect_size`` | training samples in replay | SAC, 25000 for DDPG/ + | | buffer when training starts. | TD3. + 4 | ``model.policy_`` int 256 | Linear layer size for policy | + | ``embedding_size`` | network. | + 5 | ``model.soft_q_`` int 256 | Linear layer size for soft q | + | ``embedding_size`` | network. | + 6 | ``model.emsemble`` int 10 | Number of Q-ensemble network | + | ``_num`` | | + | | | is False. + 7 | ``learn.learning`` float 3e-4 | Learning rate for soft q | Defalut to 1e-3, when + | ``_rate_q`` | network. | model.value_network + | | | is True. + 8 | ``learn.learning`` float 3e-4 | Learning rate for policy | Defalut to 1e-3, when + | ``_rate_policy`` | network. | model.value_network + | | | is True. + 9 | ``learn.learning`` float 3e-4 | Learning rate for policy | Defalut to None when + | ``_rate_value`` | network. | model.value_network + | | | is False. + 10 | ``learn.alpha`` float 1.0 | Entropy regularization | alpha is initiali- + | | coefficient. | zation for auto + | | | `alpha`, when + | | | auto_alpha is True + 11 | ``learn.eta`` bool True | Parameter of EDAC algorithm | Defalut to 1.0 + 12 | ``learn.`` bool True | Determine whether to use | Temperature parameter + | ``auto_alpha`` | auto temperature parameter | determines the + | | `alpha`. | relative importance + | | | of the entropy term + | | | against the reward. + 13 | ``learn.-`` bool False | Determine whether to ignore | Use ignore_done only + | ``ignore_done`` | done flag. | in halfcheetah env. + 14 | ``learn.-`` float 0.005 | Used for soft update of the | aka. Interpolation + | ``target_theta`` | target network. | factor in polyak aver + | | | aging for target + | | | networks. + == ==================== ======== ============= ================================= ======================= + """ + config = dict( + # (str) RL policy register name + type='edac', + cuda=False, + on_policy=False, + multi_agent=False, + priority=False, + priority_IS_weight=False, + random_collect_size=10000, + model=dict( + # (bool type) ensemble_num:num of Q-network. + ensemble_num=10, + # (bool type) value_network: Determine whether to use value network as the + # original EDAC paper (arXiv 2110.01548). + # using value_network needs to set learning_rate_value, learning_rate_q, + # and learning_rate_policy in `cfg.policy.learn`. + # Default to False. + # value_network=False, + + # (int) Hidden size for actor network head. + actor_head_hidden_size=256, + + # (int) Hidden size for critic network head. + critic_head_hidden_size=256, + ), + learn=dict( + multi_gpu=False, + update_per_collect=1, + batch_size=256, + learning_rate_q=3e-4, + learning_rate_policy=3e-4, + learning_rate_value=3e-4, + learning_rate_alpha=3e-4, + target_theta=0.005, + discount_factor=0.99, + alpha=1, + auto_alpha=True, + # (bool type) log_space: Determine whether to use auto `\alpha` in log space. + log_space=True, + # (bool) Whether ignore done(usually for max step termination env. e.g. pendulum) + # Note: Gym wraps the MuJoCo envs by default with TimeLimit environment wrappers. + # These limit HalfCheetah, and several other MuJoCo envs, to max length of 1000. + # However, interaction with HalfCheetah always gets done with done is False, + # Since we inplace done==True with done==False to keep + # TD-error accurate computation(``gamma * (1 - done) * next_v + reward``), + # when the episode step is greater than max episode step. + ignore_done=False, + # (float) Weight uniform initialization range in the last output layer + init_w=3e-3, + # (float) Loss weight for conservative item. + min_q_weight=1.0, + # (bool) Whether to use entropy in target q. + with_q_entropy=False, + eta=0.1, + ), + collect=dict( + # (int) Cut trajectories into pieces with length "unroll_len". + unroll_len=1, + ), + eval=dict(), + other=dict( + replay_buffer=dict( + # (int type) replay_buffer_size: Max size of replay buffer. + replay_buffer_size=1000000, + # (int type) max_use: Max use times of one data in the buffer. + # Data will be removed once used for too many times. + # Default to infinite. + # max_use=256, + ), + ), + ) + + def default_model(self) -> Tuple[str, List[str]]: + """ + Overview: + Return this algorithm default neural network model setting for demonstration. ``__init__`` method will \ + automatically call this method to get the default model setting and create model. + Returns: + - model_info (:obj:`Tuple[str, List[str]]`): The registered model name and model's import_names. + """ + return 'edac', ['ding.model.template.edac'] + + def _init_learn(self) -> None: + """ + Overview: + Initialize the learn mode of policy, including related attributes and modules. For EDAC, in addition \ + to the things that need to be initialized in SAC, it is also necessary to additionally define \ + eta/with_q_entropy/forward_learn_cnt. \ + This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. + """ + super()._init_learn() + # EDAC special implementation + self._eta = self._cfg.learn.eta + self._with_q_entropy = self._cfg.learn.with_q_entropy + self._forward_learn_cnt = 0 + + def _forward_learn(self, data: List[Dict[int, Any]]) -> Dict[str, Any]: + """ + Overview: + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data from the replay buffer and then returns the output \ + result, including various training information such as loss, action, priority. + Arguments: + - data (:obj:`List[Dict[int, Any]]`): The input data used for policy forward, including a batch of \ + training samples. For each element in list, the key of the dict is the name of data items and the \ + value is the corresponding data. Usually, the value is torch.Tensor or np.ndarray or there dict/list \ + combinations. In the ``_forward_learn`` method, data often need to first be stacked in the batch \ + dimension by some utility functions such as ``default_preprocess_learn``. \ + For EDAC, each element in list is a dict containing at least the following keys: ``obs``, ``action``, \ + ``logit``, ``reward``, ``next_obs``, ``done``. Sometimes, it also contains other keys like ``weight``. + Returns: + - info_dict (:obj:`Dict[str, Any]`): The information dict that indicated training result, which will be \ + recorded in text log and tensorboard, values must be python scalar or a list of scalars. For the \ + detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for EDACPolicy: \ + ``ding.policy.tests.test_edac``. + """ + loss_dict = {} + data = default_preprocess_learn( + data, + use_priority=self._priority, + use_priority_IS_weight=self._cfg.priority_IS_weight, + ignore_done=self._cfg.learn.ignore_done, + use_nstep=False + ) + if len(data.get('action').shape) == 1: + data['action'] = data['action'].reshape(-1, 1) + + if self._cuda: + data = to_device(data, self._device) + + self._learn_model.train() + self._target_model.train() + obs = data['obs'] + next_obs = data['next_obs'] + reward = data['reward'] + done = data['done'] + acs = data['action'] + + # 1. predict q value + q_value = self._learn_model.forward(data, mode='compute_critic')['q_value'] + with torch.no_grad(): + (mu, sigma) = self._learn_model.forward(next_obs, mode='compute_actor')['logit'] + + dist = Independent(Normal(mu, sigma), 1) + pred = dist.rsample() + next_action = torch.tanh(pred) + y = 1 - next_action.pow(2) + 1e-6 + next_log_prob = dist.log_prob(pred).unsqueeze(-1) + next_log_prob = next_log_prob - torch.log(y).sum(-1, keepdim=True) + + next_data = {'obs': next_obs, 'action': next_action} + target_q_value = self._target_model.forward(next_data, mode='compute_critic')['q_value'] + # the value of a policy according to the maximum entropy objective + + target_q_value, _ = torch.min(target_q_value, dim=0) + if self._with_q_entropy: + target_q_value -= self._alpha * next_log_prob.squeeze(-1) + target_q_value = self._gamma * (1 - done) * target_q_value + reward + + weight = data['weight'] + if weight is None: + weight = torch.ones_like(q_value) + td_error_per_sample = nn.MSELoss(reduction='none')(q_value, target_q_value).mean(dim=1).sum() + loss_dict['critic_loss'] = (td_error_per_sample * weight).mean() + + # penalty term of EDAC + if self._eta > 0: + # [batch_size,dim] -> [Ensemble_num,batch_size,dim] + pre_obs = obs.unsqueeze(0).repeat_interleave(self._cfg.model.ensemble_num, dim=0) + pre_acs = acs.unsqueeze(0).repeat_interleave(self._cfg.model.ensemble_num, dim=0).requires_grad_(True) + + # [Ensemble_num,batch_size] + q_pred_tile = self._learn_model.forward({ + 'obs': pre_obs, + 'action': pre_acs + }, mode='compute_critic')['q_value'].requires_grad_(True) + + q_pred_grads = torch.autograd.grad(q_pred_tile.sum(), pre_acs, retain_graph=True, create_graph=True)[0] + q_pred_grads = q_pred_grads / (torch.norm(q_pred_grads, p=2, dim=2).unsqueeze(-1) + 1e-10) + # [Ensemble_num,batch_size,act_dim] -> [batch_size,Ensemble_num,act_dim] + q_pred_grads = q_pred_grads.transpose(0, 1) + + q_pred_grads = q_pred_grads @ q_pred_grads.permute(0, 2, 1) + masks = torch.eye( + self._cfg.model.ensemble_num, device=obs.device + ).unsqueeze(dim=0).repeat(q_pred_grads.size(0), 1, 1) + q_pred_grads = (1 - masks) * q_pred_grads + grad_loss = torch.mean(torch.sum(q_pred_grads, dim=(1, 2))) / (self._cfg.model.ensemble_num - 1) + loss_dict['critic_loss'] += grad_loss * self._eta + + self._optimizer_q.zero_grad() + loss_dict['critic_loss'].backward() + self._optimizer_q.step() + + (mu, sigma) = self._learn_model.forward(data['obs'], mode='compute_actor')['logit'] + dist = Independent(Normal(mu, sigma), 1) + pred = dist.rsample() + action = torch.tanh(pred) + y = 1 - action.pow(2) + 1e-6 + log_prob = dist.log_prob(pred).unsqueeze(-1) + log_prob = log_prob - torch.log(y).sum(-1, keepdim=True) + + eval_data = {'obs': obs, 'action': action} + new_q_value = self._learn_model.forward(eval_data, mode='compute_critic')['q_value'] + new_q_value, _ = torch.min(new_q_value, dim=0) + + # 8. compute policy loss + policy_loss = (self._alpha * log_prob - new_q_value.unsqueeze(-1)).mean() + + loss_dict['policy_loss'] = policy_loss + + # 9. update policy network + self._optimizer_policy.zero_grad() + loss_dict['policy_loss'].backward() + self._optimizer_policy.step() + + # 10. compute alpha loss + if self._auto_alpha: + if self._log_space: + log_prob = log_prob + self._target_entropy + loss_dict['alpha_loss'] = -(self._log_alpha * log_prob.detach()).mean() + + self._alpha_optim.zero_grad() + loss_dict['alpha_loss'].backward() + self._alpha_optim.step() + self._alpha = self._log_alpha.detach().exp() + else: + log_prob = log_prob + self._target_entropy + loss_dict['alpha_loss'] = -(self._alpha * log_prob.detach()).mean() + + self._alpha_optim.zero_grad() + loss_dict['alpha_loss'].backward() + self._alpha_optim.step() + self._alpha = max(0, self._alpha) + + loss_dict['total_loss'] = sum(loss_dict.values()) + + # ============= + # after update + # ============= + self._forward_learn_cnt += 1 + # target update + self._target_model.update(self._learn_model.state_dict()) + return { + 'cur_lr_q': self._optimizer_q.defaults['lr'], + 'cur_lr_p': self._optimizer_policy.defaults['lr'], + 'priority': td_error_per_sample.abs().tolist(), + 'td_error': td_error_per_sample.detach().mean().item(), + 'alpha': self._alpha.item(), + 'target_q_value': target_q_value.detach().mean().item(), + **loss_dict + } diff --git a/ding/policy/fqf.py b/ding/policy/fqf.py index d98b157cf5..fae697f85d 100644 --- a/ding/policy/fqf.py +++ b/ding/policy/fqf.py @@ -1,22 +1,34 @@ -from typing import List, Dict, Any, Tuple, Union import copy +from typing import List, Dict, Any, Tuple + import torch -from ding.torch_utils import Adam, RMSprop, to_device -from ding.rl_utils import fqf_nstep_td_data, fqf_nstep_td_error, fqf_calculate_fraction_loss, \ - get_train_sample, get_nstep_return_data from ding.model import model_wrap +from ding.rl_utils import fqf_nstep_td_data, fqf_nstep_td_error, fqf_calculate_fraction_loss +from ding.torch_utils import Adam, RMSprop, to_device from ding.utils import POLICY_REGISTRY -from ding.utils.data import default_collate, default_decollate -from .dqn import DQNPolicy from .common_utils import default_preprocess_learn +from .dqn import DQNPolicy + + +def compute_grad_norm(model): + """ + Overview: + Compute grad norm of a network's parameters. + Arguments: + - model (:obj:`nn.Module`): The network to compute grad norm. + Returns: + - grad_norm (:obj:`torch.Tensor`): The grad norm of the network's parameters. + """ + return torch.norm(torch.stack([torch.norm(p.grad.detach(), 2.0) for p in model.parameters()]), 2.0) @POLICY_REGISTRY.register('fqf') class FQFPolicy(DQNPolicy): - r""" + """ Overview: - Policy class of FQF algorithm. + Policy class of FQF (Fully Parameterized Quantile Function) algorithm, proposed in + https://arxiv.org/pdf/1911.02140.pdf. Config: == ==================== ======== ============== ======================================== ======================= @@ -46,71 +58,100 @@ class FQFPolicy(DQNPolicy): """ config = dict( - # (str) RL policy register name (refer to function "POLICY_REGISTRY"). + # (str) Name of the RL policy registered in "POLICY_REGISTRY" function. type='fqf', - # (bool) Whether to use cuda for network. + # (bool) Flag to enable/disable CUDA for network computation. cuda=False, - # (bool) Whether the RL algorithm is on-policy or off-policy. + # (bool) Indicator of the RL algorithm's policy type (True for on-policy algorithms). on_policy=False, - # (bool) Whether use priority(priority sample, IS weight, update priority) + # (bool) Toggle for using prioritized experience replay (priority sampling and updating). priority=False, - # (float) Reward's future discount factor, aka. gamma. + # (float) Discount factor (gamma) for calculating the future reward. discount_factor=0.97, - # (int) N-step reward for target q_value estimation + # (int) Number of steps to consider for calculating n-step returns. nstep=1, learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, - # How many updates(iterations) to train after collector's one collection. - # Bigger "update_per_collect" means bigger off-policy. - # collect data -> update policy-> collect data -> ... + # (int) Number of training iterations per data collection from the environment. update_per_collect=3, + # (int) Size of minibatch for each update. batch_size=64, + # (float) Fractional learning rate for the fraction proposal network. learning_rate_fraction=2.5e-9, + # (float) Learning rate for the quantile regression network. learning_rate_quantile=0.00005, # ============================================================== - # The following configs are algorithm-specific + # Algorithm-specific configurations # ============================================================== - # (int) Frequence of target network update. + # (int) Frequency of target network updates. target_update_freq=100, - # (float) Threshold of Huber loss. In the FQF paper, this is denoted by kappa. Default to 1.0. + # (float) Huber loss threshold (kappa in the FQF paper). kappa=1.0, - # (float) Coefficient of entropy_loss. + # (float) Coefficient for the entropy loss term. ent_coef=0, - # (bool) Whether ignore done(usually for max step termination env) + # (bool) If set to True, the 'done' signals that indicate the end of an episode due to environment time + # limits are disregarded. By default, this is set to False. This setting is particularly useful for tasks + # that have a predetermined episode length, such as HalfCheetah and various other MuJoCo environments, + # where the maximum length is capped at 1000 steps. When enabled, any 'done' signal triggered by reaching + # the maximum episode steps will be overridden to 'False'. This ensures the accurate calculation of the + # Temporal Difference (TD) error, using the formula `gamma * (1 - done) * next_v + reward`, + # even when the episode surpasses the predefined step limit. ignore_done=False, ), - # collect_mode config collect=dict( - # (int) Only one of [n_sample, n_step, n_episode] shoule be set + # (int) Specify one of [n_sample, n_step, n_episode] for data collection. # n_sample=8, - # (int) Cut trajectories into pieces with length "unroll_len". + # (int) Length of trajectory segments for processing. unroll_len=1, ), eval=dict(), - # other config other=dict( - # Epsilon greedy with decay. + # Epsilon-greedy strategy with a decay mechanism. eps=dict( - # (str) Decay type. Support ['exp', 'linear']. + # (str) Type of decay mechanism ['exp' for exponential, 'linear']. type='exp', + # (float) Initial value of epsilon in epsilon-greedy exploration. start=0.95, + # (float) Final value of epsilon after decay. end=0.1, - # (int) Decay length(env step) + # (int) Number of environment steps over which epsilon is decayed. decay=10000, ), - replay_buffer=dict(replay_buffer_size=10000, ) + replay_buffer=dict( + # (int) Size of the replay buffer. + replay_buffer_size=10000, + ), ), ) def default_model(self) -> Tuple[str, List[str]]: + """ + Overview: + Returns the default model configuration used by the FQF algorithm. ``__init__`` method will \ + automatically call this method to get the default model setting and create model. + + Returns: + - model_info (:obj:`Tuple[str, List[str]]`): \ + Tuple containing the registered model name and model's import_names. + """ return 'fqf', ['ding.model.template.q_learning'] def _init_learn(self) -> None: - r""" + """ Overview: - Learn mode init method. Called by ``self.__init__``. - Init the optimizer, algorithm config, main and target models. + Initialize the learn mode of policy, including related attributes and modules. For FQF, it mainly \ + contains optimizer, algorithm-specific arguments such as gamma, nstep, kappa ent_coef, main and \ + target model. This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. """ self._priority = self._cfg.priority # Optimizer @@ -144,15 +185,32 @@ def _init_learn(self) -> None: self._learn_model.reset() self._target_model.reset() - def _forward_learn(self, data: dict) -> Dict[str, Any]: - r""" + def _forward_learn(self, data: List[Dict[str, Any]]) -> Dict[str, Any]: + """ Overview: - Forward and backward function of learn mode. + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data from the replay buffer and then returns the output \ + result, including various training information such as policy_loss, value_loss, entropy_loss. Arguments: - - data (:obj:`dict`): Dict type data, including at least ['obs', 'action', 'reward', 'next_obs'] + - data (:obj:`List[Dict[int, Any]]`): The input data used for policy forward, including a batch of \ + training samples. For each element in list, the key of the dict is the name of data items and the \ + value is the corresponding data. Usually, the value is torch.Tensor or np.ndarray or there dict/list \ + combinations. In the ``_forward_learn`` method, data often need to first be stacked in the batch \ + dimension by some utility functions such as ``default_preprocess_learn``. \ + For FQF, each element in list is a dict containing at least the following keys: \ + ['obs', 'action', 'reward', 'next_obs']. Returns: - - info_dict (:obj:`Dict[str, Any]`): Including current lr and loss. + - info_dict (:obj:`Dict[str, Any]`): The information dict that indicated training result, which will be \ + recorded in text log and tensorboard, values must be python scalar or a list of scalars. For the \ + detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement your own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. """ + # Data preprocessing operations, such as stack data, cpu to cuda device data = default_preprocess_learn( data, use_priority=self._priority, ignore_done=self._cfg.learn.ignore_done, use_nstep=True ) @@ -183,25 +241,18 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: data['weight'] ) value_gamma = data.get('value_gamma') - entropy_loss = -self._ent_coef * entropies.mean() - fraction_loss = fqf_calculate_fraction_loss(q_tau_i.detach(), q_value, quantiles, data['action']) + entropy_loss - quantile_loss, td_error_per_sample = fqf_nstep_td_error( data_n, self._gamma, nstep=self._nstep, kappa=self._kappa, value_gamma=value_gamma ) - # compute grad norm of a network's parameters - def compute_grad_norm(model): - return torch.norm(torch.stack([torch.norm(p.grad.detach(), 2.0) for p in model.parameters()]), 2.0) - # ==================== # fraction_proposal network update # ==================== self._fraction_loss_optimizer.zero_grad() fraction_loss.backward(retain_graph=True) - if self._cfg.learn.multi_gpu: + if self._cfg.multi_gpu: self.sync_gradients(self._learn_model) with torch.no_grad(): total_norm_quantiles_proposal = compute_grad_norm(self._model.head.quantiles_proposal) @@ -212,7 +263,7 @@ def compute_grad_norm(model): # ==================== self._quantile_loss_optimizer.zero_grad() quantile_loss.backward() - if self._cfg.learn.multi_gpu: + if self._cfg.multi_gpu: self.sync_gradients(self._learn_model) with torch.no_grad(): total_norm_Q = compute_grad_norm(self._model.head.Q) @@ -241,12 +292,25 @@ def compute_grad_norm(model): } def _monitor_vars_learn(self) -> List[str]: + """ + Overview: + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. + Returns: + - necessary_keys (:obj:`List[str]`): The list of the necessary keys to be logged. + """ return [ 'cur_lr_fraction_loss', 'cur_lr_quantile_loss', 'logit', 'fraction_loss', 'quantile_loss', 'total_norm_quantiles_proposal', 'total_norm_Q', 'total_norm_fqf_fc', 'total_norm_encoder' ] def _state_dict_learn(self) -> Dict[str, Any]: + """ + Overview: + Return the state_dict of learn mode, usually including model and optimizer. + Returns: + - state_dict (:obj:`Dict[str, Any]`): The dict of current policy learn state, for saving and restoring. + """ return { 'model': self._learn_model.state_dict(), 'target_model': self._target_model.state_dict(), @@ -255,6 +319,17 @@ def _state_dict_learn(self) -> Dict[str, Any]: } def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: + """ + Overview: + Load the state_dict variable into policy learn mode. + Arguments: + - state_dict (:obj:`Dict[str, Any]`): The dict of policy learn state saved before. + + .. tip:: + If you want to only load some parts of model, you can simply set the ``strict`` argument in \ + load_state_dict to ``False``, or refer to ``ding.torch_utils.checkpoint_helper`` for more \ + complicated operation. + """ self._learn_model.load_state_dict(state_dict['model']) self._target_model.load_state_dict(state_dict['target_model']) self._fraction_loss_optimizer.load_state_dict(state_dict['optimizer_fraction_loss']) diff --git a/ding/policy/happo.py b/ding/policy/happo.py new file mode 100644 index 0000000000..4cbd38324b --- /dev/null +++ b/ding/policy/happo.py @@ -0,0 +1,734 @@ +from typing import List, Dict, Any, Tuple, Union +from collections import namedtuple +import torch +import copy +import numpy as np +from torch.distributions import Independent, Normal + +from ding.torch_utils import Adam, to_device, to_dtype, unsqueeze, ContrastiveLoss +from ding.rl_utils import happo_data, happo_error, happo_policy_error, happo_policy_data, \ + v_nstep_td_data, v_nstep_td_error, get_train_sample, gae, gae_data, happo_error_continuous, \ + get_gae +from ding.model import model_wrap +from ding.utils import POLICY_REGISTRY, split_data_generator, RunningMeanStd +from ding.utils.data import default_collate, default_decollate +from .base_policy import Policy +from .common_utils import default_preprocess_learn + + +@POLICY_REGISTRY.register('happo') +class HAPPOPolicy(Policy): + """ + Overview: + Policy class of on policy version HAPPO algorithm. Paper link: https://arxiv.org/abs/2109.11251. + """ + config = dict( + # (str) RL policy register name (refer to function "POLICY_REGISTRY"). + type='happo', + # (bool) Whether to use cuda for network. + cuda=False, + # (bool) Whether the RL algorithm is on-policy or off-policy. (Note: in practice PPO can be off-policy used) + on_policy=True, + # (bool) Whether to use priority(priority sample, IS weight, update priority) + priority=False, + # (bool) Whether to use Importance Sampling Weight to correct biased update due to priority. + # If True, priority must be True. + priority_IS_weight=False, + # (bool) Whether to recompurete advantages in each iteration of on-policy PPO + recompute_adv=True, + # (str) Which kind of action space used in PPOPolicy, ['discrete', 'continuous', 'hybrid'] + action_space='discrete', + # (bool) Whether to use nstep return to calculate value target, otherwise, use return = adv + value + nstep_return=False, + # (bool) Whether to enable multi-agent training, i.e.: MAPPO + multi_agent=False, + # (bool) Whether to need policy data in process transition + transition_with_policy_data=True, + learn=dict( + epoch_per_collect=10, + batch_size=64, + learning_rate=3e-4, + # ============================================================== + # The following configs is algorithm-specific + # ============================================================== + # (float) The loss weight of value network, policy network weight is set to 1 + value_weight=0.5, + # (float) The loss weight of entropy regularization, policy network weight is set to 1 + entropy_weight=0.0, + # (float) PPO clip ratio, defaults to 0.2 + clip_ratio=0.2, + # (bool) Whether to use advantage norm in a whole training batch + adv_norm=True, + value_norm=True, + ppo_param_init=True, + grad_clip_type='clip_norm', + grad_clip_value=0.5, + ignore_done=False, + ), + collect=dict( + # (int) Only one of [n_sample, n_episode] shoule be set + # n_sample=64, + # (int) Cut trajectories into pieces with length "unroll_len". + unroll_len=1, + # ============================================================== + # The following configs is algorithm-specific + # ============================================================== + # (float) Reward's future discount factor, aka. gamma. + discount_factor=0.99, + # (float) GAE lambda factor for the balance of bias and variance(1-step td and mc) + gae_lambda=0.95, + ), + eval=dict(), + ) + + def _init_learn(self) -> None: + """ + Overview: + Initialize the learn mode of policy, including related attributes and modules. For HAPPO, it mainly \ + contains optimizer, algorithm-specific arguments such as loss weight, clip_ratio and recompute_adv. This \ + method also executes some special network initializations and prepares running mean/std monitor for value. + This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. + """ + self._priority = self._cfg.priority + self._priority_IS_weight = self._cfg.priority_IS_weight + assert not self._priority and not self._priority_IS_weight, "Priority is not implemented in PPO" + + assert self._cfg.action_space in ["continuous", "discrete"] + self._action_space = self._cfg.action_space + if self._cfg.learn.ppo_param_init: + for n, m in self._model.named_modules(): + if isinstance(m, torch.nn.Linear): + torch.nn.init.orthogonal_(m.weight) + torch.nn.init.zeros_(m.bias) + if self._action_space in ['continuous']: + # init log sigma + for agent_id in range(self._cfg.agent_num): + # if hasattr(self._model.agent_models[agent_id].actor_head, 'log_sigma_param'): + # torch.nn.init.constant_(self._model.agent_models[agent_id].actor_head.log_sigma_param, 1) + # The above initialization step has been changed to reparameterizationHead. + for m in list(self._model.agent_models[agent_id].critic.modules()) + \ + list(self._model.agent_models[agent_id].actor.modules()): + if isinstance(m, torch.nn.Linear): + # orthogonal initialization + torch.nn.init.orthogonal_(m.weight, gain=np.sqrt(2)) + torch.nn.init.zeros_(m.bias) + # do last policy layer scaling, this will make initial actions have (close to) + # 0 mean and std, and will help boost performances, + # see https://arxiv.org/abs/2006.05990, Fig.24 for details + for m in self._model.agent_models[agent_id].actor.modules(): + if isinstance(m, torch.nn.Linear): + torch.nn.init.zeros_(m.bias) + m.weight.data.copy_(0.01 * m.weight.data) + + # Add the actor/critic parameters of each HAVACAgent in HAVAC to the parameter list of actor/critic_optimizer + actor_params = [] + critic_params = [] + for agent_idx in range(self._model.agent_num): + actor_params.append({'params': self._model.agent_models[agent_idx].actor.parameters()}) + critic_params.append({'params': self._model.agent_models[agent_idx].critic.parameters()}) + + self._actor_optimizer = Adam( + actor_params, + lr=self._cfg.learn.learning_rate, + grad_clip_type=self._cfg.learn.grad_clip_type, + clip_value=self._cfg.learn.grad_clip_value, + # eps = 1e-5, + ) + + self._critic_optimizer = Adam( + critic_params, + lr=self._cfg.learn.critic_learning_rate, + grad_clip_type=self._cfg.learn.grad_clip_type, + clip_value=self._cfg.learn.grad_clip_value, + # eps = 1e-5, + ) + + self._learn_model = model_wrap(self._model, wrapper_name='base') + # self._learn_model = model_wrap( + # self._model, + # wrapper_name='hidden_state', + # state_num=self._cfg.learn.batch_size, + # init_fn=lambda: [None for _ in range(self._cfg.model.agent_num)] + # ) + + # Algorithm config + self._value_weight = self._cfg.learn.value_weight + self._entropy_weight = self._cfg.learn.entropy_weight + self._clip_ratio = self._cfg.learn.clip_ratio + self._adv_norm = self._cfg.learn.adv_norm + self._value_norm = self._cfg.learn.value_norm + if self._value_norm: + self._running_mean_std = RunningMeanStd(epsilon=1e-4, device=self._device) + self._gamma = self._cfg.collect.discount_factor + self._gae_lambda = self._cfg.collect.gae_lambda + self._recompute_adv = self._cfg.recompute_adv + # Main model + self._learn_model.reset() + + def prepocess_data_agent(self, data: Dict[str, Any]): + """ + Overview: + Preprocess data for agent dim. This function is used in learn mode. \ + It will be called recursively to process nested dict data. \ + It will transpose the data with shape (B, agent_num, ...) to (agent_num, B, ...). \ + Arguments: + - data (:obj:`dict`): Dict type data, where each element is the data of an agent of dict type. + Returns: + - ret (:obj:`dict`): Dict type data, where each element is the data of an agent of dict type. + """ + ret = {} + for key, value in data.items(): + if isinstance(value, dict): + ret[key] = self.prepocess_data_agent(value) + elif isinstance(value, torch.Tensor) and len(value.shape) > 1: + ret[key] = value.transpose(0, 1) + else: + ret[key] = value + return ret + + def _forward_learn(self, data: Dict[str, Any]) -> Dict[str, Any]: + """ + Overview: + Forward and backward function of learn mode. + Arguments: + - data (:obj:`dict`): List type data, where each element is the data of an agent of dict type. + Returns: + - info_dict (:obj:`Dict[str, Any]`): + Including current lr, total_loss, policy_loss, value_loss, entropy_loss, \ + adv_abs_max, approx_kl, clipfrac + Overview: + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data from the replay buffer and then returns the output \ + result, including various training information such as loss, clipfrac, approx_kl. + Arguments: + - data (:obj:`List[Dict[int, Any]]`): The input data used for policy forward, including the latest \ + collected training samples for on-policy algorithms like HAPPO. For each element in list, the key of \ + dict is the name of data items and the value is the corresponding data. Usually, the value is \ + torch.Tensor or np.ndarray or there dict/list combinations. In the ``_forward_learn`` method, data \ + often need to first be stacked in the batch dimension by some utility functions such as \ + ``default_preprocess_learn``. \ + For HAPPO, each element in list is a dict containing at least the following keys: ``obs``, \ + ``action``, ``reward``, ``logit``, ``value``, ``done``. Sometimes, it also contains other keys \ + such as ``weight``. + Returns: + - return_infos (:obj:`List[Dict[str, Any]]`): The information list that indicated training result, each \ + training iteration contains append a information dict into the final list. The list will be precessed \ + and recorded in text log and tensorboard. The value of the dict must be python scalar or a list of \ + scalars. For the detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. + + .. tip:: + The training procedure of HAPPO is three for loops. The outermost loop trains each agent separately. \ + The middle loop trains all the collected training samples with ``epoch_per_collect`` epochs. The inner \ + loop splits all the data into different mini-batch with the length of ``batch_size``. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for HAPPOPolicy: ``ding.policy.tests.test_happo``. + """ + data = default_preprocess_learn(data, ignore_done=self._cfg.learn.ignore_done, use_nstep=False) + all_data_len = data['obs']['agent_state'].shape[0] + # fator is the ratio of the old and new strategies of the first m-1 agents, initialized to 1. + # Each transition has its own factor. ref: http://arxiv.org/abs/2109.11251 + factor = torch.ones(all_data_len, 1) # (B, 1) + if self._cuda: + data = to_device(data, self._device) + factor = to_device(factor, self._device) + # process agent dim + data = self.prepocess_data_agent(data) + # ==================== + # PPO forward + # ==================== + return_infos = [] + self._learn_model.train() + + for agent_id in range(self._cfg.agent_num): + agent_data = {} + for key, value in data.items(): + if value is not None: + if type(value) is dict: + agent_data[key] = {k: v[agent_id] for k, v in value.items()} # not feasible for rnn + elif len(value.shape) > 1: + agent_data[key] = data[key][agent_id] + else: + agent_data[key] = data[key] + else: + agent_data[key] = data[key] + + # update factor + agent_data['factor'] = factor + # calculate old_logits of all data in buffer for later factor + inputs = { + 'obs': agent_data['obs'], + # 'actor_prev_state': agent_data['actor_prev_state'], + # 'critic_prev_state': agent_data['critic_prev_state'], + } + old_logits = self._learn_model.forward(agent_id, inputs, mode='compute_actor')['logit'] + + for epoch in range(self._cfg.learn.epoch_per_collect): + if self._recompute_adv: # calculate new value using the new updated value network + with torch.no_grad(): + inputs['obs'] = agent_data['obs'] + # value = self._learn_model.forward(agent_id, agent_data['obs'], mode='compute_critic')['value'] + value = self._learn_model.forward(agent_id, inputs, mode='compute_critic')['value'] + inputs['obs'] = agent_data['next_obs'] + next_value = self._learn_model.forward(agent_id, inputs, mode='compute_critic')['value'] + if self._value_norm: + value *= self._running_mean_std.std + next_value *= self._running_mean_std.std + + traj_flag = agent_data.get('traj_flag', None) # traj_flag indicates termination of trajectory + compute_adv_data = gae_data( + value, next_value, agent_data['reward'], agent_data['done'], traj_flag + ) + agent_data['adv'] = gae(compute_adv_data, self._gamma, self._gae_lambda) + + unnormalized_returns = value + agent_data['adv'] + + if self._value_norm: + agent_data['value'] = value / self._running_mean_std.std + agent_data['return'] = unnormalized_returns / self._running_mean_std.std + self._running_mean_std.update(unnormalized_returns.cpu().numpy()) + else: + agent_data['value'] = value + agent_data['return'] = unnormalized_returns + + else: # don't recompute adv + if self._value_norm: + unnormalized_return = agent_data['adv'] + agent_data['value'] * self._running_mean_std.std + agent_data['return'] = unnormalized_return / self._running_mean_std.std + self._running_mean_std.update(unnormalized_return.cpu().numpy()) + else: + agent_data['return'] = agent_data['adv'] + agent_data['value'] + + for batch in split_data_generator(agent_data, self._cfg.learn.batch_size, shuffle=True): + inputs = { + 'obs': batch['obs'], + # 'actor_prev_state': batch['actor_prev_state'], + # 'critic_prev_state': batch['critic_prev_state'], + } + output = self._learn_model.forward(agent_id, inputs, mode='compute_actor_critic') + adv = batch['adv'] + if self._adv_norm: + # Normalize advantage in a train_batch + adv = (adv - adv.mean()) / (adv.std() + 1e-8) + + # Calculate happo error + if self._action_space == 'continuous': + happo_batch = happo_data( + output['logit'], batch['logit'], batch['action'], output['value'], batch['value'], adv, + batch['return'], batch['weight'], batch['factor'] + ) + happo_loss, happo_info = happo_error_continuous(happo_batch, self._clip_ratio) + elif self._action_space == 'discrete': + happo_batch = happo_data( + output['logit'], batch['logit'], batch['action'], output['value'], batch['value'], adv, + batch['return'], batch['weight'], batch['factor'] + ) + happo_loss, happo_info = happo_error(happo_batch, self._clip_ratio) + wv, we = self._value_weight, self._entropy_weight + total_loss = happo_loss.policy_loss + wv * happo_loss.value_loss - we * happo_loss.entropy_loss + + # actor update + # critic update + self._actor_optimizer.zero_grad() + self._critic_optimizer.zero_grad() + total_loss.backward() + self._actor_optimizer.step() + self._critic_optimizer.step() + + return_info = { + 'agent{}_cur_lr'.format(agent_id): self._actor_optimizer.defaults['lr'], + 'agent{}_total_loss'.format(agent_id): total_loss.item(), + 'agent{}_policy_loss'.format(agent_id): happo_loss.policy_loss.item(), + 'agent{}_value_loss'.format(agent_id): happo_loss.value_loss.item(), + 'agent{}_entropy_loss'.format(agent_id): happo_loss.entropy_loss.item(), + 'agent{}_adv_max'.format(agent_id): adv.max().item(), + 'agent{}_adv_mean'.format(agent_id): adv.mean().item(), + 'agent{}_value_mean'.format(agent_id): output['value'].mean().item(), + 'agent{}_value_max'.format(agent_id): output['value'].max().item(), + 'agent{}_approx_kl'.format(agent_id): happo_info.approx_kl, + 'agent{}_clipfrac'.format(agent_id): happo_info.clipfrac, + } + if self._action_space == 'continuous': + return_info.update( + { + 'agent{}_act'.format(agent_id): batch['action'].float().mean().item(), + 'agent{}_mu_mean'.format(agent_id): output['logit']['mu'].mean().item(), + 'agent{}_sigma_mean'.format(agent_id): output['logit']['sigma'].mean().item(), + } + ) + return_infos.append(return_info) + # calculate the factor + inputs = { + 'obs': agent_data['obs'], + # 'actor_prev_state': agent_data['actor_prev_state'], + } + new_logits = self._learn_model.forward(agent_id, inputs, mode='compute_actor')['logit'] + if self._cfg.action_space == 'discrete': + dist_new = torch.distributions.categorical.Categorical(logits=new_logits) + dist_old = torch.distributions.categorical.Categorical(logits=old_logits) + elif self._cfg.action_space == 'continuous': + dist_new = Normal(new_logits['mu'], new_logits['sigma']) + dist_old = Normal(old_logits['mu'], old_logits['sigma']) + logp_new = dist_new.log_prob(agent_data['action']) + logp_old = dist_old.log_prob(agent_data['action']) + if len(logp_new.shape) > 1: + # for logp with shape(B, action_shape), we need to calculate the product of all action dimensions. + factor = factor * torch.prod( + torch.exp(logp_new - logp_old), dim=-1 + ).reshape(all_data_len, 1).detach() # attention the shape + else: + # for logp with shape(B, ), directly calculate factor + factor = factor * torch.exp(logp_new - logp_old).reshape(all_data_len, 1).detach() + return return_infos + + def _state_dict_learn(self) -> Dict[str, Any]: + """ + Overview: + Return the state_dict of learn mode optimizer and model. + Returns: + - state_dict (:obj:`Dict[str, Any]`): The dict of current policy learn mode. It contains the \ + state_dict of current policy network and optimizer. + """ + return { + 'model': self._learn_model.state_dict(), + 'actor_optimizer': self._actor_optimizer.state_dict(), + 'critic_optimizer': self._critic_optimizer.state_dict(), + } + + def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: + """ + Overview: + Load the state_dict of learn mode optimizer and model. + Arguments: + - state_dict (:obj:`Dict[str, Any]`): The dict of policy learn mode. It contains the state_dict \ + of current policy network and optimizer. + """ + self._learn_model.load_state_dict(state_dict['model']) + self._actor_optimizer.load_state_dict(state_dict['actor_optimizer']) + self._critic_optimizer.load_state_dict(state_dict['critic_optimizer']) + + def _init_collect(self) -> None: + """ + Overview: + Initialize the collect mode of policy, including related attributes and modules. For HAPPO, it contains \ + the collect_model to balance the exploration and exploitation (e.g. the multinomial sample mechanism in \ + discrete action space), and other algorithm-specific arguments such as unroll_len and gae_lambda. + This method will be called in ``__init__`` method if ``collect`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_collect`` method, you'd better name them \ + with prefix ``_collect_`` to avoid conflict with other modes, such as ``self._collect_attr1``. + + .. tip:: + Some variables need to initialize independently in different modes, such as gamma and gae_lambda in PPO. \ + This design is for the convenience of parallel execution of different policy modes. + """ + self._unroll_len = self._cfg.collect.unroll_len + assert self._cfg.action_space in ["continuous", "discrete"] + self._action_space = self._cfg.action_space + if self._action_space == 'continuous': + self._collect_model = model_wrap(self._model, wrapper_name='reparam_sample') + elif self._action_space == 'discrete': + self._collect_model = model_wrap(self._model, wrapper_name='multinomial_sample') + self._collect_model.reset() + self._gamma = self._cfg.collect.discount_factor + self._gae_lambda = self._cfg.collect.gae_lambda + self._recompute_adv = self._cfg.recompute_adv + + def _forward_collect(self, data: Dict[int, Any]) -> dict: + """ + Overview: + Policy forward function of collect mode (collecting training data by interacting with envs). Forward means \ + that the policy gets some necessary data (mainly observation) from the envs and then returns the output \ + data, such as the action to interact with the envs. + Arguments: + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. + Returns: + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action and \ + other necessary data (action logit and value) for learn mode defined in ``self._process_transition`` \ + method. The key of the dict is the same as the input data, i.e. environment id. + + .. tip:: + If you want to add more tricks on this policy, like temperature factor in multinomial sample, you can pass \ + related data as extra keyword arguments of this method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for HAPPOPolicy: ``ding.policy.tests.test_happo``. + """ + data_id = list(data.keys()) + data = default_collate(list(data.values())) + if self._cuda: + data = to_device(data, self._device) + data = {k: v.transpose(0, 1) for k, v in data.items()} # not feasible for rnn + self._collect_model.eval() + with torch.no_grad(): + outputs = [] + for agent_id in range(self._cfg.agent_num): + # output = self._collect_model.forward(agent_id, data, mode='compute_actor_critic') + single_agent_obs = {k: v[agent_id] for k, v in data.items()} + input = { + 'obs': single_agent_obs, + } + output = self._collect_model.forward(agent_id, input, mode='compute_actor_critic') + outputs.append(output) + # transfer data from (M, B, N)->(B, M, N) + result = {} + for key in outputs[0].keys(): + if isinstance(outputs[0][key], dict): + subkeys = outputs[0][key].keys() + stacked_subvalues = {} + for subkey in subkeys: + stacked_subvalues[subkey] = \ + torch.stack([output[key][subkey] for output in outputs], dim=0).transpose(0, 1) + result[key] = stacked_subvalues + else: + # If Value is tensor, stack it directly + if isinstance(outputs[0][key], torch.Tensor): + result[key] = torch.stack([output[key] for output in outputs], dim=0).transpose(0, 1) + else: + # If it is not tensor, assume that it is a non-stackable data type \ + # (such as int, float, etc.), and directly retain the original value + result[key] = [output[key] for output in outputs] + output = result + if self._cuda: + output = to_device(output, 'cpu') + output = default_decollate(output) + return {i: d for i, d in zip(data_id, output)} + + def _process_transition(self, obs: Any, model_output: dict, timestep: namedtuple) -> dict: + """ + Overview: + Process and pack one timestep transition data into a dict, which can be directly used for training and \ + saved in replay buffer. For HAPPO, it contains obs, next_obs, action, reward, done, logit, value. + Arguments: + - obs (:obj:`torch.Tensor`): The env observation of current timestep. + - policy_output (:obj:`Dict[str, torch.Tensor]`): The output of the policy network with the observation \ + as input. For PPO, it contains the state value, action and the logit of the action. + - timestep (:obj:`namedtuple`): The execution result namedtuple returned by the environment step method, \ + except all the elements have been transformed into tensor data. Usually, it contains the next obs, \ + reward, done, info, etc. + Returns: + - transition (:obj:`Dict[str, torch.Tensor]`): The processed transition data of the current timestep. + + .. note:: + ``next_obs`` is used to calculate nstep return when necessary, so we place in into transition by default. \ + You can delete this field to save memory occupancy if you do not need nstep return. + """ + transition = { + 'obs': obs, + 'next_obs': timestep.obs, + 'action': model_output['action'], + 'logit': model_output['logit'], + 'value': model_output['value'], + 'reward': timestep.reward, + 'done': timestep.done, + } + return transition + + def _get_train_sample(self, data: list) -> Union[None, List[Any]]: + """ + Overview: + For a given trajectory (transitions, a list of transition) data, process it into a list of sample that \ + can be used for training directly. In HAPPO, a train sample is a processed transition with new computed \ + ``traj_flag`` and ``adv`` field. This method is usually used in collectors to execute necessary \ + RL data preprocessing before training, which can help learner amortize revelant time consumption. \ + In addition, you can also implement this method as an identity function and do the data processing \ + in ``self._forward_learn`` method. + Arguments: + - transitions (:obj:`List[Dict[str, Any]`): The trajectory data (a list of transition), each element is \ + the same format as the return value of ``self._process_transition`` method. + Returns: + - samples (:obj:`List[Dict[str, Any]]`): The processed train samples, each element is the similar format \ + as input transitions, but may contain more data for training, such as GAE advantage. + """ + data = to_device(data, self._device) + for transition in data: + transition['traj_flag'] = copy.deepcopy(transition['done']) + data[-1]['traj_flag'] = True + + if self._cfg.learn.ignore_done: + data[-1]['done'] = False + + if data[-1]['done']: + last_value = torch.zeros_like(data[-1]['value']) + else: + with torch.no_grad(): + last_values = [] + for agent_id in range(self._cfg.agent_num): + inputs = {'obs': {k: unsqueeze(v[agent_id], 0) for k, v in data[-1]['next_obs'].items()}} + last_value = self._collect_model.forward(agent_id, inputs, mode='compute_actor_critic')['value'] + last_values.append(last_value) + last_value = torch.cat(last_values) + if len(last_value.shape) == 2: # multi_agent case: + last_value = last_value.squeeze(0) + if self._value_norm: + last_value *= self._running_mean_std.std + for i in range(len(data)): + data[i]['value'] *= self._running_mean_std.std + data = get_gae( + data, + to_device(last_value, self._device), + gamma=self._gamma, + gae_lambda=self._gae_lambda, + cuda=False, + ) + if self._value_norm: + for i in range(len(data)): + data[i]['value'] /= self._running_mean_std.std + + # remove next_obs for save memory when not recompute adv + if not self._recompute_adv: + for i in range(len(data)): + data[i].pop('next_obs') + return get_train_sample(data, self._unroll_len) + + def _init_eval(self) -> None: + """ + Overview: + Initialize the eval mode of policy, including related attributes and modules. For PPO, it contains the \ + eval model to select optimial action (e.g. greedily select action with argmax mechanism in discrete action). + This method will be called in ``__init__`` method if ``eval`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_eval`` method, you'd better name them \ + with prefix ``_eval_`` to avoid conflict with other modes, such as ``self._eval_attr1``. + """ + assert self._cfg.action_space in ["continuous", "discrete"] + self._action_space = self._cfg.action_space + if self._action_space == 'continuous': + self._eval_model = model_wrap(self._model, wrapper_name='deterministic_sample') + elif self._action_space == 'discrete': + self._eval_model = model_wrap(self._model, wrapper_name='argmax_sample') + self._eval_model.reset() + + def _forward_eval(self, data: dict) -> dict: + """ + Overview: + Policy forward function of eval mode (evaluation policy performance by interacting with envs). Forward \ + means that the policy gets some necessary data (mainly observation) from the envs and then returns the \ + action to interact with the envs. ``_forward_eval`` in HAPPO often uses deterministic sample method to \ + get actions while ``_forward_collect`` usually uses stochastic sample method for balance exploration and \ + exploitation. + Arguments: + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. + Returns: + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action. The \ + key of the dict is the same as the input data, i.e. environment id. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for HAPPOPolicy: ``ding.policy.tests.test_happo``. + """ + data_id = list(data.keys()) + data = default_collate(list(data.values())) + if self._cuda: + data = to_device(data, self._device) + # transfer data from (B, M, N)->(M, B, N) + data = {k: v.transpose(0, 1) for k, v in data.items()} # not feasible for rnn + self._eval_model.eval() + with torch.no_grad(): + outputs = [] + for agent_id in range(self._cfg.agent_num): + single_agent_obs = {k: v[agent_id] for k, v in data.items()} + input = { + 'obs': single_agent_obs, + } + output = self._eval_model.forward(agent_id, input, mode='compute_actor') + outputs.append(output) + output = self.revert_agent_data(outputs) + if self._cuda: + output = to_device(output, 'cpu') + output = default_decollate(output) + return {i: d for i, d in zip(data_id, output)} + + def default_model(self) -> Tuple[str, List[str]]: + """ + Overview: + Return this algorithm default neural network model setting for demonstration. ``__init__`` method will \ + automatically call this method to get the default model setting and create model. + Returns: + - model_info (:obj:`Tuple[str, List[str]]`): The registered model name and model's import_names. + + .. note:: + The user can define and use customized network model but must obey the same inferface definition indicated \ + by import_names path. For example about HAPPO, its registered name is ``happo`` and the import_names is \ + ``ding.model.template.havac``. + """ + return 'havac', ['ding.model.template.havac'] + + def _monitor_vars_learn(self) -> List[str]: + """ + Overview: + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. + Returns: + - necessary_keys (:obj:`List[str]`): The list of the necessary keys to be logged. + """ + variables = super()._monitor_vars_learn() + [ + 'policy_loss', + 'value_loss', + 'entropy_loss', + 'adv_max', + 'adv_mean', + 'approx_kl', + 'clipfrac', + 'value_max', + 'value_mean', + ] + if self._action_space == 'continuous': + variables += ['mu_mean', 'sigma_mean', 'sigma_grad', 'act'] + prefixes = [f'agent{i}_' for i in range(self._cfg.agent_num)] + variables = [prefix + var for prefix in prefixes for var in variables] + return variables + + def revert_agent_data(self, data: list): + """ + Overview: + Revert the data of each agent to the original data format. + Arguments: + - data (:obj:`list`): List type data, where each element is the data of an agent of dict type. + Returns: + - ret (:obj:`dict`): Dict type data, where each element is the data of an agent of dict type. + """ + ret = {} + # Traverse all keys of the first output + for key in data[0].keys(): + if isinstance(data[0][key], torch.Tensor): + # If the value corresponding to the current key is tensor, stack N tensors + stacked_tensor = torch.stack([output[key] for output in data], dim=0) + ret[key] = stacked_tensor.transpose(0, 1) + elif isinstance(data[0][key], dict): + # If the value corresponding to the current key is a dictionary, recursively \ + # call the function to process the contents inside the dictionary. + ret[key] = self.revert_agent_data([output[key] for output in data]) + return ret diff --git a/ding/policy/ibc.py b/ding/policy/ibc.py index 9c9536d987..887edf298d 100644 --- a/ding/policy/ibc.py +++ b/ding/policy/ibc.py @@ -20,40 +20,89 @@ class IBCPolicy(BehaviourCloningPolicy): r""" Overview: - Implicit Behavior Cloning - https://arxiv.org/abs/2109.00137.pdf + Policy class of IBC (Implicit Behavior Cloning), proposed in https://arxiv.org/abs/2109.00137.pdf. .. note:: - The code is adapted from the pytorch version of IBC https://github.com/kevinzakka/ibc, - which only supports the derivative-free optimization (dfo) variants. - This implementation moves a step forward and supports all variants of energy-based model - mentioned in the paper (dfo, autoregressive dfo, and mcmc). + The code is adapted from the pytorch version of IBC https://github.com/kevinzakka/ibc, which only supports the \ + derivative-free optimization (dfo) variants. This implementation moves a step forward and supports all \ + variants of energy-based model mentioned in the paper (dfo, autoregressive dfo, and mcmc). """ config = dict( + # (str) The policy type. 'ibc' refers to Implicit Behavior Cloning. type='ibc', + # (bool) Whether to use CUDA for training. False means CPU will be used. cuda=False, + # (bool) If True, the policy will operate on-policy. Here it's False, indicating off-policy. on_policy=False, + # (bool) Whether the action space is continuous. True for continuous action space. continuous=True, - model=dict(stochastic_optim=dict(type='mcmc', )), + # (dict) Configuration for the model, including stochastic optimization settings. + model=dict( + # (dict) Configuration for the stochastic optimization, specifying the type of optimizer. + stochastic_optim=dict( + # (str) The type of stochastic optimizer. 'mcmc' refers to Markov Chain Monte Carlo methods. + type='mcmc', + ), + ), + # (dict) Configuration for the learning process. learn=dict( + # (int) The number of training epochs. train_epoch=30, + # (int) The size of batches used during training. batch_size=256, - multi_gpu=False, + # (dict) Configuration for the optimizer used during training. optim=dict( + # (float) The learning rate for the optimizer. learning_rate=1e-5, + # (float) The weight decay regularization term for the optimizer. weight_decay=0.0, + # (float) The beta1 hyperparameter for the AdamW optimizer. beta1=0.9, + # (float) The beta2 hyperparameter for the AdamW optimizer. beta2=0.999, ), ), - eval=dict(evaluator=dict(eval_freq=10000, )), + # (dict) Configuration for the evaluation process. + eval=dict( + # (dict) Configuration for the evaluator. + evaluator=dict( + # (int) The frequency of evaluations during training, in terms of number of training steps. + eval_freq=10000, + ), + ), ) def default_model(self) -> Tuple[str, List[str]]: + """ + Overview: + Returns the default model configuration used by the IBC algorithm. ``__init__`` method will \ + automatically call this method to get the default model setting and create model. + + Returns: + - model_info (:obj:`Tuple[str, List[str]]`): \ + Tuple containing the registered model name and model's import_names. + """ return 'ebm', ['ding.model.template.ebm'] - def _init_learn(self): + def _init_learn(self) -> None: + """ + Overview: + Initialize the learn mode of policy, including related attributes and modules. For IBC, it mainly \ + contains optimizer and main model. \ + This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. + """ self._timer = EasyTimer(cuda=self._cfg.cuda) self._sync_timer = EasyTimer(cuda=self._cfg.cuda) optim_cfg = self._cfg.learn.optim @@ -68,7 +117,31 @@ def _init_learn(self): self._learn_model = model_wrap(self._model, 'base') self._learn_model.reset() - def _forward_learn(self, data): + def _forward_learn(self, data: List[Dict[str, Any]]) -> Dict[str, Any]: + """ + Overview: + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data from the replay buffer and then returns the output \ + result, including various training information such as policy_loss, value_loss, entropy_loss. + Arguments: + - data (:obj:`List[Dict[int, Any]]`): The input data used for policy forward, including a batch of \ + training samples. For each element in list, the key of the dict is the name of data items and the \ + value is the corresponding data. Usually, the value is torch.Tensor or np.ndarray or there dict/list \ + combinations. In the ``_forward_learn`` method, data often need to first be stacked in the batch \ + dimension by some utility functions such as ``default_preprocess_learn``. \ + For IBC, each element in list is a dict containing at least the following keys: \ + ['obs', 'action']. + Returns: + - info_dict (:obj:`Dict[str, Any]`): The information dict that indicated training result, which will be \ + recorded in text log and tensorboard, values must be python scalar or a list of scalars. For the \ + detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement your own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + """ with self._timer: data = default_collate(data) if self._cuda: @@ -82,7 +155,7 @@ def _forward_learn(self, data): obs, action = data['obs'], data['action'] # When action/observation space is 1, the action/observation dimension will # be squeezed in the first place, therefore unsqueeze there to make the data - # compatiable with the ibc pipeline. + # compatible with the ibc pipeline. if len(obs.shape) == 1: obs = obs.unsqueeze(-1) if len(action.shape) == 1: @@ -124,7 +197,7 @@ def _forward_learn(self, data): self._optimizer.zero_grad() loss.backward() with self._sync_timer: - if self._cfg.learn.multi_gpu: + if self._cfg.multi_gpu: self.sync_gradients(self._learn_model) sync_time = self._sync_timer.value self._optimizer.step() @@ -137,17 +210,51 @@ def _forward_learn(self, data): **loss_dict, } - def _monitor_vars_learn(self): + def _monitor_vars_learn(self) -> List[str]: + """ + Overview: + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. + Returns: + - necessary_keys (:obj:`List[str]`): The list of the necessary keys to be logged. + """ if isinstance(self._stochastic_optimizer, MCMC): return ['total_loss', 'ebm_loss', 'grad_penalty', 'total_time', 'sync_time'] else: return ['total_loss', 'ebm_loss', 'total_time', 'sync_time'] - def _init_eval(self): + def _init_eval(self) -> None: + """ + Overview: + Initialize the eval mode of policy, including related attributes and modules. + This method will be called in ``__init__`` method if ``eval`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_eval`` method, you'd better name them \ + with prefix ``_eval_`` to avoid conflict with other modes, such as ``self._eval_attr1``. + """ self._eval_model = model_wrap(self._model, wrapper_name='base') self._eval_model.reset() - def _forward_eval(self, data: dict) -> dict: + def _forward_eval(self, data: Dict[int, Any]) -> Dict[int, Any]: + """ + Overview: + Policy forward function of eval mode (evaluation policy performance by interacting with envs). Forward \ + means that the policy gets some necessary data (mainly observation) from the envs and then returns the \ + action to interact with the envs. + Arguments: + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. + Returns: + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action. The \ + key of the dict is the same as the input data, i.e., environment id. + + .. note:: + The input value can be ``torch.Tensor`` or dict/list combinations, current policy supports all of them. \ + For the data type that is not supported, the main reason is that the corresponding model does not \ + support it. You can implement your own model rather than use the default model. For more information, \ + please raise an issue in GitHub repo, and we will continue to follow up. + """ tensor_input = isinstance(data, torch.Tensor) if not tensor_input: data_id = list(data.keys()) @@ -169,6 +276,13 @@ def _forward_eval(self, data: dict) -> dict: return {i: d for i, d in zip(data_id, output)} def set_statistic(self, statistics: EasyDict) -> None: + """ + Overview: + Set the statistics of the environment, including the action space and the observation space. + Arguments: + - statistics (:obj:`EasyDict`): The statistics of the environment. For IBC, it contains at least the \ + following keys: ['action_bounds']. + """ self._stochastic_optimizer.set_action_bounds(statistics.action_bounds) # =================================================================== # diff --git a/ding/policy/il.py b/ding/policy/il.py index 72bafa3687..77989facec 100644 --- a/ding/policy/il.py +++ b/ding/policy/il.py @@ -30,7 +30,7 @@ class ILPolicy(Policy): # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. priority_IS_weight=False, learn=dict( - multi_gpu=False, + # (int) collect n_episode data, train model n_iteration time update_per_collect=20, # (int) the number of data for a train iteration diff --git a/ding/policy/impala.py b/ding/policy/impala.py index 7267b8af98..46adeb1204 100644 --- a/ding/policy/impala.py +++ b/ding/policy/impala.py @@ -2,20 +2,21 @@ from typing import List, Dict, Any, Tuple import torch +import treetensor.torch as ttorch from ding.model import model_wrap -from ding.rl_utils import vtrace_data, vtrace_error, get_train_sample +from ding.rl_utils import vtrace_data, vtrace_error_discrete_action, vtrace_error_continuous_action, get_train_sample from ding.torch_utils import Adam, RMSprop, to_device from ding.utils import POLICY_REGISTRY -from ding.utils.data import default_collate, default_decollate +from ding.utils.data import default_collate, default_decollate, ttorch_collate from ding.policy.base_policy import Policy @POLICY_REGISTRY.register('impala') class IMPALAPolicy(Policy): - r""" + """ Overview: - Policy class of IMPALA algorithm. + Policy class of IMPALA algorithm. Paper link: https://arxiv.org/abs/1802.01561. Config: == ==================== ======== ============== ======================================== ======================= @@ -40,78 +41,117 @@ class IMPALAPolicy(Policy): == ==================== ======== ============== ======================================== ======================= """ config = dict( + # (str) RL policy register name (refer to function "POLICY_REGISTRY"). type='impala', + # (bool) Whether to use cuda in policy. cuda=False, - # (bool) whether use on-policy training pipeline(behaviour policy and training policy are the same) - # here we follow ppo serial pipeline, the original is False + # (bool) Whether learning policy is the same as collecting data policy(on-policy). on_policy=False, + # (bool) Whether to enable priority experience sample. priority=False, # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. priority_IS_weight=False, - # (int) the trajectory length to calculate v-trace target + # (str) Which kind of action space used in IMPALAPolicy, ['discrete', 'continuous']. + action_space='discrete', + # (int) the trajectory length to calculate v-trace target. unroll_len=32, - # (bool) Whether to need policy data in process transition + # (bool) Whether to need policy data in process transition. transition_with_policy_data=True, + # learn_mode config learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, - # (int) collect n_sample data, train model update_per_collect times - # here we follow ppo serial pipeline + # (int) collect n_sample data, train model update_per_collect times. update_per_collect=4, - # (int) the number of data for a train iteration + # (int) the number of data for a train iteration. batch_size=16, + # (float) The step size of gradient descent. learning_rate=0.0005, - # (float) loss weight of the value network, the weight of policy network is set to 1 + # (float) loss weight of the value network, the weight of policy network is set to 1. value_weight=0.5, - # (float) loss weight of the entropy regularization, the weight of policy network is set to 1 + # (float) loss weight of the entropy regularization, the weight of policy network is set to 1. entropy_weight=0.0001, - # (float) discount factor for future reward, defaults int [0, 1] + # (float) discount factor for future reward, defaults int [0, 1]. discount_factor=0.99, - # (float) additional discounting parameter + # (float) additional discounting parameter. lambda_=0.95, - # (float) clip ratio of importance weights + # (float) clip ratio of importance weights. rho_clip_ratio=1.0, - # (float) clip ratio of importance weights + # (float) clip ratio of importance weights. c_clip_ratio=1.0, - # (float) clip ratio of importance sampling + # (float) clip ratio of importance sampling. rho_pg_clip_ratio=1.0, + # (str) The gradient clip operation type used in IMPALA, ['clip_norm', clip_value', 'clip_momentum_norm']. + grad_clip_type=None, + # (float) The gradient clip target value used in IMPALA. + # If ``grad_clip_type`` is 'clip_norm', then the maximum of gradient will be normalized to this value. + clip_value=0.5, + # (str) Optimizer used to train the network, ['adam', 'rmsprop']. + optim='adam', ), + # collect_mode config collect=dict( - # (int) collect n_sample data, train model n_iteration times + # (int) How many training samples collected in one collection procedure. + # Only one of [n_sample, n_episode] shoule be set. # n_sample=16, - collector=dict(collect_print_freq=1000, ), ), - eval=dict(evaluator=dict(eval_freq=1000, ), ), - other=dict(replay_buffer=dict( - replay_buffer_size=1000, - max_use=16, - ), ), + eval=dict(), # for compatibility + other=dict( + replay_buffer=dict( + # (int) Maximum size of replay buffer. Usually, larger buffer size is better. + replay_buffer_size=1000, + # (int) Maximum use times for a sample in buffer. If reaches this value, the sample will be removed. + max_use=16, + ), + ), ) def default_model(self) -> Tuple[str, List[str]]: + """ + Overview: + Return this algorithm default neural network model setting for demonstration. ``__init__`` method will \ + automatically call this method to get the default model setting and create model. + Returns: + - model_info (:obj:`Tuple[str, List[str]]`): The registered model name and model's import_names. + + .. note:: + The user can define and use customized network model but must obey the same inferface definition indicated \ + by import_names path. For example about IMPALA , its registered name is ``vac`` and the import_names is \ + ``ding.model.template.vac``. + """ return 'vac', ['ding.model.template.vac'] def _init_learn(self) -> None: - r""" + """ Overview: - Learn mode init method. Called by ``self.__init__``. - Initialize the optimizer, algorithm config and main model. + Initialize the learn mode of policy, including related attributes and modules. For IMPALA, it mainly \ + contains optimizer, algorithm-specific arguments such as loss weight and gamma, main (learn) model. + This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. """ + assert self._cfg.action_space in ["continuous", "discrete"], self._cfg.action_space + self._action_space = self._cfg.action_space # Optimizer - grad_clip_type = self._cfg.learn.get("grad_clip_type", None) - clip_value = self._cfg.learn.get("clip_value", None) - optim_type = self._cfg.learn.get("optim", "adam") + optim_type = self._cfg.learn.optim if optim_type == 'rmsprop': self._optimizer = RMSprop(self._model.parameters(), lr=self._cfg.learn.learning_rate) elif optim_type == 'adam': self._optimizer = Adam( self._model.parameters(), - grad_clip_type=grad_clip_type, - clip_value=clip_value, + grad_clip_type=self._cfg.learn.grad_clip_type, + clip_value=self._cfg.learn.clip_value, lr=self._cfg.learn.learning_rate ) else: - raise NotImplementedError + raise NotImplementedError("Now only support rmsprop and adam, but input is {}".format(optim_type)) self._learn_model = model_wrap(self._model, wrapper_name='base') self._action_shape = self._cfg.model.action_shape @@ -138,8 +178,8 @@ def _data_preprocess_learn(self, data: List[Dict[str, Any]]): Convert list trajectory data to to trajectory data, which is a dict of tensors. Arguments: - data (:obj:`List[Dict[str, Any]]`): List type data, a list of data for training. Each list element is a \ - dict, whose values are torch.Tensor or np.ndarray or dict/list combinations, keys include at least 'obs',\ - 'next_obs', 'logit', 'action', 'reward', 'done' + dict, whose values are torch.Tensor or np.ndarray or dict/list combinations, keys include at least \ + 'obs', 'next_obs', 'logit', 'action', 'reward', 'done' Returns: - data (:obj:`dict`): Dict type data. Values are torch.Tensor or np.ndarray or dict/list combinations. \ ReturnsKeys: @@ -155,7 +195,13 @@ def _data_preprocess_learn(self, data: List[Dict[str, Any]]): - done (:obj:`torch.FloatTensor`): :math:`(T, B)` - weight (:obj:`torch.FloatTensor`): :math:`(T, B)` """ - data = default_collate(data) + elem = data[0] + if isinstance(elem, dict): # old pipeline + data = default_collate(data) + elif isinstance(elem, list): # new task pipeline + data = default_collate(default_collate(data)) + else: + raise TypeError("not support element type ({}) in IMPALA".format(type(elem))) if self._cuda: data = to_device(data, self._device) if self._priority_IS_weight: @@ -164,47 +210,59 @@ def _data_preprocess_learn(self, data: List[Dict[str, Any]]): data['weight'] = data['IS'] else: data['weight'] = data.get('weight', None) - data['obs_plus_1'] = torch.cat((data['obs'] + data['next_obs'][-1:]), dim=0) # shape (T+1)*B,env_obs_shape - data['logit'] = torch.cat( - data['logit'], dim=0 - ).reshape(self._unroll_len, -1, self._action_shape) # shape T,B,env_action_shape - data['action'] = torch.cat(data['action'], dim=0).reshape(self._unroll_len, -1) # shape T,B, - data['done'] = torch.cat(data['done'], dim=0).reshape(self._unroll_len, -1).float() # shape T,B, - data['reward'] = torch.cat(data['reward'], dim=0).reshape(self._unroll_len, -1) # shape T,B, - data['weight'] = torch.cat( - data['weight'], dim=0 - ).reshape(self._unroll_len, -1) if data['weight'] else None # shape T,B + if isinstance(elem, dict): # old pipeline + for k in data: + if isinstance(data[k], list): + data[k] = default_collate(data[k]) + data['obs_plus_1'] = torch.cat([data['obs'], data['next_obs'][-1:]], dim=0) # shape (T+1)*B,env_obs_shape return data def _forward_learn(self, data: List[Dict[str, Any]]) -> Dict[str, Any]: - r""" + """ Overview: - Forward computation graph of learn mode(updating policy). + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data from the replay buffer and then returns the output \ + result, including various training information such as loss and current learning rate. Arguments: - - data (:obj:`List[Dict[str, Any]]`): List type data, a list of data for training. Each list element is a \ - dict, whose values are torch.Tensor or np.ndarray or dict/list combinations, keys include at least 'obs',\ - 'next_obs', 'logit', 'action', 'reward', 'done' + - data (:obj:`List[Dict[int, Any]]`): The input data used for policy forward, including a batch of \ + training samples. For each element in list, the key of the dict is the name of data items and the \ + value is the corresponding data. Usually, the value is torch.Tensor or np.ndarray or there dict/list \ + combinations. In the ``_forward_learn`` method, data often need to first be stacked in the batch \ + dimension by some utility functions such as ``default_preprocess_learn``. \ + For IMPALA, each element in list is a dict containing at least the following keys: ``obs``, \ + ``action``, ``logit``, ``reward``, ``next_obs``, ``done``. Sometimes, it also contains other keys such \ + as ``weight``. Returns: - - info_dict (:obj:`Dict[str, Any]`): Dict type data, a info dict indicated training result, which will be \ - recorded in text log and tensorboard, values are python scalar or a list of scalars. - ArgumentsKeys: - - necessary: ``obs``, ``action``, ``reward``, ``next_obs``, ``done`` - - optional: 'collect_iter', 'replay_unique_id', 'replay_buffer_idx', 'priority', 'staleness', 'use', 'IS' - ReturnsKeys: - - necessary: ``cur_lr``, ``total_loss``, ``policy_loss`,``value_loss``,``entropy_loss`` - - optional: ``priority`` + - info_dict (:obj:`Dict[str, Any]`): The information dict that indicated training result, which will be \ + recorded in text log and tensorboard, values must be python scalar or a list of scalars. For the \ + detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to unittest for IMPALAPolicy: ``ding.policy.tests.test_impala``. """ data = self._data_preprocess_learn(data) # ==================== # IMPALA forward # ==================== self._learn_model.train() - output = self._learn_model.forward(data['obs_plus_1'], mode='compute_actor_critic') + output = self._learn_model.forward( + data['obs_plus_1'].view((-1, ) + data['obs_plus_1'].shape[2:]), mode='compute_actor_critic' + ) target_logit, behaviour_logit, actions, values, rewards, weights = self._reshape_data(output, data) # Calculate vtrace error data = vtrace_data(target_logit, behaviour_logit, actions, values, rewards, weights) g, l, r, c, rg = self._gamma, self._lambda, self._rho_clip_ratio, self._c_clip_ratio, self._rho_pg_clip_ratio - vtrace_loss = vtrace_error(data, g, l, r, c, rg) + if self._action_space == 'continuous': + vtrace_loss = vtrace_error_continuous_action(data, g, l, r, c, rg) + elif self._action_space == 'discrete': + vtrace_loss = vtrace_error_discrete_action(data, g, l, r, c, rg) + wv, we = self._value_weight, self._entropy_weight total_loss = vtrace_loss.policy_loss + wv * vtrace_loss.value_loss - we * vtrace_loss.entropy_loss # ==================== @@ -221,89 +279,95 @@ def _forward_learn(self, data: List[Dict[str, Any]]) -> Dict[str, Any]: 'entropy_loss': vtrace_loss.entropy_loss.item(), } - def _reshape_data(self, output: Dict[str, Any], data: Dict[str, Any]) -> Tuple[Any, Any, Any, Any, Any, Any]: - r""" + def _reshape_data(self, output: Dict[str, Any], data: Dict[str, Any]) -> Tuple: + """ Overview: - Obtain weights for loss calculating, where should be 0 for done positions - Update values and rewards with the weight + Obtain weights for loss calculating, where should be 0 for done positions. Update values and rewards with \ + the weight. Arguments: - output (:obj:`Dict[int, Any]`): Dict type data, output of learn_model forward. \ - Values are torch.Tensor or np.ndarray or dict/list combinations,keys are value, logit. - - data (:obj:`Dict[int, Any]`): Dict type data, input of policy._forward_learn \ - Values are torch.Tensor or np.ndarray or dict/list combinations. Keys includes at \ - least ['logit', 'action', 'reward', 'done',] + Values are torch.Tensor or np.ndarray or dict/list combinations,keys are value, logit. + - data (:obj:`Dict[int, Any]`): Dict type data, input of policy._forward_learn Values are torch.Tensor or \ + np.ndarray or dict/list combinations. Keys includes at least ['logit', 'action', 'reward', 'done']. Returns: - - data (:obj:`Tuple[Any]`): Tuple of target_logit, behaviour_logit, actions, \ - values, rewards, weights + - data (:obj:`Tuple[Any]`): Tuple of target_logit, behaviour_logit, actions, values, rewards, weights. ReturnsShapes: - target_logit (:obj:`torch.FloatTensor`): :math:`((T+1), B, Obs_Shape)`, where T is timestep,\ - B is batch size and Obs_Shape is the shape of single env observation. + B is batch size and Obs_Shape is the shape of single env observation. - behaviour_logit (:obj:`torch.FloatTensor`): :math:`(T, B, N)`, where N is action dim. - actions (:obj:`torch.LongTensor`): :math:`(T, B)` - values (:obj:`torch.FloatTensor`): :math:`(T+1, B)` - rewards (:obj:`torch.FloatTensor`): :math:`(T, B)` - weights (:obj:`torch.FloatTensor`): :math:`(T, B)` """ - target_logit = output['logit'].reshape(self._unroll_len + 1, -1, - self._action_shape)[:-1] # shape (T+1),B,env_obs_shape + if self._action_space == 'continuous': + target_logit = {} + target_logit['mu'] = output['logit']['mu'].reshape(self._unroll_len + 1, -1, + self._action_shape)[:-1 + ] # shape (T+1),B,env_action_shape + target_logit['sigma'] = output['logit']['sigma'].reshape(self._unroll_len + 1, -1, self._action_shape + )[:-1] # shape (T+1),B,env_action_shape + elif self._action_space == 'discrete': + target_logit = output['logit'].reshape(self._unroll_len + 1, -1, + self._action_shape)[:-1] # shape (T+1),B,env_action_shape behaviour_logit = data['logit'] # shape T,B - actions = data['action'] # shape T,B + actions = data['action'] # shape T,B for discrete # shape T,B,env_action_shape for continuous values = output['value'].reshape(self._unroll_len + 1, -1) # shape T+1,B,env_action_shape rewards = data['reward'] # shape T,B - weights_ = 1 - data['done'] # shape T,B + weights_ = 1 - data['done'].float() # shape T,B weights = torch.ones_like(rewards) # shape T,B values[1:] = values[1:] * weights_ weights[1:] = weights_[:-1] rewards = rewards * weights # shape T,B return target_logit, behaviour_logit, actions, values, rewards, weights - def _state_dict_learn(self) -> Dict[str, Any]: - r""" - Overview: - Return the state_dict of learn mode, usually including model and optimizer. - Returns: - - state_dict (:obj:`Dict[str, Any]`): the dict of current policy learn state, for saving and restoring. + def _init_collect(self) -> None: """ - return { - 'model': self._learn_model.state_dict(), - 'optimizer': self._optimizer.state_dict(), - } - - def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: - r""" Overview: - Load the state_dict variable into policy learn mode. - Arguments: - - state_dict (:obj:`Dict[str, Any]`): the dict of policy learn state saved before. - .. tip:: - If you want to only load some parts of model, you can simply set the ``strict`` argument in \ - load_state_dict to ``False``, or refer to ``ding.torch_utils.checkpoint_helper`` for more \ - complicated operation. - """ - self._learn_model.load_state_dict(state_dict['model']) - self._optimizer.load_state_dict(state_dict['optimizer']) + Initialize the collect mode of policy, including related attributes and modules. For IMPALA, it contains \ + the collect_model to balance the exploration and exploitation (e.g. the multinomial sample mechanism in \ + discrete action space), and other algorithm-specific arguments such as unroll_len. + This method will be called in ``__init__`` method if ``collect`` field is in ``enable_field``. - def _init_collect(self) -> None: - r""" - Overview: - Collect mode init method. Called by ``self.__init__``, initialize algorithm arguments and collect_model. - Use multinomial_sample to choose action. + .. note:: + If you want to set some spacial member variables in ``_init_collect`` method, you'd better name them \ + with prefix ``_collect_`` to avoid conflict with other modes, such as ``self._collect_attr1``. """ - self._collect_model = model_wrap(self._model, wrapper_name='multinomial_sample') + assert self._cfg.action_space in ["continuous", "discrete"] + self._action_space = self._cfg.action_space + if self._action_space == 'continuous': + self._collect_model = model_wrap(self._model, wrapper_name='reparam_sample') + elif self._action_space == 'discrete': + self._collect_model = model_wrap(self._model, wrapper_name='multinomial_sample') + self._collect_model.reset() - def _forward_collect(self, data: Dict[int, Any]) -> Dict[int, Dict[str, Any]]: - r""" + def _forward_collect(self, data: Dict[int, Any]) -> Dict[int, Any]: + """ Overview: - Forward computation graph of collect mode(collect training data). + Policy forward function of collect mode (collecting training data by interacting with envs). Forward means \ + that the policy gets some necessary data (mainly observation) from the envs and then returns the output \ + data, such as the action to interact with the envs. Arguments: - - data (:obj:`Dict[int, Any]`): Dict type data, stacked env data for predicting \ - action, values are torch.Tensor or np.ndarray or dict/list combinations,keys \ - are env_id indicated by integer. + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. Returns: - - output (:obj:`Dict[int, Dict[str,Any]]`): Dict of predicting policy_output(logit, action) for each env. - ReturnsKeys - - necessary: ``logit``, ``action`` + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action and \ + other necessary data (action logit and value) for learn mode defined in ``self._process_transition`` \ + method. The key of the dict is the same as the input data, i.e. environment id. + + .. tip:: + If you want to add more tricks on this policy, like temperature factor in multinomial sample, you can pass \ + related data as extra keyword arguments of this method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to unittest for IMPALAPolicy: ``ding.policy.tests.test_impala``. """ data_id = list(data.keys()) data = default_collate(list(data.values())) @@ -319,34 +383,34 @@ def _forward_collect(self, data: Dict[int, Any]) -> Dict[int, Dict[str, Any]]: return output def _get_train_sample(self, data: List[Dict[str, Any]]) -> List[Dict[str, Any]]: - r""" + """ Overview: - For a given trajectory(transitions, a list of transition) data, process it into a list of sample that \ - can be used for training directly. + For a given trajectory (transitions, a list of transition) data, process it into a list of sample that \ + can be used for training. In IMPALA, a train sample is processed transitions with unroll_len length. Arguments: - - data (:obj:`List[Dict[str, Any]`): The trajectory data(a list of transition), each element is the same \ - format as the return value of ``self._process_transition`` method. + - transitions (:obj:`List[Dict[str, Any]`): The trajectory data (a list of transition), each element is \ + the same format as the return value of ``self._process_transition`` method. Returns: - - samples (:obj:`dict`): List of training samples. - .. note:: - We will vectorize ``process_transition`` and ``get_train_sample`` method in the following release version. \ - And the user can customize the this data processing procedure by overriding this two methods and collector \ - itself. + - samples (:obj:`List[Dict[str, Any]]`): The processed train samples, each element is the similar format \ + as input transitions, but may contain more data for training. """ return get_train_sample(data, self._unroll_len) - def _process_transition(self, obs: Any, policy_output: Dict[str, Any], timestep: namedtuple) -> Dict[str, Any]: - r""" + def _process_transition(self, obs: torch.Tensor, policy_output: Dict[str, torch.Tensor], + timestep: namedtuple) -> Dict[str, torch.Tensor]: + """ Overview: - Generate dict type transition data from inputs. + Process and pack one timestep transition data into a dict, which can be directly used for training and \ + saved in replay buffer. For IMPALA, it contains obs, next_obs, action, reward, done, logit. Arguments: - - obs (:obj:`Any`): Env observation,can be torch.Tensor or np.ndarray or dict/list combinations. - - model_output (:obj:`dict`): Output of collect model, including ['logit','action'] - - timestep (:obj:`namedtuple`): Output after env step, including at least ['obs', 'reward', 'done']\ - (here 'obs' indicates obs after env step). + - obs (:obj:`torch.Tensor`): The env observation of current timestep, such as stacked 2D image in Atari. + - policy_output (:obj:`Dict[str, torch.Tensor]`): The output of the policy network with the observation \ + as input. For IMPALA, it contains the action and the logit of the action. + - timestep (:obj:`namedtuple`): The execution result namedtuple returned by the environment step method, \ + except all the elements have been transformed into tensor data. Usually, it contains the next obs, \ + reward, done, info, etc. Returns: - - transition (:obj:`dict`): Dict type transition data, including at least ['obs','next_obs', 'logit',\ - 'action','reward', 'done'] + - transition (:obj:`Dict[str, torch.Tensor]`): The processed transition data of the current timestep. """ transition = { 'obs': obs, @@ -359,28 +423,48 @@ def _process_transition(self, obs: Any, policy_output: Dict[str, Any], timestep: return transition def _init_eval(self) -> None: - r""" + """ Overview: - Evaluate mode init method. Called by ``self.__init__``, initialize eval_model, - and use argmax_sample to choose action. + Initialize the eval mode of policy, including related attributes and modules. For IMPALA, it contains the \ + eval model to select optimial action (e.g. greedily select action with argmax mechanism in discrete action). + This method will be called in ``__init__`` method if ``eval`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_eval`` method, you'd better name them \ + with prefix ``_eval_`` to avoid conflict with other modes, such as ``self._eval_attr1``. """ - self._eval_model = model_wrap(self._model, wrapper_name='argmax_sample') + assert self._cfg.action_space in ["continuous", "discrete"], self._cfg.action_space + self._action_space = self._cfg.action_space + if self._action_space == 'continuous': + self._eval_model = model_wrap(self._model, wrapper_name='deterministic_sample') + elif self._action_space == 'discrete': + self._eval_model = model_wrap(self._model, wrapper_name='argmax_sample') + self._eval_model.reset() def _forward_eval(self, data: Dict[int, Any]) -> Dict[int, Any]: - r""" + """ Overview: - Forward computation graph of eval mode(evaluate policy performance), at most cases, it is similar to \ - ``self._forward_collect``. + Policy forward function of eval mode (evaluation policy performance by interacting with envs). Forward \ + means that the policy gets some necessary data (mainly observation) from the envs and then returns the \ + action to interact with the envs. ``_forward_eval`` in IMPALA often uses deterministic sample to get \ + actions while ``_forward_collect`` usually uses stochastic sample method for balance exploration and \ + exploitation. Arguments: - - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \ - values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer. + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. Returns: - - output (:obj:`Dict[int, Any]`): The dict of predicting action for the interaction with env. - ReturnsKeys - - necessary: ``action`` - - optional: ``logit`` + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action. The \ + key of the dict is the same as the input data, i.e. environment id. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + .. note:: + For more detailed examples, please refer to unittest for IMPALAPolicy: ``ding.policy.tests.test_impala``. """ data_id = list(data.keys()) data = default_collate(list(data.values())) @@ -396,13 +480,11 @@ def _forward_eval(self, data: Dict[int, Any]) -> Dict[int, Any]: return output def _monitor_vars_learn(self) -> List[str]: - r""" + """ Overview: - Return this algorithm default model setting for demonstration. + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. Returns: - - model_info (:obj:`Tuple[str, List[str]]`): model name and mode import_names - .. note:: - The user can define and use customized network model but must obey the same interface definition indicated \ - by import_names path. For IMPALA, ``ding.model.interface.IMPALA`` + - necessary_keys (:obj:`List[str]`): The list of the necessary keys to be logged. """ return super()._monitor_vars_learn() + ['policy_loss', 'value_loss', 'entropy_loss'] diff --git a/ding/policy/iql.py b/ding/policy/iql.py new file mode 100644 index 0000000000..567fa31214 --- /dev/null +++ b/ding/policy/iql.py @@ -0,0 +1,647 @@ +from typing import List, Dict, Any, Tuple, Union +import copy +from collections import namedtuple +import numpy as np +import torch +import torch.nn.functional as F +from torch.distributions import Normal, Independent, TransformedDistribution +from torch.distributions.transforms import TanhTransform, AffineTransform + +from ding.torch_utils import Adam, to_device +from ding.rl_utils import v_1step_td_data, v_1step_td_error, get_train_sample, \ + qrdqn_nstep_td_data, qrdqn_nstep_td_error, get_nstep_return_data +from ding.model import model_wrap +from ding.utils import POLICY_REGISTRY +from ding.utils.data import default_collate, default_decollate +from .base_policy import Policy +from .common_utils import default_preprocess_learn + + +def asymmetric_l2_loss(u, tau): + return torch.mean(torch.abs(tau - (u < 0).float()) * u ** 2) + + +@POLICY_REGISTRY.register('iql') +class IQLPolicy(Policy): + """ + Overview: + Policy class of Implicit Q-Learning (IQL) algorithm for continuous control. + Paper link: https://arxiv.org/abs/2110.06169. + + Config: + == ==================== ======== ============= ================================= ======================= + ID Symbol Type Default Value Description Other(Shape) + == ==================== ======== ============= ================================= ======================= + 1 ``type`` str iql | RL policy register name, refer | this arg is optional, + | to registry ``POLICY_REGISTRY`` | a placeholder + 2 ``cuda`` bool True | Whether to use cuda for network | + 3 | ``random_`` int 10000 | Number of randomly collected | Default to 10000 for + | ``collect_size`` | training samples in replay | SAC, 25000 for DDPG/ + | | buffer when training starts. | TD3. + 4 | ``model.policy_`` int 256 | Linear layer size for policy | + | ``embedding_size`` | network. | + 5 | ``model.soft_q_`` int 256 | Linear layer size for soft q | + | ``embedding_size`` | network. | + 6 | ``model.value_`` int 256 | Linear layer size for value | Defalut to None when + | ``embedding_size`` | network. | model.value_network + | | | is False. + 7 | ``learn.learning`` float 3e-4 | Learning rate for soft q | Defalut to 1e-3, when + | ``_rate_q`` | network. | model.value_network + | | | is True. + 8 | ``learn.learning`` float 3e-4 | Learning rate for policy | Defalut to 1e-3, when + | ``_rate_policy`` | network. | model.value_network + | | | is True. + 9 | ``learn.learning`` float 3e-4 | Learning rate for policy | Defalut to None when + | ``_rate_value`` | network. | model.value_network + | | | is False. + 10 | ``learn.alpha`` float 0.2 | Entropy regularization | alpha is initiali- + | | coefficient. | zation for auto + | | | `alpha`, when + | | | auto_alpha is True + 11 | ``learn.repara_`` bool True | Determine whether to use | + | ``meterization`` | reparameterization trick. | + 12 | ``learn.`` bool False | Determine whether to use | Temperature parameter + | ``auto_alpha`` | auto temperature parameter | determines the + | | `alpha`. | relative importance + | | | of the entropy term + | | | against the reward. + 13 | ``learn.-`` bool False | Determine whether to ignore | Use ignore_done only + | ``ignore_done`` | done flag. | in halfcheetah env. + 14 | ``learn.-`` float 0.005 | Used for soft update of the | aka. Interpolation + | ``target_theta`` | target network. | factor in polyak aver + | | | aging for target + | | | networks. + == ==================== ======== ============= ================================= ======================= + """ + + config = dict( + # (str) RL policy register name (refer to function "POLICY_REGISTRY"). + type='iql', + # (bool) Whether to use cuda for policy. + cuda=False, + # (bool) on_policy: Determine whether on-policy or off-policy. + # on-policy setting influences the behaviour of buffer. + on_policy=False, + # (bool) priority: Determine whether to use priority in buffer sample. + priority=False, + # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. + priority_IS_weight=False, + # (int) Number of training samples(randomly collected) in replay buffer when training starts. + random_collect_size=10000, + model=dict( + # (bool type) twin_critic: Determine whether to use double-soft-q-net for target q computation. + # Please refer to TD3 about Clipped Double-Q Learning trick, which learns two Q-functions instead of one . + # Default to True. + twin_critic=True, + # (str type) action_space: Use reparameterization trick for continous action + action_space='reparameterization', + # (int) Hidden size for actor network head. + actor_head_hidden_size=512, + actor_head_layer_num=3, + # (int) Hidden size for critic network head. + critic_head_hidden_size=512, + critic_head_layer_num=2, + ), + # learn_mode config + learn=dict( + # (int) How many updates (iterations) to train after collector's one collection. + # Bigger "update_per_collect" means bigger off-policy. + update_per_collect=1, + # (int) Minibatch size for gradient descent. + batch_size=256, + # (float) learning_rate_q: Learning rate for soft q network. + learning_rate_q=3e-4, + # (float) learning_rate_policy: Learning rate for policy network. + learning_rate_policy=3e-4, + # (float) learning_rate_alpha: Learning rate for auto temperature parameter ``alpha``. + learning_rate_alpha=3e-4, + # (float) target_theta: Used for soft update of the target network, + # aka. Interpolation factor in polyak averaging for target networks. + target_theta=0.005, + # (float) discount factor for the discounted sum of rewards, aka. gamma. + discount_factor=0.99, + # (float) alpha: Entropy regularization coefficient. + # Please check out the original SAC paper (arXiv 1801.01290): Eq 1 for more details. + # If auto_alpha is set to `True`, alpha is initialization for auto `\alpha`. + # Default to 0.2. + alpha=0.2, + # (bool) auto_alpha: Determine whether to use auto temperature parameter `\alpha` . + # Temperature parameter determines the relative importance of the entropy term against the reward. + # Please check out the original SAC paper (arXiv 1801.01290): Eq 1 for more details. + # Default to False. + # Note that: Using auto alpha needs to set learning_rate_alpha in `cfg.policy.learn`. + auto_alpha=True, + # (bool) log_space: Determine whether to use auto `\alpha` in log space. + log_space=True, + # (bool) Whether ignore done(usually for max step termination env. e.g. pendulum) + # Note: Gym wraps the MuJoCo envs by default with TimeLimit environment wrappers. + # These limit HalfCheetah, and several other MuJoCo envs, to max length of 1000. + # However, interaction with HalfCheetah always gets done with done is False, + # Since we inplace done==True with done==False to keep + # TD-error accurate computation(``gamma * (1 - done) * next_v + reward``), + # when the episode step is greater than max episode step. + ignore_done=False, + # (float) Weight uniform initialization range in the last output layer. + init_w=3e-3, + # (int) The numbers of action sample each at every state s from a uniform-at-random. + num_actions=10, + # (bool) Whether use lagrange multiplier in q value loss. + with_lagrange=False, + # (float) The threshold for difference in Q-values. + lagrange_thresh=-1, + # (float) Loss weight for conservative item. + min_q_weight=1.0, + # (float) coefficient for the asymmetric loss, range from [0.5, 1.0], default to 0.70. + tau=0.7, + # (float) temperature coefficient for Advantage Weighted Regression loss, default to 1.0. + beta=1.0, + ), + eval=dict(), # for compatibility + ) + + def default_model(self) -> Tuple[str, List[str]]: + """ + Overview: + Return this algorithm default neural network model setting for demonstration. ``__init__`` method will \ + automatically call this method to get the default model setting and create model. + Returns: + - model_info (:obj:`Tuple[str, List[str]]`): The registered model name and model's import_names. + """ + + return 'continuous_qvac', ['ding.model.template.qvac'] + + def _init_learn(self) -> None: + """ + Overview: + Initialize the learn mode of policy, including related attributes and modules. For SAC, it mainly \ + contains three optimizers, algorithm-specific arguments such as gamma, min_q_weight, with_lagrange, \ + main and target model. Especially, the ``auto_alpha`` mechanism for balancing max entropy \ + target is also initialized here. + This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. + """ + self._priority = self._cfg.priority + self._priority_IS_weight = self._cfg.priority_IS_weight + self._twin_critic = self._cfg.model.twin_critic + self._num_actions = self._cfg.learn.num_actions + + self._min_q_version = 3 + self._min_q_weight = self._cfg.learn.min_q_weight + self._with_lagrange = self._cfg.learn.with_lagrange and (self._lagrange_thresh > 0) + self._lagrange_thresh = self._cfg.learn.lagrange_thresh + if self._with_lagrange: + self.target_action_gap = self._lagrange_thresh + self.log_alpha_prime = torch.tensor(0.).to(self._device).requires_grad_() + self.alpha_prime_optimizer = Adam( + [self.log_alpha_prime], + lr=self._cfg.learn.learning_rate_q, + ) + + # Weight Init + init_w = self._cfg.learn.init_w + self._model.actor_head[-1].mu.weight.data.uniform_(-init_w, init_w) + self._model.actor_head[-1].mu.bias.data.uniform_(-init_w, init_w) + # self._model.actor_head[-1].log_sigma_layer.weight.data.uniform_(-init_w, init_w) + # self._model.actor_head[-1].log_sigma_layer.bias.data.uniform_(-init_w, init_w) + if self._twin_critic: + self._model.critic_q_head[0][-1].last.weight.data.uniform_(-init_w, init_w) + self._model.critic_q_head[0][-1].last.bias.data.uniform_(-init_w, init_w) + self._model.critic_q_head[1][-1].last.weight.data.uniform_(-init_w, init_w) + self._model.critic_q_head[1][-1].last.bias.data.uniform_(-init_w, init_w) + else: + self._model.critic_q_head[2].last.weight.data.uniform_(-init_w, init_w) + self._model.critic_q_head[-1].last.bias.data.uniform_(-init_w, init_w) + self._model.critic_v_head[2].last.weight.data.uniform_(-init_w, init_w) + self._model.critic_v_head[-1].last.bias.data.uniform_(-init_w, init_w) + + # Optimizers + self._optimizer_q = Adam( + self._model.critic.parameters(), + lr=self._cfg.learn.learning_rate_q, + ) + self._optimizer_policy = Adam( + self._model.actor.parameters(), + lr=self._cfg.learn.learning_rate_policy, + ) + + # Algorithm config + self._gamma = self._cfg.learn.discount_factor + + self._learn_model = model_wrap(self._model, wrapper_name='base') + self._learn_model.reset() + + self._forward_learn_cnt = 0 + + self._tau = self._cfg.learn.tau + self._beta = self._cfg.learn.beta + self._policy_start_training_counter = 10000 # 300000 + + def _forward_learn(self, data: List[Dict[str, Any]]) -> Dict[str, Any]: + """ + Overview: + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data from the offline dataset and then returns the output \ + result, including various training information such as loss, action, priority. + Arguments: + - data (:obj:`List[Dict[int, Any]]`): The input data used for policy forward, including a batch of \ + training samples. For each element in list, the key of the dict is the name of data items and the \ + value is the corresponding data. Usually, the value is torch.Tensor or np.ndarray or there dict/list \ + combinations. In the ``_forward_learn`` method, data often need to first be stacked in the batch \ + dimension by some utility functions such as ``default_preprocess_learn``. \ + For IQL, each element in list is a dict containing at least the following keys: ``obs``, ``action``, \ + ``reward``, ``next_obs``, ``done``. Sometimes, it also contains other keys such as ``weight``. + Returns: + - info_dict (:obj:`Dict[str, Any]`): The information dict that indicated training result, which will be \ + recorded in text log and tensorboard, values must be python scalar or a list of scalars. For the \ + detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + """ + loss_dict = {} + data = default_preprocess_learn( + data, + use_priority=self._priority, + use_priority_IS_weight=self._cfg.priority_IS_weight, + ignore_done=self._cfg.learn.ignore_done, + use_nstep=False + ) + if len(data.get('action').shape) == 1: + data['action'] = data['action'].reshape(-1, 1) + + if self._cuda: + data = to_device(data, self._device) + + self._learn_model.train() + obs = data['obs'] + next_obs = data['next_obs'] + reward = data['reward'] + done = data['done'] + + # 1. predict q and v value + value = self._learn_model.forward(data, mode='compute_critic') + q_value, v_value = value['q_value'], value['v_value'] + + # 2. predict target value + with torch.no_grad(): + (mu, sigma) = self._learn_model.forward(next_obs, mode='compute_actor')['logit'] + + next_obs_dist = TransformedDistribution( + Independent(Normal(mu, sigma), 1), + transforms=[TanhTransform(cache_size=1), + AffineTransform(loc=0.0, scale=1.05)] + ) + next_action = next_obs_dist.rsample() + next_log_prob = next_obs_dist.log_prob(next_action) + + next_data = {'obs': next_obs, 'action': next_action} + next_value = self._learn_model.forward(next_data, mode='compute_critic') + next_q_value, next_v_value = next_value['q_value'], next_value['v_value'] + + # the value of a policy according to the maximum entropy objective + if self._twin_critic: + next_q_value = torch.min(next_q_value[0], next_q_value[1]) + + # 3. compute v loss + if self._twin_critic: + q_value_min = torch.min(q_value[0], q_value[1]).detach() + v_loss = asymmetric_l2_loss(q_value_min - v_value, self._tau) + else: + advantage = q_value.detach() - v_value + v_loss = asymmetric_l2_loss(advantage, self._tau) + + # 4. compute q loss + if self._twin_critic: + q_data0 = v_1step_td_data(q_value[0], next_v_value, reward, done, data['weight']) + loss_dict['critic_q_loss'], td_error_per_sample0 = v_1step_td_error(q_data0, self._gamma) + q_data1 = v_1step_td_data(q_value[1], next_v_value, reward, done, data['weight']) + loss_dict['twin_critic_q_loss'], td_error_per_sample1 = v_1step_td_error(q_data1, self._gamma) + q_loss = (loss_dict['critic_q_loss'] + loss_dict['twin_critic_q_loss']) / 2 + else: + q_data = v_1step_td_data(q_value, next_v_value, reward, done, data['weight']) + loss_dict['critic_q_loss'], td_error_per_sample = v_1step_td_error(q_data, self._gamma) + q_loss = loss_dict['critic_q_loss'] + + # 5. update q and v network + self._optimizer_q.zero_grad() + v_loss.backward() + q_loss.backward() + self._optimizer_q.step() + + # 6. evaluate to get action distribution + (mu, sigma) = self._learn_model.forward(data['obs'], mode='compute_actor')['logit'] + + dist = TransformedDistribution( + Independent(Normal(mu, sigma), 1), + transforms=[TanhTransform(cache_size=1), AffineTransform(loc=0.0, scale=1.05)] + ) + action = data['action'] + log_prob = dist.log_prob(action) + + eval_data = {'obs': obs, 'action': action} + new_value = self._learn_model.forward(eval_data, mode='compute_critic') + new_q_value, new_v_value = new_value['q_value'], new_value['v_value'] + if self._twin_critic: + new_q_value = torch.min(new_q_value[0], new_q_value[1]) + new_advantage = new_q_value - new_v_value + + # 8. compute policy loss + policy_loss = (-log_prob * torch.exp(new_advantage.detach() / self._beta).clamp(max=20.0)).mean() + self._policy_start_training_counter -= 1 + + loss_dict['policy_loss'] = policy_loss + + # 9. update policy network + self._optimizer_policy.zero_grad() + policy_loss.backward() + policy_grad_norm = torch.nn.utils.clip_grad_norm_(self._model.actor.parameters(), 1) + self._optimizer_policy.step() + + loss_dict['total_loss'] = sum(loss_dict.values()) + + # ============= + # after update + # ============= + self._forward_learn_cnt += 1 + + return { + 'cur_lr_q': self._optimizer_q.defaults['lr'], + 'cur_lr_p': self._optimizer_policy.defaults['lr'], + 'priority': q_loss.abs().tolist(), + 'q_loss': q_loss.detach().mean().item(), + 'v_loss': v_loss.detach().mean().item(), + 'log_prob': log_prob.detach().mean().item(), + 'next_q_value': next_q_value.detach().mean().item(), + 'next_v_value': next_v_value.detach().mean().item(), + 'policy_loss': policy_loss.detach().mean().item(), + 'total_loss': loss_dict['total_loss'].detach().item(), + 'advantage_max': new_advantage.max().detach().item(), + 'new_q_value': new_q_value.detach().mean().item(), + 'new_v_value': new_v_value.detach().mean().item(), + 'policy_grad_norm': policy_grad_norm, + } + + def _get_policy_actions(self, data: Dict, num_actions: int = 10, epsilon: float = 1e-6) -> List: + # evaluate to get action distribution + obs = data['obs'] + obs = obs.unsqueeze(1).repeat(1, num_actions, 1).view(obs.shape[0] * num_actions, obs.shape[1]) + (mu, sigma) = self._learn_model.forward(obs, mode='compute_actor')['logit'] + dist = Independent(Normal(mu, sigma), 1) + pred = dist.rsample() + action = torch.tanh(pred) + + # evaluate action log prob depending on Jacobi determinant. + y = 1 - action.pow(2) + epsilon + log_prob = dist.log_prob(pred).unsqueeze(-1) + log_prob = log_prob - torch.log(y).sum(-1, keepdim=True) + + return action, log_prob.view(-1, num_actions, 1) + + def _get_q_value(self, data: Dict, keep: bool = True) -> torch.Tensor: + new_q_value = self._learn_model.forward(data, mode='compute_critic')['q_value'] + if self._twin_critic: + new_q_value = [value.view(-1, self._num_actions, 1) for value in new_q_value] + else: + new_q_value = new_q_value.view(-1, self._num_actions, 1) + if self._twin_critic and not keep: + new_q_value = torch.min(new_q_value[0], new_q_value[1]) + return new_q_value + + def _get_v_value(self, data: Dict, keep: bool = True) -> torch.Tensor: + new_v_value = self._learn_model.forward(data, mode='compute_critic')['v_value'] + if self._twin_critic: + new_v_value = [value.view(-1, self._num_actions, 1) for value in new_v_value] + else: + new_v_value = new_v_value.view(-1, self._num_actions, 1) + if self._twin_critic and not keep: + new_v_value = torch.min(new_v_value[0], new_v_value[1]) + return new_v_value + + def _init_collect(self) -> None: + """ + Overview: + Initialize the collect mode of policy, including related attributes and modules. For SAC, it contains the \ + collect_model other algorithm-specific arguments such as unroll_len. \ + This method will be called in ``__init__`` method if ``collect`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_collect`` method, you'd better name them \ + with prefix ``_collect_`` to avoid conflict with other modes, such as ``self._collect_attr1``. + """ + self._unroll_len = self._cfg.collect.unroll_len + self._collect_model = model_wrap(self._model, wrapper_name='base') + self._collect_model.reset() + + def _forward_collect(self, data: Dict[int, Any], **kwargs) -> Dict[int, Any]: + """ + Overview: + Policy forward function of collect mode (collecting training data by interacting with envs). Forward means \ + that the policy gets some necessary data (mainly observation) from the envs and then returns the output \ + data, such as the action to interact with the envs. + Arguments: + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. + Returns: + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action and \ + other necessary data for learn mode defined in ``self._process_transition`` method. The key of the \ + dict is the same as the input data, i.e. environment id. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + ``logit`` in SAC means the mu and sigma of Gaussioan distribution. Here we use this name for consistency. + + .. note:: + For more detailed examples, please refer to our unittest for SACPolicy: ``ding.policy.tests.test_sac``. + """ + data_id = list(data.keys()) + data = default_collate(list(data.values())) + if self._cuda: + data = to_device(data, self._device) + self._collect_model.eval() + with torch.no_grad(): + (mu, sigma) = self._collect_model.forward(data, mode='compute_actor')['logit'] + dist = Independent(Normal(mu, sigma), 1) + action = torch.tanh(dist.rsample()) + output = {'logit': (mu, sigma), 'action': action} + if self._cuda: + output = to_device(output, 'cpu') + output = default_decollate(output) + return {i: d for i, d in zip(data_id, output)} + + def _process_transition(self, obs: torch.Tensor, policy_output: Dict[str, torch.Tensor], + timestep: namedtuple) -> Dict[str, torch.Tensor]: + """ + Overview: + Process and pack one timestep transition data into a dict, which can be directly used for training and \ + saved in replay buffer. For continuous SAC, it contains obs, next_obs, action, reward, done. The logit \ + will be also added when ``collector_logit`` is True. + Arguments: + - obs (:obj:`torch.Tensor`): The env observation of current timestep, such as stacked 2D image in Atari. + - policy_output (:obj:`Dict[str, torch.Tensor]`): The output of the policy network with the observation \ + as input. For continuous SAC, it contains the action and the logit (mu and sigma) of the action. + - timestep (:obj:`namedtuple`): The execution result namedtuple returned by the environment step method, \ + except all the elements have been transformed into tensor data. Usually, it contains the next obs, \ + reward, done, info, etc. + Returns: + - transition (:obj:`Dict[str, torch.Tensor]`): The processed transition data of the current timestep. + """ + if self._cfg.collect.collector_logit: + transition = { + 'obs': obs, + 'next_obs': timestep.obs, + 'logit': policy_output['logit'], + 'action': policy_output['action'], + 'reward': timestep.reward, + 'done': timestep.done, + } + else: + transition = { + 'obs': obs, + 'next_obs': timestep.obs, + 'action': policy_output['action'], + 'reward': timestep.reward, + 'done': timestep.done, + } + return transition + + def _get_train_sample(self, transitions: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Overview: + For a given trajectory (transitions, a list of transition) data, process it into a list of sample that \ + can be used for training directly. In continuous SAC, a train sample is a processed transition \ + (unroll_len=1). + Arguments: + - transitions (:obj:`List[Dict[str, Any]`): The trajectory data (a list of transition), each element is \ + the same format as the return value of ``self._process_transition`` method. + Returns: + - samples (:obj:`List[Dict[str, Any]]`): The processed train samples, each element is the similar format \ + as input transitions, but may contain more data for training. + """ + return get_train_sample(transitions, self._unroll_len) + + def _init_eval(self) -> None: + """ + Overview: + Initialize the eval mode of policy, including related attributes and modules. For SAC, it contains the \ + eval model, which is equipped with ``base`` model wrapper to ensure compability. + This method will be called in ``__init__`` method if ``eval`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_eval`` method, you'd better name them \ + with prefix ``_eval_`` to avoid conflict with other modes, such as ``self._eval_attr1``. + """ + self._eval_model = model_wrap(self._model, wrapper_name='base') + self._eval_model.reset() + + def _forward_eval(self, data: Dict[int, Any]) -> Dict[int, Any]: + """ + Overview: + Policy forward function of eval mode (evaluation policy performance by interacting with envs). Forward \ + means that the policy gets some necessary data (mainly observation) from the envs and then returns the \ + action to interact with the envs. + Arguments: + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. + Returns: + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action. The \ + key of the dict is the same as the input data, i.e. environment id. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + ``logit`` in SAC means the mu and sigma of Gaussioan distribution. Here we use this name for consistency. + + .. note:: + For more detailed examples, please refer to our unittest for SACPolicy: ``ding.policy.tests.test_sac``. + """ + data_id = list(data.keys()) + data = default_collate(list(data.values())) + if self._cuda: + data = to_device(data, self._device) + self._eval_model.eval() + with torch.no_grad(): + (mu, sigma) = self._eval_model.forward(data, mode='compute_actor')['logit'] + action = torch.tanh(mu) / 1.05 # deterministic_eval + output = {'action': action} + if self._cuda: + output = to_device(output, 'cpu') + output = default_decollate(output) + return {i: d for i, d in zip(data_id, output)} + + def _monitor_vars_learn(self) -> List[str]: + """ + Overview: + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. + Returns: + - necessary_keys (:obj:`List[str]`): The list of the necessary keys to be logged. + """ + twin_critic = ['twin_critic_loss'] if self._twin_critic else [] + return [ + 'cur_lr_q', + 'cur_lr_p', + 'value_loss' + 'policy_loss', + 'q_loss', + 'v_loss', + 'policy_loss', + 'log_prob', + 'total_loss', + 'advantage_max', + 'next_q_value', + 'next_v_value', + 'new_q_value', + 'new_v_value', + 'policy_grad_norm', + ] + twin_critic + + def _state_dict_learn(self) -> Dict[str, Any]: + """ + Overview: + Return the state_dict of learn mode, usually including model and optimizer. + Returns: + - state_dict (:obj:`Dict[str, Any]`): The dict of current policy learn state, for saving and restoring. + """ + return { + 'model': self._learn_model.state_dict(), + 'optimizer_q': self._optimizer_q.state_dict(), + 'optimizer_policy': self._optimizer_policy.state_dict(), + } + + def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: + """ + Overview: + Load the state_dict variable into policy learn mode. + Arguments: + - state_dict (:obj:`Dict[str, Any]`): The dict of policy learn state saved before. + + .. tip:: + If you want to only load some parts of model, you can simply set the ``strict`` argument in \ + load_state_dict to ``False``, or refer to ``ding.torch_utils.checkpoint_helper`` for more \ + complicated operation. + """ + self._learn_model.load_state_dict(state_dict['model']) + self._optimizer_q.load_state_dict(state_dict['optimizer_q']) + self._optimizer_policy.load_state_dict(state_dict['optimizer_policy']) diff --git a/ding/policy/iqn.py b/ding/policy/iqn.py index 24785c41b0..7ff69d0528 100644 --- a/ding/policy/iqn.py +++ b/ding/policy/iqn.py @@ -13,9 +13,13 @@ @POLICY_REGISTRY.register('iqn') class IQNPolicy(DQNPolicy): - r""" + """ Overview: - Policy class of IQN algorithm. + Policy class of IQN algorithm. Paper link: https://arxiv.org/pdf/1806.06923.pdf. \ + Distrbutional RL is a new direction of RL, which is more stable than the traditional RL algorithm. \ + The core idea of distributional RL is to estimate the distribution of action value instead of the \ + expectation. The difference between IQN and DQN is that IQN uses quantile regression to estimate the \ + quantile value of the action distribution, while DQN uses the expectation of the action distribution. \ Config: == ==================== ======== ============== ======================================== ======================= @@ -58,8 +62,6 @@ class IQNPolicy(DQNPolicy): # (int) N-step reward for target q_value estimation nstep=1, learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, # How many updates(iterations) to train after collector's one collection. # Bigger "update_per_collect" means bigger off-policy. # collect data -> update policy-> collect data -> ... @@ -100,13 +102,37 @@ class IQNPolicy(DQNPolicy): ) def default_model(self) -> Tuple[str, List[str]]: + """ + Overview: + Return this algorithm default neural network model setting for demonstration. ``__init__`` method will \ + automatically call this method to get the default model setting and create model. + Returns: + - model_info (:obj:`Tuple[str, List[str]]`): The registered model name and model's import_names. + + .. note:: + The user can define and use customized network model but must obey the same inferface definition indicated \ + by import_names path. For example about IQN, its registered name is ``iqn`` and the import_names is \ + ``ding.model.template.q_learning``. + """ return 'iqn', ['ding.model.template.q_learning'] def _init_learn(self) -> None: - r""" + """ Overview: - Learn mode init method. Called by ``self.__init__``. - Init the optimizer, algorithm config, main and target models. + Initialize the learn mode of policy, including related attributes and modules. For IQN, it mainly contains \ + optimizer, algorithm-specific arguments such as nstep, kappa and gamma, main and target model. + This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. """ self._priority = self._cfg.priority # Optimizer @@ -128,14 +154,34 @@ def _init_learn(self) -> None: self._learn_model.reset() self._target_model.reset() - def _forward_learn(self, data: dict) -> Dict[str, Any]: - r""" + def _forward_learn(self, data: List[Dict[int, Any]]) -> Dict[str, Any]: + """ Overview: - Forward and backward function of learn mode. + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data from the replay buffer and then returns the output \ + result, including various training information such as loss, priority. Arguments: - - data (:obj:`dict`): Dict type data, including at least ['obs', 'action', 'reward', 'next_obs'] + - data (:obj:`List[Dict[int, Any]]`): The input data used for policy forward, including a batch of \ + training samples. For each element in list, the key of the dict is the name of data items and the \ + value is the corresponding data. Usually, the value is torch.Tensor or np.ndarray or there dict/list \ + combinations. In the ``_forward_learn`` method, data often need to first be stacked in the batch \ + dimension by some utility functions such as ``default_preprocess_learn``. \ + For IQN, each element in list is a dict containing at least the following keys: ``obs``, ``action``, \ + ``reward``, ``next_obs``, ``done``. Sometimes, it also contains other keys such as ``weight`` \ + and ``value_gamma``. Returns: - - info_dict (:obj:`Dict[str, Any]`): Including current lr and loss. + - info_dict (:obj:`Dict[str, Any]`): The information dict that indicated training result, which will be \ + recorded in text log and tensorboard, values must be python scalar or a list of scalars. For the \ + detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for IQNPolicy: ``ding.policy.tests.test_iqn``. """ data = default_preprocess_learn( data, use_priority=self._priority, ignore_done=self._cfg.learn.ignore_done, use_nstep=True @@ -171,7 +217,7 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: # ==================== self._optimizer.zero_grad() loss.backward() - if self._cfg.learn.multi_gpu: + if self._cfg.multi_gpu: self.sync_gradients(self._learn_model) self._optimizer.step() @@ -188,6 +234,12 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: } def _state_dict_learn(self) -> Dict[str, Any]: + """ + Overview: + Return the state_dict of learn mode, usually including model, target_model and optimizer. + Returns: + - state_dict (:obj:`Dict[str, Any]`): The dict of current policy learn state, for saving and restoring. + """ return { 'model': self._learn_model.state_dict(), 'target_model': self._target_model.state_dict(), @@ -195,6 +247,17 @@ def _state_dict_learn(self) -> Dict[str, Any]: } def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: + """ + Overview: + Load the state_dict variable into policy learn mode. + Arguments: + - state_dict (:obj:`Dict[str, Any]`): The dict of policy learn state saved before. + + .. tip:: + If you want to only load some parts of model, you can simply set the ``strict`` argument in \ + load_state_dict to ``False``, or refer to ``ding.torch_utils.checkpoint_helper`` for more \ + complicated operation. + """ self._learn_model.load_state_dict(state_dict['model']) self._target_model.load_state_dict(state_dict['target_model']) self._optimizer.load_state_dict(state_dict['optimizer']) diff --git a/ding/policy/madqn.py b/ding/policy/madqn.py new file mode 100644 index 0000000000..185992af64 --- /dev/null +++ b/ding/policy/madqn.py @@ -0,0 +1,354 @@ +from typing import List, Dict, Any, Tuple, Union, Optional +from collections import namedtuple +import torch +import copy + +from ding.torch_utils import RMSprop, to_device +from ding.rl_utils import v_1step_td_data, v_1step_td_error, get_train_sample, \ + v_nstep_td_data, v_nstep_td_error, get_nstep_return_data +from ding.model import model_wrap +from ding.utils import POLICY_REGISTRY +from ding.utils.data import timestep_collate, default_collate, default_decollate +from .qmix import QMIXPolicy + + +@POLICY_REGISTRY.register('madqn') +class MADQNPolicy(QMIXPolicy): + config = dict( + # (str) RL policy register name (refer to function "POLICY_REGISTRY"). + type='madqn', + # (bool) Whether to use cuda for network. + cuda=True, + # (bool) Whether the RL algorithm is on-policy or off-policy. + on_policy=False, + # (bool) Whether use priority(priority sample, IS weight, update priority) + priority=False, + # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. + priority_IS_weight=False, + nstep=3, + learn=dict( + update_per_collect=20, + batch_size=32, + learning_rate=0.0005, + clip_value=100, + # ============================================================== + # The following configs is algorithm-specific + # ============================================================== + # (float) Target network update momentum parameter. + # in [0, 1]. + target_update_theta=0.008, + # (float) The discount factor for future rewards, + # in [0, 1]. + discount_factor=0.99, + # (bool) Whether to use double DQN mechanism(target q for surpassing over estimation) + double_q=False, + weight_decay=1e-5, + ), + collect=dict( + # (int) Only one of [n_sample, n_episode] shoule be set + n_episode=32, + # (int) Cut trajectories into pieces with length "unroll_len", the length of timesteps + # in each forward when training. In qmix, it is greater than 1 because there is RNN. + unroll_len=10, + ), + eval=dict(), + other=dict( + eps=dict( + # (str) Type of epsilon decay + type='exp', + # (float) Start value for epsilon decay, in [0, 1]. + # 0 means not use epsilon decay. + start=1, + # (float) Start value for epsilon decay, in [0, 1]. + end=0.05, + # (int) Decay length(env step) + decay=50000, + ), + replay_buffer=dict( + replay_buffer_size=5000, + # (int) The maximum reuse times of each data + max_reuse=1e+9, + max_staleness=1e+9, + ), + ), + ) + + def default_model(self) -> Tuple[str, List[str]]: + """ + Overview: + Return this algorithm default model setting for demonstration. + Returns: + - model_info (:obj:`Tuple[str, List[str]]`): model name and mode import_names + """ + return 'madqn', ['ding.model.template.madqn'] + + def _init_learn(self) -> None: + self._priority = self._cfg.priority + self._priority_IS_weight = self._cfg.priority_IS_weight + assert not self._priority and not self._priority_IS_weight, "Priority is not implemented in QMIX" + self._optimizer_current = RMSprop( + params=self._model.current.parameters(), + lr=self._cfg.learn.learning_rate, + alpha=0.99, + eps=0.00001, + weight_decay=self._cfg.learn.weight_decay + ) + self._optimizer_cooperation = RMSprop( + params=self._model.cooperation.parameters(), + lr=self._cfg.learn.learning_rate, + alpha=0.99, + eps=0.00001, + weight_decay=self._cfg.learn.weight_decay + ) + self._gamma = self._cfg.learn.discount_factor + self._nstep = self._cfg.nstep + self._target_model = copy.deepcopy(self._model) + self._target_model = model_wrap( + self._target_model, + wrapper_name='target', + update_type='momentum', + update_kwargs={'theta': self._cfg.learn.target_update_theta} + ) + self._target_model = model_wrap( + self._target_model, + wrapper_name='hidden_state', + state_num=self._cfg.learn.batch_size, + init_fn=lambda: [None for _ in range(self._cfg.model.agent_num)] + ) + self._learn_model = model_wrap( + self._model, + wrapper_name='hidden_state', + state_num=self._cfg.learn.batch_size, + init_fn=lambda: [None for _ in range(self._cfg.model.agent_num)] + ) + self._learn_model.reset() + self._target_model.reset() + + def _data_preprocess_learn(self, data: List[Any]) -> dict: + r""" + Overview: + Preprocess the data to fit the required data format for learning + Arguments: + - data (:obj:`List[Dict[str, Any]]`): the data collected from collect function + Returns: + - data (:obj:`Dict[str, Any]`): the processed data, from \ + [len=B, ele={dict_key: [len=T, ele=Tensor(any_dims)]}] -> {dict_key: Tensor([T, B, any_dims])} + """ + # data preprocess + data = timestep_collate(data) + if self._cuda: + data = to_device(data, self._device) + data['weight'] = data.get('weight', None) + data['done'] = data['done'].float() + return data + + def _forward_learn(self, data: dict) -> Dict[str, Any]: + r""" + Overview: + Forward and backward function of learn mode. + Arguments: + - data (:obj:`Dict[str, Any]`): Dict type data, a batch of data for training, values are torch.Tensor or \ + np.ndarray or dict/list combinations. + Returns: + - info_dict (:obj:`Dict[str, Any]`): Dict type data, a info dict indicated training result, which will be \ + recorded in text log and tensorboard, values are python scalar or a list of scalars. + ArgumentsKeys: + - necessary: ``obs``, ``next_obs``, ``action``, ``reward``, ``weight``, ``prev_state``, ``done`` + ReturnsKeys: + - necessary: ``cur_lr``, ``total_loss`` + - cur_lr (:obj:`float`): Current learning rate + - total_loss (:obj:`float`): The calculated loss + """ + data = self._data_preprocess_learn(data) + # ==================== + # Q-mix forward + # ==================== + self._learn_model.train() + self._target_model.train() + # for hidden_state plugin, we need to reset the main model and target model + self._learn_model.reset(state=data['prev_state'][0]) + self._target_model.reset(state=data['prev_state'][0]) + inputs = {'obs': data['obs'], 'action': data['action']} + + total_q = self._learn_model.forward(inputs, single_step=False)['total_q'] + + if self._cfg.learn.double_q: + next_inputs = {'obs': data['next_obs']} + self._learn_model.reset(state=data['prev_state'][1]) + logit_detach = self._learn_model.forward(next_inputs, single_step=False)['logit'].clone().detach() + next_inputs = {'obs': data['next_obs'], 'action': logit_detach.argmax(dim=-1)} + else: + next_inputs = {'obs': data['next_obs']} + with torch.no_grad(): + target_total_q = self._target_model.forward(next_inputs, cooperation=True, single_step=False)['total_q'] + + if self._nstep == 1: + + v_data = v_1step_td_data(total_q, target_total_q, data['reward'], data['done'], data['weight']) + loss, td_error_per_sample = v_1step_td_error(v_data, self._gamma) + # for visualization + with torch.no_grad(): + if data['done'] is not None: + target_v = self._gamma * (1 - data['done']) * target_total_q + data['reward'] + else: + target_v = self._gamma * target_total_q + data['reward'] + else: + data['reward'] = data['reward'].permute(0, 2, 1).contiguous() + loss = [] + td_error_per_sample = [] + for t in range(self._cfg.collect.unroll_len): + v_data = v_nstep_td_data( + total_q[t], target_total_q[t], data['reward'][t], data['done'][t], data['weight'], None + ) + # calculate v_nstep_td critic_loss + loss_i, td_error_per_sample_i = v_nstep_td_error(v_data, self._gamma, self._nstep) + loss.append(loss_i) + td_error_per_sample.append(td_error_per_sample_i) + loss = sum(loss) / (len(loss) + 1e-8) + td_error_per_sample = sum(td_error_per_sample) / (len(td_error_per_sample) + 1e-8) + + self._optimizer_current.zero_grad() + loss.backward() + grad_norm = torch.nn.utils.clip_grad_norm_(self._model.current.parameters(), self._cfg.learn.clip_value) + self._optimizer_current.step() + + # cooperation + self._learn_model.reset(state=data['prev_state'][0]) + self._target_model.reset(state=data['prev_state'][0]) + cooperation_total_q = self._learn_model.forward(inputs, cooperation=True, single_step=False)['total_q'] + next_inputs = {'obs': data['next_obs']} + with torch.no_grad(): + cooperation_target_total_q = self._target_model.forward( + next_inputs, cooperation=True, single_step=False + )['total_q'] + + if self._nstep == 1: + v_data = v_1step_td_data( + cooperation_total_q, cooperation_target_total_q, data['reward'], data['done'], data['weight'] + ) + cooperation_loss, _ = v_1step_td_error(v_data, self._gamma) + else: + cooperation_loss_all = [] + for t in range(self._cfg.collect.unroll_len): + v_data = v_nstep_td_data( + cooperation_total_q[t], + cooperation_target_total_q[t], + data['reward'][t], + data['done'][t], + data['weight'], + None, + ) + cooperation_loss, _ = v_nstep_td_error(v_data, self._gamma, self._nstep) + cooperation_loss_all.append(cooperation_loss) + cooperation_loss = sum(cooperation_loss_all) / (len(cooperation_loss_all) + 1e-8) + self._optimizer_cooperation.zero_grad() + cooperation_loss.backward() + cooperation_grad_norm = torch.nn.utils.clip_grad_norm_( + self._model.cooperation.parameters(), self._cfg.learn.clip_value + ) + self._optimizer_cooperation.step() + + # ============= + # after update + # ============= + self._target_model.update(self._learn_model.state_dict()) + return { + 'cur_lr': self._optimizer_current.defaults['lr'], + 'total_loss': loss.item(), + 'total_q': total_q.mean().item() / self._cfg.model.agent_num, + 'target_total_q': target_total_q.mean().item() / self._cfg.model.agent_num, + 'grad_norm': grad_norm, + 'cooperation_grad_norm': cooperation_grad_norm, + 'cooperation_loss': cooperation_loss.item(), + } + + def _reset_learn(self, data_id: Optional[List[int]] = None) -> None: + r""" + Overview: + Reset learn model to the state indicated by data_id + Arguments: + - data_id (:obj:`Optional[List[int]]`): The id that store the state and we will reset\ + the model state to the state indicated by data_id + """ + self._learn_model.reset(data_id=data_id) + + def _state_dict_learn(self) -> Dict[str, Any]: + r""" + Overview: + Return the state_dict of learn mode, usually including model and optimizer. + Returns: + - state_dict (:obj:`Dict[str, Any]`): the dict of current policy learn state, for saving and restoring. + """ + return { + 'model': self._learn_model.state_dict(), + 'target_model': self._target_model.state_dict(), + 'optimizer_current': self._optimizer_current.state_dict(), + 'optimizer_cooperation': self._optimizer_cooperation.state_dict(), + } + + def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: + """ + Overview: + Load the state_dict variable into policy learn mode. + Arguments: + - state_dict (:obj:`Dict[str, Any]`): the dict of policy learn state saved before. + + .. tip:: + If you want to only load some parts of model, you can simply set the ``strict`` argument in \ + load_state_dict to ``False``, or refer to ``ding.torch_utils.checkpoint_helper`` for more \ + complicated operation. + """ + self._learn_model.load_state_dict(state_dict['model']) + self._target_model.load_state_dict(state_dict['target_model']) + self._optimizer_current.load_state_dict(state_dict['optimizer_current']) + self._optimizer_cooperation.load_state_dict(state_dict['optimizer_cooperation']) + + def _process_transition(self, obs: Any, model_output: dict, timestep: namedtuple) -> dict: + r""" + Overview: + Generate dict type transition data from inputs. + Arguments: + - obs (:obj:`Any`): Env observation + - model_output (:obj:`dict`): Output of collect model, including at least ['action', 'prev_state'] + - timestep (:obj:`namedtuple`): Output after env step, including at least ['obs', 'reward', 'done']\ + (here 'obs' indicates obs after env step). + Returns: + - transition (:obj:`dict`): Dict type transition data, including 'obs', 'next_obs', 'prev_state',\ + 'action', 'reward', 'done' + """ + transition = { + 'obs': obs, + 'next_obs': timestep.obs, + 'prev_state': model_output['prev_state'], + 'action': model_output['action'], + 'reward': timestep.reward, + 'done': timestep.done, + } + return transition + + def _get_train_sample(self, data: list) -> Union[None, List[Any]]: + r""" + Overview: + Get the train sample from trajectory. + Arguments: + - data (:obj:`list`): The trajectory's cache + Returns: + - samples (:obj:`dict`): The training samples generated + """ + if self._cfg.nstep == 1: + return get_train_sample(data, self._unroll_len) + else: + data = get_nstep_return_data(data, self._nstep, gamma=self._gamma) + return get_train_sample(data, self._unroll_len) + + def _monitor_vars_learn(self) -> List[str]: + r""" + Overview: + Return variables' name if variables are to used in monitor. + Returns: + - vars (:obj:`List[str]`): Variables' name list. + """ + return [ + 'cur_lr', 'total_loss', 'total_q', 'target_total_q', 'grad_norm', 'target_reward_total_q', + 'cooperation_grad_norm', 'cooperation_loss' + ] diff --git a/ding/policy/mbpolicy/__init__.py b/ding/policy/mbpolicy/__init__.py index 7d528cd15b..e23c8d823d 100644 --- a/ding/policy/mbpolicy/__init__.py +++ b/ding/policy/mbpolicy/__init__.py @@ -1 +1,2 @@ from .mbsac import MBSACPolicy +from .dreamer import DREAMERPolicy diff --git a/ding/policy/mbpolicy/dreamer.py b/ding/policy/mbpolicy/dreamer.py new file mode 100644 index 0000000000..35287c1bb6 --- /dev/null +++ b/ding/policy/mbpolicy/dreamer.py @@ -0,0 +1,361 @@ +from typing import List, Dict, Any, Tuple, Union +from collections import namedtuple +import torch +from torch import nn +from copy import deepcopy +from ding.torch_utils import Adam, to_device +from ding.rl_utils import get_train_sample +from ding.utils import POLICY_REGISTRY, deep_merge_dicts +from ding.utils.data import default_collate, default_decollate +from ding.policy import Policy +from ding.model import model_wrap +from ding.policy.common_utils import default_preprocess_learn + +from .utils import imagine, compute_target, compute_actor_loss, RewardEMA, tensorstats + + +@POLICY_REGISTRY.register('dreamer') +class DREAMERPolicy(Policy): + config = dict( + # (str) RL policy register name (refer to function "POLICY_REGISTRY"). + type='dreamer', + # (bool) Whether to use cuda for network and loss computation. + cuda=False, + # (int) Number of training samples (randomly collected) in replay buffer when training starts. + random_collect_size=5000, + # (bool) Whether to need policy-specific data in preprocess transition. + transition_with_policy_data=False, + # (int) + imag_horizon=15, + learn=dict( + # (float) Lambda for TD-lambda return. + lambda_=0.95, + # (float) Max norm of gradients. + grad_clip=100, + learning_rate=3e-5, + batch_size=16, + batch_length=64, + imag_sample=True, + slow_value_target=True, + slow_target_update=1, + slow_target_fraction=0.02, + discount=0.997, + reward_EMA=True, + actor_entropy=3e-4, + actor_state_entropy=0.0, + value_decay=0.0, + ), + ) + + def default_model(self) -> Tuple[str, List[str]]: + return 'dreamervac', ['ding.model.template.vac'] + + def _init_learn(self) -> None: + r""" + Overview: + Learn mode init method. Called by ``self.__init__``. + Init the optimizer, algorithm config, main and target models. + """ + # Algorithm config + self._lambda = self._cfg.learn.lambda_ + self._grad_clip = self._cfg.learn.grad_clip + + self._critic = self._model.critic + self._actor = self._model.actor + + if self._cfg.learn.slow_value_target: + self._slow_value = deepcopy(self._critic) + self._updates = 0 + + # Optimizer + self._optimizer_value = Adam( + self._critic.parameters(), + lr=self._cfg.learn.learning_rate, + ) + self._optimizer_actor = Adam( + self._actor.parameters(), + lr=self._cfg.learn.learning_rate, + ) + + self._learn_model = model_wrap(self._model, wrapper_name='base') + self._learn_model.reset() + + self._forward_learn_cnt = 0 + + if self._cfg.learn.reward_EMA: + self.reward_ema = RewardEMA(device=self._device) + + def _forward_learn(self, start: dict, world_model, envstep) -> Dict[str, Any]: + # log dict + log_vars = {} + self._learn_model.train() + self._update_slow_target() + + self._actor.requires_grad_(requires_grad=True) + # start is dict of {stoch, deter, logit} + if self._cuda: + start = to_device(start, self._device) + + # train self._actor + imag_feat, imag_state, imag_action = imagine( + self._cfg.learn, world_model, start, self._actor, self._cfg.imag_horizon + ) + reward = world_model.heads["reward"](world_model.dynamics.get_feat(imag_state)).mode() + actor_ent = self._actor(imag_feat).entropy() + state_ent = world_model.dynamics.get_dist(imag_state).entropy() + # this target is not scaled + # slow is flag to indicate whether slow_target is used for lambda-return + target, weights, base = compute_target( + self._cfg.learn, world_model, self._critic, imag_feat, imag_state, reward, actor_ent, state_ent + ) + actor_loss, mets = compute_actor_loss( + self._cfg.learn, + self._actor, + self.reward_ema, + imag_feat, + imag_state, + imag_action, + target, + actor_ent, + state_ent, + weights, + base, + ) + log_vars.update(mets) + value_input = imag_feat + self._actor.requires_grad_(requires_grad=False) + + self._critic.requires_grad_(requires_grad=True) + value = self._critic(value_input[:-1].detach()) + # to do + # target = torch.stack(target, dim=1) + # (time, batch, 1), (time, batch, 1) -> (time, batch) + value_loss = -value.log_prob(target.detach()) + slow_target = self._slow_value(value_input[:-1].detach()) + if self._cfg.learn.slow_value_target: + value_loss = value_loss - value.log_prob(slow_target.mode().detach()) + if self._cfg.learn.value_decay: + value_loss += self._cfg.learn.value_decay * value.mode() + # (time, batch, 1), (time, batch, 1) -> (1,) + value_loss = torch.mean(weights[:-1] * value_loss[:, :, None]) + self._critic.requires_grad_(requires_grad=False) + + log_vars.update(tensorstats(value.mode(), "value")) + log_vars.update(tensorstats(target, "target")) + log_vars.update(tensorstats(reward, "imag_reward")) + log_vars.update(tensorstats(imag_action, "imag_action")) + log_vars["actor_ent"] = torch.mean(actor_ent).detach().cpu().numpy().item() + # ==================== + # actor-critic update + # ==================== + self._model.requires_grad_(requires_grad=True) + world_model.requires_grad_(requires_grad=True) + + loss_dict = { + 'critic_loss': value_loss, + 'actor_loss': actor_loss, + } + + norm_dict = self._update(loss_dict) + + self._model.requires_grad_(requires_grad=False) + world_model.requires_grad_(requires_grad=False) + # ============= + # after update + # ============= + self._forward_learn_cnt += 1 + + return { + **log_vars, + **norm_dict, + **loss_dict, + } + + def _update(self, loss_dict): + # update actor + self._optimizer_actor.zero_grad() + loss_dict['actor_loss'].backward() + actor_norm = nn.utils.clip_grad_norm_(self._model.actor.parameters(), self._grad_clip) + self._optimizer_actor.step() + # update critic + self._optimizer_value.zero_grad() + loss_dict['critic_loss'].backward() + critic_norm = nn.utils.clip_grad_norm_(self._model.critic.parameters(), self._grad_clip) + self._optimizer_value.step() + return {'actor_grad_norm': actor_norm, 'critic_grad_norm': critic_norm} + + def _update_slow_target(self): + if self._cfg.learn.slow_value_target: + if self._updates % self._cfg.learn.slow_target_update == 0: + mix = self._cfg.learn.slow_target_fraction + for s, d in zip(self._critic.parameters(), self._slow_value.parameters()): + d.data = mix * s.data + (1 - mix) * d.data + self._updates += 1 + + def _state_dict_learn(self) -> Dict[str, Any]: + ret = { + 'model': self._learn_model.state_dict(), + 'optimizer_value': self._optimizer_value.state_dict(), + 'optimizer_actor': self._optimizer_actor.state_dict(), + } + return ret + + def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: + self._learn_model.load_state_dict(state_dict['model']) + self._optimizer_value.load_state_dict(state_dict['optimizer_value']) + self._optimizer_actor.load_state_dict(state_dict['optimizer_actor']) + + def _init_collect(self) -> None: + self._unroll_len = self._cfg.collect.unroll_len + self._collect_model = model_wrap(self._model, wrapper_name='base') + self._collect_model.reset() + + def _forward_collect(self, data: dict, world_model, envstep, reset=None, state=None) -> dict: + data_id = list(data.keys()) + data = default_collate(list(data.values())) + if self._cuda: + data = to_device(data, self._device) + self._collect_model.eval() + + if state is None: + batch_size = len(data_id) + latent = world_model.dynamics.initial(batch_size) # {logit, stoch, deter} + action = torch.zeros((batch_size, self._cfg.collect.action_size)).to(self._device) + else: + #state = default_collate(list(state.values())) + latent = to_device(default_collate(list(zip(*state))[0]), self._device) + action = to_device(default_collate(list(zip(*state))[1]), self._device) + if len(action.shape) == 1: + action = action.unsqueeze(-1) + if reset.any(): + mask = 1 - reset + for key in latent.keys(): + for i in range(latent[key].shape[0]): + latent[key][i] *= mask[i] + for i in range(len(action)): + action[i] *= mask[i] + assert world_model.obs_type == 'vector' or world_model.obs_type == 'RGB', \ + "action type must be vector or RGB" + # normalize RGB image input + if world_model.obs_type == 'RGB': + data = data - 0.5 + embed = world_model.encoder(data) + latent, _ = world_model.dynamics.obs_step(latent, action, embed, self._cfg.collect.collect_dyn_sample) + feat = world_model.dynamics.get_feat(latent) + + actor = self._actor(feat) + action = actor.sample() + logprob = actor.log_prob(action) + latent = {k: v.detach() for k, v in latent.items()} + action = action.detach() + + state = (latent, action) + assert world_model.action_type == 'discrete' or world_model.action_type == 'continuous', \ + "action type must be continuous or discrete" + if world_model.action_type == 'discrete': + action = torch.where(action == 1)[1] + output = {"action": action, "logprob": logprob, "state": state} + + if self._cuda: + output = to_device(output, 'cpu') + output = default_decollate(output) + if world_model.action_type == 'discrete': + for l in range(len(output)): + output[l]['action'] = output[l]['action'].squeeze(0) + return {i: d for i, d in zip(data_id, output)} + + def _process_transition(self, obs: Any, model_output: dict, timestep: namedtuple) -> dict: + r""" + Overview: + Generate dict type transition data from inputs. + Arguments: + - obs (:obj:`Any`): Env observation + - model_output (:obj:`dict`): Output of collect model, including at least ['action'] + - timestep (:obj:`namedtuple`): Output after env step, including at least ['obs', 'reward', 'done'] \ + (here 'obs' indicates obs after env step). + Returns: + - transition (:obj:`dict`): Dict type transition data. + """ + transition = { + 'obs': obs, + 'action': model_output['action'], + # TODO(zp) random_collect just have action + #'logprob': model_output['logprob'], + 'reward': timestep.reward, + 'discount': 1. - timestep.done, # timestep.info['discount'], + 'done': timestep.done, + } + return transition + + def _get_train_sample(self, data: list) -> Union[None, List[Any]]: + return get_train_sample(data, self._unroll_len) + + def _init_eval(self) -> None: + self._eval_model = model_wrap(self._model, wrapper_name='base') + self._eval_model.reset() + + def _forward_eval(self, data: dict, world_model, reset=None, state=None) -> dict: + data_id = list(data.keys()) + data = default_collate(list(data.values())) + if self._cuda: + data = to_device(data, self._device) + self._eval_model.eval() + + if state is None: + batch_size = len(data_id) + latent = world_model.dynamics.initial(batch_size) # {logit, stoch, deter} + action = torch.zeros((batch_size, self._cfg.collect.action_size)).to(self._device) + else: + #state = default_collate(list(state.values())) + latent = to_device(default_collate(list(zip(*state))[0]), self._device) + action = to_device(default_collate(list(zip(*state))[1]), self._device) + if len(action.shape) == 1: + action = action.unsqueeze(-1) + if reset.any(): + mask = 1 - reset + for key in latent.keys(): + for i in range(latent[key].shape[0]): + latent[key][i] *= mask[i] + for i in range(len(action)): + action[i] *= mask[i] + + # normalize RGB image input + if world_model.obs_type == 'RGB': + data = data - 0.5 + embed = world_model.encoder(data) + latent, _ = world_model.dynamics.obs_step(latent, action, embed, self._cfg.collect.collect_dyn_sample) + feat = world_model.dynamics.get_feat(latent) + + actor = self._actor(feat) + action = actor.mode() + logprob = actor.log_prob(action) + latent = {k: v.detach() for k, v in latent.items()} + action = action.detach() + + state = (latent, action) + if world_model.action_type == 'discrete': + action = torch.where(action == 1)[1] + output = {"action": action, "logprob": logprob, "state": state} + + if self._cuda: + output = to_device(output, 'cpu') + output = default_decollate(output) + if world_model.action_type == 'discrete': + for l in range(len(output)): + output[l]['action'] = output[l]['action'].squeeze(0) + return {i: d for i, d in zip(data_id, output)} + + def _monitor_vars_learn(self) -> List[str]: + r""" + Overview: + Return variables' name if variables are to used in monitor. + Returns: + - vars (:obj:`List[str]`): Variables' name list. + """ + return [ + 'normed_target_mean', 'normed_target_std', 'normed_target_min', 'normed_target_max', 'EMA_005', 'EMA_095', + 'actor_entropy', 'actor_state_entropy', 'value_mean', 'value_std', 'value_min', 'value_max', 'target_mean', + 'target_std', 'target_min', 'target_max', 'imag_reward_mean', 'imag_reward_std', 'imag_reward_min', + 'imag_reward_max', 'imag_action_mean', 'imag_action_std', 'imag_action_min', 'imag_action_max', 'actor_ent', + 'actor_loss', 'critic_loss', 'actor_grad_norm', 'critic_grad_norm' + ] diff --git a/ding/policy/mbpolicy/mbsac.py b/ding/policy/mbpolicy/mbsac.py index 65f17d0fb3..1918e161db 100644 --- a/ding/policy/mbpolicy/mbsac.py +++ b/ding/policy/mbpolicy/mbsac.py @@ -1,15 +1,13 @@ from typing import Dict, Any, List from functools import partial -import copy import torch from torch import Tensor from torch import nn -from torch.distributions import Normal, Independent, TransformedDistribution, TanhTransform -from easydict import EasyDict +from torch.distributions import Normal, Independent from ding.torch_utils import to_device, fold_batch, unfold_batch, unsqueeze_repeat -from ding.utils import POLICY_REGISTRY, deep_merge_dicts +from ding.utils import POLICY_REGISTRY from ding.policy import SACPolicy from ding.rl_utils import generalized_lambda_returns from ding.policy.common_utils import default_preprocess_learn @@ -19,27 +17,27 @@ @POLICY_REGISTRY.register('mbsac') class MBSACPolicy(SACPolicy): - r""" - Overview: - Model based SAC with value expansion (arXiv: 1803.00101) - and value gradient (arXiv: 1510.09142) w.r.t lambda-return. - - https://arxiv.org/pdf/1803.00101.pdf - https://arxiv.org/pdf/1510.09142.pdf - - Config: - == ==================== ======== ============= ================================== - ID Symbol Type Default Value Description - == ==================== ======== ============= ================================== - 1 ``learn._lambda`` float 0.8 | Lambda for TD-lambda return. - 2 ``learn.grad_clip` float 100.0 | Max norm of gradients. - 3 ``learn.sample_state`` bool True | Whether to sample states or tra- - | nsitions from environment buffer. - == ==================== ======== ============= ================================== - - .. note:: - For other configs, please refer to ding.policy.sac.SACPolicy. - """ + """ + Overview: + Model based SAC with value expansion (arXiv: 1803.00101) + and value gradient (arXiv: 1510.09142) w.r.t lambda-return. + + https://arxiv.org/pdf/1803.00101.pdf + https://arxiv.org/pdf/1510.09142.pdf + + Config: + == ==================== ======== ============= ================================== + ID Symbol Type Default Value Description + == ==================== ======== ============= ================================== + 1 ``learn._lambda`` float 0.8 | Lambda for TD-lambda return. + 2 ``learn.grad_clip` float 100.0 | Max norm of gradients. + 3 | ``learn.sample`` bool True | Whether to sample states or + | ``_state`` | transitions from env buffer. + == ==================== ======== ============= ================================== + + .. note:: + For other configs, please refer to ding.policy.sac.SACPolicy. + """ config = dict( learn=dict( @@ -52,13 +50,6 @@ class MBSACPolicy(SACPolicy): ) ) - @classmethod - def default_config(cls: type) -> EasyDict: - cfg = copy.deepcopy(cls.config) - cfg = EasyDict(deep_merge_dicts(super().config, cfg)) - cfg.cfg_type = cls.__name__ + 'Dict' - return cfg - def _init_learn(self) -> None: super()._init_learn() self._target_model.requires_grad_(False) @@ -98,6 +89,7 @@ def critic_fn(obss: Tensor, actions: Tensor, model: nn.Module): return q_values self._critic_fn = critic_fn + self._forward_learn_cnt = 0 def _forward_learn(self, data: dict, world_model, envstep) -> Dict[str, Any]: # preprocess data @@ -255,13 +247,6 @@ class STEVESACPolicy(SACPolicy): ) ) - @classmethod - def default_config(cls: type) -> EasyDict: - cfg = copy.deepcopy(cls.config) - cfg = EasyDict(deep_merge_dicts(super().config, cfg)) - cfg.cfg_type = cls.__name__ + 'Dict' - return cfg - def _init_learn(self) -> None: super()._init_learn() self._target_model.requires_grad_(False) @@ -294,6 +279,7 @@ def critic_fn(obss: Tensor, actions: Tensor, model: nn.Module): return q_values self._critic_fn = critic_fn + self._forward_learn_cnt = 0 def _forward_learn(self, data: dict, world_model, envstep) -> Dict[str, Any]: # preprocess data diff --git a/ding/policy/mbpolicy/utils.py b/ding/policy/mbpolicy/utils.py index bd1d77d3f9..b17c36e47f 100644 --- a/ding/policy/mbpolicy/utils.py +++ b/ding/policy/mbpolicy/utils.py @@ -1,7 +1,9 @@ from typing import Callable, Tuple, Union - +import torch from torch import Tensor from ding.torch_utils import fold_batch, unfold_batch +from ding.rl_utils import generalized_lambda_returns +from ding.torch_utils.network.dreamer import static_scan def q_evaluation(obss: Tensor, actions: Tensor, q_critic_fn: Callable[[Tensor, Tensor], @@ -36,3 +38,111 @@ def q_evaluation(obss: Tensor, actions: Tensor, q_critic_fn: Callable[[Tensor, T if isinstance(q_values, list): return [unfold_batch(q_values[0], dim), unfold_batch(q_values[1], dim)] return unfold_batch(q_values, dim) + + +def imagine(cfg, world_model, start, actor, horizon, repeats=None): + dynamics = world_model.dynamics + flatten = lambda x: x.reshape([-1] + list(x.shape[2:])) + start = {k: flatten(v) for k, v in start.items()} + + def step(prev, _): + state, _, _ = prev + feat = dynamics.get_feat(state) + inp = feat.detach() + action = actor(inp).sample() + succ = dynamics.img_step(state, action, sample=cfg.imag_sample) + return succ, feat, action + + succ, feats, actions = static_scan(step, [torch.arange(horizon)], (start, None, None)) + states = {k: torch.cat([start[k][None], v[:-1]], 0) for k, v in succ.items()} + + return feats, states, actions + + +def compute_target(cfg, world_model, critic, imag_feat, imag_state, reward, actor_ent, state_ent): + if "discount" in world_model.heads: + inp = world_model.dynamics.get_feat(imag_state) + discount = cfg.discount * world_model.heads["discount"](inp).mean + # TODO whether to detach + discount = discount.detach() + else: + discount = cfg.discount * torch.ones_like(reward) + + value = critic(imag_feat).mode() + # value(imag_horizon, 16*64, 1) + # action(imag_horizon, 16*64, ch) + # discount(imag_horizon, 16*64, 1) + target = generalized_lambda_returns(value, reward[:-1], discount[:-1], cfg.lambda_) + weights = torch.cumprod(torch.cat([torch.ones_like(discount[:1]), discount[:-1]], 0), 0).detach() + return target, weights, value[:-1] + + +def compute_actor_loss( + cfg, + actor, + reward_ema, + imag_feat, + imag_state, + imag_action, + target, + actor_ent, + state_ent, + weights, + base, +): + metrics = {} + inp = imag_feat.detach() + policy = actor(inp) + actor_ent = policy.entropy() + # Q-val for actor is not transformed using symlog + if cfg.reward_EMA: + offset, scale = reward_ema(target) + normed_target = (target - offset) / scale + normed_base = (base - offset) / scale + adv = normed_target - normed_base + metrics.update(tensorstats(normed_target, "normed_target")) + values = reward_ema.values + metrics["EMA_005"] = values[0].detach().cpu().numpy().item() + metrics["EMA_095"] = values[1].detach().cpu().numpy().item() + + actor_target = adv + if cfg.actor_entropy > 0: + actor_entropy = cfg.actor_entropy * actor_ent[:-1][:, :, None] + actor_target += actor_entropy + metrics["actor_entropy"] = torch.mean(actor_entropy).detach().cpu().numpy().item() + if cfg.actor_state_entropy > 0: + state_entropy = cfg.actor_state_entropy * state_ent[:-1] + actor_target += state_entropy + metrics["actor_state_entropy"] = torch.mean(state_entropy).detach().cpu().numpy().item() + actor_loss = -torch.mean(weights[:-1] * actor_target) + return actor_loss, metrics + + +class RewardEMA(object): + """running mean and std""" + + def __init__(self, device, alpha=1e-2): + self.device = device + self.values = torch.zeros((2, )).to(device) + self.alpha = alpha + self.range = torch.tensor([0.05, 0.95]).to(device) + + def __call__(self, x): + flat_x = torch.flatten(x.detach()) + x_quantile = torch.quantile(input=flat_x, q=self.range) + self.values = self.alpha * x_quantile + (1 - self.alpha) * self.values + scale = torch.clip(self.values[1] - self.values[0], min=1.0) + offset = self.values[0] + return offset.detach(), scale.detach() + + +def tensorstats(tensor, prefix=None): + metrics = { + 'mean': torch.mean(tensor).detach().cpu().numpy(), + 'std': torch.std(tensor).detach().cpu().numpy(), + 'min': torch.min(tensor).detach().cpu().numpy(), + 'max': torch.max(tensor).detach().cpu().numpy(), + } + if prefix: + metrics = {f'{prefix}_{k}': v.item() for k, v in metrics.items()} + return metrics diff --git a/ding/policy/mdqn.py b/ding/policy/mdqn.py new file mode 100644 index 0000000000..8842c11102 --- /dev/null +++ b/ding/policy/mdqn.py @@ -0,0 +1,281 @@ +from typing import List, Dict, Any +import copy +import torch + +from ding.torch_utils import Adam, to_device +from ding.rl_utils import m_q_1step_td_data, m_q_1step_td_error +from ding.model import model_wrap +from ding.utils import POLICY_REGISTRY + +from .dqn import DQNPolicy +from .common_utils import default_preprocess_learn + + +@POLICY_REGISTRY.register('mdqn') +class MDQNPolicy(DQNPolicy): + """ + Overview: + Policy class of Munchausen DQN algorithm, extended by auxiliary objectives. + Paper link: https://arxiv.org/abs/2007.14430. + Config: + == ==================== ======== ============== ======================================== ======================= + ID Symbol Type Default Value Description Other(Shape) + == ==================== ======== ============== ======================================== ======================= + 1 ``type`` str mdqn | RL policy register name, refer to | This arg is optional, + | registry ``POLICY_REGISTRY`` | a placeholder + 2 ``cuda`` bool False | Whether to use cuda for network | This arg can be diff- + | erent from modes + 3 ``on_policy`` bool False | Whether the RL algorithm is on-policy + | or off-policy + 4 ``priority`` bool False | Whether use priority(PER) | Priority sample, + | update priority + 5 | ``priority_IS`` bool False | Whether use Importance Sampling Weight + | ``_weight`` | to correct biased update. If True, + | priority must be True. + 6 | ``discount_`` float 0.97, | Reward's future discount factor, aka. | May be 1 when sparse + | ``factor`` [0.95, 0.999] | gamma | reward env + 7 ``nstep`` int 1, | N-step reward discount sum for target + [3, 5] | q_value estimation + 8 | ``learn.update`` int 1 | How many updates(iterations) to train | This args can be vary + | ``per_collect`` | after collector's one collection. Only | from envs. Bigger val + | valid in serial training | means more off-policy + | ``_gpu`` + 10 | ``learn.batch_`` int 32 | The number of samples of an iteration + | ``size`` + 11 | ``learn.learning`` float 0.001 | Gradient step length of an iteration. + | ``_rate`` + 12 | ``learn.target_`` int 2000 | Frequence of target network update. | Hard(assign) update + | ``update_freq`` + 13 | ``learn.ignore_`` bool False | Whether ignore done for target value | Enable it for some + | ``done`` | calculation. | fake termination env + 14 ``collect.n_sample`` int 4 | The number of training samples of a | It varies from + | call of collector. | different envs + 15 | ``collect.unroll`` int 1 | unroll length of an iteration | In RNN, unroll_len>1 + | ``_len`` + 16 | ``other.eps.type`` str exp | exploration rate decay type | Support ['exp', + | 'linear']. + 17 | ``other.eps.`` float 0.01 | start value of exploration rate | [0,1] + | ``start`` + 18 | ``other.eps.`` float 0.001 | end value of exploration rate | [0,1] + | ``end`` + 19 | ``other.eps.`` int 250000 | decay length of exploration | greater than 0. set + | ``decay`` | decay=250000 means + | the exploration rate + | decay from start + | value to end value + | during decay length. + 20 | ``entropy_tau`` float 0.003 | the ration of entropy in TD loss + 21 | ``alpha`` float 0.9 | the ration of Munchausen term to the + | TD loss + == ==================== ======== ============== ======================================== ======================= + """ + config = dict( + # (str) RL policy register name (refer to function "POLICY_REGISTRY"). + type='mdqn', + # (bool) Whether to use cuda in policy. + cuda=False, + # (bool) Whether learning policy is the same as collecting data policy(on-policy). + on_policy=False, + # (bool) Whether to enable priority experience sample. + priority=False, + # (bool) Whether to use Importance Sampling Weight to correct biased update. If True, priority must be True. + priority_IS_weight=False, + # (float) Discount factor(gamma) for returns. + discount_factor=0.97, + # (float) Entropy factor (tau) for Munchausen DQN. + entropy_tau=0.03, + # (float) Discount factor (alpha) for Munchausen term. + m_alpha=0.9, + # (int) The number of step for calculating target q_value. + nstep=1, + # learn_mode config + learn=dict( + # (int) How many updates(iterations) to train after collector's one collection. + # Bigger "update_per_collect" means bigger off-policy. + # collect data -> update policy-> collect data -> ... + update_per_collect=3, + # (int) How many samples in a training batch + batch_size=64, + # (float) The step size of gradient descent + learning_rate=0.001, + # (int) Frequence of target network update. + target_update_freq=100, + # (bool) Whether ignore done(usually for max step termination env). + # Note: Gym wraps the MuJoCo envs by default with TimeLimit environment wrappers. + # These limit HalfCheetah, and several other MuJoCo envs, to max length of 1000. + # However, interaction with HalfCheetah always gets done with done is False, + # Since we inplace done==True with done==False to keep + # TD-error accurate computation(``gamma * (1 - done) * next_v + reward``), + # when the episode step is greater than max episode step. + ignore_done=False, + ), + # collect_mode config + collect=dict( + # (int) How many training samples collected in one collection procedure. + # Only one of [n_sample, n_episode] shoule be set. + n_sample=4, + # (int) Split episodes or trajectories into pieces with length `unroll_len`. + unroll_len=1, + ), + eval=dict(), # for compability + # other config + other=dict( + # Epsilon greedy with decay. + eps=dict( + # (str) Decay type. Support ['exp', 'linear']. + type='exp', + # (float) Epsilon start value. + start=0.95, + # (float) Epsilon end value. + end=0.1, + # (int) Decay length(env step). + decay=10000, + ), + replay_buffer=dict( + # (int) Maximum size of replay buffer. Usually, larger buffer size is better. + replay_buffer_size=10000, + ), + ), + ) + + def _init_learn(self) -> None: + """ + Overview: + Initialize the learn mode of policy, including related attributes and modules. For MDQN, it contains \ + optimizer, algorithm-specific arguments such as entropy_tau, m_alpha and nstep, main and target model. + This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. + """ + self._priority = self._cfg.priority + self._priority_IS_weight = self._cfg.priority_IS_weight + # Optimizer + # set eps in order to consistent with the original paper implementation + self._optimizer = Adam(self._model.parameters(), lr=self._cfg.learn.learning_rate, eps=0.0003125) + + self._gamma = self._cfg.discount_factor + self._nstep = self._cfg.nstep + self._entropy_tau = self._cfg.entropy_tau + self._m_alpha = self._cfg.m_alpha + + # use model_wrapper for specialized demands of different modes + self._target_model = copy.deepcopy(self._model) + if 'target_update_freq' in self._cfg.learn: + self._target_model = model_wrap( + self._target_model, + wrapper_name='target', + update_type='assign', + update_kwargs={'freq': self._cfg.learn.target_update_freq} + ) + elif 'target_theta' in self._cfg.learn: + self._target_model = model_wrap( + self._target_model, + wrapper_name='target', + update_type='momentum', + update_kwargs={'theta': self._cfg.learn.target_theta} + ) + else: + raise RuntimeError("DQN needs target network, please either indicate target_update_freq or target_theta") + self._learn_model = model_wrap(self._model, wrapper_name='argmax_sample') + self._learn_model.reset() + self._target_model.reset() + + def _forward_learn(self, data: Dict[str, Any]) -> Dict[str, Any]: + """ + Overview: + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data from the replay buffer and then returns the output \ + result, including various training information such as loss, action_gap, clip_frac, priority. + Arguments: + - data (:obj:`List[Dict[int, Any]]`): The input data used for policy forward, including a batch of \ + training samples. For each element in list, the key of the dict is the name of data items and the \ + value is the corresponding data. Usually, the value is torch.Tensor or np.ndarray or there dict/list \ + combinations. In the ``_forward_learn`` method, data often need to first be stacked in the batch \ + dimension by some utility functions such as ``default_preprocess_learn``. \ + For MDQN, each element in list is a dict containing at least the following keys: ``obs``, ``action``, \ + ``reward``, ``next_obs``, ``done``. Sometimes, it also contains other keys such as ``weight`` \ + and ``value_gamma``. + Returns: + - info_dict (:obj:`Dict[str, Any]`): The information dict that indicated training result, which will be \ + recorded in text log and tensorboard, values must be python scalar or a list of scalars. For the \ + detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for MDQNPolicy: ``ding.policy.tests.test_mdqn``. + """ + data = default_preprocess_learn( + data, + use_priority=self._priority, + use_priority_IS_weight=self._cfg.priority_IS_weight, + ignore_done=self._cfg.learn.ignore_done, + use_nstep=True + ) + if self._cuda: + data = to_device(data, self._device) + # ==================== + # Q-learning forward + # ==================== + self._learn_model.train() + self._target_model.train() + # Current q value (main model) + q_value = self._learn_model.forward(data['obs'])['logit'] + # Target q value + with torch.no_grad(): + target_q_value_current = self._target_model.forward(data['obs'])['logit'] + target_q_value = self._target_model.forward(data['next_obs'])['logit'] + + data_m = m_q_1step_td_data( + q_value, target_q_value_current, target_q_value, data['action'], data['reward'].squeeze(0), data['done'], + data['weight'] + ) + + loss, td_error_per_sample, action_gap, clipfrac = m_q_1step_td_error( + data_m, self._gamma, self._entropy_tau, self._m_alpha + ) + # ==================== + # Q-learning update + # ==================== + self._optimizer.zero_grad() + loss.backward() + if self._cfg.multi_gpu: + self.sync_gradients(self._learn_model) + self._optimizer.step() + + # ============= + # after update + # ============= + self._target_model.update(self._learn_model.state_dict()) + return { + 'cur_lr': self._optimizer.defaults['lr'], + 'total_loss': loss.item(), + 'q_value': q_value.mean().item(), + 'target_q_value': target_q_value.mean().item(), + 'priority': td_error_per_sample.abs().tolist(), + 'action_gap': action_gap.item(), + 'clip_frac': clipfrac.mean().item(), + } + + def _monitor_vars_learn(self) -> List[str]: + """ + Overview: + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. + Returns: + - necessary_keys (:obj:`List[str]`): The list of the necessary keys to be logged. + """ + return ['cur_lr', 'total_loss', 'q_value', 'action_gap', 'clip_frac'] diff --git a/ding/policy/ngu.py b/ding/policy/ngu.py index 0d594cfa4c..51fabda807 100644 --- a/ding/policy/ngu.py +++ b/ding/policy/ngu.py @@ -85,8 +85,6 @@ class NGUPolicy(Policy): # i.e., = = + learn_unroll_len=80, # set this key according to the episode length learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, update_per_collect=1, batch_size=64, learning_rate=0.0001, @@ -292,7 +290,6 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: 'action': data['burnin_nstep_action'], 'reward': data['burnin_nstep_reward'], 'beta': data['burnin_nstep_beta'], - 'enable_fast_timestep': True } tmp = self._learn_model.forward( inputs, saved_state_timesteps=[self._burnin_step, self._burnin_step + self._nstep] @@ -306,7 +303,6 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: 'action': data['main_action'], 'reward': data['main_reward'], 'beta': data['main_beta'], - 'enable_fast_timestep': True } self._learn_model.reset(data_id=None, state=tmp['saved_state'][0]) q_value = self._learn_model.forward(inputs)['logit'] @@ -319,7 +315,6 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: 'action': data['target_action'], 'reward': data['target_reward'], 'beta': data['target_beta'], - 'enable_fast_timestep': True } with torch.no_grad(): target_q_value = self._target_model.forward(next_inputs)['logit'] diff --git a/ding/policy/offppo_collect_traj.py b/ding/policy/offppo_collect_traj.py index d165cf5a74..219d582c83 100644 --- a/ding/policy/offppo_collect_traj.py +++ b/ding/policy/offppo_collect_traj.py @@ -37,8 +37,7 @@ class OffPPOCollectTrajPolicy(Policy): nstep_return=False, nstep=3, learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, + # How many updates(iterations) to train after collector's one collection. # Bigger "update_per_collect" means bigger off-policy. # collect data -> update policy-> collect data -> ... diff --git a/ding/policy/pc.py b/ding/policy/pc.py new file mode 100644 index 0000000000..7c472462b0 --- /dev/null +++ b/ding/policy/pc.py @@ -0,0 +1,186 @@ +import math +from typing import List, Dict, Any, Tuple +from collections import namedtuple + +import torch +import torch.nn as nn +from torch.optim import Adam, SGD, AdamW +from torch.optim.lr_scheduler import LambdaLR + +from ding.policy import Policy +from ding.model import model_wrap +from ding.torch_utils import to_device +from ding.utils import EasyTimer +from ding.utils import POLICY_REGISTRY + + +@POLICY_REGISTRY.register('pc_bfs') +class ProcedureCloningBFSPolicy(Policy): + + def default_model(self) -> Tuple[str, List[str]]: + return 'pc_bfs', ['ding.model.template.procedure_cloning'] + + config = dict( + type='pc', + cuda=False, + on_policy=False, + continuous=False, + max_bfs_steps=100, + learn=dict( + update_per_collect=1, + batch_size=32, + learning_rate=1e-5, + lr_decay=False, + decay_epoch=30, + decay_rate=0.1, + warmup_lr=1e-4, + warmup_epoch=3, + optimizer='SGD', + momentum=0.9, + weight_decay=1e-4, + ), + collect=dict( + unroll_len=1, + noise=False, + noise_sigma=0.2, + noise_range=dict( + min=-0.5, + max=0.5, + ), + ), + eval=dict(), + other=dict(replay_buffer=dict(replay_buffer_size=10000)), + ) + + def _init_learn(self): + assert self._cfg.learn.optimizer in ['SGD', 'Adam'] + if self._cfg.learn.optimizer == 'SGD': + self._optimizer = SGD( + self._model.parameters(), + lr=self._cfg.learn.learning_rate, + weight_decay=self._cfg.learn.weight_decay, + momentum=self._cfg.learn.momentum + ) + elif self._cfg.learn.optimizer == 'Adam': + if self._cfg.learn.weight_decay is None: + self._optimizer = Adam( + self._model.parameters(), + lr=self._cfg.learn.learning_rate, + ) + else: + self._optimizer = AdamW( + self._model.parameters(), + lr=self._cfg.learn.learning_rate, + weight_decay=self._cfg.learn.weight_decay + ) + if self._cfg.learn.lr_decay: + + def lr_scheduler_fn(epoch): + if epoch <= self._cfg.learn.warmup_epoch: + return self._cfg.learn.warmup_lr / self._cfg.learn.learning_rate + else: + ratio = (epoch - self._cfg.learn.warmup_epoch) // self._cfg.learn.decay_epoch + return math.pow(self._cfg.learn.decay_rate, ratio) + + self._lr_scheduler = LambdaLR(self._optimizer, lr_scheduler_fn) + self._timer = EasyTimer(cuda=True) + self._learn_model = model_wrap(self._model, 'base') + self._learn_model.reset() + self._max_bfs_steps = self._cfg.max_bfs_steps + self._maze_size = self._cfg.maze_size + self._num_actions = self._cfg.num_actions + + self._loss = nn.CrossEntropyLoss() + + def process_states(self, observations, maze_maps): + """Returns [B, W, W, 3] binary values. Channels are (wall; goal; obs)""" + loc = torch.nn.functional.one_hot( + (observations[:, 0] * self._maze_size + observations[:, 1]).long(), + self._maze_size * self._maze_size, + ).long() + loc = torch.reshape(loc, [observations.shape[0], self._maze_size, self._maze_size]) + states = torch.cat([maze_maps, loc], dim=-1).long() + return states + + def _forward_learn(self, data): + if self._cuda: + collated_data = to_device(data, self._device) + else: + collated_data = data + observations = collated_data['obs'], + bfs_input_maps, bfs_output_maps = collated_data['bfs_in'].long(), collated_data['bfs_out'].long() + states = observations + bfs_input_onehot = torch.nn.functional.one_hot(bfs_input_maps, self._num_actions + 1).float() + + bfs_states = torch.cat([ + states, + bfs_input_onehot, + ], dim=-1) + logits = self._model(bfs_states)['logit'] + logits = logits.flatten(0, -2) + labels = bfs_output_maps.flatten(0, -1) + + loss = self._loss(logits, labels) + preds = torch.argmax(logits, dim=-1) + acc = torch.sum((preds == labels)) / preds.shape[0] + + self._optimizer.zero_grad() + loss.backward() + self._optimizer.step() + pred_loss = loss.item() + + cur_lr = [param_group['lr'] for param_group in self._optimizer.param_groups] + cur_lr = sum(cur_lr) / len(cur_lr) + return {'cur_lr': cur_lr, 'total_loss': pred_loss, 'acc': acc} + + def _monitor_vars_learn(self): + return ['cur_lr', 'total_loss', 'acc'] + + def _init_eval(self): + self._eval_model = model_wrap(self._model, wrapper_name='base') + self._eval_model.reset() + + def _forward_eval(self, data): + if self._cuda: + data = to_device(data, self._device) + max_len = self._max_bfs_steps + data_id = list(data.keys()) + output = {} + + for ii in data_id: + states = data[ii].unsqueeze(0) + bfs_input_maps = self._num_actions * torch.ones([1, self._maze_size, self._maze_size]).long() + if self._cuda: + bfs_input_maps = to_device(bfs_input_maps, self._device) + xy = torch.where(states[:, :, :, -1] == 1) + observation = (xy[1][0].item(), xy[2][0].item()) + + i = 0 + while bfs_input_maps[0, observation[0], observation[1]].item() == self._num_actions and i < max_len: + bfs_input_onehot = torch.nn.functional.one_hot(bfs_input_maps, self._num_actions + 1).long() + + bfs_states = torch.cat([ + states, + bfs_input_onehot, + ], dim=-1) + logits = self._model(bfs_states)['logit'] + bfs_input_maps = torch.argmax(logits, dim=-1) + i += 1 + output[ii] = bfs_input_maps[0, observation[0], observation[1]] + if self._cuda: + output[ii] = {'action': to_device(output[ii], 'cpu'), 'info': {}} + if output[ii]['action'].item() == self._num_actions: + output[ii]['action'] = torch.randint(low=0, high=self._num_actions, size=[1])[0] + return output + + def _init_collect(self) -> None: + raise NotImplementedError + + def _forward_collect(self, data: Dict[int, Any], **kwargs) -> Dict[int, Any]: + raise NotImplementedError + + def _process_transition(self, obs: Any, model_output: dict, timestep: namedtuple) -> dict: + raise NotImplementedError + + def _get_train_sample(self, data: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + raise NotImplementedError diff --git a/ding/policy/pdqn.py b/ding/policy/pdqn.py index f62e028e3b..6b66e263ab 100644 --- a/ding/policy/pdqn.py +++ b/ding/policy/pdqn.py @@ -14,9 +14,10 @@ @POLICY_REGISTRY.register('pdqn') class PDQNPolicy(Policy): - r""": + """ Overview: Policy class of PDQN algorithm, which extends the DQN algorithm on discrete-continuous hybrid action spaces. + Paper link: https://arxiv.org/abs/1810.06394. Config: == ==================== ======== ============== ======================================== ======================= @@ -43,7 +44,6 @@ class PDQNPolicy(Policy): | valid in serial training | means more off-policy 9 | ``learn.batch_`` int 64 | The number of samples of an iteration | ``size`` - 10 | ``learn.multi`` bool False | whether to use multi gpu during | ``_gpu`` 11 | ``learn.learning`` float 0.001 | Gradient step length of an iteration. | ``_rate`` @@ -59,12 +59,12 @@ class PDQNPolicy(Policy): | ``_sigma`` | during collection 17 | ``other.eps.type`` str exp | exploration rate decay type | Support ['exp', | 'linear']. - 18 | ``other.eps. float 0.95 | start value of exploration rate | [0,1] - | start`` - 19 | ``other.eps. float 0.05 | end value of exploration rate | [0,1] - | end`` - 20 | ``other.eps. int 10000 | decay length of exploration | greater than 0. set - | decay`` | decay=10000 means + 18 | ``other.eps.`` float 0.95 | start value of exploration rate | [0,1] + | ``start`` + 19 | ``other.eps.`` float 0.05 | end value of exploration rate | [0,1] + | ``end`` + 20 | ``other.eps.`` int 10000 | decay length of exploration | greater than 0. set + | ``decay`` | decay=10000 means | the exploration rate | decay from start | value to end value @@ -72,74 +72,104 @@ class PDQNPolicy(Policy): == ==================== ======== ============== ======================================== ======================= """ config = dict( + # (str) RL policy register name (refer to function "POLICY_REGISTRY"). type='pdqn', + # (bool) Whether to use cuda in policy. cuda=False, + # (bool) Whether learning policy is the same as collecting data policy(on-policy). on_policy=False, + # (bool) Whether to enable priority experience sample. priority=False, - # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. + # (bool) Whether to use Importance Sampling Weight to correct biased update. If True, priority must be True. priority_IS_weight=False, + # (float) Discount factor(gamma) for returns. discount_factor=0.97, + # (int) The number of step for calculating target q_value. nstep=1, + # learn_mode config learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, - # How many updates(iterations) to train after collector's one collection. + # (int) How many updates(iterations) to train after collector's one collection. # Bigger "update_per_collect" means bigger off-policy. # collect data -> update policy-> collect data -> ... update_per_collect=3, + # (int) How many samples in a training batch. batch_size=64, + # (float) The step size of gradient descent. learning_rate=0.001, - # ============================================================== - # The following configs are algorithm-specific - # ============================================================== # (int) Frequence of target network update. target_theta=0.005, - # (bool) Whether ignore done(usually for max step termination env) + # (bool) Whether ignore done(usually for max step termination env). + # Note: Gym wraps the MuJoCo envs by default with TimeLimit environment wrappers. + # These limit HalfCheetah, and several other MuJoCo envs, to max length of 1000. + # However, interaction with HalfCheetah always gets done with done is False, + # Since we inplace done==True with done==False to keep + # TD-error accurate computation(``gamma * (1 - done) * next_v + reward``), + # when the episode step is greater than max episode step. ignore_done=False, ), # collect_mode config collect=dict( - # (int) Only one of [n_sample, n_episode] shoule be set + # (int) How many training samples collected in one collection procedure. + # Only one of [n_sample, n_episode] shoule be set. # n_sample=8, - # (int) Cut trajectories into pieces with length "unroll_len". + # (int) Split episodes or trajectories into pieces with length `unroll_len`. unroll_len=1, - # (float) It is a must to add noise during collection. So here omits "noise" and only set "noise_sigma". + # (float) It is a must to add noise during collection. So here omits noise and only set ``noise_sigma``. noise_sigma=0.1, ), - eval=dict(), + eval=dict(), # for compatibility # other config other=dict( # Epsilon greedy with decay. eps=dict( # (str) Decay type. Support ['exp', 'linear']. type='exp', + # (float) Epsilon start value. start=0.95, + # (float) Epsilon end value. end=0.1, # (int) Decay length(env step) decay=10000, ), - replay_buffer=dict(replay_buffer_size=10000, ), + replay_buffer=dict( + # (int) Maximum size of replay buffer. Usually, larger buffer size is better. + replay_buffer_size=10000, + ), ), ) def default_model(self) -> Tuple[str, List[str]]: """ Overview: - Return this algorithm default model setting for demonstration. + Return this algorithm default neural network model setting for demonstration. ``__init__`` method will \ + automatically call this method to get the default model setting and create model. Returns: - - model_info (:obj:`Tuple[str, List[str]]`): model name and mode import_names + - model_info (:obj:`Tuple[str, List[str]]`): The registered model name and model's import_names. .. note:: The user can define and use customized network model but must obey the same inferface definition indicated \ - by import_names path. For PDQN, ``ding.model.template.pdqn.PDQN`` + by import_names path. For example about PDQN, its registered name is ``pdqn`` and the import_names is \ + ``ding.model.template.pdqn``. """ return 'pdqn', ['ding.model.template.pdqn'] def _init_learn(self) -> None: """ Overview: - Learn mode init method. Called by ``self.__init__``, initialize the optimizer, algorithm arguments, main \ - and target model. + Initialize the learn mode of policy, including related attributes and modules. For PDQN, it mainly \ + contains two optimizers, algorithm-specific arguments such as nstep and gamma, main and target model. + This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. """ self._priority = self._cfg.priority self._priority_IS_weight = self._cfg.priority_IS_weight @@ -172,19 +202,31 @@ def _init_learn(self) -> None: def _forward_learn(self, data: Dict[str, Any]) -> Dict[str, Any]: """ Overview: - Forward computation graph of learn mode(updating policy). + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data from the replay buffer and then returns the output \ + result, including various training information such as loss, q value, target_q_value, priority. Arguments: - - data (:obj:`Dict[str, Any]`): Dict type data, a batch of data for training, values are torch.Tensor or \ - np.ndarray or dict/list combinations. + - data (:obj:`List[Dict[int, Any]]`): The input data used for policy forward, including a batch of \ + training samples. For each element in list, the key of the dict is the name of data items and the \ + value is the corresponding data. Usually, the value is torch.Tensor or np.ndarray or there dict/list \ + combinations. In the ``_forward_learn`` method, data often need to first be stacked in the batch \ + dimension by some utility functions such as ``default_preprocess_learn``. \ + For PDQN, each element in list is a dict containing at least the following keys: ``obs``, ``action``, \ + ``reward``, ``next_obs``, ``done``. Sometimes, it also contains other keys such as ``weight`` \ + and ``value_gamma``. Returns: - - info_dict (:obj:`Dict[str, Any]`): Dict type data, a info dict indicated training result, which will be \ - recorded in text log and tensorboard, values are python scalar or a list of scalars. - ArgumentsKeys: - - necessary: ``obs``, ``action``, ``reward``, ``next_obs``, ``done`` - - optional: ``value_gamma``, ``IS`` - ReturnsKeys: - - necessary: ``cur_lr``, ``q_loss``, ``continuous_loss``, - ``q_value``, ``priority``, ``reward``, ``target_q_value`` + - info_dict (:obj:`Dict[str, Any]`): The information dict that indicated training result, which will be \ + recorded in text log and tensorboard, values must be python scalar or a list of scalars. For the \ + detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for PDQNPolicy: ``ding.policy.tests.test_pdqn``. """ data = default_preprocess_learn( data, @@ -267,7 +309,7 @@ def _forward_learn(self, data: Dict[str, Any]) -> Dict[str, Any]: return { 'cur_lr': self._dis_optimizer.defaults['lr'], 'q_loss': dis_loss.item(), - 'total_loss': (cont_loss + dis_loss).item(), + 'total_loss': cont_loss.item() + dis_loss.item(), 'continuous_loss': cont_loss.item(), 'q_value': q_pi_action_value.mean().item(), 'priority': td_error_per_sample.abs().tolist(), @@ -278,7 +320,8 @@ def _forward_learn(self, data: Dict[str, Any]) -> Dict[str, Any]: def _state_dict_learn(self) -> Dict[str, Any]: """ Overview: - Return the state_dict of learn mode, usually including model and optimizer. + Return the state_dict of learn mode, usually including model, target model, discrete part optimizer, and \ + continuous part optimizer. Returns: - state_dict (:obj:`Dict[str, Any]`): the dict of current policy learn state, for saving and restoring. """ @@ -309,8 +352,19 @@ def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: def _init_collect(self) -> None: """ Overview: - Collect mode init method. Called by ``self.__init__``, initialize algorithm arguments and collect_model, \ - enable the eps_greedy_sample for exploration. + Initialize the collect mode of policy, including related attributes and modules. For PDQN, it contains the \ + collect_model to balance the exploration and exploitation with epsilon-greedy sample mechanism and \ + continuous action mechanism, besides, other algorithm-specific arguments such as unroll_len and nstep are \ + also initialized here. + This method will be called in ``__init__`` method if ``collect`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_collect`` method, you'd better name them \ + with prefix ``_collect_`` to avoid conflict with other modes, such as ``self._collect_attr1``. + + .. tip:: + Some variables need to initialize independently in different modes, such as gamma and nstep in PDQN. This \ + design is for the convenience of parallel execution of different policy modes. """ self._unroll_len = self._cfg.collect.unroll_len self._gamma = self._cfg.discount_factor # necessary for parallel @@ -331,18 +385,27 @@ def _init_collect(self) -> None: def _forward_collect(self, data: Dict[int, Any], eps: float) -> Dict[int, Any]: """ Overview: - Forward computation graph of collect mode(collect training data), with eps_greedy for exploration. + Policy forward function of collect mode (collecting training data by interacting with envs). Forward means \ + that the policy gets some necessary data (mainly observation) from the envs and then returns the output \ + data, such as the action to interact with the envs. Besides, this policy also needs ``eps`` argument for \ + exploration, i.e., classic epsilon-greedy exploration strategy. Arguments: - - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \ - values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer. - - eps (:obj:`float`): epsilon value for exploration, which is decayed by collected env step. + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. + - eps (:obj:`float`): The epsilon value for exploration. Returns: - - output (:obj:`Dict[int, Any]`): The dict of predicting policy_output(action) for the interaction with \ - env and the constructing of transition. - ArgumentsKeys: - - necessary: ``obs`` - ReturnsKeys - - necessary: ``logit``, ``action`` + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action and \ + other necessary data for learn mode defined in ``self._process_transition`` method. The key of the \ + dict is the same as the input data, i.e. environment id. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for PDQNPolicy: ``ding.policy.tests.test_pdqn``. """ data_id = list(data.keys()) data = default_collate(list(data.values())) @@ -358,69 +421,86 @@ def _forward_collect(self, data: Dict[int, Any], eps: float) -> Dict[int, Any]: output = default_decollate(output) return {i: d for i, d in zip(data_id, output)} - def _get_train_sample(self, data: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + def _get_train_sample(self, transitions: List[Dict[str, Any]]) -> List[Dict[str, Any]]: """ Overview: - For a given trajectory(transitions, a list of transition) data, process it into a list of sample that \ - can be used for training directly. A train sample can be a processed transition(DQN with nstep TD) \ - or some continuous transitions(DRQN). + For a given trajectory (transitions, a list of transition) data, process it into a list of sample that \ + can be used for training directly. In PDQN, a train sample is a processed transition. \ + This method is usually used in collectors to execute necessary \ + RL data preprocessing before training, which can help learner amortize revelant time consumption. \ + In addition, you can also implement this method as an identity function and do the data processing \ + in ``self._forward_learn`` method. Arguments: - - data (:obj:`List[Dict[str, Any]`): The trajectory data(a list of transition), each element is the same \ - format as the return value of ``self._process_transition`` method. + - transitions (:obj:`List[Dict[str, Any]`): The trajectory data (a list of transition), each element is \ + the same format as the return value of ``self._process_transition`` method. Returns: - - samples (:obj:`dict`): The list of training samples. - - .. note:: - We will vectorize ``process_transition`` and ``get_train_sample`` method in the following release version. \ - And the user can customize the this data processing procecure by overriding this two methods and collector \ - itself. + - samples (:obj:`List[Dict[str, Any]]`): The processed train samples, each element is the similar format \ + as input transitions, but may contain more data for training, such as nstep reward and target obs. """ - data = get_nstep_return_data(data, self._nstep, gamma=self._gamma) - return get_train_sample(data, self._unroll_len) + transitions = get_nstep_return_data(transitions, self._nstep, gamma=self._gamma) + return get_train_sample(transitions, self._unroll_len) - def _process_transition(self, obs: Any, model_output: Dict[str, Any], timestep: namedtuple) -> Dict[str, Any]: + def _process_transition(self, obs: torch.Tensor, policy_output: Dict[str, torch.Tensor], + timestep: namedtuple) -> Dict[str, torch.Tensor]: """ Overview: - Generate a transition(e.g.: ) for this algorithm training. + Process and pack one timestep transition data into a dict, which can be directly used for training and \ + saved in replay buffer. For PDQN, it contains obs, next_obs, action, reward, done and logit. Arguments: - - obs (:obj:`Any`): Env observation. - - policy_output (:obj:`Dict[str, Any]`): The output of policy collect mode(``self._forward_collect``),\ - including at least ``action``. - - timestep (:obj:`namedtuple`): The output after env step(execute policy output action), including at \ - least ``obs``, ``reward``, ``done``, (here obs indicates obs after env step). + - obs (:obj:`torch.Tensor`): The env observation of current timestep, such as stacked 2D image in Atari. + - policy_output (:obj:`Dict[str, torch.Tensor]`): The output of the policy network with the observation \ + as input. For PDQN, it contains the hybrid action and the logit (discrete part q_value) of the action. + - timestep (:obj:`namedtuple`): The execution result namedtuple returned by the environment step method, \ + except all the elements have been transformed into tensor data. Usually, it contains the next obs, \ + reward, done, info, etc. Returns: - - transition (:obj:`dict`): Dict type transition data. + - transition (:obj:`Dict[str, torch.Tensor]`): The processed transition data of the current timestep. """ transition = { 'obs': obs, 'next_obs': timestep.obs, - 'action': model_output['action'], - 'logit': model_output['logit'], + 'action': policy_output['action'], + 'logit': policy_output['logit'], 'reward': timestep.reward, 'done': timestep.done, } return transition def _init_eval(self) -> None: - r""" + """ Overview: - Evaluate mode init method. Called by ``self.__init__``, initialize eval_model. + Initialize the eval mode of policy, including related attributes and modules. For PDQN, it contains the \ + eval model to greedily select action with argmax q_value mechanism. + This method will be called in ``__init__`` method if ``eval`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_eval`` method, you'd better name them \ + with prefix ``_eval_`` to avoid conflict with other modes, such as ``self._eval_attr1``. """ self._eval_model = model_wrap(self._model, wrapper_name='hybrid_argmax_sample') self._eval_model.reset() - def _forward_eval(self, data: dict) -> dict: - r""" + def _forward_eval(self, data: Dict[int, Any]) -> Dict[int, Any]: + """ Overview: - Forward function of eval mode, similar to ``self._forward_collect``. + Policy forward function of eval mode (evaluation policy performance by interacting with envs). Forward \ + means that the policy gets some necessary data (mainly observation) from the envs and then returns the \ + action to interact with the envs. Arguments: - - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \ - values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer. + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. Returns: - - output (:obj:`Dict[int, Any]`): The dict of predicting action for the interaction with env. - ReturnsKeys - - necessary: ``action`` - - optional: ``logit`` + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action. The \ + key of the dict is the same as the input data, i.e. environment id. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for PDQNPolicy: ``ding.policy.tests.test_pdqn``. """ data_id = list(data.keys()) data = default_collate(list(data.values())) @@ -437,10 +517,11 @@ def _forward_eval(self, data: dict) -> dict: return {i: d for i, d in zip(data_id, output)} def _monitor_vars_learn(self) -> List[str]: - r""" + """ Overview: - Return variables' names if variables are to used in monitor. + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. Returns: - - vars (:obj:`List[str]`): Variables' name list. + - necessary_keys (:obj:`List[str]`): The list of the necessary keys to be logged. """ return ['cur_lr', 'total_loss', 'q_loss', 'continuous_loss', 'q_value', 'reward', 'target_q_value'] diff --git a/ding/policy/pg.py b/ding/policy/pg.py new file mode 100644 index 0000000000..9ae0c87f25 --- /dev/null +++ b/ding/policy/pg.py @@ -0,0 +1,332 @@ +from typing import List, Dict, Any, Tuple, Union +from collections import namedtuple +import torch +import treetensor as ttorch + +from ding.rl_utils import get_gae_with_default_last_value, get_train_sample +from ding.torch_utils import Adam, to_device +from ding.utils import POLICY_REGISTRY, split_data_generator +from ding.utils.data import default_collate, default_decollate +from .base_policy import Policy +from .common_utils import default_preprocess_learn + + +@POLICY_REGISTRY.register('pg') +class PGPolicy(Policy): + """ + Overview: + Policy class of Policy Gradient (REINFORCE) algorithm. Paper link: \ + https://proceedings.neurips.cc/paper_files/paper/1999/file/464d828b85b0bed98e80ade0a5c43b0f-Paper.pdf + """ + config = dict( + # (string) RL policy register name (refer to function "register_policy"). + type='pg', + # (bool) whether to use cuda for network. + cuda=False, + # (bool) whether use on-policy training pipeline(behaviour policy and training policy are the same) + on_policy=True, # for pg strictly on policy algorithm, this line should not be modified by users + # (str) action space type: ['discrete', 'continuous'] + action_space='discrete', + # (bool) whether to use deterministic action for evaluation. + deterministic_eval=True, + learn=dict( + + # (int) the number of samples for one update. + batch_size=64, + # (float) the step size of one gradient descend. + learning_rate=0.001, + # ============================================================== + # The following configs is algorithm-specific + # ============================================================== + # (float) loss weight of the entropy regularization, the weight of policy network is set to 1 + entropy_weight=0.01, + # (float) max grad norm value. + grad_norm=5, + # (bool) whether to ignore done signal for non-termination env. + ignore_done=False, + ), + collect=dict( + # (int) collect n_sample data, train model n_iteration times + # n_episode=8, + # (int) trajectory unroll length + unroll_len=1, + # ============================================================== + # The following configs is algorithm-specific + # ============================================================== + # (float) discount factor for future reward, defaults int [0, 1] + discount_factor=0.99, + collector=dict(get_train_sample=True), + ), + eval=dict(), + ) + + def default_model(self) -> Tuple[str, List[str]]: + """ + Overview: + Return this algorithm default neural network model setting for demonstration. ``__init__`` method will \ + automatically call this method to get the default model setting and create model. + Returns: + - model_info (:obj:`Tuple[str, List[str]]`): The registered model name and model's import_names. + """ + return 'pg', ['ding.model.template.pg'] + + def _init_learn(self) -> None: + """ + Overview: + Initialize the learn mode of policy, including related attributes and modules. For PG, it mainly \ + contains optimizer, algorithm-specific arguments such as entropy weight and grad norm. This method \ + also executes some special network initializations. + This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. + """ + # Optimizer + self._optimizer = Adam(self._model.parameters(), lr=self._cfg.learn.learning_rate) + + self._entropy_weight = self._cfg.learn.entropy_weight + self._grad_norm = self._cfg.learn.grad_norm + self._learn_model = self._model # for compatibility + + def _forward_learn(self, data: List[Dict[int, Any]]) -> Dict[str, Any]: + """ + Overview: + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data from the replay buffer and then returns the output \ + result, including various training information such as loss, clipfrac, approx_kl. + Arguments: + - data (:obj:`List[Dict[int, Any]]`): The input data used for policy forward, including the latest \ + collected training samples for on-policy algorithms like PG. For each element in list, the key of the \ + dict is the name of data items and the value is the corresponding data. Usually, the value is \ + torch.Tensor or np.ndarray or there dict/list combinations. In the ``_forward_learn`` method, data \ + often need to first be stacked in the batch dimension by some utility functions such as \ + ``default_preprocess_learn``. \ + For PG, each element in list is a dict containing at least the following keys: ``obs``, ``action``, \ + ``return``. + Returns: + - return_infos (:obj:`List[Dict[str, Any]]`): The information list that indicated training result, each \ + training iteration contains append a information dict into the final list. The list will be precessed \ + and recorded in text log and tensorboard. The value of the dict must be python scalar or a list of \ + scalars. For the detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + """ + data = default_preprocess_learn(data, ignore_done=self._cfg.learn.ignore_done, use_nstep=False) + if self._cuda: + data = to_device(data, self._device) + self._model.train() + + return_infos = [] + for batch in split_data_generator(data, self._cfg.learn.batch_size, shuffle=True): + # forward + output = self._learn_model.forward(batch['obs']) + return_ = batch['return'] + dist = output['dist'] + # calculate PG loss + log_prob = dist.log_prob(batch['action']) + policy_loss = -(log_prob * return_).mean() + entropy_loss = -self._cfg.learn.entropy_weight * dist.entropy().mean() + total_loss = policy_loss + entropy_loss + + # update + self._optimizer.zero_grad() + total_loss.backward() + + grad_norm = torch.nn.utils.clip_grad_norm_( + list(self._learn_model.parameters()), + max_norm=self._grad_norm, + ) + self._optimizer.step() + + # only record last updates information in logger + return_info = { + 'cur_lr': self._optimizer.param_groups[0]['lr'], + 'total_loss': total_loss.item(), + 'policy_loss': policy_loss.item(), + 'entropy_loss': entropy_loss.item(), + 'return_abs_max': return_.abs().max().item(), + 'grad_norm': grad_norm, + } + return_infos.append(return_info) + return return_infos + + def _init_collect(self) -> None: + """ + Overview: + Initialize the collect mode of policy, including related attributes and modules. For PPG, it contains \ + algorithm-specific arguments such as unroll_len and gamma. + This method will be called in ``__init__`` method if ``collect`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_collect`` method, you'd better name them \ + with prefix ``_collect_`` to avoid conflict with other modes, such as ``self._collect_attr1``. + """ + self._unroll_len = self._cfg.collect.unroll_len + self._gamma = self._cfg.collect.discount_factor + + def _forward_collect(self, data: Dict[int, Any]) -> dict: + """ + Overview: + Policy forward function of collect mode (collecting training data by interacting with envs). Forward means \ + that the policy gets some necessary data (mainly observation) from the envs and then returns the output \ + data, such as the action to interact with the envs. + Arguments: + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. + Returns: + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action and \ + other necessary data (action logit) for learn mode defined in ``self._process_transition`` \ + method. The key of the dict is the same as the input data, i.e. environment id. + + .. tip:: + If you want to add more tricks on this policy, like temperature factor in multinomial sample, you can pass \ + related data as extra keyword arguments of this method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + """ + data_id = list(data.keys()) + data = default_collate(list(data.values())) + if self._cuda: + data = to_device(data, self._device) + self._model.eval() + with torch.no_grad(): + output = self._model.forward(data) + output['action'] = output['dist'].sample() + if self._cuda: + output = to_device(output, 'cpu') + output = default_decollate(output) + return {i: d for i, d in zip(data_id, output)} + + def _process_transition(self, obs: Any, model_output: Dict[str, torch.Tensor], timestep: namedtuple) -> dict: + """ + Overview: + Process and pack one timestep transition data into a dict, which can be directly used for training and \ + saved in replay buffer. For PG, it contains obs, action, reward, done. + Arguments: + - obs (:obj:`torch.Tensor`): The env observation of current timestep, such as stacked 2D image in Atari. + - model_output (:obj:`Dict[str, torch.Tensor]`): The output of the policy network with the observation \ + as input. For PG, it contains the action. + - timestep (:obj:`namedtuple`): The execution result namedtuple returned by the environment step method, \ + except all the elements have been transformed into tensor data. Usually, it contains the next obs, \ + reward, done, info, etc. + Returns: + - transition (:obj:`Dict[str, torch.Tensor]`): The processed transition data of the current timestep. + """ + return { + 'obs': obs, + 'action': model_output['action'], + 'reward': timestep.reward, + 'done': timestep.done, + } + + def _get_train_sample(self, data: List[Dict[str, Any]]) -> Union[None, List[Any]]: + """ + Overview: + For a given entire episode data (a list of transition), process it into a list of sample that \ + can be used for training directly. In PG, a train sample is a processed transition with new computed \ + ``return`` field. This method is usually used in collectors to execute necessary \ + RL data preprocessing before training, which can help learner amortize revelant time consumption. \ + In addition, you can also implement this method as an identity function and do the data processing \ + in ``self._forward_learn`` method. + Arguments: + - data (:obj:`List[Dict[str, Any]`): The episode data (a list of transition), each element is \ + the same format as the return value of ``self._process_transition`` method. Note that PG needs \ + a complete epsiode + Returns: + - samples (:obj:`List[Dict[str, Any]]`): The processed train samples, each element is the similar format \ + as input transitions, but may contain more data for training, such as discounted episode return. + """ + assert data[-1]['done'], "PG needs a complete epsiode" + + if self._cfg.learn.ignore_done: + raise NotImplementedError + + R = 0. + if isinstance(data, list): + for i in reversed(range(len(data))): + R = self._gamma * R + data[i]['reward'] + data[i]['return'] = R + return get_train_sample(data, self._unroll_len) + elif isinstance(data, ttorch.Tensor): + data_size = data['done'].shape[0] + data['return'] = ttorch.torch.zeros(data_size) + for i in reversed(range(data_size)): + R = self._gamma * R + data['reward'][i] + data['return'][i] = R + return get_train_sample(data, self._unroll_len) + else: + raise ValueError + + def _init_eval(self) -> None: + pass + + def _forward_eval(self, data: Dict[int, Any]) -> dict: + """ + Overview: + Policy forward function of eval mode (evaluation policy performance by interacting with envs). Forward \ + means that the policy gets some necessary data (mainly observation) from the envs and then returns the \ + action to interact with the envs. ``_forward_eval`` in PG often uses deterministic sample method to get \ + actions while ``_forward_collect`` usually uses stochastic sample method for balance exploration and \ + exploitation. + Arguments: + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. + Returns: + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action. The \ + key of the dict is the same as the input data, i.e. environment id. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for PGPGPolicy: ``ding.policy.tests.test_pg``. + """ + data_id = list(data.keys()) + data = default_collate(list(data.values())) + if self._cuda: + data = to_device(data, self._device) + self._model.eval() + with torch.no_grad(): + output = self._model.forward(data) + if self._cfg.deterministic_eval: + if self._cfg.action_space == 'discrete': + output['action'] = output['logit'].argmax(dim=-1) + elif self._cfg.action_space == 'continuous': + output['action'] = output['logit']['mu'] + else: + raise KeyError("invalid action_space: {}".format(self._cfg.action_space)) + else: + output['action'] = output['dist'].sample() + if self._cuda: + output = to_device(output, 'cpu') + output = default_decollate(output) + return {i: d for i, d in zip(data_id, output)} + + def _monitor_vars_learn(self) -> List[str]: + """ + Overview: + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. + Returns: + - necessary_keys (:obj:`List[str]`): The list of the necessary keys to be logged. + """ + return super()._monitor_vars_learn() + ['policy_loss', 'entropy_loss', 'return_abs_max', 'grad_norm'] diff --git a/ding/policy/plan_diffuser.py b/ding/policy/plan_diffuser.py new file mode 100755 index 0000000000..ad58546a15 --- /dev/null +++ b/ding/policy/plan_diffuser.py @@ -0,0 +1,400 @@ +from typing import List, Dict, Any, Optional, Tuple, Union +from collections import namedtuple, defaultdict +import copy +import numpy as np +import torch +import torch.nn.functional as F +from torch.distributions import Normal, Independent + +from ding.torch_utils import Adam, to_device +from ding.rl_utils import v_1step_td_data, v_1step_td_error, get_train_sample, \ + qrdqn_nstep_td_data, qrdqn_nstep_td_error, get_nstep_return_data +from ding.policy import Policy +from ding.model import model_wrap +from ding.utils import POLICY_REGISTRY, DatasetNormalizer +from ding.utils.data import default_collate, default_decollate +from .common_utils import default_preprocess_learn + + +@POLICY_REGISTRY.register('pd') +class PDPolicy(Policy): + r""" + Overview: + Implicit Plan Diffuser + https://arxiv.org/pdf/2205.09991.pdf + + """ + config = dict( + type='pd', + # (bool) Whether to use cuda for network. + cuda=False, + # (bool type) priority: Determine whether to use priority in buffer sample. + # Default False in SAC. + priority=False, + # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. + priority_IS_weight=False, + # (int) Number of training samples(randomly collected) in replay buffer when training starts. + # Default 10000 in SAC. + random_collect_size=10000, + nstep=1, + # normalizer type + normalizer='GaussianNormalizer', + model=dict( + diffuser_model='GaussianDiffusion', + diffuser_model_cfg=dict( + # the type of model + model='TemporalUnet', + # config of model + model_cfg=dict( + # model dim, In GaussianInvDynDiffusion, it is obs_dim. In others, it is obs_dim + action_dim + transition_dim=23, + dim=32, + dim_mults=[1, 2, 4, 8], + # whether use return as a condition + returns_condition=False, + condition_dropout=0.1, + # whether use calc energy + calc_energy=False, + kernel_size=5, + # whether use attention + attention=False, + ), + # horizon of tarjectory which generated by model + horizon=80, + # timesteps of diffusion + n_timesteps=1000, + # hidden dim of action model + # Whether predict epsilon + predict_epsilon=True, + # discount of loss + loss_discount=1.0, + # whether clip denoise + clip_denoised=False, + action_weight=10, + ), + value_model='ValueDiffusion', + value_model_cfg=dict( + # the type of model + model='TemporalValue', + # config of model + model_cfg=dict( + horizon=4, + # model dim, In GaussianInvDynDiffusion, it is obs_dim. In others, it is obs_dim + action_dim + transition_dim=23, + dim=32, + dim_mults=[1, 2, 4, 8], + # whether use calc energy + kernel_size=5, + ), + # horizon of tarjectory which generated by model + horizon=80, + # timesteps of diffusion + n_timesteps=1000, + # hidden dim of action model + predict_epsilon=True, + # discount of loss + loss_discount=1.0, + # whether clip denoise + clip_denoised=False, + action_weight=1.0, + ), + # guide_steps for p sample + n_guide_steps=2, + # scale of grad for p sample + scale=0.1, + # t of stopgrad for p sample + t_stopgrad=2, + # whether use std as a scale for grad + scale_grad_by_std=True, + ), + learn=dict( + + # How many updates(iterations) to train after collector's one collection. + # Bigger "update_per_collect" means bigger off-policy. + # collect data -> update policy-> collect data -> ... + update_per_collect=1, + # (int) Minibatch size for gradient descent. + batch_size=100, + + # (float type) learning_rate_q: Learning rate for model. + # Default to 3e-4. + # Please set to 1e-3, when model.value_network is True. + learning_rate=3e-4, + # (bool) Whether ignore done(usually for max step termination env. e.g. pendulum) + # Note: Gym wraps the MuJoCo envs by default with TimeLimit environment wrappers. + # These limit HalfCheetah, and several other MuJoCo envs, to max length of 1000. + # However, interaction with HalfCheetah always gets done with done is False, + # Since we inplace done==True with done==False to keep + # TD-error accurate computation(``gamma * (1 - done) * next_v + reward``), + # when the episode step is greater than max episode step. + ignore_done=False, + + # (float type) target_theta: Used for soft update of the target network, + # aka. Interpolation factor in polyak averaging for target networks. + # Default to 0.005. + target_theta=0.005, + # (float) discount factor for the discounted sum of rewards, aka. gamma. + discount_factor=0.99, + gradient_accumulate_every=2, + # train_epoch = train_epoch * gradient_accumulate_every + train_epoch=60000, + # batch_size of every env when eval + plan_batch_size=64, + + # step start update target model and frequence + step_start_update_target=2000, + update_target_freq=10, + # update weight of target net + target_weight=0.995, + value_step=200e3, + + # dataset weight include returns + include_returns=True, + + # (float) Weight uniform initialization range in the last output layer + init_w=3e-3, + ), + ) + + def default_model(self) -> Tuple[str, List[str]]: + return 'pd', ['ding.model.template.diffusion'] + + def _init_learn(self) -> None: + r""" + Overview: + Learn mode init method. Called by ``self.__init__``. + Init q, value and policy's optimizers, algorithm config, main and target models. + """ + # Init + self._priority = self._cfg.priority + self._priority_IS_weight = self._cfg.priority_IS_weight + self.action_dim = self._cfg.model.diffuser_model_cfg.action_dim + self.obs_dim = self._cfg.model.diffuser_model_cfg.obs_dim + self.n_timesteps = self._cfg.model.diffuser_model_cfg.n_timesteps + self.gradient_accumulate_every = self._cfg.learn.gradient_accumulate_every + self.plan_batch_size = self._cfg.learn.plan_batch_size + self.gradient_steps = 1 + self.update_target_freq = self._cfg.learn.update_target_freq + self.step_start_update_target = self._cfg.learn.step_start_update_target + self.target_weight = self._cfg.learn.target_weight + self.value_step = self._cfg.learn.value_step + self.use_target = False + self.horizon = self._cfg.model.diffuser_model_cfg.horizon + self.include_returns = self._cfg.learn.include_returns + + # Optimizers + self._plan_optimizer = Adam( + self._model.diffuser.model.parameters(), + lr=self._cfg.learn.learning_rate, + ) + if self._model.value: + self._value_optimizer = Adam( + self._model.value.model.parameters(), + lr=self._cfg.learn.learning_rate, + ) + + # Algorithm config + self._gamma = self._cfg.learn.discount_factor + + # Main and target models + self._target_model = copy.deepcopy(self._model) + # self._target_model = model_wrap( + # self._target_model, + # wrapper_name='target', + # update_type='momentum', + # update_kwargs={'theta': self._cfg.learn.target_theta} + # ) + self._learn_model = model_wrap(self._model, wrapper_name='base') + self._learn_model.reset() + # self._target_model.reset() + + self._forward_learn_cnt = 0 + + def _forward_learn(self, data: dict) -> Dict[str, Any]: + loss_dict = {} + + data = default_preprocess_learn( + data, + use_priority=self._priority, + use_priority_IS_weight=self._cfg.priority_IS_weight, + ignore_done=self._cfg.learn.ignore_done, + use_nstep=False + ) + + conds = {} + vals = data['condition_val'] + ids = data['condition_id'] + for i in range(len(ids)): + conds[ids[i][0].item()] = vals[i] + if len(ids) > 1: + self.use_target = True + data['conditions'] = conds + if 'returns' in data.keys(): + data['returns'] = data['returns'].unsqueeze(-1) + if self._cuda: + data = to_device(data, self._device) + + self._learn_model.train() + # self._target_model.train() + x = data['trajectories'] + + batch_size = len(x) + t = torch.randint(0, self.n_timesteps, (batch_size, ), device=x.device).long() + cond = data['conditions'] + if 'returns' in data.keys(): + target = data['returns'] + loss_dict['diffuse_loss'], loss_dict['a0_loss'] = self._model.diffuser_loss(x, cond, t) + loss_dict['diffuse_loss'] = loss_dict['diffuse_loss'] / self.gradient_accumulate_every + loss_dict['diffuse_loss'].backward() + if self._forward_learn_cnt < self.value_step and self._model.value: + loss_dict['value_loss'], logs = self._model.value_loss(x, cond, target, t) + loss_dict['value_loss'] = loss_dict['value_loss'] / self.gradient_accumulate_every + loss_dict['value_loss'].backward() + loss_dict.update(logs) + + if self.gradient_steps >= self.gradient_accumulate_every: + self._plan_optimizer.step() + self._plan_optimizer.zero_grad() + if self._forward_learn_cnt < self.value_step and self._model.value: + self._value_optimizer.step() + self._value_optimizer.zero_grad() + self.gradient_steps = 1 + else: + self.gradient_steps += 1 + self._forward_learn_cnt += 1 + if self._forward_learn_cnt % self.update_target_freq == 0: + if self._forward_learn_cnt < self.step_start_update_target: + self._target_model.load_state_dict(self._model.state_dict()) + else: + self.update_model_average(self._target_model, self._learn_model) + + if 'returns' in data.keys(): + loss_dict['max_return'] = target.max().item() + loss_dict['min_return'] = target.min().item() + loss_dict['mean_return'] = target.mean().item() + loss_dict['max_traj'] = x.max().item() + loss_dict['min_traj'] = x.min().item() + loss_dict['mean_traj'] = x.mean().item() + return loss_dict + + def update_model_average(self, ma_model, current_model): + for current_params, ma_params in zip(current_model.parameters(), ma_model.parameters()): + old_weight, up_weight = ma_params.data, current_params.data + if old_weight is None: + ma_params.data = up_weight + else: + old_weight * self.target_weight + (1 - self.target_weight) * up_weight + + def _monitor_vars_learn(self) -> List[str]: + return [ + 'diffuse_loss', + 'value_loss', + 'max_return', + 'min_return', + 'mean_return', + 'max_traj', + 'min_traj', + 'mean_traj', + 'mean_pred', + 'max_pred', + 'min_pred', + 'a0_loss', + ] + + def _state_dict_learn(self) -> Dict[str, Any]: + if self._model.value: + return { + 'model': self._learn_model.state_dict(), + 'target_model': self._target_model.state_dict(), + 'plan_optimizer': self._plan_optimizer.state_dict(), + 'value_optimizer': self._value_optimizer.state_dict(), + } + else: + return { + 'model': self._learn_model.state_dict(), + 'target_model': self._target_model.state_dict(), + 'plan_optimizer': self._plan_optimizer.state_dict(), + } + + def _init_eval(self): + self._eval_model = model_wrap(self._target_model, wrapper_name='base') + self._eval_model.reset() + if self.use_target: + self._plan_seq = [] + + def init_data_normalizer(self, normalizer: DatasetNormalizer = None): + self.normalizer = normalizer + + def _forward_eval(self, data: dict) -> Dict[str, Any]: + data_id = list(data.keys()) + data = default_collate(list(data.values())) + + self._eval_model.eval() + if self.use_target: + cur_obs = self.normalizer.normalize(data[:, :self.obs_dim], 'observations') + target_obs = self.normalizer.normalize(data[:, self.obs_dim:], 'observations') + else: + obs = self.normalizer.normalize(data, 'observations') + with torch.no_grad(): + if self.use_target: + cur_obs = torch.tensor(cur_obs) + target_obs = torch.tensor(target_obs) + if self._cuda: + cur_obs = to_device(cur_obs, self._device) + target_obs = to_device(target_obs, self._device) + conditions = {0: cur_obs, self.horizon - 1: target_obs} + else: + obs = torch.tensor(obs) + if self._cuda: + obs = to_device(obs, self._device) + conditions = {0: obs} + + if self.use_target: + if self._plan_seq == [] or 0 in self._eval_t: + plan_traj = self._eval_model.get_eval(conditions, self.plan_batch_size) + plan_traj = to_device(plan_traj, 'cpu').numpy() + if self._plan_seq == []: + self._plan_seq = plan_traj + self._eval_t = [0] * len(data_id) + else: + for id in data_id: + if self._eval_t[id] == 0: + self._plan_seq[id] = plan_traj[id] + action = [] + for id in data_id: + if self._eval_t[id] < len(self._plan_seq[id]) - 1: + next_waypoint = self._plan_seq[id][self._eval_t[id] + 1] + else: + next_waypoint = self._plan_seq[id][-1].copy() + next_waypoint[2:] = 0 + cur_ob = cur_obs[id] + cur_ob = to_device(cur_ob, 'cpu').numpy() + act = next_waypoint[:2] - cur_ob[:2] + (next_waypoint[2:] - cur_ob[2:]) + action.append(act) + self._eval_t[id] += 1 + else: + action = self._eval_model.get_eval(conditions, self.plan_batch_size) + if self._cuda: + action = to_device(action, 'cpu') + action = self.normalizer.unnormalize(action, 'actions') + action = torch.tensor(action).to('cpu') + output = {'action': action} + output = default_decollate(output) + return {i: d for i, d in zip(data_id, output)} + + def _reset_eval(self, data_id: Optional[List[int]] = None) -> None: + if self.use_target and data_id: + for id in data_id: + self._eval_t[id] = 0 + + def _init_collect(self) -> None: + pass + + def _forward_collect(self, data: dict, **kwargs) -> dict: + pass + + def _process_transition(self, obs: Any, model_output: dict, timestep: namedtuple) -> dict: + pass + + def _get_train_sample(self, data: list) -> Union[None, List[Any]]: + pass diff --git a/ding/policy/policy_factory.py b/ding/policy/policy_factory.py index 3ccdccbb9b..09bdcf840a 100644 --- a/ding/policy/policy_factory.py +++ b/ding/policy/policy_factory.py @@ -1,24 +1,40 @@ from typing import Dict, Any, Callable from collections import namedtuple from easydict import EasyDict +import gym +import gymnasium import torch from ding.torch_utils import to_device -import gym class PolicyFactory: - r""" + """ Overview: - Pure random policy. Only used for initial sample collecting if `cfg.policy.random_collect_size` > 0. + Policy factory class, used to generate different policies for general purpose. Such as random action policy, \ + which is used for initial sample collecting for better exploration when ``random_collect_size`` > 0. + Interfaces: + ``get_random_policy`` """ @staticmethod def get_random_policy( - policy: 'BasePolicy', # noqa + policy: 'Policy.collect_mode', # noqa action_space: 'gym.spaces.Space' = None, # noqa forward_fn: Callable = None, - ) -> None: + ) -> 'Policy.collect_mode': # noqa + """ + Overview: + According to the given action space, define the forward function of the random policy, then pack it with \ + other interfaces of the given policy, and return the final collect mode interfaces of policy. + Arguments: + - policy (:obj:`Policy.collect_mode`): The collect mode interfaces of the policy. + - action_space (:obj:`gym.spaces.Space`): The action space of the environment, gym-style. + - forward_fn (:obj:`Callable`): It action space is too complex, you can define your own forward function \ + and pass it to this function, note you should set ``action_space`` to ``None`` in this case. + Returns: + - random_policy (:obj:`Policy.collect_mode`): The collect mode intefaces of the random policy. + """ assert not (action_space is None and forward_fn is None) random_collect_function = namedtuple( 'random_collect_function', [ @@ -34,23 +50,35 @@ def forward(data: Dict[int, Any], *args, **kwargs) -> Dict[int, Any]: actions = {} for env_id in data: - if not isinstance(action_space, list): - action = torch.as_tensor(action_space.sample()) - if isinstance(action_space, gym.spaces.MultiDiscrete): - action = [torch.LongTensor([v]) for v in action] - actions[env_id] = {'action': action} - elif 'global_state' in data[env_id].keys(): - # for smac - logit = torch.ones_like(data[env_id]['action_mask']) - logit[data[env_id]['action_mask'] == 0.0] = -1e8 - dist = torch.distributions.categorical.Categorical(logits=torch.Tensor(logit)) - actions[env_id] = {'action': dist.sample(), 'logit': torch.as_tensor(logit)} - else: - # for gfootball + if isinstance(action_space, list): + if 'global_state' in data[env_id].keys(): + # for smac + logit = torch.ones_like(data[env_id]['action_mask']) + logit[data[env_id]['action_mask'] == 0.0] = -1e8 + dist = torch.distributions.categorical.Categorical(logits=torch.Tensor(logit)) + actions[env_id] = {'action': dist.sample(), 'logit': torch.as_tensor(logit)} + else: + # for gfootball + actions[env_id] = { + 'action': torch.as_tensor( + [action_space_agent.sample() for action_space_agent in action_space] + ), + 'logit': torch.ones([len(action_space), action_space[0].n]) + } + elif isinstance(action_space, gymnasium.spaces.Dict): # pettingzoo actions[env_id] = { - 'action': torch.as_tensor([action_space_agent.sample() for action_space_agent in action_space]), - 'logit': torch.ones([len(action_space), action_space[0].n]) + 'action': torch.as_tensor( + [action_space_agent.sample() for action_space_agent in action_space.values()] + ) } + else: + if isinstance(action_space, gym.spaces.Discrete): + action = torch.LongTensor([action_space.sample()]) + elif isinstance(action_space, gym.spaces.MultiDiscrete): + action = [torch.LongTensor([v]) for v in action_space.sample()] + else: + action = torch.as_tensor(action_space.sample()) + actions[env_id] = {'action': action} return actions def reset(*args, **kwargs) -> None: @@ -66,7 +94,23 @@ def reset(*args, **kwargs) -> None: ) -def get_random_policy(cfg: EasyDict, policy: 'Policy.collect_mode', env: 'BaseEnvManager'): # noqa +def get_random_policy( + cfg: EasyDict, + policy: 'Policy.collect_mode', # noqa + env: 'BaseEnvManager' # noqa +) -> 'Policy.collect_mode': # noqa + """ + Overview: + The entry function to get the corresponding random policy. If a policy needs special data items in a \ + transition, then return itself, otherwise, we will use ``PolicyFactory`` to return a general random policy. + Arguments: + - cfg (:obj:`EasyDict`): The EasyDict-type dict configuration. + - policy (:obj:`Policy.collect_mode`): The collect mode interfaces of the policy. + - env (:obj:`BaseEnvManager`): The env manager instance, which is used to get the action space for random \ + action generation. + Returns: + - random_policy (:obj:`Policy.collect_mode`): The collect mode intefaces of the random policy. + """ if cfg.policy.get('transition_with_policy_data', False): return policy else: diff --git a/ding/policy/ppg.py b/ding/policy/ppg.py index 436fef788f..e9732b46c6 100644 --- a/ding/policy/ppg.py +++ b/ding/policy/ppg.py @@ -14,8 +14,20 @@ class ExperienceDataset(Dataset): + """ + Overview: + A dataset class for storing and accessing experience data. + + Interface: + ``__init__``, ``__len__``, ``__getitem__``. + """ def __init__(self, data): + """ + Arguments: + - data (:obj:`dict`): A dictionary containing the experience data, where the keys represent the data types \ + and the values are the corresponding data arrays. + """ super().__init__() self.data = data @@ -38,11 +50,16 @@ def create_shuffled_dataloader(data, batch_size): class PPGPolicy(Policy): """ Overview: - Policy class of PPG algorithm. + Policy class of PPG algorithm. PPG is a policy gradient algorithm with auxiliary phase training. \ + The auxiliary phase training is proposed to distill the value into the policy network, \ + while making sure the policy network does not change the action predictions (kl div loss). \ + Paper link: https://arxiv.org/abs/2009.04416. + Interface: - _init_learn, _data_preprocess_learn, _forward_learn, _state_dict_learn, _load_state_dict_learn\ - _init_collect, _forward_collect, _process_transition, _get_train_sample, _get_batch_size, _init_eval,\ - _forward_eval, default_model, _monitor_vars_learn, learn_aux + ``_init_learn``, ``_data_preprocess_learn``, ``_forward_learn``, ``_state_dict_learn``, \ + ``_load_state_dict_learn``, ``_init_collect``, ``_forward_collect``, ``_process_transition``, \ + ``_get_train_sample``, ``_get_batch_size``, ``_init_eval``, ``_forward_eval``, ``default_model``, \ + ``_monitor_vars_learn``, ``learn_aux``. Config: == ==================== ======== ============== ======================================== ======================= ID Symbol Type Default Value Description Other(Shape) @@ -91,8 +108,6 @@ class PPGPolicy(Policy): # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. priority_IS_weight=False, learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, actor_epoch_per_collect=1, critic_epoch_per_collect=1, batch_size=64, @@ -136,27 +151,32 @@ class PPGPolicy(Policy): def default_model(self) -> Tuple[str, List[str]]: """ Overview: - Return this algorithm default model setting for demonstration. - Returns: - - model_info (:obj:`Tuple[str, List[str]]`): model name and mode import_names + Return this algorithm default neural network model setting for demonstration. ``__init__`` method will \ + automatically call this method to get the default model setting and create model. - .. note:: - The user can define and use customized network model but must obey the same inferface definition indicated \ - by import_names path. + Returns: + - model_info (:obj:`Tuple[str, List[str]]`): The registered model name and model's import_names. """ return 'ppg', ['ding.model.template.ppg'] def _init_learn(self) -> None: - r""" + """ Overview: - Learn mode init method. Called by ``self.__init__``. - Init the optimizer, algorithm config and the main model. - Arguments: - .. note:: + Initialize the learn mode of policy, including related attributes and modules. For PPG, it mainly \ + contains optimizer, algorithm-specific arguments such as aux_bc_weight and aux_train_epoch. This method \ + also executes some special network initializations and prepares running mean/std monitor for value. \ + This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. - The _init_learn method takes the argument from the self._cfg.learn in the config file + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. - - learning_rate (:obj:`float`): The learning rate fo the optimizer + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. """ # Optimizer self._optimizer_ac = Adam(self._model.actor_critic.parameters(), lr=self._cfg.learn.learning_rate) @@ -191,9 +211,9 @@ def _data_preprocess_learn(self, data: List[Any]) -> dict: collate(stack data into batch), ignore done(in some fake terminate env),\ prepare loss weight per training sample, and cpu tensor to cuda. Arguments: - - data (:obj:`List[Dict[str, Any]]`): the data collected from collect function + - data (:obj:`List[Dict[str, Any]]`): The data collected from collect function. Returns: - - data (:obj:`Dict[str, Any]`): the processed data, including at least ['done', 'weight'] + - data (:obj:`Dict[str, Any]`): The processed data, including at least ['done', 'weight']. """ # data preprocess data = default_collate(data) @@ -212,30 +232,27 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: Overview: Forward and backward function of learn mode. Arguments: - - data (:obj:`Dict[str, Any]`): Dict type data, a batch of data for training, values are torch.Tensor or \ - np.ndarray or dict/list combinations. + - data (:obj:`Dict[str, Any]`): Input data used for policy forward, including the \ + collected training samples from replay buffer. For each element in dict, the key of the \ + dict is the name of data items and the value is the corresponding data. Usually, the value is \ + torch.Tensor or np.ndarray or there dict/list combinations. In the ``_forward_learn`` method, data \ + often need to first be stacked in the batch dimension by some utility functions such as \ + ``default_preprocess_learn``. \ + For PPG, each element in list is a dict containing at least the following keys: ``obs``, ``action``, \ + ``reward``, ``logit``, ``value``, ``done``. Sometimes, it also contains other keys such as ``weight``. Returns: - info_dict (:obj:`Dict[str, Any]`): Dict type data, a info dict indicated training result, which will be \ - recorded in text log and tensorboard, values are python scalar or a list of scalars. - ArgumentsKeys: - - necessary: 'obs', 'logit', 'action', 'value', 'reward', 'done' - ReturnsKeys: - - necessary: current lr, total_loss, policy_loss, value_loss, entropy_loss, \ - adv_abs_max, approx_kl, clipfrac\ - aux_value_loss, auxiliary_loss, behavioral_cloning_loss - - - current_lr (:obj:`float`): Current learning rate - - total_loss (:obj:`float`): The calculated loss - - policy_loss (:obj:`float`): The policy(actor) loss of ppg - - value_loss (:obj:`float`): The value(critic) loss of ppg - - entropy_loss (:obj:`float`): The entropy loss - - auxiliary_loss (:obj:`float`): The auxiliary loss, we use the value function loss \ - as the auxiliary objective, thereby sharing features between the policy and value function\ - while minimizing distortions to the policy - - aux_value_loss (:obj:`float`): The auxiliary value loss, we need to train the value network extra \ - during the auxiliary phase, it's the value loss we train the value network during auxiliary phase - - behavioral_cloning_loss (:obj:`float`): The behavioral cloning loss, used to optimize the auxiliary\ - objective while otherwise preserving the original policy + recorded in text log and tensorboard, values are python scalar or a list of scalars. \ + For the detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for PPGPolicy: ``ding.policy.tests.test_ppgs``. """ data = self._data_preprocess_learn(data) # ==================== @@ -260,7 +277,7 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: policy_output = self._learn_model.forward(policy_data['obs'], mode='compute_actor') policy_error_data = ppo_policy_data( policy_output['logit'], policy_data['logit'], policy_data['action'], policy_adv, - policy_data['weight'] + policy_data['weight'], None ) ppo_policy_loss, ppo_info = ppo_policy_error(policy_error_data, self._clip_ratio) policy_loss = ppo_policy_loss.policy_loss - self._entropy_weight * ppo_policy_loss.entropy_loss @@ -327,7 +344,7 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: } def _state_dict_learn(self) -> Dict[str, Any]: - r""" + """ Overview: Return the state_dict of learn mode, usually including model and optimizer. Returns: @@ -359,10 +376,16 @@ def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: self._optimizer_aux_critic.load_state_dict(state_dict['optimizer_aux_critic']) def _init_collect(self) -> None: - r""" + """ Overview: - Collect mode init method. Called by ``self.__init__``. - Init unroll length, collect model. + Initialize the collect mode of policy, including related attributes and modules. For PPG, it contains the \ + collect_model to balance the exploration and exploitation (e.g. the multinomial sample mechanism in \ + discrete action space), and other algorithm-specific arguments such as unroll_len and gae_lambda. + This method will be called in ``__init__`` method if ``collect`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_collect`` method, you'd better name them \ + with prefix ``_collect_`` to avoid conflict with other modes, such as ``self._collect_attr1``. """ self._unroll_len = self._cfg.collect.unroll_len self._collect_model = model_wrap(self._model, wrapper_name='multinomial_sample') @@ -372,16 +395,34 @@ def _init_collect(self) -> None: self._gae_lambda = self._cfg.collect.gae_lambda def _forward_collect(self, data: dict) -> dict: - r""" + """ Overview: - Forward function of collect mode. + Policy forward function of collect mode (collecting training data by interacting with envs). Forward means \ + that the policy gets some necessary data (mainly observation) from the envs and then returns the output \ + data, such as the action to interact with the envs. + Arguments: - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \ values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer. + Returns: - - output (:obj:`Dict[int, Any]`): Dict type data, including at least inferred action according to input obs. - ReturnsKeys - - necessary: ``action`` + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action and \ + other necessary data (action logit and value) for learn mode defined in \ + ``self._process_transition`` method. The key of the dict is the same as the input data, \ + i.e. environment id. + + .. tip:: + If you want to add more tricks on this policy, like temperature factor in multinomial sample, you can pass \ + related data as extra keyword arguments of this method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for PPGPolicy: ``ding.policy.tests.test_ppg``. """ data_id = list(data.keys()) data = default_collate(list(data.values())) @@ -398,14 +439,21 @@ def _forward_collect(self, data: dict) -> dict: def _process_transition(self, obs: Any, model_output: dict, timestep: namedtuple) -> dict: """ Overview: - Generate dict type transition data from inputs. + Process and pack one timestep transition data into a dict, which can be directly used for training and \ + saved in replay buffer. For PPG, it contains obs, next_obs, action, reward, done, logit, value. Arguments: - - obs (:obj:`Any`): Env observation - - model_output (:obj:`dict`): Output of collect model, including at least ['action'] - - timestep (:obj:`namedtuple`): Output after env step, including at least ['obs', 'reward', 'done']\ - (here 'obs' indicates obs after env step). + - obs (:obj:`Any`): Env observation + - model_output (:obj:`dict`): The output of the policy network with the observation \ + as input. For PPG, it contains the state value, action and the logit of the action. + - timestep (:obj:`namedtuple`): The execution result namedtuple returned by the environment step \ + method, except all the elements have been transformed into tensor data. Usually, it contains the next \ + obs, reward, done, info, etc. Returns: - - transition (:obj:`dict`): Dict type transition data. + - transition (:obj:`dict`): The processed transition data of the current timestep. + + .. note:: + ``next_obs`` is used to calculate nstep return when necessary, so we place in into transition by default. \ + You can delete this field to save memory occupancy if you do not need nstep return. """ transition = { 'obs': obs, @@ -418,14 +466,21 @@ def _process_transition(self, obs: Any, model_output: dict, timestep: namedtuple } return transition - def _get_train_sample(self, data: list) -> Union[None, List[Any]]: - r""" + def _get_train_sample(self, data: List[Dict[str, Any]]) -> Union[None, List[Any]]: + """ Overview: - Get the trajectory and calculate GAE, return one data to cache for next time calculation + For a given trajectory (transitions, a list of transition) data, process it into a list of sample that \ + can be used for training directly. In PPG, a train sample is a processed transition with new computed \ + ``adv`` field. This method is usually used in collectors to execute necessary. \ + RL data preprocessing before training, which can help learner amortize revelant time consumption. \ + In addition, you can also implement this method as an identity function and do the data processing \ + in ``self._forward_learn`` method. Arguments: - - data (:obj:`list`): The trajectory's cache + - data (:obj:`List[Dict[str, Any]]`): The trajectory data (a list of transition), each element is \ + the same format as the return value of ``self._process_transition`` method. Returns: - - samples (:obj:`dict`): The training samples generated + - samples (:obj:`dict`): The processed train samples, each element is the similar format \ + as input transitions, but may contain more data for training, such as GAE advantage. """ data = to_device(data, self._device) if self._cfg.learn.ignore_done: @@ -468,25 +523,43 @@ def _get_batch_size(self) -> Dict[str, int]: return {'policy': bs, 'value': bs} def _init_eval(self) -> None: - r""" + """ Overview: - Evaluate mode init method. Called by ``self.__init__``. - Init eval model with argmax strategy. + Initialize the eval mode of policy, including related attributes and modules. For PPG, it contains the \ + eval model to select optimial action (e.g. greedily select action with argmax mechanism in discrete \ + action). This method will be called in ``__init__`` method if ``eval`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_eval`` method, you'd better name them \ + with prefix ``_eval_`` to avoid conflict with other modes, such as ``self._eval_attr1``. """ self._eval_model = model_wrap(self._model, wrapper_name='argmax_sample') self._eval_model.reset() def _forward_eval(self, data: dict) -> dict: - r""" + """ Overview: - Forward function of eval mode, similar to ``self._forward_collect``. + Policy forward function of eval mode (evaluation policy performance by interacting with envs). Forward \ + means that the policy gets some necessary data (mainly observation) from the envs and then returns the \ + action to interact with the envs. ``_forward_eval`` in PPG often uses deterministic sample method to get \ + actions while ``_forward_collect`` usually uses stochastic sample method for balance exploration and \ + exploitation. Arguments: - - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \ - values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer. + - data (:obj:`Dict[str, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. + Returns: - - output (:obj:`Dict[int, Any]`): The dict of predicting action for the interaction with env. - ReturnsKeys - - necessary: ``action`` + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action. The \ + key of the dict is the same as the input data, i.e. environment id. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for PPGPolicy: ``ding.policy.tests.test_ppg``. """ data_id = list(data.keys()) data = default_collate(list(data.values())) @@ -501,11 +574,12 @@ def _forward_eval(self, data: dict) -> dict: return {i: d for i, d in zip(data_id, output)} def _monitor_vars_learn(self) -> List[str]: - r""" + """ Overview: - Return variables' name if variables are to used in monitor. + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. Returns: - - vars (:obj:`List[str]`): Variables' name list. + - vars (:obj:`List[str]`): The list of the necessary keys to be logged. """ return [ 'policy_cur_lr', @@ -524,10 +598,13 @@ def _monitor_vars_learn(self) -> List[str]: def learn_aux(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ Overview: - The auxiliary phase training, where the value is distilled into the policy network + The auxiliary phase training, where the value is distilled into the policy network. In PPG algorithm, \ + we use the value function loss as the auxiliary objective, thereby sharing features between the policy \ + and value function while minimizing distortions to the policy. We also use behavioral cloning loss to \ + optimize the auxiliary objective while otherwise preserving the original policy. Returns: - - aux_loss (:obj:`Tuple[torch.Tensor, torch.Tensor, torch.Tensor]`): including average auxiliary loss\ - average behavioral cloning loss, and average auxiliary value loss + - aux_loss (:obj:`Tuple[torch.Tensor, torch.Tensor, torch.Tensor]`): Including average auxiliary loss\ + average behavioral cloning loss, and average auxiliary value loss. """ aux_memories = self._aux_memories # gather states and target values into one tensor @@ -617,11 +694,15 @@ def learn_aux(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: class PPGOffPolicy(Policy): """ Overview: - Policy class of PPG algorithm. + Policy class of PPG algorithm with off-policy training mode. Off-policy PPG contains two different data \ + max_use buffers. The policy buffer offers data for policy phase , while the value buffer provides auxiliary \ + phase's data. The whole training procedure is similar to off-policy PPO but execute additional auxiliary \ + phase with a fixed frequency. Interface: - _init_learn, _data_preprocess_learn, _forward_learn, _state_dict_learn, _load_state_dict_learn\ - _init_collect, _forward_collect, _process_transition, _get_train_sample, _get_batch_size, _init_eval,\ - _forward_eval, default_model, _monitor_vars_learn, learn_aux + ``_init_learn``, ``_data_preprocess_learn``, ``_forward_learn``, ``_state_dict_learn``, \ + ``_load_state_dict_learn``, ``_init_collect``, ``_forward_collect``, ``_process_transition``, \ + ``_get_train_sample``, ``_get_batch_size``, ``_init_eval``, ``_forward_eval``, ``default_model``, \ + ``_monitor_vars_learn``, ``learn_aux``. Config: == ==================== ======== ============== ======================================== ======================= ID Symbol Type Default Value Description Other(Shape) @@ -672,8 +753,6 @@ class PPGOffPolicy(Policy): # (bool) Whether to need policy data in process transition transition_with_policy_data=True, learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, update_per_collect=5, batch_size=64, learning_rate=0.001, @@ -721,9 +800,11 @@ class PPGOffPolicy(Policy): def default_model(self) -> Tuple[str, List[str]]: """ Overview: - Return this algorithm default model setting for demonstration. + Return this algorithm default neural network model setting for demonstration. ``__init__`` method will \ + automatically call this method to get the default model setting and create model. + Returns: - - model_info (:obj:`Tuple[str, List[str]]`): model name and mode import_names + - model_info (:obj:`Tuple[str, List[str]]`): The registered model name and model's import_names. .. note:: The user can define and use customized network model but must obey the same inferface definition indicated \ @@ -732,16 +813,23 @@ def default_model(self) -> Tuple[str, List[str]]: return 'ppg', ['ding.model.template.ppg'] def _init_learn(self) -> None: - r""" + """ Overview: - Learn mode init method. Called by ``self.__init__``. - Init the optimizer, algorithm config and the main model. - Arguments: - .. note:: + Initialize the learn mode of policy, including related attributes and modules. For PPG, it mainly \ + contains optimizer, algorithm-specific arguments such as aux_bc_weight and aux_train_epoch. This method \ + also executes some special network initializations and prepares running mean/std monitor for value. \ + This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. - The _init_learn method takes the argument from the self._cfg.learn in the config file + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. - - learning_rate (:obj:`float`): The learning rate fo the optimizer + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. """ # Optimizer self._optimizer_ac = Adam(self._model.actor_critic.parameters(), lr=self._cfg.learn.learning_rate) @@ -773,9 +861,9 @@ def _data_preprocess_learn(self, data: List[Any]) -> dict: collate(stack data into batch), ignore done(in some fake terminate env),\ prepare loss weight per training sample, and cpu tensor to cuda. Arguments: - - data (:obj:`List[Dict[str, Any]]`): the data collected from collect function + - data (:obj:`List[Dict[str, Any]]`): The data collected from collect function. Returns: - - data (:obj:`Dict[str, Any]`): the processed data, including at least ['done', 'weight'] + - data (:obj:`Dict[str, Any]`): The processed data, including at least ['done', 'weight']. """ # data preprocess for k, data_item in data.items(): @@ -796,30 +884,38 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: Overview: Forward and backward function of learn mode. Arguments: - - data (:obj:`Dict[str, Any]`): Dict type data, a batch of data for training, values are torch.Tensor or \ - np.ndarray or dict/list combinations. + - data (:obj:`Dict[str, Any]`): Input data used for policy forward, including the \ + collected training samples from replay buffer. For each element in dict, the key of the \ + dict is the name of data items and the value is the corresponding data. Usually, \ + the class type of value is either torch.Tensor or np.ndarray, or a dict/list containing \ + either torch.Tensor or np.ndarray items In the ``_forward_learn`` method, data \ + often need to first be stacked in the batch dimension by some utility functions such as \ + ``default_preprocess_learn``. \ + For PPGOff, each element in list is a dict containing at least the following keys: ``obs``, \ + ``action``, ``reward``, ``logit``, ``value``, ``done``. Sometimes, it also contains other keys \ + such as ``weight``. Returns: - info_dict (:obj:`Dict[str, Any]`): Dict type data, a info dict indicated training result, which will be \ - recorded in text log and tensorboard, values are python scalar or a list of scalars. - ArgumentsKeys: - - necessary: 'obs', 'logit', 'action', 'value', 'reward', 'done' + recorded in text log and tensorboard, values are python scalar or a list of scalars. \ + For the detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. + ReturnsKeys: - - necessary: current lr, total_loss, policy_loss, value_loss, entropy_loss, \ - adv_abs_max, approx_kl, clipfrac\ - aux_value_loss, auxiliary_loss, behavioral_cloning_loss - - - current_lr (:obj:`float`): Current learning rate - - total_loss (:obj:`float`): The calculated loss - - policy_loss (:obj:`float`): The policy(actor) loss of ppg - - value_loss (:obj:`float`): The value(critic) loss of ppg - - entropy_loss (:obj:`float`): The entropy loss + - necessary: "current lr", "total_loss", "policy_loss", "value_loss", "entropy_loss", \ + "adv_abs_max", "approx_kl", "clipfrac", \ + "aux_value_loss", "auxiliary_loss", "behavioral_cloning_loss". + + - current_lr (:obj:`float`): Current learning rate. + - total_loss (:obj:`float`): The calculated loss. + - policy_loss (:obj:`float`): The policy(actor) loss of ppg. + - value_loss (:obj:`float`): The value(critic) loss of ppg. + - entropy_loss (:obj:`float`): The entropy loss. - auxiliary_loss (:obj:`float`): The auxiliary loss, we use the value function loss \ as the auxiliary objective, thereby sharing features between the policy and value function\ - while minimizing distortions to the policy + while minimizing distortions to the policy. - aux_value_loss (:obj:`float`): The auxiliary value loss, we need to train the value network extra \ - during the auxiliary phase, it's the value loss we train the value network during auxiliary phase + during the auxiliary phase, it's the value loss we train the value network during auxiliary phase. - behavioral_cloning_loss (:obj:`float`): The behavioral cloning loss, used to optimize the auxiliary\ - objective while otherwise preserving the original policy + objective while otherwise preserving the original policy. """ data = self._data_preprocess_learn(data) # ==================== @@ -836,7 +932,7 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: # Policy Phase(Policy) policy_output = self._learn_model.forward(policy_data['obs'], mode='compute_actor') policy_error_data = ppo_policy_data( - policy_output['logit'], policy_data['logit'], policy_data['action'], policy_adv, policy_data['weight'] + policy_output['logit'], policy_data['logit'], policy_data['action'], policy_adv, policy_data['weight'], None ) ppo_policy_loss, ppo_info = ppo_policy_error(policy_error_data, self._clip_ratio) policy_loss = ppo_policy_loss.policy_loss - self._entropy_weight * ppo_policy_loss.entropy_loss @@ -896,7 +992,7 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: } def _state_dict_learn(self) -> Dict[str, Any]: - r""" + """ Overview: Return the state_dict of learn mode, usually including model and optimizer. Returns: @@ -928,10 +1024,16 @@ def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: self._optimizer_aux_critic.load_state_dict(state_dict['optimizer_aux_critic']) def _init_collect(self) -> None: - r""" + """ Overview: - Collect mode init method. Called by ``self.__init__``. - Init unroll length, collect model. + Initialize the collect mode of policy, including related attributes and modules. For PPO, it contains the \ + collect_model to balance the exploration and exploitation (e.g. the multinomial sample mechanism in \ + discrete action space), and other algorithm-specific arguments such as unroll_len and gae_lambda. + This method will be called in ``__init__`` method if ``collect`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_collect`` method, you'd better name them \ + with prefix ``_collect_`` to avoid conflict with other modes, such as ``self._collect_attr1``. """ self._unroll_len = self._cfg.collect.unroll_len self._collect_model = model_wrap(self._model, wrapper_name='multinomial_sample') @@ -941,16 +1043,34 @@ def _init_collect(self) -> None: self._gae_lambda = self._cfg.collect.gae_lambda def _forward_collect(self, data: dict) -> dict: - r""" + """ Overview: - Forward function of collect mode. + Policy forward function of collect mode (collecting training data by interacting with envs). Forward means \ + that the policy gets some necessary data (mainly observation) from the envs and then returns the output \ + data, such as the action to interact with the envs. + Arguments: - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \ values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer. + Returns: - - output (:obj:`Dict[int, Any]`): Dict type data, including at least inferred action according to input obs. - ReturnsKeys - - necessary: ``action`` + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action and \ + other necessary data (action logit and value) for learn mode defined in \ + ``self._process_transition`` method. The key of the dict is the same as the input data, \ + i.e. environment id. + + .. tip:: + If you want to add more tricks on this policy, like temperature factor in multinomial sample, you can pass \ + related data as extra keyword arguments of this method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for PPGOffPolicy: ``ding.policy.tests.test_ppg``. """ data_id = list(data.keys()) data = default_collate(list(data.values())) @@ -967,14 +1087,21 @@ def _forward_collect(self, data: dict) -> dict: def _process_transition(self, obs: Any, model_output: dict, timestep: namedtuple) -> dict: """ Overview: - Generate dict type transition data from inputs. + Process and pack one timestep transition data into a dict, which can be directly used for training and \ + saved in replay buffer. For PPG, it contains obs, next_obs, action, reward, done, logit, value. Arguments: - obs (:obj:`Any`): Env observation - - model_output (:obj:`dict`): Output of collect model, including at least ['action'] - - timestep (:obj:`namedtuple`): Output after env step, including at least ['obs', 'reward', 'done']\ - (here 'obs' indicates obs after env step). + - model_output (:obj:`dict`): The output of the policy network with the observation \ + as input. For PPG, it contains the state value, action and the logit of the action. + - timestep (:obj:`namedtuple`): The execution result namedtuple returned by the environment step \ + method, except all the elements have been transformed into tensor data. Usually, it contains the next \ + obs, reward, done, info, etc. Returns: - - transition (:obj:`dict`): Dict type transition data. + - transition (:obj:`dict`): The processed transition data of the current timestep. + + .. note:: + ``next_obs`` is used to calculate nstep return when necessary, so we place in into transition by default. \ + You can delete this field to save memory occupancy if you do not need nstep return. """ transition = { 'obs': obs, @@ -988,13 +1115,20 @@ def _process_transition(self, obs: Any, model_output: dict, timestep: namedtuple return transition def _get_train_sample(self, data: list) -> Union[None, List[Any]]: - r""" + """ Overview: - Get the trajectory and calculate GAE, return one data to cache for next time calculation + For a given trajectory (transitions, a list of transition) data, process it into a list of sample that \ + can be used for training directly. In PPG, a train sample is a processed transition with new computed \ + ``adv`` field. This method is usually used in collectors to execute necessary. \ + RL data preprocessing before training, which can help learner amortize revelant time consumption. \ + In addition, you can also implement this method as an identity function and do the data processing \ + in ``self._forward_learn`` method. Arguments: - - data (:obj:`list`): The trajectory's cache + - data (:obj:`list`): The trajectory data (a list of transition), each element is \ + the same format as the return value of ``self._process_transition`` method. Returns: - - samples (:obj:`dict`): The training samples generated + - samples (:obj:`dict`): The processed train samples, each element is the similar format \ + as input transitions, but may contain more data for training, such as GAE advantage. """ data = get_gae_with_default_last_value( data, @@ -1021,10 +1155,15 @@ def _get_batch_size(self) -> Dict[str, int]: return {'policy': bs, 'value': bs} def _init_eval(self) -> None: - r""" + """ Overview: - Evaluate mode init method. Called by ``self.__init__``. - Init eval model with argmax strategy. + Initialize the eval mode of policy, including related attributes and modules. For PPG, it contains the \ + eval model to select optimial action (e.g. greedily select action with argmax mechanism in discrete \ + action). This method will be called in ``__init__`` method if ``eval`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_eval`` method, you'd better name them \ + with prefix ``_eval_`` to avoid conflict with other modes, such as ``self._eval_attr1``. """ self._eval_model = model_wrap(self._model, wrapper_name='argmax_sample') self._eval_model.reset() @@ -1032,14 +1171,27 @@ def _init_eval(self) -> None: def _forward_eval(self, data: dict) -> dict: r""" Overview: - Forward function of eval mode, similar to ``self._forward_collect``. + Policy forward function of eval mode (evaluation policy performance by interacting with envs). Forward \ + means that the policy gets some necessary data (mainly observation) from the envs and then returns the \ + action to interact with the envs. ``_forward_eval`` in PPG often uses deterministic sample method to get \ + actions while ``_forward_collect`` usually uses stochastic sample method for balance exploration and \ + exploitation. Arguments: - - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \ - values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer. + - data (:obj:`Dict[str, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. + Returns: - - output (:obj:`Dict[int, Any]`): The dict of predicting action for the interaction with env. - ReturnsKeys - - necessary: ``action`` + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action. The \ + key of the dict is the same as the input data, i.e. environment id. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for PPGOffPolicy: ``ding.policy.tests.test_ppg``. """ data_id = list(data.keys()) data = default_collate(list(data.values())) @@ -1054,11 +1206,12 @@ def _forward_eval(self, data: dict) -> dict: return {i: d for i, d in zip(data_id, output)} def _monitor_vars_learn(self) -> List[str]: - r""" + """ Overview: - Return variables' name if variables are to used in monitor. + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. Returns: - - vars (:obj:`List[str]`): Variables' name list. + - vars (:obj:`List[str]`): The list of the necessary keys to be logged. """ return [ 'policy_cur_lr', @@ -1077,10 +1230,13 @@ def _monitor_vars_learn(self) -> List[str]: def learn_aux(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ Overview: - The auxiliary phase training, where the value is distilled into the policy network + The auxiliary phase training, where the value is distilled into the policy network. In PPG algorithm, \ + we use the value function loss as the auxiliary objective, thereby sharing features between the policy \ + and value function while minimizing distortions to the policy. We also use behavioral cloning loss to \ + optimize the auxiliary objective while otherwise preserving the original policy. Returns: - - aux_loss (:obj:`Tuple[torch.Tensor, torch.Tensor, torch.Tensor]`): including average auxiliary loss\ - average behavioral cloning loss, and average auxiliary value loss + - aux_loss (:obj:`Tuple[torch.Tensor, torch.Tensor, torch.Tensor]`): Including average auxiliary loss\ + average behavioral cloning loss, and average auxiliary value loss. """ aux_memories = self._aux_memories # gather states and target values into one tensor diff --git a/ding/policy/ppo.py b/ding/policy/ppo.py index 25746ecfe9..006df6961a 100644 --- a/ding/policy/ppo.py +++ b/ding/policy/ppo.py @@ -4,7 +4,7 @@ import copy import numpy as np -from ding.torch_utils import Adam, to_device, unsqueeze, ContrastiveLoss +from ding.torch_utils import Adam, to_device, to_dtype, unsqueeze, ContrastiveLoss from ding.rl_utils import ppo_data, ppo_error, ppo_policy_error, ppo_policy_data, get_gae_with_default_last_value, \ v_nstep_td_data, v_nstep_td_error, get_nstep_return_data, get_train_sample, gae, gae_data, ppo_error_continuous, \ get_gae, ppo_policy_error_continuous @@ -17,9 +17,9 @@ @POLICY_REGISTRY.register('ppo') class PPOPolicy(Policy): - r""" + """ Overview: - Policy class of on policy version PPO algorithm. + Policy class of on-policy version PPO algorithm. Paper link: https://arxiv.org/abs/1707.06347. """ config = dict( # (str) RL policy register name (refer to function "POLICY_REGISTRY"). @@ -28,70 +28,125 @@ class PPOPolicy(Policy): cuda=False, # (bool) Whether the RL algorithm is on-policy or off-policy. (Note: in practice PPO can be off-policy used) on_policy=True, - # (bool) Whether to use priority(priority sample, IS weight, update priority) + # (bool) Whether to use priority (priority sample, IS weight, update priority). priority=False, # (bool) Whether to use Importance Sampling Weight to correct biased update due to priority. # If True, priority must be True. priority_IS_weight=False, - # (bool) Whether to recompurete advantages in each iteration of on-policy PPO + # (bool) Whether to recompurete advantages in each iteration of on-policy PPO. recompute_adv=True, # (str) Which kind of action space used in PPOPolicy, ['discrete', 'continuous', 'hybrid'] action_space='discrete', - # (bool) Whether to use nstep return to calculate value target, otherwise, use return = adv + value + # (bool) Whether to use nstep return to calculate value target, otherwise, use return = adv + value. nstep_return=False, - # (bool) Whether to enable multi-agent training, i.e.: MAPPO + # (bool) Whether to enable multi-agent training, i.e.: MAPPO. multi_agent=False, - # (bool) Whether to need policy data in process transition + # (bool) Whether to need policy ``_forward_collect`` output data in process transition. transition_with_policy_data=True, + # learn_mode config learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, + # (int) After collecting n_sample/n_episode data, how many epoches to train models. + # Each epoch means the one entire passing of training data. epoch_per_collect=10, + # (int) How many samples in a training batch. batch_size=64, + # (float) The step size of gradient descent. learning_rate=3e-4, - # ============================================================== - # The following configs is algorithm-specific - # ============================================================== - # (float) The loss weight of value network, policy network weight is set to 1 + # (dict or None) The learning rate decay. + # If not None, should contain key 'epoch_num' and 'min_lr_lambda'. + # where 'epoch_num' is the total epoch num to decay the learning rate to min value, + # 'min_lr_lambda' is the final decayed learning rate. + lr_scheduler=None, + # (float) The loss weight of value network, policy network weight is set to 1. value_weight=0.5, - # (float) The loss weight of entropy regularization, policy network weight is set to 1 + # (float) The loss weight of entropy regularization, policy network weight is set to 1. entropy_weight=0.0, - # (float) PPO clip ratio, defaults to 0.2 + # (float) PPO clip ratio, defaults to 0.2. clip_ratio=0.2, - # (bool) Whether to use advantage norm in a whole training batch + # (bool) Whether to use advantage norm in a whole training batch. adv_norm=True, + # (bool) Whether to use value norm with running mean and std in the whole training process. value_norm=True, + # (bool) Whether to enable special network parameters initialization scheme in PPO, such as orthogonal init. ppo_param_init=True, + # (str) The gradient clip operation type used in PPO, ['clip_norm', clip_value', 'clip_momentum_norm']. grad_clip_type='clip_norm', + # (float) The gradient clip target value used in PPO. + # If ``grad_clip_type`` is 'clip_norm', then the maximum of gradient will be normalized to this value. grad_clip_value=0.5, + # (bool) Whether ignore done (usually for max step termination env). ignore_done=False, + # (str) The type of KL divergence loss between current policy and pretrained policy, ['k1', 'k2', 'k3']. + # Reference: http://joschu.net/blog/kl-approx.html + kl_type='k1', + # (float) The weight of KL divergence loss. + kl_beta=0.0, + # (Optional[str]) The path of pretrained model checkpoint. + # If provided, KL regularizer will be calculated between current policy and pretrained policy. + # Default to None, which means KL is not calculated. + pretrained_model_path=None, ), + # collect_mode config collect=dict( - # (int) Only one of [n_sample, n_episode] shoule be set + # (int) How many training samples collected in one collection procedure. + # Only one of [n_sample, n_episode] should be set. # n_sample=64, - # (int) Cut trajectories into pieces with length "unroll_len". + # (int) Split episodes or trajectories into pieces with length `unroll_len`. unroll_len=1, - # ============================================================== - # The following configs is algorithm-specific - # ============================================================== # (float) Reward's future discount factor, aka. gamma. discount_factor=0.99, # (float) GAE lambda factor for the balance of bias and variance(1-step td and mc) gae_lambda=0.95, ), - eval=dict(), + eval=dict(), # for compability ) + def default_model(self) -> Tuple[str, List[str]]: + """ + Overview: + Return this algorithm default neural network model setting for demonstration. ``__init__`` method will \ + automatically call this method to get the default model setting and create model. + Returns: + - model_info (:obj:`Tuple[str, List[str]]`): The registered model name and model's import_names. + + .. note:: + The user can define and use customized network model but must obey the same inferface definition indicated \ + by import_names path. For example about PPO, its registered name is ``ppo`` and the import_names is \ + ``ding.model.template.vac``. + + .. note:: + Because now PPO supports both single-agent and multi-agent usages, so we can implement these functions \ + with the same policy and two different default models, which is controled by ``self._cfg.multi_agent``. + """ + if self._cfg.multi_agent: + return 'mavac', ['ding.model.template.mavac'] + else: + return 'vac', ['ding.model.template.vac'] + def _init_learn(self) -> None: - r""" + """ Overview: - Learn mode init method. Called by ``self.__init__``. - Init the optimizer, algorithm config and the main model. + Initialize the learn mode of policy, including related attributes and modules. For PPO, it mainly contains \ + optimizer, algorithm-specific arguments such as loss weight, clip_ratio and recompute_adv. This method \ + also executes some special network initializations and prepares running mean/std monitor for value. + This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. """ self._priority = self._cfg.priority self._priority_IS_weight = self._cfg.priority_IS_weight assert not self._priority and not self._priority_IS_weight, "Priority is not implemented in PPO" + assert self._cfg.action_space in ["continuous", "discrete", "hybrid"] self._action_space = self._cfg.action_space if self._cfg.learn.ppo_param_init: for n, m in self._model.named_modules(): @@ -128,14 +183,35 @@ def _init_learn(self) -> None: clip_value=self._cfg.learn.grad_clip_value ) + # Define linear lr scheduler + if self._cfg.learn.lr_scheduler is not None: + epoch_num = self._cfg.learn.lr_scheduler['epoch_num'] + min_lr_lambda = self._cfg.learn.lr_scheduler['min_lr_lambda'] + + self._lr_scheduler = torch.optim.lr_scheduler.LambdaLR( + self._optimizer, + lr_lambda=lambda epoch: max(1.0 - epoch * (1.0 - min_lr_lambda) / epoch_num, min_lr_lambda) + ) + self._learn_model = model_wrap(self._model, wrapper_name='base') + # load pretrained model + if self._cfg.learn.pretrained_model_path is not None: + self._pretrained_model = copy.deepcopy(self._model) + state_dict = torch.load(self._cfg.learn.pretrained_model_path, map_location='cpu') + self._pretrained_model.load_state_dict(state_dict) + self._pretrained_model.eval() + else: + self._pretrained_model = None + # Algorithm config self._value_weight = self._cfg.learn.value_weight self._entropy_weight = self._cfg.learn.entropy_weight self._clip_ratio = self._cfg.learn.clip_ratio self._adv_norm = self._cfg.learn.adv_norm self._value_norm = self._cfg.learn.value_norm + self._kl_type = self._cfg.learn.kl_type + self._kl_beta = self._cfg.learn.kl_beta if self._value_norm: self._running_mean_std = RunningMeanStd(epsilon=1e-4, device=self._device) self._gamma = self._cfg.collect.discount_factor @@ -144,20 +220,47 @@ def _init_learn(self) -> None: # Main model self._learn_model.reset() - def _forward_learn(self, data: Dict[str, Any]) -> Dict[str, Any]: - r""" + def _forward_learn(self, data: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ Overview: - Forward and backward function of learn mode. + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data from the replay buffer and then returns the output \ + result, including various training information such as loss, clipfrac, approx_kl. Arguments: - - data (:obj:`dict`): Dict type data + - data (:obj:`List[Dict[int, Any]]`): The input data used for policy forward, including the latest \ + collected training samples for on-policy algorithms like PPO. For each element in list, the key of the \ + dict is the name of data items and the value is the corresponding data. Usually, the value is \ + torch.Tensor or np.ndarray or there dict/list combinations. In the ``_forward_learn`` method, data \ + often need to first be stacked in the batch dimension by some utility functions such as \ + ``default_preprocess_learn``. \ + For PPO, each element in list is a dict containing at least the following keys: ``obs``, ``action``, \ + ``reward``, ``logit``, ``value``, ``done``. Sometimes, it also contains other keys such as ``weight``. Returns: - - info_dict (:obj:`Dict[str, Any]`): - Including current lr, total_loss, policy_loss, value_loss, entropy_loss, \ - adv_abs_max, approx_kl, clipfrac + - return_infos (:obj:`List[Dict[str, Any]]`): The information list that indicated training result, each \ + training iteration contains append a information dict into the final list. The list will be precessed \ + and recorded in text log and tensorboard. The value of the dict must be python scalar or a list of \ + scalars. For the detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. + + .. tip:: + The training procedure of PPO is two for loops. The outer loop trains all the collected training samples \ + with ``epoch_per_collect`` epochs. The inner loop splits all the data into different mini-batch with \ + the length of ``batch_size``. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for PPOPolicy: ``ding.policy.tests.test_ppo``. """ data = default_preprocess_learn(data, ignore_done=self._cfg.learn.ignore_done, use_nstep=False) if self._cuda: data = to_device(data, self._device) + data['obs'] = to_dtype(data['obs'], torch.float32) + if 'next_obs' in data: + data['next_obs'] = to_dtype(data['next_obs'], torch.float32) # ==================== # PPO forward # ==================== @@ -202,52 +305,69 @@ def _forward_learn(self, data: Dict[str, Any]) -> Dict[str, Any]: # Normalize advantage in a train_batch adv = (adv - adv.mean()) / (adv.std() + 1e-8) + if self._pretrained_model is not None: + with torch.no_grad(): + logit_pretrained = self._pretrained_model.forward(batch['obs'], mode='compute_actor')['logit'] + else: + logit_pretrained = None + # Calculate ppo error if self._action_space == 'continuous': ppo_batch = ppo_data( output['logit'], batch['logit'], batch['action'], output['value'], batch['value'], adv, - batch['return'], batch['weight'] + batch['return'], batch['weight'], logit_pretrained ) - ppo_loss, ppo_info = ppo_error_continuous(ppo_batch, self._clip_ratio) + ppo_loss, ppo_info = ppo_error_continuous(ppo_batch, self._clip_ratio, kl_type=self._kl_type) elif self._action_space == 'discrete': ppo_batch = ppo_data( output['logit'], batch['logit'], batch['action'], output['value'], batch['value'], adv, - batch['return'], batch['weight'] + batch['return'], batch['weight'], logit_pretrained ) - ppo_loss, ppo_info = ppo_error(ppo_batch, self._clip_ratio) + ppo_loss, ppo_info = ppo_error(ppo_batch, self._clip_ratio, kl_type=self._kl_type) elif self._action_space == 'hybrid': # discrete part (discrete policy loss and entropy loss) ppo_discrete_batch = ppo_policy_data( output['logit']['action_type'], batch['logit']['action_type'], batch['action']['action_type'], - adv, batch['weight'] + adv, batch['weight'], logit_pretrained + ) + ppo_discrete_loss, ppo_discrete_info = ppo_policy_error( + ppo_discrete_batch, self._clip_ratio, kl_type=self._kl_type ) - ppo_discrete_loss, ppo_discrete_info = ppo_policy_error(ppo_discrete_batch, self._clip_ratio) # continuous part (continuous policy loss and entropy loss, value loss) ppo_continuous_batch = ppo_data( output['logit']['action_args'], batch['logit']['action_args'], batch['action']['action_args'], - output['value'], batch['value'], adv, batch['return'], batch['weight'] + output['value'], batch['value'], adv, batch['return'], batch['weight'], None ) ppo_continuous_loss, ppo_continuous_info = ppo_error_continuous( - ppo_continuous_batch, self._clip_ratio + ppo_continuous_batch, self._clip_ratio, kl_type=self._kl_type ) # sum discrete and continuous loss ppo_loss = type(ppo_continuous_loss)( ppo_continuous_loss.policy_loss + ppo_discrete_loss.policy_loss, ppo_continuous_loss.value_loss, - ppo_continuous_loss.entropy_loss + ppo_discrete_loss.entropy_loss + ppo_continuous_loss.entropy_loss + ppo_discrete_loss.entropy_loss, ppo_continuous_loss.kl_div ) ppo_info = type(ppo_continuous_info)( max(ppo_continuous_info.approx_kl, ppo_discrete_info.approx_kl), max(ppo_continuous_info.clipfrac, ppo_discrete_info.clipfrac) ) wv, we = self._value_weight, self._entropy_weight - total_loss = ppo_loss.policy_loss + wv * ppo_loss.value_loss - we * ppo_loss.entropy_loss + kl_div = ppo_loss.kl_div + total_loss = ( + ppo_loss.policy_loss + wv * ppo_loss.value_loss - we * ppo_loss.entropy_loss + + self._kl_beta * kl_div + ) self._optimizer.zero_grad() total_loss.backward() self._optimizer.step() + if self._cfg.learn.lr_scheduler is not None: + cur_lr = sum(self._lr_scheduler.get_last_lr()) / len(self._lr_scheduler.get_last_lr()) + else: + cur_lr = self._optimizer.defaults['lr'] + return_info = { - 'cur_lr': self._optimizer.defaults['lr'], + 'cur_lr': cur_lr, 'total_loss': total_loss.item(), 'policy_loss': ppo_loss.policy_loss.item(), 'value_loss': ppo_loss.value_loss.item(), @@ -258,6 +378,7 @@ def _forward_learn(self, data: Dict[str, Any]) -> Dict[str, Any]: 'value_max': output['value'].max().item(), 'approx_kl': ppo_info.approx_kl, 'clipfrac': ppo_info.clipfrac, + 'kl_div': kl_div.item(), } if self._action_space == 'continuous': return_info.update( @@ -268,25 +389,30 @@ def _forward_learn(self, data: Dict[str, Any]) -> Dict[str, Any]: } ) return_infos.append(return_info) - return return_infos - def _state_dict_learn(self) -> Dict[str, Any]: - return { - 'model': self._learn_model.state_dict(), - 'optimizer': self._optimizer.state_dict(), - } + if self._cfg.learn.lr_scheduler is not None: + self._lr_scheduler.step() - def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: - self._learn_model.load_state_dict(state_dict['model']) - self._optimizer.load_state_dict(state_dict['optimizer']) + return return_infos def _init_collect(self) -> None: - r""" + """ Overview: - Collect mode init method. Called by ``self.__init__``. - Init traj and unroll length, collect model. + Initialize the collect mode of policy, including related attributes and modules. For PPO, it contains the \ + collect_model to balance the exploration and exploitation (e.g. the multinomial sample mechanism in \ + discrete action space), and other algorithm-specific arguments such as unroll_len and gae_lambda. + This method will be called in ``__init__`` method if ``collect`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_collect`` method, you'd better name them \ + with prefix ``_collect_`` to avoid conflict with other modes, such as ``self._collect_attr1``. + + .. tip:: + Some variables need to initialize independently in different modes, such as gamma and gae_lambda in PPO. \ + This design is for the convenience of parallel execution of different policy modes. """ self._unroll_len = self._cfg.collect.unroll_len + assert self._cfg.action_space in ["continuous", "discrete", "hybrid"], self._cfg.action_space self._action_space = self._cfg.action_space if self._action_space == 'continuous': self._collect_model = model_wrap(self._model, wrapper_name='reparam_sample') @@ -299,17 +425,32 @@ def _init_collect(self) -> None: self._gae_lambda = self._cfg.collect.gae_lambda self._recompute_adv = self._cfg.recompute_adv - def _forward_collect(self, data: dict) -> dict: - r""" + def _forward_collect(self, data: Dict[int, Any]) -> Dict[int, Any]: + """ Overview: - Forward function of collect mode. + Policy forward function of collect mode (collecting training data by interacting with envs). Forward means \ + that the policy gets some necessary data (mainly observation) from the envs and then returns the output \ + data, such as the action to interact with the envs. Arguments: - - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \ - values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer. + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. Returns: - - output (:obj:`Dict[int, Any]`): Dict type data, including at least inferred action according to input obs. - ReturnsKeys - - necessary: ``action`` + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action and \ + other necessary data (action logit and value) for learn mode defined in ``self._process_transition`` \ + method. The key of the dict is the same as the input data, i.e. environment id. + + .. tip:: + If you want to add more tricks on this policy, like temperature factor in multinomial sample, you can pass \ + related data as extra keyword arguments of this method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for PPOPolicy: ``ding.policy.tests.test_ppo``. """ data_id = list(data.keys()) data = default_collate(list(data.values())) @@ -323,38 +464,54 @@ def _forward_collect(self, data: dict) -> dict: output = default_decollate(output) return {i: d for i, d in zip(data_id, output)} - def _process_transition(self, obs: Any, model_output: dict, timestep: namedtuple) -> dict: + def _process_transition(self, obs: torch.Tensor, policy_output: Dict[str, torch.Tensor], + timestep: namedtuple) -> Dict[str, torch.Tensor]: """ Overview: - Generate dict type transition data from inputs. + Process and pack one timestep transition data into a dict, which can be directly used for training and \ + saved in replay buffer. For PPO, it contains obs, next_obs, action, reward, done, logit, value. Arguments: - - obs (:obj:`Any`): Env observation - - model_output (:obj:`dict`): Output of collect model, including at least ['action'] - - timestep (:obj:`namedtuple`): Output after env step, including at least ['obs', 'reward', 'done']\ - (here 'obs' indicates obs after env step). + - obs (:obj:`torch.Tensor`): The env observation of current timestep, such as stacked 2D image in Atari. + - policy_output (:obj:`Dict[str, torch.Tensor]`): The output of the policy network with the observation \ + as input. For PPO, it contains the state value, action and the logit of the action. + - timestep (:obj:`namedtuple`): The execution result namedtuple returned by the environment step method, \ + except all the elements have been transformed into tensor data. Usually, it contains the next obs, \ + reward, done, info, etc. Returns: - - transition (:obj:`dict`): Dict type transition data. + - transition (:obj:`Dict[str, torch.Tensor]`): The processed transition data of the current timestep. + + .. note:: + ``next_obs`` is used to calculate nstep return when necessary, so we place in into transition by default. \ + You can delete this field to save memory occupancy if you do not need nstep return. """ transition = { 'obs': obs, 'next_obs': timestep.obs, - 'action': model_output['action'], - 'logit': model_output['logit'], - 'value': model_output['value'], + 'action': policy_output['action'], + 'logit': policy_output['logit'], + 'value': policy_output['value'], 'reward': timestep.reward, 'done': timestep.done, } return transition - def _get_train_sample(self, data: list) -> Union[None, List[Any]]: - r""" + def _get_train_sample(self, transitions: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ Overview: - Get the trajectory and calculate GAE, return one data to cache for next time calculation + For a given trajectory (transitions, a list of transition) data, process it into a list of sample that \ + can be used for training directly. In PPO, a train sample is a processed transition with new computed \ + ``traj_flag`` and ``adv`` field. This method is usually used in collectors to execute necessary \ + RL data preprocessing before training, which can help learner amortize revelant time consumption. \ + In addition, you can also implement this method as an identity function and do the data processing \ + in ``self._forward_learn`` method. Arguments: - - data (:obj:`list`): The trajectory's cache + - transitions (:obj:`List[Dict[str, Any]`): The trajectory data (a list of transition), each element is \ + the same format as the return value of ``self._process_transition`` method. Returns: - - samples (:obj:`dict`): The training samples generated + - samples (:obj:`List[Dict[str, Any]]`): The processed train samples, each element is the similar format \ + as input transitions, but may contain more data for training, such as GAE advantage. """ + data = transitions data = to_device(data, self._device) for transition in data: transition['traj_flag'] = copy.deepcopy(transition['done']) @@ -394,31 +551,50 @@ def _get_train_sample(self, data: list) -> Union[None, List[Any]]: return get_train_sample(data, self._unroll_len) def _init_eval(self) -> None: - r""" + """ Overview: - Evaluate mode init method. Called by ``self.__init__``. - Init eval model with argmax strategy. + Initialize the eval mode of policy, including related attributes and modules. For PPO, it contains the \ + eval model to select optimial action (e.g. greedily select action with argmax mechanism in discrete action). + This method will be called in ``__init__`` method if ``eval`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_eval`` method, you'd better name them \ + with prefix ``_eval_`` to avoid conflict with other modes, such as ``self._eval_attr1``. """ + assert self._cfg.action_space in ["continuous", "discrete", "hybrid"] self._action_space = self._cfg.action_space if self._action_space == 'continuous': self._eval_model = model_wrap(self._model, wrapper_name='deterministic_sample') elif self._action_space == 'discrete': self._eval_model = model_wrap(self._model, wrapper_name='argmax_sample') elif self._action_space == 'hybrid': - self._eval_model = model_wrap(self._model, wrapper_name='hybrid_deterministic_argmax_sample') + self._eval_model = model_wrap(self._model, wrapper_name='hybrid_reparam_multinomial_sample') + self._eval_model.reset() - def _forward_eval(self, data: dict) -> dict: - r""" + def _forward_eval(self, data: Dict[int, Any]) -> Dict[int, Any]: + """ Overview: - Forward function of eval mode, similar to ``self._forward_collect``. + Policy forward function of eval mode (evaluation policy performance by interacting with envs). Forward \ + means that the policy gets some necessary data (mainly observation) from the envs and then returns the \ + action to interact with the envs. ``_forward_eval`` in PPO often uses deterministic sample method to get \ + actions while ``_forward_collect`` usually uses stochastic sample method for balance exploration and \ + exploitation. Arguments: - - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \ - values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer. + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. Returns: - - output (:obj:`Dict[int, Any]`): The dict of predicting action for the interaction with env. - ReturnsKeys - - necessary: ``action`` + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action. The \ + key of the dict is the same as the input data, i.e. environment id. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for PPOPolicy: ``ding.policy.tests.test_ppo``. """ data_id = list(data.keys()) data = default_collate(list(data.values())) @@ -432,13 +608,14 @@ def _forward_eval(self, data: dict) -> dict: output = default_decollate(output) return {i: d for i, d in zip(data_id, output)} - def default_model(self) -> Tuple[str, List[str]]: - if self._cfg.multi_agent: - return 'mavac', ['ding.model.template.mavac'] - else: - return 'vac', ['ding.model.template.vac'] - def _monitor_vars_learn(self) -> List[str]: + """ + Overview: + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. + Returns: + - necessary_keys (:obj:`List[str]`): The list of the necessary keys to be logged. + """ variables = super()._monitor_vars_learn() + [ 'policy_loss', 'value_loss', @@ -450,6 +627,8 @@ def _monitor_vars_learn(self) -> List[str]: 'value_max', 'value_mean', ] + if self._pretrained_model is not None: + variables += ['kl_div'] if self._action_space == 'continuous': variables += ['mu_mean', 'sigma_mean', 'sigma_grad', 'act'] return variables @@ -457,9 +636,10 @@ def _monitor_vars_learn(self) -> List[str]: @POLICY_REGISTRY.register('ppo_pg') class PPOPGPolicy(Policy): - r""" + """ Overview: - Policy class of on policy version PPO algorithm (pure policy gradient). + Policy class of on policy version PPO algorithm (pure policy gradient without value network). + Paper link: https://arxiv.org/abs/1707.06347. """ config = dict( # (str) RL policy register name (refer to function "POLICY_REGISTRY"). @@ -470,53 +650,75 @@ class PPOPGPolicy(Policy): on_policy=True, # (str) Which kind of action space used in PPOPolicy, ['discrete', 'continuous', 'hybrid'] action_space='discrete', - # (bool) Whether to enable multi-agent training, i.e.: MAPPO + # (bool) Whether to enable multi-agent training, i.e.: MAPPO. multi_agent=False, - # (bool) Whether to need policy data in process transition + # (bool) Whether to need policy data in process transition. transition_with_policy_data=True, + # learn_mode config learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, + # (int) After collecting n_sample/n_episode data, how many epoches to train models. + # Each epoch means the one entire passing of training data. epoch_per_collect=10, + # (int) How many samples in a training batch. batch_size=64, + # (float) The step size of gradient descent. learning_rate=3e-4, - # ============================================================== - # The following configs is algorithm-specific - # ============================================================== - # (float) The loss weight of entropy regularization, policy network weight is set to 1 + # (float) The loss weight of entropy regularization, policy network weight is set to 1. entropy_weight=0.0, - # (float) PPO clip ratio, defaults to 0.2 + # (float) PPO clip ratio, defaults to 0.2. clip_ratio=0.2, - # (bool) Whether to use advantage norm in a whole training batch + # (bool) Whether to enable special network parameters initialization scheme in PPO, such as orthogonal init. ppo_param_init=True, + # (str) The gradient clip operation type used in PPO, ['clip_norm', clip_value', 'clip_momentum_norm']. grad_clip_type='clip_norm', + # (float) The gradient clip target value used in PPO. + # If ``grad_clip_type`` is 'clip_norm', then the maximum of gradient will be normalized to this value. grad_clip_value=0.5, + # (bool) Whether ignore done (usually for max step termination env). ignore_done=False, ), + # collect_mode config collect=dict( - # (int) Only one of [n_sample, n_episode] shoule be set - # n_sample=64, + # (int) How many training episodes collected in one collection process. Only one of n_episode shoule be set. + # n_episode=8, # (int) Cut trajectories into pieces with length "unroll_len". unroll_len=1, - # ============================================================== - # The following configs is algorithm-specific - # ============================================================== # (float) Reward's future discount factor, aka. gamma. discount_factor=0.99, ), - eval=dict(), + eval=dict(), # for compability ) def default_model(self) -> Tuple[str, List[str]]: - if self._cfg.action_space == 'discrete': - return 'discrete_bc', ['ding.model.template.bc'] - else: - return RuntimeError( - "PPOPGPolicy doesn't have default_model, please define your own model and pass it " + - "into policy, you can refer to " + "dizoo/box2d/bipedalwalker/config/bipedalwalker_ppopg_config.py" - ) + """ + Overview: + Return this algorithm default neural network model setting for demonstration. ``__init__`` method will \ + automatically call this method to get the default model setting and create model. + Returns: + - model_info (:obj:`Tuple[str, List[str]]`): The registered model name and model's import_names. + """ + return 'pg', ['ding.model.template.pg'] def _init_learn(self) -> None: + """ + Overview: + Initialize the learn mode of policy, including related attributes and modules. For PPOPG, it mainly \ + contains optimizer, algorithm-specific arguments such as loss weight and clip_ratio. This method \ + also executes some special network initializations. + This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. + """ + assert self._cfg.action_space in ["continuous", "discrete"] self._action_space = self._cfg.action_space if self._cfg.learn.ppo_param_init: for n, m in self._model.named_modules(): @@ -548,7 +750,39 @@ def _init_learn(self) -> None: # Main model self._learn_model.reset() - def _forward_learn(self, data: Dict[str, Any]) -> Dict[str, Any]: + def _forward_learn(self, data: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Overview: + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data from the replay buffer and then returns the output \ + result, including various training information such as loss, clipfrac, approx_kl. + Arguments: + - data (:obj:`List[Dict[int, Any]]`): The input data used for policy forward, including the latest \ + collected training samples for on-policy algorithms like PPO. For each element in list, the key of the \ + dict is the name of data items and the value is the corresponding data. Usually, the value is \ + torch.Tensor or np.ndarray or there dict/list combinations. In the ``_forward_learn`` method, data \ + often need to first be stacked in the batch dimension by some utility functions such as \ + ``default_preprocess_learn``. \ + For PPOPG, each element in list is a dict containing at least the following keys: ``obs``, ``action``, \ + ``return``, ``logit``, ``done``. Sometimes, it also contains other keys such as ``weight``. + Returns: + - return_infos (:obj:`List[Dict[str, Any]]`): The information list that indicated training result, each \ + training iteration contains append a information dict into the final list. The list will be precessed \ + and recorded in text log and tensorboard. The value of the dict must be python scalar or a list of \ + scalars. For the detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. + + .. tip:: + The training procedure of PPOPG is two for loops. The outer loop trains all the collected training samples \ + with ``epoch_per_collect`` epochs. The inner loop splits all the data into different mini-batch with \ + the length of ``batch_size``. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + """ + data = default_preprocess_learn(data) if self._cuda: data = to_device(data, self._device) @@ -560,7 +794,7 @@ def _forward_learn(self, data: Dict[str, Any]) -> Dict[str, Any]: output = self._learn_model.forward(batch['obs']) ppo_batch = ppo_policy_data( - output['logit'], batch['logit'], batch['action'], batch['return'], batch['weight'] + output['logit'], batch['logit'], batch['action'], batch['return'], batch['weight'], None ) if self._action_space == 'continuous': ppo_loss, ppo_info = ppo_policy_error_continuous(ppo_batch, self._clip_ratio) @@ -591,17 +825,23 @@ def _forward_learn(self, data: Dict[str, Any]) -> Dict[str, Any]: return_infos.append(return_info) return return_infos - def _state_dict_learn(self) -> Dict[str, Any]: - return { - 'model': self._learn_model.state_dict(), - 'optimizer': self._optimizer.state_dict(), - } - - def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: - self._learn_model.load_state_dict(state_dict['model']) - self._optimizer.load_state_dict(state_dict['optimizer']) - def _init_collect(self) -> None: + """ + Overview: + Initialize the collect mode of policy, including related attributes and modules. For PPOPG, it contains \ + the collect_model to balance the exploration and exploitation (e.g. the multinomial sample mechanism in \ + discrete action space), and other algorithm-specific arguments such as unroll_len and gae_lambda. + This method will be called in ``__init__`` method if ``collect`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_collect`` method, you'd better name them \ + with prefix ``_collect_`` to avoid conflict with other modes, such as ``self._collect_attr1``. + + .. tip:: + Some variables need to initialize independently in different modes, such as gamma and gae_lambda in PPO. \ + This design is for the convenience of parallel execution of different policy modes. + """ + assert self._cfg.action_space in ["continuous", "discrete"], self._cfg.action_space self._action_space = self._cfg.action_space self._unroll_len = self._cfg.collect.unroll_len if self._action_space == 'continuous': @@ -611,7 +851,30 @@ def _init_collect(self) -> None: self._collect_model.reset() self._gamma = self._cfg.collect.discount_factor - def _forward_collect(self, data: dict) -> dict: + def _forward_collect(self, data: Dict[int, Any]) -> Dict[int, Any]: + """ + Overview: + Policy forward function of collect mode (collecting training data by interacting with envs). Forward means \ + that the policy gets some necessary data (mainly observation) from the envs and then returns the output \ + data, such as the action to interact with the envs. + Arguments: + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. + Returns: + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action and \ + other necessary data (action logit) for learn mode defined in ``self._process_transition`` \ + method. The key of the dict is the same as the input data, i.e. environment id. + + .. tip:: + If you want to add more tricks on this policy, like temperature factor in multinomial sample, you can pass \ + related data as extra keyword arguments of this method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + """ data_id = list(data.keys()) data = default_collate(list(data.values())) if self._cuda: @@ -624,17 +887,47 @@ def _forward_collect(self, data: dict) -> dict: output = default_decollate(output) return {i: d for i, d in zip(data_id, output)} - def _process_transition(self, obs: Any, model_output: dict, timestep: namedtuple) -> dict: + def _process_transition(self, obs: torch.Tensor, policy_output: Dict[str, torch.Tensor], + timestep: namedtuple) -> Dict[str, torch.Tensor]: + """ + Overview: + Process and pack one timestep transition data into a dict, which can be directly used for training and \ + saved in replay buffer. For PPOPG, it contains obs, action, reward, done, logit. + Arguments: + - obs (:obj:`torch.Tensor`): The env observation of current timestep, such as stacked 2D image in Atari. + - policy_output (:obj:`Dict[str, torch.Tensor]`): The output of the policy network with the observation \ + as input. For PPOPG, it contains the action and the logit of the action. + - timestep (:obj:`namedtuple`): The execution result namedtuple returned by the environment step method, \ + except all the elements have been transformed into tensor data. Usually, it contains the next obs, \ + reward, done, info, etc. + Returns: + - transition (:obj:`Dict[str, torch.Tensor]`): The processed transition data of the current timestep. + """ transition = { 'obs': obs, - 'action': model_output['action'], - 'logit': model_output['logit'], + 'action': policy_output['action'], + 'logit': policy_output['logit'], 'reward': timestep.reward, 'done': timestep.done, } return transition - def _get_train_sample(self, data: list) -> Union[None, List[Any]]: + def _get_train_sample(self, data: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Overview: + For a given entire episode data (a list of transition), process it into a list of sample that \ + can be used for training directly. In PPOPG, a train sample is a processed transition with new computed \ + ``return`` field. This method is usually used in collectors to execute necessary \ + RL data preprocessing before training, which can help learner amortize revelant time consumption. \ + In addition, you can also implement this method as an identity function and do the data processing \ + in ``self._forward_learn`` method. + Arguments: + - data (:obj:`List[Dict[str, Any]`): The episode data (a list of transition), each element is \ + the same format as the return value of ``self._process_transition`` method. + Returns: + - samples (:obj:`List[Dict[str, Any]]`): The processed train samples, each element is the similar format \ + as input transitions, but may contain more data for training, such as discounted episode return. + """ assert data[-1]['done'] is True, "PPO-PG needs a complete epsiode" if self._cfg.learn.ignore_done: @@ -648,6 +941,17 @@ def _get_train_sample(self, data: list) -> Union[None, List[Any]]: return get_train_sample(data, self._unroll_len) def _init_eval(self) -> None: + """ + Overview: + Initialize the eval mode of policy, including related attributes and modules. For PPOPG, it contains the \ + eval model to select optimial action (e.g. greedily select action with argmax mechanism in discrete action). + This method will be called in ``__init__`` method if ``eval`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_eval`` method, you'd better name them \ + with prefix ``_eval_`` to avoid conflict with other modes, such as ``self._eval_attr1``. + """ + assert self._cfg.action_space in ["continuous", "discrete"] self._action_space = self._cfg.action_space if self._action_space == 'continuous': self._eval_model = model_wrap(self._model, wrapper_name='deterministic_sample') @@ -655,7 +959,30 @@ def _init_eval(self) -> None: self._eval_model = model_wrap(self._model, wrapper_name='argmax_sample') self._eval_model.reset() - def _forward_eval(self, data: dict) -> dict: + def _forward_eval(self, data: Dict[int, Any]) -> Dict[int, Any]: + """ + Overview: + Policy forward function of eval mode (evaluation policy performance by interacting with envs). Forward \ + means that the policy gets some necessary data (mainly observation) from the envs and then returns the \ + action to interact with the envs. ``_forward_eval`` in PPO often uses deterministic sample method to get \ + actions while ``_forward_collect`` usually uses stochastic sample method for balance exploration and \ + exploitation. + Arguments: + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. + Returns: + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action. The \ + key of the dict is the same as the input data, i.e. environment id. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for PPOPGPolicy: ``ding.policy.tests.test_ppo``. + """ data_id = list(data.keys()) data = default_collate(list(data.values())) if self._cuda: @@ -669,6 +996,13 @@ def _forward_eval(self, data: dict) -> dict: return {i: d for i, d in zip(data_id, output)} def _monitor_vars_learn(self) -> List[str]: + """ + Overview: + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. + Returns: + - necessary_keys (:obj:`List[str]`): The list of the necessary keys to be logged. + """ return super()._monitor_vars_learn() + [ 'policy_loss', 'entropy_loss', @@ -679,9 +1013,10 @@ def _monitor_vars_learn(self) -> List[str]: @POLICY_REGISTRY.register('ppo_offpolicy') class PPOOffPolicy(Policy): - r""" + """ Overview: - Policy class of PPO algorithm. + Policy class of off-policy version PPO algorithm. Paper link: https://arxiv.org/abs/1707.06347. + This version is more suitable for large-scale distributed training. """ config = dict( # (str) RL policy register name (refer to function "POLICY_REGISTRY"). @@ -689,71 +1024,142 @@ class PPOOffPolicy(Policy): # (bool) Whether to use cuda for network. cuda=False, on_policy=False, - # (bool) Whether to use priority(priority sample, IS weight, update priority) + # (bool) Whether to use priority (priority sample, IS weight, update priority). priority=False, # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. priority_IS_weight=False, - # (bool) Whether to use nstep_return for value loss + # (str) Which kind of action space used in PPOPolicy, ["continuous", "discrete", "hybrid"]. + action_space='discrete', + # (bool) Whether to use nstep_return for value loss. nstep_return=False, + # (int) The timestep of TD (temporal-difference) loss. nstep=3, - # (bool) Whether to need policy data in process transition + # (bool) Whether to need policy data in process transition. transition_with_policy_data=True, + # learn_mode config learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, - # How many updates(iterations) to train after collector's one collection. + # (int) How many updates(iterations) to train after collector's one collection. # Bigger "update_per_collect" means bigger off-policy. # collect data -> update policy-> collect data -> ... update_per_collect=5, + # (int) How many samples in a training batch. batch_size=64, + # (float) The step size of gradient descent. learning_rate=0.001, - # ============================================================== - # The following configs is algorithm-specific - # ============================================================== - # (float) The loss weight of value network, policy network weight is set to 1 + # (float) The loss weight of value network, policy network weight is set to 1. value_weight=0.5, - # (float) The loss weight of entropy regularization, policy network weight is set to 1 + # (float) The loss weight of entropy regularization, policy network weight is set to 1. entropy_weight=0.01, - # (float) PPO clip ratio, defaults to 0.2 + # (float) PPO clip ratio, defaults to 0.2. clip_ratio=0.2, - # (bool) Whether to use advantage norm in a whole training batch + # (bool) Whether to use advantage norm in a whole training batch. adv_norm=False, + # (bool) Whether to use value norm with running mean and std in the whole training process. + value_norm=True, + # (bool) Whether to enable special network parameters initialization scheme in PPO, such as orthogonal init. + ppo_param_init=True, + # (str) The gradient clip operation type used in PPO, ['clip_norm', clip_value', 'clip_momentum_norm']. + grad_clip_type='clip_norm', + # (float) The gradient clip target value used in PPO. + # If ``grad_clip_type`` is 'clip_norm', then the maximum of gradient will be normalized to this value. + grad_clip_value=0.5, + # (bool) Whether ignore done (usually for max step termination env). ignore_done=False, + # (float) The weight decay (L2 regularization) loss weight, defaults to 0.0. + weight_decay=0.0, ), + # collect_mode config collect=dict( - # (int) Only one of [n_sample, n_episode] shoule be set + # (int) How many training samples collected in one collection procedure. + # Only one of [n_sample, n_episode] shoule be set. # n_sample=64, # (int) Cut trajectories into pieces with length "unroll_len". unroll_len=1, - # ============================================================== - # The following configs is algorithm-specific - # ============================================================== # (float) Reward's future discount factor, aka. gamma. discount_factor=0.99, - # (float) GAE lambda factor for the balance of bias and variance(1-step td and mc) + # (float) GAE lambda factor for the balance of bias and variance (1-step td and mc). gae_lambda=0.95, ), - eval=dict(), - other=dict(replay_buffer=dict(replay_buffer_size=10000, ), ), + eval=dict(), # for compability + other=dict( + replay_buffer=dict( + # (int) Maximum size of replay buffer. Usually, larger buffer size is better. + replay_buffer_size=10000, + ), + ), ) + def default_model(self) -> Tuple[str, List[str]]: + """ + Overview: + Return this algorithm default neural network model setting for demonstration. ``__init__`` method will \ + automatically call this method to get the default model setting and create model. + Returns: + - model_info (:obj:`Tuple[str, List[str]]`): The registered model name and model's import_names. + """ + return 'vac', ['ding.model.template.vac'] + def _init_learn(self) -> None: - r""" + """ Overview: - Learn mode init method. Called by ``self.__init__``. - Init the optimizer, algorithm config and the main model. + Initialize the learn mode of policy, including related attributes and modules. For PPOOff, it mainly \ + contains optimizer, algorithm-specific arguments such as loss weight and clip_ratio. This method \ + also executes some special network initializations and prepares running mean/std monitor for value. + This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. """ self._priority = self._cfg.priority self._priority_IS_weight = self._cfg.priority_IS_weight - assert not self._priority and not self._priority_IS_weight, "Priority is not implemented in PPO" - # Orthogonal init - for m in self._model.modules(): - if isinstance(m, torch.nn.Conv2d): - torch.nn.init.orthogonal_(m.weight) - if isinstance(m, torch.nn.Linear): - torch.nn.init.orthogonal_(m.weight) + assert not self._priority and not self._priority_IS_weight, "Priority is not implemented in PPOOff" + + assert self._cfg.action_space in ["continuous", "discrete", "hybrid"] + self._action_space = self._cfg.action_space + + if self._cfg.learn.ppo_param_init: + for n, m in self._model.named_modules(): + if isinstance(m, torch.nn.Linear): + torch.nn.init.orthogonal_(m.weight) + torch.nn.init.zeros_(m.bias) + if self._action_space in ['continuous', 'hybrid']: + # init log sigma + if self._action_space == 'continuous': + if hasattr(self._model.actor_head, 'log_sigma_param'): + torch.nn.init.constant_(self._model.actor_head.log_sigma_param, -2.0) + elif self._action_space == 'hybrid': # actor_head[1]: ReparameterizationHead, for action_args + if hasattr(self._model.actor_head[1], 'log_sigma_param'): + torch.nn.init.constant_(self._model.actor_head[1].log_sigma_param, -0.5) + + for m in list(self._model.critic.modules()) + list(self._model.actor.modules()): + if isinstance(m, torch.nn.Linear): + # orthogonal initialization + torch.nn.init.orthogonal_(m.weight, gain=np.sqrt(2)) + torch.nn.init.zeros_(m.bias) + # do last policy layer scaling, this will make initial actions have (close to) + # 0 mean and std, and will help boost performances, + # see https://arxiv.org/abs/2006.05990, Fig.24 for details + for m in self._model.actor.modules(): + if isinstance(m, torch.nn.Linear): + torch.nn.init.zeros_(m.bias) + m.weight.data.copy_(0.01 * m.weight.data) + # Optimizer - self._optimizer = Adam(self._model.parameters(), lr=self._cfg.learn.learning_rate) + self._optimizer = Adam( + self._model.parameters(), + lr=self._cfg.learn.learning_rate, + grad_clip_type=self._cfg.learn.grad_clip_type, + clip_value=self._cfg.learn.grad_clip_value + ) + self._learn_model = model_wrap(self._model, wrapper_name='base') # Algorithm config @@ -761,44 +1167,106 @@ def _init_learn(self) -> None: self._entropy_weight = self._cfg.learn.entropy_weight self._clip_ratio = self._cfg.learn.clip_ratio self._adv_norm = self._cfg.learn.adv_norm + self._value_norm = self._cfg.learn.value_norm + if self._value_norm: + self._running_mean_std = RunningMeanStd(epsilon=1e-4, device=self._device) + self._gamma = self._cfg.collect.discount_factor + self._gae_lambda = self._cfg.collect.gae_lambda self._nstep = self._cfg.nstep self._nstep_return = self._cfg.nstep_return # Main model self._learn_model.reset() - def _forward_learn(self, data: dict) -> Dict[str, Any]: - r""" + def _forward_learn(self, data: List[Dict[str, Any]]) -> Dict[str, Any]: + """ Overview: - Forward and backward function of learn mode. + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data from the replay buffer and then returns the output \ + result, including various training information such as loss, clipfrac and approx_kl. Arguments: - - data (:obj:`dict`): Dict type data + - data (:obj:`List[Dict[int, Any]]`): The input data used for policy forward, including a batch of \ + training samples. For each element in list, the key of the dict is the name of data items and the \ + value is the corresponding data. Usually, the value is torch.Tensor or np.ndarray or there dict/list \ + combinations. In the ``_forward_learn`` method, data often need to first be stacked in the batch \ + dimension by some utility functions such as ``default_preprocess_learn``. \ + For PPOOff, each element in list is a dict containing at least the following keys: ``obs``, ``adv``, \ + ``action``, ``logit``, ``value``, ``done``. Sometimes, it also contains other keys such as ``weight`` \ + and ``value_gamma``. Returns: - - info_dict (:obj:`Dict[str, Any]`): - Including current lr, total_loss, policy_loss, value_loss, entropy_loss, \ - adv_abs_max, approx_kl, clipfrac + - info_dict (:obj:`Dict[str, Any]`): The information dict that indicated training result, which will be \ + recorded in text log and tensorboard, values must be python scalar or a list of scalars. For the \ + detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. """ data = default_preprocess_learn(data, ignore_done=self._cfg.learn.ignore_done, use_nstep=self._nstep_return) if self._cuda: data = to_device(data, self._device) + data['obs'] = to_dtype(data['obs'], torch.float32) + if 'next_obs' in data: + data['next_obs'] = to_dtype(data['next_obs'], torch.float32) # ==================== # PPO forward # ==================== self._learn_model.train() + + with torch.no_grad(): + if self._value_norm: + unnormalized_return = data['adv'] + data['value'] * self._running_mean_std.std + data['return'] = unnormalized_return / self._running_mean_std.std + self._running_mean_std.update(unnormalized_return.cpu().numpy()) + else: + data['return'] = data['adv'] + data['value'] + # normal ppo if not self._nstep_return: output = self._learn_model.forward(data['obs'], mode='compute_actor_critic') adv = data['adv'] - return_ = data['value'] + adv + if self._adv_norm: # Normalize advantage in a total train_batch adv = (adv - adv.mean()) / (adv.std() + 1e-8) # Calculate ppo loss - ppodata = ppo_data( - output['logit'], data['logit'], data['action'], output['value'], data['value'], adv, return_, - data['weight'] - ) - ppo_loss, ppo_info = ppo_error(ppodata, self._clip_ratio) + if self._action_space == 'continuous': + ppodata = ppo_data( + output['logit'], data['logit'], data['action'], output['value'], data['value'], adv, data['return'], + data['weight'], None + ) + ppo_loss, ppo_info = ppo_error_continuous(ppodata, self._clip_ratio) + elif self._action_space == 'discrete': + ppodata = ppo_data( + output['logit'], data['logit'], data['action'], output['value'], data['value'], adv, data['return'], + data['weight'], None + ) + ppo_loss, ppo_info = ppo_error(ppodata, self._clip_ratio) + elif self._action_space == 'hybrid': + # discrete part (discrete policy loss and entropy loss) + ppo_discrete_batch = ppo_policy_data( + output['logit']['action_type'], data['logit']['action_type'], data['action']['action_type'], adv, + data['weight'], None + ) + ppo_discrete_loss, ppo_discrete_info = ppo_policy_error(ppo_discrete_batch, self._clip_ratio) + # continuous part (continuous policy loss and entropy loss, value loss) + ppo_continuous_batch = ppo_data( + output['logit']['action_args'], data['logit']['action_args'], data['action']['action_args'], + output['value'], data['value'], adv, data['return'], data['weight'], None + ) + ppo_continuous_loss, ppo_continuous_info = ppo_error_continuous(ppo_continuous_batch, self._clip_ratio) + # sum discrete and continuous loss + ppo_loss = type(ppo_continuous_loss)( + ppo_continuous_loss.policy_loss + ppo_discrete_loss.policy_loss, ppo_continuous_loss.value_loss, + ppo_continuous_loss.entropy_loss + ppo_discrete_loss.entropy_loss + ) + ppo_info = type(ppo_continuous_info)( + max(ppo_continuous_info.approx_kl, ppo_discrete_info.approx_kl), + max(ppo_continuous_info.clipfrac, ppo_discrete_info.clipfrac) + ) + wv, we = self._value_weight, self._entropy_weight total_loss = ppo_loss.policy_loss + wv * ppo_loss.value_loss - we * ppo_loss.entropy_loss @@ -810,8 +1278,37 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: adv = (adv - adv.mean()) / (adv.std() + 1e-8) # Calculate ppo loss - ppodata = ppo_policy_data(output['logit'], data['logit'], data['action'], adv, data['weight']) - ppo_policy_loss, ppo_info = ppo_policy_error(ppodata, self._clip_ratio) + if self._action_space == 'continuous': + ppodata = ppo_policy_data(output['logit'], data['logit'], data['action'], adv, data['weight'], None) + ppo_policy_loss, ppo_info = ppo_policy_error_continuous(ppodata, self._clip_ratio) + elif self._action_space == 'discrete': + ppodata = ppo_policy_data(output['logit'], data['logit'], data['action'], adv, data['weight'], None) + ppo_policy_loss, ppo_info = ppo_policy_error(ppodata, self._clip_ratio) + elif self._action_space == 'hybrid': + # discrete part (discrete policy loss and entropy loss) + ppo_discrete_data = ppo_policy_data( + output['logit']['action_type'], data['logit']['action_type'], data['action']['action_type'], adv, + data['weight'], None + ) + ppo_discrete_loss, ppo_discrete_info = ppo_policy_error(ppo_discrete_data, self._clip_ratio) + # continuous part (continuous policy loss and entropy loss, value loss) + ppo_continuous_data = ppo_policy_data( + output['logit']['action_args'], data['logit']['action_args'], data['action']['action_args'], adv, + data['weight'], None + ) + ppo_continuous_loss, ppo_continuous_info = ppo_policy_error_continuous( + ppo_continuous_data, self._clip_ratio + ) + # sum discrete and continuous loss + ppo_policy_loss = type(ppo_continuous_loss)( + ppo_continuous_loss.policy_loss + ppo_discrete_loss.policy_loss, + ppo_continuous_loss.entropy_loss + ppo_discrete_loss.entropy_loss + ) + ppo_info = type(ppo_continuous_info)( + max(ppo_continuous_info.approx_kl, ppo_discrete_info.approx_kl), + max(ppo_continuous_info.clipfrac, ppo_discrete_info.clipfrac) + ) + wv, we = self._value_weight, self._entropy_weight next_obs = data.get('next_obs') value_gamma = data.get('value_gamma') @@ -838,52 +1335,87 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: self._optimizer.zero_grad() total_loss.backward() self._optimizer.step() - return { + return_info = { 'cur_lr': self._optimizer.defaults['lr'], 'total_loss': total_loss.item(), 'policy_loss': ppo_loss.policy_loss.item(), + 'value': data['value'].mean().item(), 'value_loss': ppo_loss.value_loss.item(), 'entropy_loss': ppo_loss.entropy_loss.item(), 'adv_abs_max': adv.abs().max().item(), 'approx_kl': ppo_info.approx_kl, 'clipfrac': ppo_info.clipfrac, } - - def _state_dict_learn(self) -> Dict[str, Any]: - return { - 'model': self._learn_model.state_dict(), - 'optimizer': self._optimizer.state_dict(), - } - - def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: - self._learn_model.load_state_dict(state_dict['model']) - self._optimizer.load_state_dict(state_dict['optimizer']) + if self._action_space == 'continuous': + return_info.update( + { + 'act': data['action'].float().mean().item(), + 'mu_mean': output['logit']['mu'].mean().item(), + 'sigma_mean': output['logit']['sigma'].mean().item(), + } + ) + return return_info def _init_collect(self) -> None: - r""" + """ Overview: - Collect mode init method. Called by ``self.__init__``. - Init traj and unroll length, collect model. + Initialize the collect mode of policy, including related attributes and modules. For PPOOff, it contains \ + collect_model to balance the exploration and exploitation (e.g. the multinomial sample mechanism in \ + discrete action space), and other algorithm-specific arguments such as unroll_len and gae_lambda. + This method will be called in ``__init__`` method if ``collect`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_collect`` method, you'd better name them \ + with prefix ``_collect_`` to avoid conflict with other modes, such as ``self._collect_attr1``. + + .. tip:: + Some variables need to initialize independently in different modes, such as gamma and gae_lambda in PPOOff. + This design is for the convenience of parallel execution of different policy modes. """ self._unroll_len = self._cfg.collect.unroll_len - self._collect_model = model_wrap(self._model, wrapper_name='multinomial_sample') + assert self._cfg.action_space in ["continuous", "discrete", "hybrid"] + self._action_space = self._cfg.action_space + if self._action_space == 'continuous': + self._collect_model = model_wrap(self._model, wrapper_name='reparam_sample') + elif self._action_space == 'discrete': + self._collect_model = model_wrap(self._model, wrapper_name='multinomial_sample') + elif self._action_space == 'hybrid': + self._collect_model = model_wrap(self._model, wrapper_name='hybrid_reparam_multinomial_sample') self._collect_model.reset() self._gamma = self._cfg.collect.discount_factor self._gae_lambda = self._cfg.collect.gae_lambda self._nstep = self._cfg.nstep self._nstep_return = self._cfg.nstep_return + self._value_norm = self._cfg.learn.value_norm + if self._value_norm: + self._running_mean_std = RunningMeanStd(epsilon=1e-4, device=self._device) - def _forward_collect(self, data: dict) -> dict: - r""" + def _forward_collect(self, data: Dict[int, Any]) -> Dict[int, Any]: + """ Overview: - Forward function of collect mode. + Policy forward function of collect mode (collecting training data by interacting with envs). Forward means \ + that the policy gets some necessary data (mainly observation) from the envs and then returns the output \ + data, such as the action to interact with the envs. Arguments: - - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \ - values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer. + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. Returns: - - output (:obj:`Dict[int, Any]`): Dict type data, including at least inferred action according to input obs. - ReturnsKeys - - necessary: ``action`` + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action and \ + other necessary data (action logit and value) for learn mode defined in ``self._process_transition`` \ + method. The key of the dict is the same as the input data, i.e. environment id. + + .. tip:: + If you want to add more tricks on this policy, like temperature factor in multinomial sample, you can pass \ + related data as extra keyword arguments of this method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for PPOOffPolicy: ``ding.policy.tests.test_ppo``. """ data_id = list(data.keys()) data = default_collate(list(data.values())) @@ -897,71 +1429,136 @@ def _forward_collect(self, data: dict) -> dict: output = default_decollate(output) return {i: d for i, d in zip(data_id, output)} - def _process_transition(self, obs: Any, model_output: dict, timestep: namedtuple) -> dict: + def _process_transition(self, obs: torch.Tensor, policy_output: Dict[str, torch.Tensor], + timestep: namedtuple) -> Dict[str, torch.Tensor]: """ Overview: - Generate dict type transition data from inputs. + Process and pack one timestep transition data into a dict, which can be directly used for training and \ + saved in replay buffer. For PPO, it contains obs, next_obs, action, reward, done, logit, value. Arguments: - - obs (:obj:`Any`): Env observation - - model_output (:obj:`dict`): Output of collect model, including at least ['action'] - - timestep (:obj:`namedtuple`): Output after env step, including at least ['obs', 'reward', 'done']\ - (here 'obs' indicates obs after env step). + - obs (:obj:`torch.Tensor`): The env observation of current timestep, such as stacked 2D image in Atari. + - policy_output (:obj:`Dict[str, torch.Tensor]`): The output of the policy network with the observation \ + as input. For PPO, it contains the state value, action and the logit of the action. + - timestep (:obj:`namedtuple`): The execution result namedtuple returned by the environment step method, \ + except all the elements have been transformed into tensor data. Usually, it contains the next obs, \ + reward, done, info, etc. Returns: - - transition (:obj:`dict`): Dict type transition data. + - transition (:obj:`Dict[str, torch.Tensor]`): The processed transition data of the current timestep. + + .. note:: + ``next_obs`` is used to calculate nstep return when necessary, so we place in into transition by default. \ + You can delete this field to save memory occupancy if you do not need nstep return. """ transition = { 'obs': obs, 'next_obs': timestep.obs, - 'logit': model_output['logit'], - 'action': model_output['action'], - 'value': model_output['value'], + 'logit': policy_output['logit'], + 'action': policy_output['action'], + 'value': policy_output['value'], 'reward': timestep.reward, 'done': timestep.done, } return transition - def _get_train_sample(self, data: list) -> Union[None, List[Any]]: - r""" + def _get_train_sample(self, transitions: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ Overview: - Get the trajectory and calculate GAE, return one data to cache for next time calculation + For a given trajectory (transitions, a list of transition) data, process it into a list of sample that \ + can be used for training directly. In PPO, a train sample is a processed transition with new computed \ + ``traj_flag`` and ``adv`` field. This method is usually used in collectors to execute necessary \ + RL data preprocessing before training, which can help learner amortize revelant time consumption. \ + In addition, you can also implement this method as an identity function and do the data processing \ + in ``self._forward_learn`` method. Arguments: - - data (:obj:`list`): The trajectory's cache + - transitions (:obj:`List[Dict[str, Any]`): The trajectory data (a list of transition), each element is \ + the same format as the return value of ``self._process_transition`` method. Returns: - - samples (:obj:`dict`): The training samples generated + - samples (:obj:`List[Dict[str, Any]]`): The processed train samples, each element is the similar format \ + as input transitions, but may contain more data for training, such as GAE advantage. """ - data = get_gae_with_default_last_value( + data = transitions + data = to_device(data, self._device) + for transition in data: + transition['traj_flag'] = copy.deepcopy(transition['done']) + data[-1]['traj_flag'] = True + + if self._cfg.learn.ignore_done: + data[-1]['done'] = False + + if data[-1]['done']: + last_value = torch.zeros_like(data[-1]['value']) + else: + with torch.no_grad(): + last_value = self._collect_model.forward( + unsqueeze(data[-1]['next_obs'], 0), mode='compute_actor_critic' + )['value'] + if len(last_value.shape) == 2: # multi_agent case: + last_value = last_value.squeeze(0) + if self._value_norm: + last_value *= self._running_mean_std.std + for i in range(len(data)): + data[i]['value'] *= self._running_mean_std.std + data = get_gae( data, - data[-1]['done'], + to_device(last_value, self._device), gamma=self._gamma, gae_lambda=self._gae_lambda, cuda=False, ) + if self._value_norm: + for i in range(len(data)): + data[i]['value'] /= self._running_mean_std.std + if not self._nstep_return: return get_train_sample(data, self._unroll_len) else: return get_nstep_return_data(data, self._nstep) def _init_eval(self) -> None: - r""" + """ Overview: - Evaluate mode init method. Called by ``self.__init__``. - Init eval model with argmax strategy. + Initialize the eval mode of policy, including related attributes and modules. For PPOOff, it contains the \ + eval model to select optimial action (e.g. greedily select action with argmax mechanism in discrete action). + This method will be called in ``__init__`` method if ``eval`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_eval`` method, you'd better name them \ + with prefix ``_eval_`` to avoid conflict with other modes, such as ``self._eval_attr1``. """ - self._eval_model = model_wrap(self._model, wrapper_name='argmax_sample') + assert self._cfg.action_space in ["continuous", "discrete", "hybrid"] + self._action_space = self._cfg.action_space + if self._action_space == 'continuous': + self._eval_model = model_wrap(self._model, wrapper_name='deterministic_sample') + elif self._action_space == 'discrete': + self._eval_model = model_wrap(self._model, wrapper_name='argmax_sample') + elif self._action_space == 'hybrid': + self._eval_model = model_wrap(self._model, wrapper_name='hybrid_deterministic_argmax_sample') self._eval_model.reset() - def _forward_eval(self, data: dict) -> dict: - r""" + def _forward_eval(self, data: Dict[int, Any]) -> Dict[int, Any]: + """ Overview: - Forward function of eval mode, similar to ``self._forward_collect``. + Policy forward function of eval mode (evaluation policy performance by interacting with envs). Forward \ + means that the policy gets some necessary data (mainly observation) from the envs and then returns the \ + action to interact with the envs. ``_forward_eval`` in PPO often uses deterministic sample method to get \ + actions while ``_forward_collect`` usually uses stochastic sample method for balance exploration and \ + exploitation. Arguments: - - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \ - values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer. + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. Returns: - - output (:obj:`Dict[int, Any]`): The dict of predicting action for the interaction with env. - ReturnsKeys - - necessary: ``action`` + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action. The \ + key of the dict is the same as the input data, i.e. environment id. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for PPOOffPolicy: ``ding.policy.tests.test_ppo``. """ data_id = list(data.keys()) data = default_collate(list(data.values())) @@ -975,13 +1572,20 @@ def _forward_eval(self, data: dict) -> dict: output = default_decollate(output) return {i: d for i, d in zip(data_id, output)} - def default_model(self) -> Tuple[str, List[str]]: - return 'vac', ['ding.model.template.vac'] - def _monitor_vars_learn(self) -> List[str]: - return super()._monitor_vars_learn() + [ - 'policy_loss', 'value_loss', 'entropy_loss', 'adv_abs_max', 'approx_kl', 'clipfrac' + """ + Overview: + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. + Returns: + - necessary_keys (:obj:`List[str]`): The list of the necessary keys to be logged. + """ + variables = super()._monitor_vars_learn() + [ + 'policy_loss', 'value', 'value_loss', 'entropy_loss', 'adv_abs_max', 'approx_kl', 'clipfrac' ] + if self._action_space == 'continuous': + variables += ['mu_mean', 'sigma_mean', 'sigma_grad', 'act'] + return variables @POLICY_REGISTRY.register('ppo_stdim') @@ -989,6 +1593,8 @@ class PPOSTDIMPolicy(PPOPolicy): """ Overview: Policy class of on policy version PPO algorithm with ST-DIM auxiliary model. + PPO paper link: https://arxiv.org/abs/1707.06347. + ST-DIM paper link: https://arxiv.org/abs/1906.08226. """ config = dict( # (str) RL policy register name (refer to function "POLICY_REGISTRY"). @@ -1012,6 +1618,8 @@ class PPOSTDIMPolicy(PPOPolicy): multi_agent=False, # (bool) Whether to need policy data in process transition transition_with_policy_data=True, + # (float) The loss weight of the auxiliary model to the main loss. + aux_loss_weight=0.001, aux_model=dict( # (int) the encoding size (of each head) to apply contrastive loss. encode_shape=64, @@ -1022,44 +1630,48 @@ class PPOSTDIMPolicy(PPOPolicy): # (float) a parameter to adjust the polarity between positive and negative samples. temperature=1.0, ), + # learn_mode config learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, + # (int) After collecting n_sample/n_episode data, how many epoches to train models. + # Each epoch means the one entire passing of training data. epoch_per_collect=10, + # (int) How many samples in a training batch. batch_size=64, + # (float) The step size of gradient descent. learning_rate=3e-4, - # ============================================================== - # The following configs is algorithm-specific - # ============================================================== - # (float) The loss weight of value network, policy network weight is set to 1 + # (float) The loss weight of value network, policy network weight is set to 1. value_weight=0.5, - # (float) The loss weight of entropy regularization, policy network weight is set to 1 + # (float) The loss weight of entropy regularization, policy network weight is set to 1. entropy_weight=0.0, - # (float) PPO clip ratio, defaults to 0.2 + # (float) PPO clip ratio, defaults to 0.2. clip_ratio=0.2, - # (bool) Whether to use advantage norm in a whole training batch + # (bool) Whether to use advantage norm in a whole training batch. adv_norm=True, + # (bool) Whether to use value norm with running mean and std in the whole training process. value_norm=True, + # (bool) Whether to enable special network parameters initialization scheme in PPO, such as orthogonal init. ppo_param_init=True, + # (str) The gradient clip operation type used in PPO, ['clip_norm', clip_value', 'clip_momentum_norm']. grad_clip_type='clip_norm', + # (float) The gradient clip target value used in PPO. + # If ``grad_clip_type`` is 'clip_norm', then the maximum of gradient will be normalized to this value. grad_clip_value=0.5, + # (bool) Whether ignore done (usually for max step termination env). ignore_done=False, ), + # collect_mode config collect=dict( - # (int) Only one of [n_sample, n_episode] shoule be set + # (int) How many training samples collected in one collection procedure. + # Only one of [n_sample, n_episode] shoule be set. # n_sample=64, # (int) Cut trajectories into pieces with length "unroll_len". unroll_len=1, - # ============================================================== - # The following configs is algorithm-specific - # ============================================================== # (float) Reward's future discount factor, aka. gamma. discount_factor=0.99, - # (float) GAE lambda factor for the balance of bias and variance(1-step td and mc) + # (float) GAE lambda factor for the balance of bias and variance (1-step td and mc). gae_lambda=0.95, ), - eval=dict(), - aux_loss_weight=0.001, + eval=dict(), # for compability ) def _init_learn(self) -> None: @@ -1178,7 +1790,7 @@ def _forward_learn(self, data: Dict[str, Any]) -> Dict[str, Any]: # the BP process of the auxiliary network self._aux_optimizer.zero_grad() aux_loss_learn.backward() - if self._cfg.learn.multi_gpu: + if self._cfg.multi_gpu: self.sync_gradients(self._aux_model) self._aux_optimizer.step() @@ -1192,13 +1804,13 @@ def _forward_learn(self, data: Dict[str, Any]) -> Dict[str, Any]: if self._action_space == 'continuous': ppo_batch = ppo_data( output['logit'], batch['logit'], batch['action'], output['value'], batch['value'], adv, - batch['return'], batch['weight'] + batch['return'], batch['weight'], None ) ppo_loss, ppo_info = ppo_error_continuous(ppo_batch, self._clip_ratio) elif self._action_space == 'discrete': ppo_batch = ppo_data( output['logit'], batch['logit'], batch['action'], output['value'], batch['value'], adv, - batch['return'], batch['weight'] + batch['return'], batch['weight'], None ) ppo_loss, ppo_info = ppo_error(ppo_batch, self._clip_ratio) @@ -1248,6 +1860,13 @@ def _forward_learn(self, data: Dict[str, Any]) -> Dict[str, Any]: return return_infos def _state_dict_learn(self) -> Dict[str, Any]: + """ + Overview: + Return the state_dict of learn mode, usually including model, optimizer and aux_optimizer for \ + representation learning. + Returns: + - state_dict (:obj:`Dict[str, Any]`): The dict of current policy learn state, for saving and restoring. + """ return { 'model': self._learn_model.state_dict(), 'optimizer': self._optimizer.state_dict(), @@ -1255,9 +1874,27 @@ def _state_dict_learn(self) -> Dict[str, Any]: } def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: + """ + Overview: + Load the state_dict variable into policy learn mode. + Arguments: + - state_dict (:obj:`Dict[str, Any]`): The dict of policy learn state saved before. + + .. tip:: + If you want to only load some parts of model, you can simply set the ``strict`` argument in \ + load_state_dict to ``False``, or refer to ``ding.torch_utils.checkpoint_helper`` for more \ + complicated operation. + """ self._learn_model.load_state_dict(state_dict['model']) self._optimizer.load_state_dict(state_dict['optimizer']) self._aux_optimizer.load_state_dict(state_dict['aux_optimizer']) def _monitor_vars_learn(self) -> List[str]: + """ + Overview: + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. + Returns: + - necessary_keys (:obj:`List[str]`): The list of the necessary keys to be logged. + """ return super()._monitor_vars_learn() + ["aux_loss_learn", "aux_loss_eval"] diff --git a/ding/policy/ppof.py b/ding/policy/ppof.py new file mode 100644 index 0000000000..81e605384c --- /dev/null +++ b/ding/policy/ppof.py @@ -0,0 +1,359 @@ +from typing import List, Dict, Any, Tuple, Union, Callable, Optional +from collections import namedtuple +from easydict import EasyDict +import copy +import random +import numpy as np +import torch +import treetensor.torch as ttorch +from torch.optim import AdamW + +from ding.rl_utils import ppo_data, ppo_error, ppo_policy_error, ppo_policy_data, gae, gae_data, ppo_error_continuous, \ + get_gae, ppo_policy_error_continuous, ArgmaxSampler, MultinomialSampler, ReparameterizationSampler, MuSampler, \ + HybridStochasticSampler, HybridDeterminsticSampler, value_transform, value_inv_transform, symlog, inv_symlog +from ding.utils import POLICY_REGISTRY, RunningMeanStd + + +@POLICY_REGISTRY.register('ppof') +class PPOFPolicy: + config = dict( + type='ppo', + on_policy=True, + cuda=True, + action_space='discrete', + discount_factor=0.99, + gae_lambda=0.95, + # learn + epoch_per_collect=10, + batch_size=64, + learning_rate=3e-4, + # learningrate scheduler, which the format is (10000, 0.1) + lr_scheduler=None, + weight_decay=0, + value_weight=0.5, + entropy_weight=0.01, + clip_ratio=0.2, + adv_norm=True, + value_norm='baseline', + ppo_param_init=True, + grad_norm=0.5, + # collect + n_sample=128, + unroll_len=1, + # eval + deterministic_eval=True, + # model + model=dict(), + ) + mode = ['learn', 'collect', 'eval'] + + @classmethod + def default_config(cls: type) -> EasyDict: + cfg = EasyDict(copy.deepcopy(cls.config)) + cfg.cfg_type = cls.__name__ + 'Dict' + return cfg + + @classmethod + def default_model(cls: type) -> Callable: + from .model import PPOFModel + return PPOFModel + + def __init__(self, cfg: "EasyDict", model: torch.nn.Module, enable_mode: List[str] = None) -> None: + self._cfg = cfg + if model is None: + self._model = self.default_model() + else: + self._model = model + if self._cfg.cuda and torch.cuda.is_available(): + self._device = 'cuda' + self._model.cuda() + else: + self._device = 'cpu' + assert self._cfg.action_space in ["continuous", "discrete", "hybrid", 'multi_discrete'] + self._action_space = self._cfg.action_space + if self._cfg.ppo_param_init: + self._model_param_init() + + if enable_mode is None: + enable_mode = self.mode + self.enable_mode = enable_mode + if 'learn' in enable_mode: + self._optimizer = AdamW( + self._model.parameters(), + lr=self._cfg.learning_rate, + weight_decay=self._cfg.weight_decay, + ) + # define linear lr scheduler + if self._cfg.lr_scheduler is not None: + epoch_num, min_lr_lambda = self._cfg.lr_scheduler + + self._lr_scheduler = torch.optim.lr_scheduler.LambdaLR( + self._optimizer, + lr_lambda=lambda epoch: max(1.0 - epoch * (1.0 - min_lr_lambda) / epoch_num, min_lr_lambda) + ) + + if self._cfg.value_norm: + self._running_mean_std = RunningMeanStd(epsilon=1e-4, device=self._device) + if 'collect' in enable_mode: + if self._action_space == 'discrete': + self._collect_sampler = MultinomialSampler() + elif self._action_space == 'continuous': + self._collect_sampler = ReparameterizationSampler() + elif self._action_space == 'hybrid': + self._collect_sampler = HybridStochasticSampler() + if 'eval' in enable_mode: + if self._action_space == 'discrete': + if self._cfg.deterministic_eval: + self._eval_sampler = ArgmaxSampler() + else: + self._eval_sampler = MultinomialSampler() + elif self._action_space == 'continuous': + if self._cfg.deterministic_eval: + self._eval_sampler = MuSampler() + else: + self._eval_sampler = ReparameterizationSampler() + elif self._action_space == 'hybrid': + if self._cfg.deterministic_eval: + self._eval_sampler = HybridDeterminsticSampler() + else: + self._eval_sampler = HybridStochasticSampler() + # for compatibility + self.learn_mode = self + self.collect_mode = self + self.eval_mode = self + + def _model_param_init(self): + for n, m in self._model.named_modules(): + if isinstance(m, torch.nn.Linear): + torch.nn.init.orthogonal_(m.weight) + torch.nn.init.zeros_(m.bias) + if self._action_space in ['continuous', 'hybrid']: + for m in list(self._model.critic.modules()) + list(self._model.actor.modules()): + if isinstance(m, torch.nn.Linear): + # orthogonal initialization + torch.nn.init.orthogonal_(m.weight, gain=np.sqrt(2)) + torch.nn.init.zeros_(m.bias) + # init log sigma + if self._action_space == 'continuous': + torch.nn.init.constant_(self._model.actor_head.log_sigma_param, -0.5) + for m in self._model.actor_head.mu.modules(): + if isinstance(m, torch.nn.Linear): + torch.nn.init.zeros_(m.bias) + m.weight.data.copy_(0.01 * m.weight.data) + elif self._action_space == 'hybrid': # actor_head[1]: ReparameterizationHead, for action_args + if hasattr(self._model.actor_head[1], 'log_sigma_param'): + torch.nn.init.constant_(self._model.actor_head[1].log_sigma_param, -0.5) + for m in self._model.actor_head[1].mu.modules(): + if isinstance(m, torch.nn.Linear): + torch.nn.init.zeros_(m.bias) + m.weight.data.copy_(0.01 * m.weight.data) + + def forward(self, data: ttorch.Tensor) -> Dict[str, Any]: + return_infos = [] + self._model.train() + bs = self._cfg.batch_size + data = data[:self._cfg.n_sample // bs * bs] # rounding + + # outer training loop + for epoch in range(self._cfg.epoch_per_collect): + # recompute adv + with torch.no_grad(): + # get the value dictionary + # In popart, the dictionary has two keys: 'pred' and 'unnormalized_pred' + value = self._model.compute_critic(data.obs) + next_value = self._model.compute_critic(data.next_obs) + reward = data.reward + + assert self._cfg.value_norm in ['popart', 'value_rescale', 'symlog', 'baseline'],\ + 'Not supported value normalization! Value normalization supported: \ + popart, value rescale, symlog, baseline' + + if self._cfg.value_norm == 'popart': + unnormalized_value = value['unnormalized_pred'] + unnormalized_next_value = value['unnormalized_pred'] + + mu = self._model.critic_head.popart.mu + sigma = self._model.critic_head.popart.sigma + reward = (reward - mu) / sigma + + value = value['pred'] + next_value = next_value['pred'] + elif self._cfg.value_norm == 'value_rescale': + value = value_inv_transform(value['pred']) + next_value = value_inv_transform(next_value['pred']) + elif self._cfg.value_norm == 'symlog': + value = inv_symlog(value['pred']) + next_value = inv_symlog(next_value['pred']) + elif self._cfg.value_norm == 'baseline': + value = value['pred'] * self._running_mean_std.std + next_value = next_value['pred'] * self._running_mean_std.std + + traj_flag = data.get('traj_flag', None) # traj_flag indicates termination of trajectory + adv_data = gae_data(value, next_value, reward, data.done, traj_flag) + data.adv = gae(adv_data, self._cfg.discount_factor, self._cfg.gae_lambda) + + unnormalized_returns = value + data.adv # In popart, this return is normalized + + if self._cfg.value_norm == 'popart': + self._model.critic_head.popart.update_parameters((data.reward).unsqueeze(1)) + elif self._cfg.value_norm == 'value_rescale': + value = value_transform(value) + unnormalized_returns = value_transform(unnormalized_returns) + elif self._cfg.value_norm == 'symlog': + value = symlog(value) + unnormalized_returns = symlog(unnormalized_returns) + elif self._cfg.value_norm == 'baseline': + value /= self._running_mean_std.std + unnormalized_returns /= self._running_mean_std.std + self._running_mean_std.update(unnormalized_returns.cpu().numpy()) + data.value = value + data.return_ = unnormalized_returns + + # inner training loop + split_data = ttorch.split(data, self._cfg.batch_size) + random.shuffle(list(split_data)) + for batch in split_data: + output = self._model.compute_actor_critic(batch.obs) + adv = batch.adv + if self._cfg.adv_norm: + # Normalize advantage in a train_batch + adv = (adv - adv.mean()) / (adv.std() + 1e-8) + + # Calculate ppo error + if self._action_space == 'continuous': + ppo_batch = ppo_data( + output.logit, batch.logit, batch.action, output.value, batch.value, adv, batch.return_, None + ) + ppo_loss, ppo_info = ppo_error_continuous(ppo_batch, self._cfg.clip_ratio) + elif self._action_space == 'discrete': + ppo_batch = ppo_data( + output.logit, batch.logit, batch.action, output.value, batch.value, adv, batch.return_, None + ) + ppo_loss, ppo_info = ppo_error(ppo_batch, self._cfg.clip_ratio) + elif self._action_space == 'hybrid': + # discrete part (discrete policy loss and entropy loss) + ppo_discrete_batch = ppo_policy_data( + output.logit.action_type, batch.logit.action_type, batch.action.action_type, adv, None + ) + ppo_discrete_loss, ppo_discrete_info = ppo_policy_error(ppo_discrete_batch, self._cfg.clip_ratio) + # continuous part (continuous policy loss and entropy loss, value loss) + ppo_continuous_batch = ppo_data( + output.logit.action_args, batch.logit.action_args, batch.action.action_args, output.value, + batch.value, adv, batch.return_, None + ) + ppo_continuous_loss, ppo_continuous_info = ppo_error_continuous( + ppo_continuous_batch, self._cfg.clip_ratio + ) + # sum discrete and continuous loss + ppo_loss = type(ppo_continuous_loss)( + ppo_continuous_loss.policy_loss + ppo_discrete_loss.policy_loss, ppo_continuous_loss.value_loss, + ppo_continuous_loss.entropy_loss + ppo_discrete_loss.entropy_loss + ) + ppo_info = type(ppo_continuous_info)( + max(ppo_continuous_info.approx_kl, ppo_discrete_info.approx_kl), + max(ppo_continuous_info.clipfrac, ppo_discrete_info.clipfrac) + ) + wv, we = self._cfg.value_weight, self._cfg.entropy_weight + total_loss = ppo_loss.policy_loss + wv * ppo_loss.value_loss - we * ppo_loss.entropy_loss + + self._optimizer.zero_grad() + total_loss.backward() + torch.nn.utils.clip_grad_norm_(self._model.parameters(), self._cfg.grad_norm) + self._optimizer.step() + + return_info = { + 'cur_lr': self._optimizer.defaults['lr'], + 'total_loss': total_loss.item(), + 'policy_loss': ppo_loss.policy_loss.item(), + 'value_loss': ppo_loss.value_loss.item(), + 'entropy_loss': ppo_loss.entropy_loss.item(), + 'adv_max': adv.max().item(), + 'adv_mean': adv.mean().item(), + 'value_mean': output.value.mean().item(), + 'value_max': output.value.max().item(), + 'approx_kl': ppo_info.approx_kl, + 'clipfrac': ppo_info.clipfrac, + } + if self._action_space == 'continuous': + return_info.update( + { + 'action': batch.action.float().mean().item(), + 'mu_mean': output.logit.mu.mean().item(), + 'sigma_mean': output.logit.sigma.mean().item(), + } + ) + elif self._action_space == 'hybrid': + return_info.update( + { + 'action': batch.action.action_args.float().mean().item(), + 'mu_mean': output.logit.action_args.mu.mean().item(), + 'sigma_mean': output.logit.action_args.sigma.mean().item(), + } + ) + return_infos.append(return_info) + + if self._cfg.lr_scheduler is not None: + self._lr_scheduler.step() + + return return_infos + + def state_dict(self) -> Dict[str, Any]: + state_dict = { + 'model': self._model.state_dict(), + } + if 'learn' in self.enable_mode: + state_dict['optimizer'] = self._optimizer.state_dict() + return state_dict + + def load_state_dict(self, state_dict: Dict[str, Any]) -> None: + self._model.load_state_dict(state_dict['model']) + if 'learn' in self.enable_mode: + self._optimizer.load_state_dict(state_dict['optimizer']) + + def collect(self, data: ttorch.Tensor) -> ttorch.Tensor: + self._model.eval() + with torch.no_grad(): + output = self._model.compute_actor_critic(data) + action = self._collect_sampler(output.logit) + output.action = action + return output + + def process_transition(self, obs: ttorch.Tensor, inference_output: dict, timestep: namedtuple) -> ttorch.Tensor: + return ttorch.as_tensor( + { + 'obs': obs, + 'next_obs': timestep.obs, + 'action': inference_output.action, + 'logit': inference_output.logit, + 'value': inference_output.value, + 'reward': timestep.reward, + 'done': timestep.done, + } + ) + + def eval(self, data: ttorch.Tensor) -> ttorch.Tensor: + self._model.eval() + with torch.no_grad(): + logit = self._model.compute_actor(data) + action = self._eval_sampler(logit) + return ttorch.as_tensor({'logit': logit, 'action': action}) + + def monitor_vars(self) -> List[str]: + variables = [ + 'cur_lr', + 'policy_loss', + 'value_loss', + 'entropy_loss', + 'adv_max', + 'adv_mean', + 'approx_kl', + 'clipfrac', + 'value_max', + 'value_mean', + ] + if self._action_space in ['action', 'mu_mean', 'sigma_mean']: + variables += ['mu_mean', 'sigma_mean', 'action'] + return variables + + def reset(self, env_id_list: Optional[List[int]] = None) -> None: + pass diff --git a/ding/policy/prompt_awr.py b/ding/policy/prompt_awr.py new file mode 100644 index 0000000000..4b39057d22 --- /dev/null +++ b/ding/policy/prompt_awr.py @@ -0,0 +1,274 @@ +from collections import namedtuple +from typing import List, Dict, Any, Tuple, Union + +import torch + +from ding.model import model_wrap +from ding.rl_utils import get_train_sample +from ding.torch_utils import Adam, to_device +from ding.utils import POLICY_REGISTRY +from ding.utils.data import default_collate, default_decollate +from .base_policy import Policy + + +@POLICY_REGISTRY.register('prompt_awr') +class PromptAWRPolicy(Policy): + """ + Overview: + Policy class of AWR (Advantage Weighted Regression) algorithm, proposed in https://arxiv.org/abs/1910.00177. + Especially, this policy is designed for training a language model policy. + In this policy, the environment's observation includes the current context, a list of optional actions + (strings). The final output of the policy is a set of optional actions with a size of ``shot_number``. + """ + config = dict( + # (str) Name of the registered RL policy (refer to the "register_policy" function). + type='prompt_awr', + # (bool) Flag to enable CUDA for model computation. + cuda=False, + # (bool) Flag for using on-policy training (training policy is the same as the behavior policy). + on_policy=False, + # (bool) Flag for enabling priority experience replay. Must be False when priority_IS_weight is False. + priority=False, + # (bool) Flag for using Importance Sampling weights to correct updates. Requires `priority` to be True. + priority_IS_weight=False, + # (str) Type of action space used in the policy, with valid options ['discrete', 'continuous']. + action_space='discrete', + # (int) The number of actions that can be done simultaneously in one timestep. + shot_number=1, + # learn_mode configuration + learn=dict( + # (int) Number of updates per data collection. A2C requires this to be set to 1. + update_per_collect=1, + # (int) Batch size for learning. + batch_size=64, + # (float) Learning rate for optimizer. + learning_rate=0.001, + # (Tuple[float, float]) Coefficients used for computing running averages of gradient and its square. + betas=(0.9, 0.999), + # (float) Term added to the denominator to improve numerical stability in optimizer. + eps=1e-8, + # (float) Maximum norm for gradients. + grad_norm=0.5, + # (float) Scaling factor for value network loss relative to policy network loss. + value_weight=0.5, + # (float) Coefficient that controls the exp scale in awr algorithm. + beta=1.0, + # (float) Weight of entropy regularization in the loss function. + entropy_weight=0.001, + # (Tuple[float, float]) The range of adv. Value that exceeds this range will be clipped. + adv_range=(-0.5, 0.5), + # (bool) If set to True, the 'done' signals that indicate the end of an episode due to environment time + # limits are disregarded. By default, this is set to False. This setting is particularly useful for tasks + # that have a predetermined episode length, such as HalfCheetah and various other MuJoCo environments, + # where the maximum length is capped at 1000 steps. When enabled, any 'done' signal triggered by reaching + # the maximum episode steps will be overridden to 'False'. This ensures the accurate calculation of the + # Temporal Difference (TD) error, using the formula `gamma * (1 - done) * next_v + reward`, + # even when the episode surpasses the predefined step limit. + ignore_done=False, + ), + # collect_mode configuration + collect=dict( + # (int) The length of rollout for data collection. + unroll_len=1, + # (float) Discount factor for calculating future rewards, typically in the range [0, 1]. + discount_factor=0.9, + # (float) Trade-off parameter for balancing TD-error and Monte Carlo error in GAE. + gae_lambda=0.95, + ), + # eval_mode configuration (kept empty for compatibility purposes) + eval=dict(), + ) + + def default_model(self) -> Tuple[str, List[str]]: + """ + Overview: + Returns the default model configuration used by the AWR algorithm. ``__init__`` method will \ + automatically call this method to get the default model setting and create model. + + Returns: + - model_info (:obj:`Tuple[str, List[str]]`): \ + Tuple containing the registered model name and model's import_names. + """ + return 'language_transformer', ['ding.model.template.language_transformer'] + + def _init_learn(self) -> None: + """ + Overview: + Initialize the learn mode of policy, including related attributes and modules. For AWR, it mainly \ + contains optimizer, algorithm-specific arguments such as value_weight, entropy_weight, adv_norm + and grad_norm, and main model. \ + This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. + """ + assert self._cfg.action_space == "discrete" + # Optimizer + self._optimizer = Adam( + self._model.parameters(), + lr=self._cfg.learn.learning_rate, + betas=self._cfg.learn.betas, + eps=self._cfg.learn.eps + ) + + # Algorithm config + self._priority = self._cfg.priority + self._priority_IS_weight = self._cfg.priority_IS_weight + self._value_weight = self._cfg.learn.value_weight + self._entropy_weight = self._cfg.learn.entropy_weight + self._adv_norm = self._cfg.learn.adv_norm + self._grad_norm = self._cfg.learn.grad_norm + + # Main and target models + self._learn_model = self._model + + def _forward_learn(self, data: List[Dict[str, Any]]) -> Dict[str, Any]: + # Data preprocessing operations, such as stack data, cpu to cuda device + self._learn_model.train() + + for i in range(0, len(data), self._cfg.learn.batch_size): + batch = default_collate(data[i:i + self._cfg.learn.batch_size]) + if self._cuda: + batch = to_device(batch, self._device) + + # Prepare train_sample (the question to be answered) and the candidate_samples (the prompts to be selected) + train_samples, cand_samples = batch["obs"]["train_sample"], batch["obs"]["candidate_samples"] + for cand_n in range(len(cand_samples)): + cand_samples[cand_n] = cand_samples[cand_n][0] + output = self._learn_model.forward(train_samples, cand_samples, mode='compute_actor_critic') + return_ = batch['return'] + + # Calculate AWR loss + real_act = batch['action'] + + # Ensure the shape of real_act is: (B, shot_number) + if len(real_act.shape) == 1: + real_act = real_act.unsqueeze(-1) + + # Calculate different parts of loss. + total_policy_loss, total_entropy_loss, total_value_loss = 0, 0, 0 + for shot_n in range(self._cfg.shot_number): + log_prob = output['dist'].log_prob(real_act[:, shot_n]) + # Clamp the adv for better stability. + adv = torch.clamp( + return_ - batch['value'], min=self._cfg.learn.norm_range[0], max=self._cfg.learn.norm_range[1] + ) + # The policy loss for AWR algorithm. + policy_loss = -(log_prob * torch.exp(adv / self._cfg.learn.beta)).mean() + total_policy_loss += policy_loss + # The value loss for AWR algorithm. + value_loss = ((return_ - output['value']) ** 2).mean() + total_value_loss += value_loss + # The entropy loss for AWR algorithm. + total_entropy_loss += -self._cfg.learn.entropy_weight * output['dist'].entropy().mean() + total_loss = total_entropy_loss + total_policy_loss + total_value_loss + + self._optimizer.zero_grad() + total_loss.backward() + + grad_norm = torch.nn.utils.clip_grad_norm_( + list(self._learn_model.parameters()), + max_norm=self._grad_norm, + ) + self._optimizer.step() + + return { + 'cur_lr': self._optimizer.param_groups[0]['lr'], + 'total_loss': total_loss.item(), + 'policy_loss': total_policy_loss.item(), + 'entropy_loss': total_entropy_loss.item(), + 'value_loss': total_value_loss.item(), + 'return_abs_max': return_.abs().max().item(), + 'grad_norm': grad_norm, + } + + def _init_collect(self) -> None: + self._unroll_len = self._cfg.collect.unroll_len + self._gamma = self._cfg.collect.discount_factor + self._collect_model = model_wrap(self._model, wrapper_name='combination_multinomial_sample') + + def _forward_collect(self, data: Dict[int, Any]) -> Dict[int, Any]: + """ + Overview: + Policy forward function of collect mode (collecting training data by interacting with envs). Forward means \ + that the policy gets some necessary data (mainly observation) from the envs and then returns the output \ + data, such as the action to interact with the envs. + Arguments: + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. + Returns: + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action and \ + other necessary data for learn mode defined in ``self._process_transition`` method. The key of the \ + dict is the same as the input data, i.e. environment id. + """ + data_id = list(data.keys()) + data = default_collate(list(data.values())) + self._model.eval() + with torch.no_grad(): + # Prepare train_sample (the question to be answered) and the candidate_samples (the prompts to be selected) + for ii in range(len(data['candidate_samples'])): + data['candidate_samples'][ii] = data['candidate_samples'][ii][0] + output = self._collect_model.forward( + self._cfg.shot_number, data['train_sample'], data['candidate_samples'], mode="compute_actor_critic" + ) + if self._cuda: + output = to_device(output, 'cpu') + output = default_decollate(output) + return {i: d for i, d in zip(data_id, output)} + + def _process_transition(self, obs: Any, policy_output: Dict[str, torch.Tensor], + timestep: namedtuple) -> Dict[str, torch.Tensor]: + return { + 'obs': obs, + 'action': policy_output['action'], + 'value': policy_output['value'], + 'reward': timestep.reward, + 'done': timestep.done, + } + + def _get_train_sample(self, data: list) -> Union[None, List[Any]]: + r""" + Overview: + Get the trajectory and the n step return data, then sample from the n_step return data + Arguments: + - data (:obj:`list`): The trajectory's buffer list + Returns: + - samples (:obj:`dict`): The training samples generated + """ + if self._cfg.learn.ignore_done: + raise NotImplementedError + + R = 0. + for i in reversed(range(len(data))): + R = self._gamma * R + data[i]['reward'] + data[i]['return'] = R + return get_train_sample(data, self._unroll_len) + + def _init_eval(self) -> None: + self._eval_model = model_wrap(self._model, wrapper_name='combination_argmax_sample') + + def _forward_eval(self, data: dict) -> dict: + data_id = list(data.keys()) + data = default_collate(list(data.values())) + self._model.eval() + with torch.no_grad(): + # Prepare train_sample (the question to be answered) and the candidate_samples (the prompts to be selected) + for ii in range(len(data['candidate_samples'])): + data['candidate_samples'][ii] = data['candidate_samples'][ii][0] + output = self._eval_model.forward(self._cfg.shot_number, data['train_sample'], data['candidate_samples']) + if self._cuda: + output = to_device(output, 'cpu') + output = default_decollate(output) + return {i: d for i, d in zip(data_id, output)} + + def _monitor_vars_learn(self) -> List[str]: + return super()._monitor_vars_learn() + \ + ['policy_loss', 'entropy_loss', 'return_abs_max', 'grad_norm', 'value_loss'] diff --git a/ding/policy/prompt_pg.py b/ding/policy/prompt_pg.py new file mode 100644 index 0000000000..a76e0e5faf --- /dev/null +++ b/ding/policy/prompt_pg.py @@ -0,0 +1,210 @@ +from typing import List, Dict, Any, Tuple, Union +from collections import namedtuple +import torch + +from ding.rl_utils import get_train_sample +from ding.torch_utils import Adam, to_device +from ding.utils import POLICY_REGISTRY, split_data_generator +from ding.utils.data import default_collate, default_decollate +from .base_policy import Policy +from ..model import model_wrap + + +@POLICY_REGISTRY.register('prompt_pg') +class PromptPGPolicy(Policy): + r""" + Overview: + Policy class of Prompt Policy Gradient (PromptPG) algorithm. + Link of the original paper: https://arxiv.org/abs/2209.14610 + """ + config = dict( + # (string) RL policy register name (refer to function "register_policy"). + type='prompt_pg', + # (bool) whether to use cuda for network. + cuda=True, + # (bool) whether use on-policy training pipeline(behaviour policy and training policy are the same) + on_policy=True, # for pg strictly on policy algorithm, this line should not be modified by users + # (bool) whether to use deterministic action for evaluation. + deterministic_eval=True, + # (int) The number of actions that can be done simultaneously in one timestep. + shot_number=1, + learn=dict( + # (int) the number of samples for one update. + batch_size=64, + # (float) the step size of one gradient descend. + learning_rate=0.001, + # ============================================================== + # The following configs is algorithm-specific + # ============================================================== + # (float) loss weight of the entropy regularization, the weight of policy network is set to 1 + entropy_weight=0.01, + # (float) max grad norm value. + grad_norm=5, + # (bool) whether to ignore done signal for non-termination env. + ignore_done=False, + ), + collect=dict( + # (int) collect n_sample data, train model n_iteration times + # n_episode=8, + # (int) trajectory unroll length + unroll_len=1, + # ============================================================== + # The following configs is algorithm-specific + # ============================================================== + # (float) discount factor for future reward, defaults int [0, 1] + discount_factor=0, + collector=dict(get_train_sample=True), + ), + eval=dict(), + ) + + def default_model(self) -> Tuple[str, List[str]]: + return 'language_transformer', ['ding.model.template.language_transformer'] + + def _init_learn(self) -> None: + r""" + Overview: + Learn mode init method. Called by ``self.__init__``. + Init the optimizer, algorithm config, main and target models. + """ + # Optimizer + self._optimizer = Adam(self._model.parameters(), lr=self._cfg.learn.learning_rate) + + self._entropy_weight = self._cfg.learn.entropy_weight + self._grad_norm = self._cfg.learn.grad_norm + self._learn_model = self._model # for compatibility + + def _forward_learn(self, data: dict) -> Dict[str, Any]: + r""" + Overview: + Forward and backward function of learn mode. + Arguments: + - data (:obj:`dict`): Dict type data, including at least ['obs', 'action', 'reward'] + Returns: + - info_dict (:obj:`Dict[str, Any]`): Including current lr and loss. + """ + self._model.train() + + return_infos = [] + for i in range(0, len(data), self._cfg.learn.batch_size): + batch = default_collate(data[i:i + self._cfg.learn.batch_size]) + if self._cuda: + batch = to_device(batch, self._device) + + # Prepare train_sample (the question to be answered) and the candidate_samples (the prompts to be selected) + train_samples, cand_samples = batch["obs"]["train_sample"], batch["obs"]["candidate_samples"] + for ii in range(len(cand_samples)): + cand_samples[ii] = cand_samples[ii][0] + output = self._learn_model.forward(train_samples, cand_samples) + return_ = batch['return'] + + # calculate PG loss + real_act = batch['action'] # shape: (B, shot_number) + if len(real_act.shape) == 1: + real_act = real_act.unsqueeze(-1) + # Calculate loss. + total_policy_loss, total_entropy_loss = 0, 0 + for ii in range(self._cfg.shot_number): + log_prob = output['dist'].log_prob(real_act[:, ii]) + policy_loss = -(log_prob * return_).mean() + total_policy_loss += policy_loss + total_entropy_loss += -self._cfg.learn.entropy_weight * output['dist'].entropy().mean() + total_loss = total_entropy_loss + total_policy_loss + + # update + self._optimizer.zero_grad() + total_loss.backward() + + grad_norm = torch.nn.utils.clip_grad_norm_( + list(self._learn_model.parameters()), + max_norm=self._grad_norm, + ) + self._optimizer.step() + + # only record last updates information in logger + return_info = { + 'cur_lr': self._optimizer.param_groups[0]['lr'], + 'total_loss': total_loss.item(), + 'policy_loss': total_policy_loss.item(), + 'entropy_loss': total_entropy_loss.item(), + 'return_abs_max': return_.abs().max().item(), + 'grad_norm': grad_norm, + } + return_infos.append(return_info) + return return_infos + + def _init_collect(self) -> None: + self._unroll_len = self._cfg.collect.unroll_len + self._gamma = self._cfg.collect.discount_factor + self._collect_model = model_wrap(self._model, wrapper_name='combination_multinomial_sample') + + def _forward_collect(self, data: dict) -> dict: + data_id = list(data.keys()) + data = default_collate(list(data.values())) + self._model.eval() + with torch.no_grad(): + # Prepare train_sample (the question to be answered) and the candidate_samples (the prompts to be selected) + for ii in range(len(data['candidate_samples'])): + data['candidate_samples'][ii] = data['candidate_samples'][ii][0] + output = self._collect_model.forward(self._cfg.shot_number, data['train_sample'], data['candidate_samples']) + if self._cuda: + output = to_device(output, 'cpu') + output = default_decollate(output) + return {i: d for i, d in zip(data_id, output)} + + def _process_transition(self, obs: Any, model_output: dict, timestep: namedtuple) -> dict: + r""" + Overview: + Generate dict type transition data from inputs. + Arguments: + - obs (:obj:`Any`): Env observation + - model_output (:obj:`dict`): Output of collect model, including at least ['action'] + - timestep (:obj:`namedtuple`): Output after env step, including at least ['obs', 'reward', 'done'] \ + (here 'obs' indicates obs after env step). + Returns: + - transition (:obj:`dict`): Dict type transition data. + """ + return { + 'obs': obs, + 'action': model_output['action'], + 'reward': timestep.reward, + 'done': timestep.done, + } + + def _get_train_sample(self, data: list) -> Union[None, List[Any]]: + r""" + Overview: + Get the trajectory and the n step return data, then sample from the n_step return data + Arguments: + - data (:obj:`list`): The trajectory's buffer list + Returns: + - samples (:obj:`dict`): The training samples generated + """ + if self._cfg.learn.ignore_done: + raise NotImplementedError + + R = 0. + for i in reversed(range(len(data))): + R = self._gamma * R + data[i]['reward'] + data[i]['return'] = R + return get_train_sample(data, self._unroll_len) + + def _init_eval(self) -> None: + self._eval_model = model_wrap(self._model, wrapper_name='combination_argmax_sample') + + def _forward_eval(self, data: dict) -> dict: + data_id = list(data.keys()) + data = default_collate(list(data.values())) + self._model.eval() + with torch.no_grad(): + # Prepare train_sample (the question to be answered) and the candidate_samples (the prompts to be selected) + for ii in range(len(data['candidate_samples'])): + data['candidate_samples'][ii] = data['candidate_samples'][ii][0] + output = self._eval_model.forward(self._cfg.shot_number, data['train_sample'], data['candidate_samples']) + if self._cuda: + output = to_device(output, 'cpu') + output = default_decollate(output) + return {i: d for i, d in zip(data_id, output)} + + def _monitor_vars_learn(self) -> List[str]: + return super()._monitor_vars_learn() + ['policy_loss', 'entropy_loss', 'return_abs_max', 'grad_norm'] diff --git a/ding/policy/qgpo.py b/ding/policy/qgpo.py new file mode 100644 index 0000000000..2bd3884e11 --- /dev/null +++ b/ding/policy/qgpo.py @@ -0,0 +1,268 @@ +############################################################# +# This QGPO model is a modification implementation from https://github.com/ChenDRAG/CEP-energy-guided-diffusion +############################################################# + +from typing import List, Dict, Any +import torch +from ding.utils import POLICY_REGISTRY +from ding.utils.data import default_collate +from ding.torch_utils import to_device +from .base_policy import Policy + + +@POLICY_REGISTRY.register('qgpo') +class QGPOPolicy(Policy): + """ + Overview: + Policy class of QGPO algorithm (https://arxiv.org/abs/2304.12824). + Contrastive Energy Prediction for Exact Energy-Guided Diffusion Sampling in Offline Reinforcement Learning + Interfaces: + ``__init__``, ``forward``, ``learn``, ``eval``, ``state_dict``, ``load_state_dict`` + """ + + config = dict( + # (str) RL policy register name (refer to function "POLICY_REGISTRY"). + type='qgpo', + # (bool) Whether to use cuda for network. + cuda=False, + # (bool type) on_policy: Determine whether on-policy or off-policy. + # on-policy setting influences the behaviour of buffer. + # Default False in QGPO. + on_policy=False, + multi_agent=False, + model=dict( + qgpo_critic=dict( + # (float) The scale of the energy guidance when training qt. + # \pi_{behavior}\exp(f(s,a)) \propto \pi_{behavior}\exp(alpha * Q(s,a)) + alpha=3, + # (float) The scale of the energy guidance when training q0. + # \mathcal{T}Q(s,a)=r(s,a)+\mathbb{E}_{s'\sim P(s'|s,a),a'\sim\pi_{support}(a'|s')}Q(s',a') + # \pi_{support} \propto \pi_{behavior}\exp(q_alpha * Q(s,a)) + q_alpha=1, + ), + device='cuda', + # obs_dim + # action_dim + ), + learn=dict( + # learning rate for behavior model training + learning_rate=1e-4, + # batch size during the training of behavior model + batch_size=4096, + # batch size during the training of q value + batch_size_q=256, + # number of fake action support + M=16, + # number of diffusion time steps + diffusion_steps=15, + # training iterations when behavior model is fixed + behavior_policy_stop_training_iter=600000, + # training iterations when energy-guided policy begin training + energy_guided_policy_begin_training_iter=600000, + # training iterations when q value stop training, default None means no limit + q_value_stop_training_iter=1100000, + ), + eval=dict( + # energy guidance scale for policy in evaluation + # \pi_{evaluation} \propto \pi_{behavior}\exp(guidance_scale * alpha * Q(s,a)) + guidance_scale=[0.0, 1.0, 2.0, 3.0, 5.0, 8.0, 10.0], + ), + ) + + def _init_learn(self) -> None: + """ + Overview: + Learn mode initialization method. For QGPO, it mainly contains the optimizer, \ + algorithm-specific arguments such as qt_update_momentum, discount, behavior_policy_stop_training_iter, \ + energy_guided_policy_begin_training_iter and q_value_stop_training_iter, etc. + This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. + """ + self.cuda = self._cfg.cuda + + self.behavior_model_optimizer = torch.optim.Adam( + self._model.score_model.parameters(), lr=self._cfg.learn.learning_rate + ) + self.q_optimizer = torch.optim.Adam(self._model.q.q0.parameters(), lr=3e-4) + self.qt_optimizer = torch.optim.Adam(self._model.q.qt.parameters(), lr=3e-4) + + self.qt_update_momentum = 0.005 + self.discount = 0.99 + + self.behavior_policy_stop_training_iter = self._cfg.learn.behavior_policy_stop_training_iter + self.energy_guided_policy_begin_training_iter = self._cfg.learn.energy_guided_policy_begin_training_iter + self.q_value_stop_training_iter = self._cfg.learn.q_value_stop_training_iter + + def _forward_learn(self, data: dict) -> Dict[str, Any]: + """ + Overview: + Forward function for learning mode. + The training of QGPO algorithm is based on contrastive energy prediction, \ + which needs true action and fake action. The true action is sampled from the dataset, and the fake action \ + is sampled from the action support generated by the behavior policy. + The training process is divided into two stages: + 1. Train the behavior model, which is modeled as a diffusion model by parameterizing the score function. + 2. Train the Q function by fake action support generated by the behavior model. + 3. Train the energy-guided policy by the Q function. + Arguments: + - data (:obj:`dict`): Dict type data. + Returns: + - result (:obj:`dict`): Dict type data of algorithm results. + """ + + if self.cuda: + data = to_device(data, self._device) + + s = data['s'] + a = data['a'] + r = data['r'] + s_ = data['s_'] + d = data['d'] + fake_a = data['fake_a'] + fake_a_ = data['fake_a_'] + + # training behavior model + if self.behavior_policy_stop_training_iter > 0: + + behavior_model_training_loss = self._model.score_model_loss_fn(a, s) + + self.behavior_model_optimizer.zero_grad() + behavior_model_training_loss.backward() + self.behavior_model_optimizer.step() + + self.behavior_policy_stop_training_iter -= 1 + behavior_model_training_loss = behavior_model_training_loss.item() + else: + behavior_model_training_loss = 0 + + # training Q function + self.energy_guided_policy_begin_training_iter -= 1 + self.q_value_stop_training_iter -= 1 + if self.energy_guided_policy_begin_training_iter < 0: + if self.q_value_stop_training_iter > 0: + q0_loss = self._model.q_loss_fn(a, s, r, s_, d, fake_a_, discount=self.discount) + + self.q_optimizer.zero_grad() + q0_loss.backward() + self.q_optimizer.step() + + # Update target + for param, target_param in zip(self._model.q.q0.parameters(), self._model.q.q0_target.parameters()): + target_param.data.copy_( + self.qt_update_momentum * param.data + (1 - self.qt_update_momentum) * target_param.data + ) + + q0_loss = q0_loss.item() + + else: + q0_loss = 0 + qt_loss = self._model.qt_loss_fn(s, fake_a) + + self.qt_optimizer.zero_grad() + qt_loss.backward() + self.qt_optimizer.step() + + qt_loss = qt_loss.item() + + else: + q0_loss = 0 + qt_loss = 0 + + total_loss = behavior_model_training_loss + q0_loss + qt_loss + + return dict( + total_loss=total_loss, + behavior_model_training_loss=behavior_model_training_loss, + q0_loss=q0_loss, + qt_loss=qt_loss, + ) + + def _init_collect(self) -> None: + """ + Overview: + Collect mode initialization method. Not supported for QGPO. + """ + pass + + def _forward_collect(self) -> None: + """ + Overview: + Forward function for collect mode. Not supported for QGPO. + """ + pass + + def _init_eval(self) -> None: + """ + Overview: + Eval mode initialization method. For QGPO, it mainly contains the guidance_scale and diffusion_steps, etc. + This method will be called in ``__init__`` method if ``eval`` field is in ``enable_field``. + """ + + self.diffusion_steps = self._cfg.eval.diffusion_steps + + def _forward_eval(self, data: dict, guidance_scale: float) -> dict: + """ + Overview: + Forward function for eval mode. The eval process is based on the energy-guided policy, \ + which is modeled as a diffusion model by parameterizing the score function. + Arguments: + - data (:obj:`dict`): Dict type data. + - guidance_scale (:obj:`float`): The scale of the energy guidance. + Returns: + - output (:obj:`dict`): Dict type data of algorithm output. + """ + + data_id = list(data.keys()) + states = default_collate(list(data.values())) + actions = self._model.select_actions( + states, diffusion_steps=self.diffusion_steps, guidance_scale=guidance_scale + ) + output = actions + + return {i: {"action": d} for i, d in zip(data_id, output)} + + def _get_train_sample(self, transitions: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Overview: + Get the train sample from the replay buffer, currently not supported for QGPO. + Arguments: + - transitions (:obj:`List[Dict[str, Any]]`): The data from replay buffer. + Returns: + - samples (:obj:`List[Dict[str, Any]]`): The data for training. + """ + pass + + def _process_transition(self) -> None: + """ + Overview: + Process the transition data, currently not supported for QGPO. + """ + pass + + def _state_dict_learn(self) -> Dict[str, Any]: + """ + Overview: + Return the state dict for saving. + Returns: + - state_dict (:obj:`Dict[str, Any]`): Dict type data of state dict. + """ + return { + 'model': self._model.state_dict(), + 'behavior_model_optimizer': self.behavior_model_optimizer.state_dict(), + } + + def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: + """ + Overview: + Load the state dict. + Arguments: + - state_dict (:obj:`Dict[str, Any]`): Dict type data of state dict. + """ + self._model.load_state_dict(state_dict['model']) + self.behavior_model_optimizer.load_state_dict(state_dict['behavior_model_optimizer']) + + def _monitor_vars_learn(self) -> List[str]: + """ + Overview: + Return the variables names to be monitored. + """ + return ['total_loss', 'behavior_model_training_loss', 'q0_loss', 'qt_loss'] diff --git a/ding/policy/qmix.py b/ding/policy/qmix.py index 4e6268eee5..ff1d66f7c8 100644 --- a/ding/policy/qmix.py +++ b/ding/policy/qmix.py @@ -1,7 +1,7 @@ -from typing import List, Dict, Any, Tuple, Union, Optional +from typing import List, Dict, Any, Tuple, Optional from collections import namedtuple -import torch import copy +import torch from ding.torch_utils import RMSprop, to_device from ding.rl_utils import v_1step_td_data, v_1step_td_error, get_train_sample @@ -15,12 +15,8 @@ class QMIXPolicy(Policy): """ Overview: - Policy class of QMIX algorithm. QMIX is a multi model reinforcement learning algorithm, \ - you can view the paper in the following link https://arxiv.org/abs/1803.11485 - Interface: - _init_learn, _data_preprocess_learn, _forward_learn, _reset_learn, _state_dict_learn, _load_state_dict_learn \ - _init_collect, _forward_collect, _reset_collect, _process_transition, _init_eval, _forward_eval \ - _reset_eval, _get_train_sample, default_model + Policy class of QMIX algorithm. QMIX is a multi-agent reinforcement learning algorithm, \ + you can view the paper in the following link https://arxiv.org/abs/1803.11485. Config: == ==================== ======== ============== ======================================== ======================= ID Symbol Type Default Value Description Other(Shape) @@ -55,50 +51,48 @@ class QMIXPolicy(Policy): priority=False, # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. priority_IS_weight=False, + # learn_mode config learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, + # (int) How many updates(iterations) to train after collector's one collection. + # Bigger "update_per_collect" means bigger off-policy. + # collect data -> update policy-> collect data -> ... update_per_collect=20, + # (int) How many samples in a training batch. batch_size=32, + # (float) The step size of gradient descent. learning_rate=0.0005, clip_value=100, - # ============================================================== - # The following configs is algorithm-specific - # ============================================================== - # (float) Target network update momentum parameter. - # in [0, 1]. + # (float) Target network update momentum parameter, in [0, 1]. target_update_theta=0.008, - # (float) The discount factor for future rewards, - # in [0, 1]. + # (float) The discount factor for future rewards, in [0, 1]. discount_factor=0.99, - # (bool) Whether to use double DQN mechanism(target q for surpassing over estimation) + # (bool) Whether to use double DQN mechanism(target q for surpassing over estimation). double_q=False, ), + # collect_mode config collect=dict( - # (int) Only one of [n_sample, n_episode] shoule be set - # n_episode=32, - # (int) Cut trajectories into pieces with length "unroll_len", the length of timesteps + # (int) How many training samples collected in one collection procedure. + # In each collect phase, we collect a total of sequence samples, a sample with length unroll_len. + # n_sample=32, + # (int) Split trajectories into pieces with length ``unroll_len``, the length of timesteps # in each forward when training. In qmix, it is greater than 1 because there is RNN. unroll_len=10, ), - eval=dict(), + eval=dict(), # for compatibility other=dict( eps=dict( - # (str) Type of epsilon decay + # (str) Type of epsilon decay. type='exp', # (float) Start value for epsilon decay, in [0, 1]. - # 0 means not use epsilon decay. start=1, # (float) Start value for epsilon decay, in [0, 1]. end=0.05, - # (int) Decay length(env step) + # (int) Decay length(env step). decay=50000, ), replay_buffer=dict( + # (int) Maximum size of replay buffer. Usually, larger buffer size is better. replay_buffer_size=5000, - # (int) The maximum reuse times of each data - max_reuse=1e+9, - max_staleness=1e+9, ), ), ) @@ -119,17 +113,25 @@ def default_model(self) -> Tuple[str, List[str]]: def _init_learn(self) -> None: """ Overview: - Learn mode init method. Called by ``self.__init__``. - Init the learner model of QMIXPolicy - Arguments: - .. note:: + Initialize the learn mode of policy, including some attributes and modules. For QMIX, it mainly contains \ + optimizer, algorithm-specific arguments such as gamma, main and target model. Because of the use of RNN, \ + all the models should be wrappered with ``hidden_state`` which needs to be initialized with proper size. + This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. - The _init_learn method takes the argument from the self._cfg.learn in the config file + .. tip:: + For multi-agent algorithm, we often need to use ``agent_num`` to initialize some necessary variables. - - learning_rate (:obj:`float`): The learning rate fo the optimizer - - gamma (:obj:`float`): The discount factor + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. - agent_num (:obj:`int`): Since this is a multi-agent algorithm, we need to input the agent num. - - batch_size (:obj:`int`): Need batch size info to init hidden_state plugins """ self._priority = self._cfg.priority self._priority_IS_weight = self._cfg.priority_IS_weight @@ -165,8 +167,8 @@ def _init_learn(self) -> None: self._learn_model.reset() self._target_model.reset() - def _data_preprocess_learn(self, data: List[Any]) -> dict: - r""" + def _data_preprocess_learn(self, data: List[Dict[str, Any]]) -> Dict[str, Any]: + """ Overview: Preprocess the data to fit the required data format for learning Arguments: @@ -183,22 +185,35 @@ def _data_preprocess_learn(self, data: List[Any]) -> dict: data['done'] = data['done'].float() return data - def _forward_learn(self, data: dict) -> Dict[str, Any]: - r""" + def _forward_learn(self, data: List[List[Dict[str, Any]]]) -> Dict[str, Any]: + """ Overview: - Forward and backward function of learn mode. + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data (trajectory for QMIX) from the replay buffer and then \ + returns the output result, including various training information such as loss, q value, grad_norm. Arguments: - - data (:obj:`Dict[str, Any]`): Dict type data, a batch of data for training, values are torch.Tensor or \ - np.ndarray or dict/list combinations. + - data (:obj:`List[List[Dict[int, Any]]]`): The input data used for policy forward, including a batch of \ + training samples. For each dict element, the key of the dict is the name of data items and the \ + value is the corresponding data. Usually, the value is torch.Tensor or np.ndarray or there dict/list \ + combinations. In the ``_forward_learn`` method, data often need to first be stacked in the time and \ + batch dimension by the utility functions ``self._data_preprocess_learn``. \ + For QMIX, each element in list is a trajectory with the length of ``unroll_len``, and the element in \ + trajectory list is a dict containing at least the following keys: ``obs``, ``action``, ``prev_state``, \ + ``reward``, ``next_obs``, ``done``. Sometimes, it also contains other keys such as ``weight`` \ + and ``value_gamma``. Returns: - - info_dict (:obj:`Dict[str, Any]`): Dict type data, a info dict indicated training result, which will be \ - recorded in text log and tensorboard, values are python scalar or a list of scalars. - ArgumentsKeys: - - necessary: ``obs``, ``next_obs``, ``action``, ``reward``, ``weight``, ``prev_state``, ``done`` - ReturnsKeys: - - necessary: ``cur_lr``, ``total_loss`` - - cur_lr (:obj:`float`): Current learning rate - - total_loss (:obj:`float`): The calculated loss + - info_dict (:obj:`Dict[str, Any]`): The information dict that indicated training result, which will be \ + recorded in text log and tensorboard, values must be python scalar or a list of scalars. For the \ + detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for QMIXPolicy: ``ding.policy.tests.test_qmix``. """ data = self._data_preprocess_learn(data) # ==================== @@ -251,21 +266,24 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: } def _reset_learn(self, data_id: Optional[List[int]] = None) -> None: - r""" + """ Overview: - Reset learn model to the state indicated by data_id + Reset some stateful variables for learn mode when necessary, such as the hidden state of RNN or the \ + memory bank of some special algortihms. If ``data_id`` is None, it means to reset all the stateful \ + varaibles. Otherwise, it will reset the stateful variables according to the ``data_id``. For example, \ + different trajectories in ``data_id`` will have different hidden state in RNN. Arguments: - - data_id (:obj:`Optional[List[int]]`): The id that store the state and we will reset\ - the model state to the state indicated by data_id + - data_id (:obj:`Optional[List[int]]`): The id of the data, which is used to reset the stateful variables \ + (i.e. RNN hidden_state in QMIX) specified by ``data_id``. """ self._learn_model.reset(data_id=data_id) def _state_dict_learn(self) -> Dict[str, Any]: - r""" + """ Overview: - Return the state_dict of learn mode, usually including model and optimizer. + Return the state_dict of learn mode, usually including model, target_model and optimizer. Returns: - - state_dict (:obj:`Dict[str, Any]`): the dict of current policy learn state, for saving and restoring. + - state_dict (:obj:`Dict[str, Any]`): The dict of current policy learn state, for saving and restoring. """ return { 'model': self._learn_model.state_dict(), @@ -278,7 +296,7 @@ def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: Overview: Load the state_dict variable into policy learn mode. Arguments: - - state_dict (:obj:`Dict[str, Any]`): the dict of policy learn state saved before. + - state_dict (:obj:`Dict[str, Any]`): The dict of policy learn state saved before. .. tip:: If you want to only load some parts of model, you can simply set the ``strict`` argument in \ @@ -290,11 +308,17 @@ def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: self._optimizer.load_state_dict(state_dict['optimizer']) def _init_collect(self) -> None: - r""" + """ Overview: - Collect mode init method. Called by ``self.__init__``. - Init traj and unroll length, collect model. - Enable the eps_greedy_sample and the hidden_state plugin. + Initialize the collect mode of policy, including related attributes and modules. For QMIX, it contains the \ + collect_model to balance the exploration and exploitation with epsilon-greedy sample mechanism and \ + maintain the hidden state of rnn. Besides, there are some initialization operations about other \ + algorithm-specific arguments such as burnin_step, unroll_len and nstep. + This method will be called in ``__init__`` method if ``collect`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_collect`` method, you'd better name them \ + with prefix ``_collect_`` to avoid conflict with other modes, such as ``self._collect_attr1``. """ self._unroll_len = self._cfg.collect.unroll_len self._collect_model = model_wrap( @@ -307,18 +331,34 @@ def _init_collect(self) -> None: self._collect_model = model_wrap(self._collect_model, wrapper_name='eps_greedy_sample') self._collect_model.reset() - def _forward_collect(self, data: dict, eps: float) -> dict: - r""" + def _forward_collect(self, data: Dict[int, Any], eps: float) -> Dict[int, Any]: + """ Overview: - Forward function for collect mode with eps_greedy + Policy forward function of collect mode (collecting training data by interacting with envs). Forward means \ + that the policy gets some necessary data (mainly observation) from the envs and then returns the output \ + data, such as the action to interact with the envs. Besides, this policy also needs ``eps`` argument for \ + exploration, i.e., classic epsilon-greedy exploration strategy. Arguments: - - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \ - values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer. - - eps (:obj:`float`): epsilon value for exploration, which is decayed by collected env step. + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. + - eps (:obj:`float`): The epsilon value for exploration. Returns: - - output (:obj:`Dict[int, Any]`): Dict type data, including at least inferred action according to input obs. - ReturnsKeys - - necessary: ``action`` + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action and \ + other necessary data (prev_state) for learn mode defined in ``self._process_transition`` method. The \ + key of the dict is the same as the input data, i.e. environment id. + + .. note:: + RNN's hidden states are maintained in the policy, so we don't need pass them into data but to reset the \ + hidden states with ``_reset_collect`` method when episode ends. Besides, the previous hidden states are \ + necessary for training, so we need to return them in ``_process_transition`` method. + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for QMIXPolicy: ``ding.policy.tests.test_qmix``. """ data_id = list(data.keys()) data = default_collate(list(data.values())) @@ -334,43 +374,73 @@ def _forward_collect(self, data: dict, eps: float) -> dict: return {i: d for i, d in zip(data_id, output)} def _reset_collect(self, data_id: Optional[List[int]] = None) -> None: - r""" + """ Overview: - Reset collect model to the state indicated by data_id + Reset some stateful variables for eval mode when necessary, such as the hidden state of RNN or the \ + memory bank of some special algortihms. If ``data_id`` is None, it means to reset all the stateful \ + varaibles. Otherwise, it will reset the stateful variables according to the ``data_id``. For example, \ + different environments/episodes in evaluation in ``data_id`` will have different hidden state in RNN. Arguments: - - data_id (:obj:`Optional[List[int]]`): The id that store the state and we will reset\ - the model state to the state indicated by data_id + - data_id (:obj:`Optional[List[int]]`): The id of the data, which is used to reset the stateful variables \ + (i.e., RNN hidden_state in QMIX) specified by ``data_id``. """ self._collect_model.reset(data_id=data_id) - def _process_transition(self, obs: Any, model_output: dict, timestep: namedtuple) -> dict: - r""" + def _process_transition(self, obs: torch.Tensor, policy_output: Dict[str, torch.Tensor], + timestep: namedtuple) -> Dict[str, torch.Tensor]: + """ Overview: - Generate dict type transition data from inputs. + Process and pack one timestep transition data into a dict, which can be directly used for training and \ + saved in replay buffer. For QMIX, it contains obs, next_obs, action, prev_state, reward, done. Arguments: - - obs (:obj:`Any`): Env observation - - model_output (:obj:`dict`): Output of collect model, including at least ['action', 'prev_state'] - - timestep (:obj:`namedtuple`): Output after env step, including at least ['obs', 'reward', 'done']\ - (here 'obs' indicates obs after env step). + - obs (:obj:`torch.Tensor`): The env observation of current timestep, usually including ``agent_obs`` \ + and ``global_obs`` in multi-agent environment like MPE and SMAC. + - policy_output (:obj:`Dict[str, torch.Tensor]`): The output of the policy network with the observation \ + as input. For QMIX, it contains the action and the prev_state of RNN. + - timestep (:obj:`namedtuple`): The execution result namedtuple returned by the environment step method, \ + except all the elements have been transformed into tensor data. Usually, it contains the next obs, \ + reward, done, info, etc. Returns: - - transition (:obj:`dict`): Dict type transition data, including 'obs', 'next_obs', 'prev_state',\ - 'action', 'reward', 'done' + - transition (:obj:`Dict[str, torch.Tensor]`): The processed transition data of the current timestep. """ transition = { 'obs': obs, 'next_obs': timestep.obs, - 'prev_state': model_output['prev_state'], - 'action': model_output['action'], + 'prev_state': policy_output['prev_state'], + 'action': policy_output['action'], 'reward': timestep.reward, 'done': timestep.done, } return transition + def _get_train_sample(self, transitions: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Overview: + For a given trajectory (transitions, a list of transition) data, process it into a list of sample that \ + can be used for training directly. In QMIX, a train sample is processed transitions with unroll_len \ + length. This method is usually used in collectors to execute necessary \ + RL data preprocessing before training, which can help learner amortize revelant time consumption. \ + In addition, you can also implement this method as an identity function and do the data processing \ + in ``self._forward_learn`` method. + Arguments: + - transitions (:obj:`List[Dict[str, Any]`): The trajectory data (a list of transition), each element is \ + the same format as the return value of ``self._process_transition`` method. + Returns: + - samples (:obj:`List[Dict[str, Any]]`): The processed train samples, each sample is a fixed-length \ + trajectory, and each element in a sample is the similar format as input transitions. + """ + return get_train_sample(transitions, self._unroll_len) + def _init_eval(self) -> None: - r""" + """ Overview: - Evaluate mode init method. Called by ``self.__init__``. - Init eval model with argmax strategy and the hidden_state plugin. + Initialize the eval mode of policy, including related attributes and modules. For QMIX, it contains the \ + eval model to greedily select action with argmax q_value mechanism and main the hidden state. + This method will be called in ``__init__`` method if ``eval`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_eval`` method, you'd better name them \ + with prefix ``_eval_`` to avoid conflict with other modes, such as ``self._eval_attr1``. """ self._eval_model = model_wrap( self._model, @@ -383,16 +453,31 @@ def _init_eval(self) -> None: self._eval_model.reset() def _forward_eval(self, data: dict) -> dict: - r""" + """ Overview: - Forward function of eval mode, similar to ``self._forward_collect``. + Policy forward function of eval mode (evaluation policy performance by interacting with envs). Forward \ + means that the policy gets some necessary data (mainly observation) from the envs and then returns the \ + action to interact with the envs. ``_forward_eval`` often use argmax sample method to get actions that \ + q_value is the highest. Arguments: - - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \ - values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer. + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. Returns: - - output (:obj:`Dict[int, Any]`): The dict of predicting action for the interaction with env. - ReturnsKeys - - necessary: ``action`` + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action. The \ + key of the dict is the same as the input data, i.e. environment id. + + .. note:: + RNN's hidden states are maintained in the policy, so we don't need pass them into data but to reset the \ + hidden states with ``_reset_eval`` method when the episode ends. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for QMIXPolicy: ``ding.policy.tests.test_qmix``. """ data_id = list(data.keys()) data = default_collate(list(data.values())) @@ -408,31 +493,24 @@ def _forward_eval(self, data: dict) -> dict: return {i: d for i, d in zip(data_id, output)} def _reset_eval(self, data_id: Optional[List[int]] = None) -> None: - r""" - Overview: - Reset eval model to the state indicated by data_id - Arguments: - - data_id (:obj:`Optional[List[int]]`): The id that store the state and we will reset\ - the model state to the state indicated by data_id """ - self._eval_model.reset(data_id=data_id) - - def _get_train_sample(self, data: list) -> Union[None, List[Any]]: - r""" Overview: - Get the train sample from trajectory. + Reset some stateful variables for eval mode when necessary, such as the hidden state of RNN or the \ + memory bank of some special algortihms. If ``data_id`` is None, it means to reset all the stateful \ + varaibles. Otherwise, it will reset the stateful variables according to the ``data_id``. For example, \ + different environments/episodes in evaluation in ``data_id`` will have different hidden state in RNN. Arguments: - - data (:obj:`list`): The trajectory's cache - Returns: - - samples (:obj:`dict`): The training samples generated + - data_id (:obj:`Optional[List[int]]`): The id of the data, which is used to reset the stateful variables \ + (i.e., RNN hidden_state in QMIX) specified by ``data_id``. """ - return get_train_sample(data, self._unroll_len) + self._eval_model.reset(data_id=data_id) def _monitor_vars_learn(self) -> List[str]: - r""" + """ Overview: - Return variables' name if variables are to used in monitor. + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. Returns: - - vars (:obj:`List[str]`): Variables' name list. + - necessary_keys (:obj:`List[str]`): The list of the necessary keys to be logged. """ return ['cur_lr', 'total_loss', 'total_q', 'target_total_q', 'grad_norm', 'target_reward_total_q'] diff --git a/ding/policy/qrdqn.py b/ding/policy/qrdqn.py index 7d81a922a3..d2ed004464 100644 --- a/ding/policy/qrdqn.py +++ b/ding/policy/qrdqn.py @@ -15,7 +15,10 @@ class QRDQNPolicy(DQNPolicy): r""" Overview: - Policy class of QRDQN algorithm. + Policy class of QRDQN algorithm. QRDQN (https://arxiv.org/pdf/1710.10044.pdf) is a distributional RL \ + algorithm, which is an extension of DQN. The main idea of QRDQN is to use quantile regression to \ + estimate the quantile of the distribution of the return value, and then use the quantile to calculate \ + the quantile loss. Config: == ==================== ======== ============== ======================================== ======================= @@ -58,8 +61,7 @@ class QRDQNPolicy(DQNPolicy): # (int) N-step reward for target q_value estimation nstep=1, learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, + # How many updates(iterations) to train after collector's one collection. # Bigger "update_per_collect" means bigger off-policy. # collect data -> update policy-> collect data -> ... @@ -98,13 +100,31 @@ class QRDQNPolicy(DQNPolicy): ) def default_model(self) -> Tuple[str, List[str]]: + """ + Overview: + Return this algorithm default neural network model setting for demonstration. ``__init__`` method will \ + automatically call this method to get the default model setting and create model. + + Returns: + - model_info (:obj:`Tuple[str, List[str]]`): The registered model name and model's import_names. + """ return 'qrdqn', ['ding.model.template.q_learning'] def _init_learn(self) -> None: - r""" + """ Overview: - Learn mode init method. Called by ``self.__init__``. - Init the optimizer, algorithm config, main and target models. + Initialize the learn mode of policy, including related attributes and modules. For QRDQN, it mainly \ + contains optimizer, algorithm-specific arguments such as nstep and gamma. This method \ + also executes some special network initializations and prepares running mean/std monitor for value. + This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. """ self._priority = self._cfg.priority # Optimizer @@ -126,14 +146,38 @@ def _init_learn(self) -> None: self._target_model.reset() def _forward_learn(self, data: dict) -> Dict[str, Any]: - r""" + """ Overview: - Forward and backward function of learn mode. + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data from the replay buffer and then returns the output \ + result, including various training information such as loss, current lr. + Arguments: - - data (:obj:`dict`): Dict type data, including at least ['obs', 'action', 'reward', 'next_obs'] + - data (:obj:`dict`): Input data used for policy forward, including the \ + collected training samples from replay buffer. For each element in dict, the key of the \ + dict is the name of data items and the value is the corresponding data. Usually, the value is \ + torch.Tensor or np.ndarray or there dict/list combinations. In the ``_forward_learn`` method, data \ + often need to first be stacked in the batch dimension by some utility functions such as \ + ``default_preprocess_learn``. \ + For QRDQN, each element in list is a dict containing at least the following keys: ``obs``, \ + ``action``, ``reward``, ``next_obs``. Sometimes, it also contains other keys such as ``weight``. + Returns: - - info_dict (:obj:`Dict[str, Any]`): Including current lr and loss. + - info_dict (:obj:`Dict[str, Any]`): The output result dict of forward learn, \ + containing current lr, total_loss and priority. When discrete action satisfying \ + len(data['action'])==1, it also could contain ``action_distribution`` which is used \ + to draw histogram on tensorboard. For more information, please refer to the :class:`DQNPolicy`. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for QRDQNPolicy: ``ding.policy.tests.test_qrdqn``. """ + data = default_preprocess_learn( data, use_priority=self._priority, ignore_done=self._cfg.learn.ignore_done, use_nstep=True ) @@ -166,7 +210,7 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: # ==================== self._optimizer.zero_grad() loss.backward() - if self._cfg.learn.multi_gpu: + if self._cfg.multi_gpu: self.sync_gradients(self._learn_model) self._optimizer.step() diff --git a/ding/policy/qtran.py b/ding/policy/qtran.py index 07e7a37636..c75b942eb8 100644 --- a/ding/policy/qtran.py +++ b/ding/policy/qtran.py @@ -54,8 +54,6 @@ class QTRANPolicy(Policy): # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. priority_IS_weight=False, learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, update_per_collect=20, batch_size=32, learning_rate=0.0005, diff --git a/ding/policy/r2d2.py b/ding/policy/r2d2.py index 4615026757..420d238cad 100644 --- a/ding/policy/r2d2.py +++ b/ding/policy/r2d2.py @@ -15,16 +15,16 @@ @POLICY_REGISTRY.register('r2d2') class R2D2Policy(Policy): - r""" + """ Overview: Policy class of R2D2, from paper `Recurrent Experience Replay in Distributed Reinforcement Learning` . - R2D2 proposes that several tricks should be used to improve upon DRQN, - namely some recurrent experience replay tricks such as burn-in. + R2D2 proposes that several tricks should be used to improve upon DRQN, namely some recurrent experience replay \ + tricks and the burn-in mechanism for off-policy training. Config: == ==================== ======== ============== ======================================== ======================= ID Symbol Type Default Value Description Other(Shape) == ==================== ======== ============== ======================================== ======================= - 1 ``type`` str dqn | RL policy register name, refer to | This arg is optional, + 1 ``type`` str r2d2 | RL policy register name, refer to | This arg is optional, | registry ``POLICY_REGISTRY`` | a placeholder 2 ``cuda`` bool False | Whether to use cuda for network | This arg can be diff- | erent from modes @@ -68,13 +68,10 @@ class R2D2Policy(Policy): cuda=False, # (bool) Whether the RL algorithm is on-policy or off-policy. on_policy=False, - # (bool) Whether use priority(priority sample, IS weight, update priority) + # (bool) Whether to use priority(priority sample, IS weight, update priority) priority=True, - # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. + # (bool) Whether to use Importance Sampling Weight to correct biased update. If True, priority must be True. priority_IS_weight=True, - # ============================================================== - # The following configs are algorithm-specific - # ============================================================== # (float) Reward's future discount factor, aka. gamma. discount_factor=0.997, # (int) N-step reward for target q_value estimation @@ -85,66 +82,103 @@ class R2D2Policy(Policy): # (int) the trajectory length to unroll the RNN network minus # the timestep of burnin operation learn_unroll_len=80, + # learn_mode config learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, + # (int) The number of training updates (iterations) to perform after each data collection by the collector. + # A larger "update_per_collect" value implies a more off-policy approach. + # The whole pipeline process follows this cycle: collect data -> update policy -> collect data -> ... update_per_collect=1, + # (int) The number of samples in a training batch. batch_size=64, + # (float) The step size of gradient descent, determining the rate of learning. learning_rate=0.0001, - # ============================================================== - # The following configs are algorithm-specific - # ============================================================== # (int) Frequence of target network update. # target_update_freq=100, target_update_theta=0.001, # (bool) whether use value_rescale function for predicted value value_rescale=True, + # (bool) Whether ignore done(usually for max step termination env). + # Note: Gym wraps the MuJoCo envs by default with TimeLimit environment wrappers. + # These limit HalfCheetah, and several other MuJoCo envs, to max length of 1000. + # However, interaction with HalfCheetah always gets done with done is False, + # Since we inplace done==True with done==False to keep + # TD-error accurate computation(``gamma * (1 - done) * next_v + reward``), + # when the episode step is greater than max episode step. ignore_done=False, ), + # collect_mode config collect=dict( - # NOTE: It is important that set key traj_len_inf=True here, + # (int) How many training samples collected in one collection procedure. + # In each collect phase, we collect a total of sequence samples. + n_sample=32, + # (bool) It is important that set key traj_len_inf=True here, # to make sure self._traj_len=INF in serial_sample_collector.py. # In R2D2 policy, for each collect_env, we want to collect data of length self._traj_len=INF # unless the episode enters the 'done' state. - # In each collect phase, we collect a total of sequence samples. - n_sample=32, traj_len_inf=True, - # `env_num` is used in hidden state, should equal to that one in env config. - # User should specify this value in user config. + # (int) `env_num` is used in hidden state, should equal to that one in env config (e.g. collector_env_num). + # User should specify this value in user config. `None` is a placeholder. env_num=None, ), + # eval_mode config eval=dict( - # `env_num` is used in hidden state, should equal to that one in env config. + # (int) `env_num` is used in hidden state, should equal to that one in env config (e.g. evaluator_env_num). # User should specify this value in user config. env_num=None, ), other=dict( + # Epsilon greedy with decay. eps=dict( + # (str) Type of decay. Supports either 'exp' (exponential) or 'linear'. type='exp', + # (float) Initial value of epsilon at the start. start=0.95, + # (float) Final value of epsilon after decay. end=0.05, + # (int) The number of environment steps over which epsilon should decay. decay=10000, ), - replay_buffer=dict(replay_buffer_size=10000, ), + replay_buffer=dict( + # (int) Maximum size of replay buffer. Usually, larger buffer size is better. + replay_buffer_size=10000, + ), ), ) def default_model(self) -> Tuple[str, List[str]]: + """ + Overview: + Return this algorithm default neural network model setting for demonstration. ``__init__`` method will \ + automatically call this method to get the default model setting and create model. + Returns: + - model_info (:obj:`Tuple[str, List[str]]`): The registered model name and model's import_names. + + .. note:: + The user can define and use customized network model but must obey the same inferface definition indicated \ + by import_names path. For example about R2D2, its registered name is ``drqn`` and the import_names is \ + ``ding.model.template.q_learning``. + """ return 'drqn', ['ding.model.template.q_learning'] def _init_learn(self) -> None: - r""" + """ Overview: - Init the learner model of R2D2Policy - Arguments: - - learning_rate (:obj:`float`): The learning rate fo the optimizer - - gamma (:obj:`float`): The discount factor - - nstep (:obj:`int`): The num of n step return - - value_rescale (:obj:`bool`): Whether to use value rescaled loss in algorithm - - burnin_step (:obj:`int`): The num of step of burnin + Initialize the learn mode of policy, including some attributes and modules. For R2D2, it mainly contains \ + optimizer, algorithm-specific arguments such as burnin_step, value_rescale and gamma, main and target \ + model. Because of the use of RNN, all the models should be wrappered with ``hidden_state`` which needs to \ + be initialized with proper size. + This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. .. note:: - The _init_learn method takes the argument from the self._cfg.learn in the config file + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. """ self._priority = self._cfg.priority self._priority_IS_weight = self._cfg.priority_IS_weight @@ -176,16 +210,15 @@ def _init_learn(self) -> None: self._learn_model.reset() self._target_model.reset() - def _data_preprocess_learn(self, data: List[Dict[str, Any]]) -> dict: - r""" + def _data_preprocess_learn(self, data: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]: + """ Overview: Preprocess the data to fit the required data format for learning Arguments: - - data (:obj:`List[Dict[str, Any]]`): the data collected from collect function + - data (:obj:`List[Dict[str, Any]]`): The data collected from collect function Returns: - - data (:obj:`Dict[str, Any]`): the processed data, including at least \ + - data (:obj:`Dict[str, torch.Tensor]`): The processed data, including at least \ ['main_obs', 'target_obs', 'burnin_obs', 'action', 'reward', 'done', 'weight'] - - data_info (:obj:`dict`): the data info, such as replay_buffer_idx, replay_unique_id """ # data preprocess data = timestep_collate(data) @@ -243,18 +276,35 @@ def _data_preprocess_learn(self, data: List[Dict[str, Any]]) -> dict: return data - def _forward_learn(self, data: dict) -> Dict[str, Any]: - r""" + def _forward_learn(self, data: List[List[Dict[str, Any]]]) -> Dict[str, Any]: + """ Overview: - Forward and backward function of learn mode. - Acquire the data, calculate the loss and optimize learner model. + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data (trajectory for R2D2) from the replay buffer and then \ + returns the output result, including various training information such as loss, q value, priority. Arguments: - - data (:obj:`dict`): Dict type data, including at least \ - ['main_obs', 'target_obs', 'burnin_obs', 'action', 'reward', 'done', 'weight'] + - data (:obj:`List[List[Dict[int, Any]]]`): The input data used for policy forward, including a batch of \ + training samples. For each dict element, the key of the dict is the name of data items and the \ + value is the corresponding data. Usually, the value is torch.Tensor or np.ndarray or there dict/list \ + combinations. In the ``_forward_learn`` method, data often need to first be stacked in the time and \ + batch dimension by the utility functions ``self._data_preprocess_learn``. \ + For R2D2, each element in list is a trajectory with the length of ``unroll_len``, and the element in \ + trajectory list is a dict containing at least the following keys: ``obs``, ``action``, ``prev_state``, \ + ``reward``, ``next_obs``, ``done``. Sometimes, it also contains other keys such as ``weight`` \ + and ``value_gamma``. Returns: - - info_dict (:obj:`Dict[str, Any]`): Including cur_lr and total_loss - - cur_lr (:obj:`float`): Current learning rate - - total_loss (:obj:`float`): The calculated loss + - info_dict (:obj:`Dict[str, Any]`): The information dict that indicated training result, which will be \ + recorded in text log and tensorboard, values must be python scalar or a list of scalars. For the \ + detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for R2D2Policy: ``ding.policy.tests.test_r2d2``. """ # forward data = self._data_preprocess_learn(data) # output datatype: Dict @@ -267,7 +317,7 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: if len(data['burnin_nstep_obs']) != 0: with torch.no_grad(): - inputs = {'obs': data['burnin_nstep_obs'], 'enable_fast_timestep': True} + inputs = {'obs': data['burnin_nstep_obs']} burnin_output = self._learn_model.forward( inputs, saved_state_timesteps=[self._burnin_step, self._burnin_step + self._nstep] ) # keys include 'logit', 'hidden_state' 'saved_state', \ @@ -277,12 +327,12 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: ) self._learn_model.reset(data_id=None, state=burnin_output['saved_state'][0]) - inputs = {'obs': data['main_obs'], 'enable_fast_timestep': True} + inputs = {'obs': data['main_obs']} q_value = self._learn_model.forward(inputs)['logit'] self._learn_model.reset(data_id=None, state=burnin_output['saved_state'][1]) self._target_model.reset(data_id=None, state=burnin_output_target['saved_state'][1]) - next_inputs = {'obs': data['target_obs'], 'enable_fast_timestep': True} + next_inputs = {'obs': data['target_obs']} with torch.no_grad(): target_q_value = self._target_model.forward(next_inputs)['logit'] # argmax_action double_dqn @@ -345,23 +395,64 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: } def _reset_learn(self, data_id: Optional[List[int]] = None) -> None: + """ + Overview: + Reset some stateful variables for learn mode when necessary, such as the hidden state of RNN or the \ + memory bank of some special algortihms. If ``data_id`` is None, it means to reset all the stateful \ + varaibles. Otherwise, it will reset the stateful variables according to the ``data_id``. For example, \ + different trajectories in ``data_id`` will have different hidden state in RNN. + Arguments: + - data_id (:obj:`Optional[List[int]]`): The id of the data, which is used to reset the stateful variables \ + (i.e. RNN hidden_state in R2D2) specified by ``data_id``. + """ + self._learn_model.reset(data_id=data_id) def _state_dict_learn(self) -> Dict[str, Any]: + """ + Overview: + Return the state_dict of learn mode, usually including model, target_model and optimizer. + Returns: + - state_dict (:obj:`Dict[str, Any]`): The dict of current policy learn state, for saving and restoring. + """ return { 'model': self._learn_model.state_dict(), + 'target_model': self._target_model.state_dict(), 'optimizer': self._optimizer.state_dict(), } def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: + """ + Overview: + Load the state_dict variable into policy learn mode. + Arguments: + - state_dict (:obj:`Dict[str, Any]`): The dict of policy learn state saved before. + + .. tip:: + If you want to only load some parts of model, you can simply set the ``strict`` argument in \ + load_state_dict to ``False``, or refer to ``ding.torch_utils.checkpoint_helper`` for more \ + complicated operation. + """ self._learn_model.load_state_dict(state_dict['model']) + self._target_model.load_state_dict(state_dict['target_model']) self._optimizer.load_state_dict(state_dict['optimizer']) def _init_collect(self) -> None: - r""" + """ Overview: - Collect mode init method. Called by ``self.__init__``. - Init traj and unroll length, collect model. + Initialize the collect mode of policy, including related attributes and modules. For R2D2, it contains the \ + collect_model to balance the exploration and exploitation with epsilon-greedy sample mechanism and \ + maintain the hidden state of rnn. Besides, there are some initialization operations about other \ + algorithm-specific arguments such as burnin_step, unroll_len and nstep. + This method will be called in ``__init__`` method if ``collect`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_collect`` method, you'd better name them \ + with prefix ``_collect_`` to avoid conflict with other modes, such as ``self._collect_attr1``. + + .. tip:: + Some variables need to initialize independently in different modes, such as gamma and nstep in R2D2. This \ + design is for the convenience of parallel execution of different policy modes. """ self._nstep = self._cfg.nstep self._burnin_step = self._cfg.burnin_step @@ -377,18 +468,34 @@ def _init_collect(self) -> None: self._collect_model = model_wrap(self._collect_model, wrapper_name='eps_greedy_sample') self._collect_model.reset() - def _forward_collect(self, data: dict, eps: float) -> dict: - r""" + def _forward_collect(self, data: Dict[int, Any], eps: float) -> Dict[int, Any]: + """ Overview: - Forward function for collect mode with eps_greedy + Policy forward function of collect mode (collecting training data by interacting with envs). Forward means \ + that the policy gets some necessary data (mainly observation) from the envs and then returns the output \ + data, such as the action to interact with the envs. Besides, this policy also needs ``eps`` argument for \ + exploration, i.e., classic epsilon-greedy exploration strategy. Arguments: - - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \ - values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer. - - eps (:obj:`float`): epsilon value for exploration, which is decayed by collected env step. + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. + - eps (:obj:`float`): The epsilon value for exploration. Returns: - - output (:obj:`Dict[int, Any]`): Dict type data, including at least inferred action according to input obs. - ReturnsKeys - - necessary: ``action`` + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action and \ + other necessary data (prev_state) for learn mode defined in ``self._process_transition`` method. The \ + key of the dict is the same as the input data, i.e. environment id. + + .. note:: + RNN's hidden states are maintained in the policy, so we don't need pass them into data but to reset the \ + hidden states with ``_reset_collect`` method when episode ends. Besides, the previous hidden states are \ + necessary for training, so we need to return them in ``_process_transition`` method. + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for R2D2Policy: ``ding.policy.tests.test_r2d2``. """ data_id = list(data.keys()) data = default_collate(list(data.values())) @@ -406,62 +513,104 @@ def _forward_collect(self, data: dict, eps: float) -> dict: return {i: d for i, d in zip(data_id, output)} def _reset_collect(self, data_id: Optional[List[int]] = None) -> None: + """ + Overview: + Reset some stateful variables for eval mode when necessary, such as the hidden state of RNN or the \ + memory bank of some special algortihms. If ``data_id`` is None, it means to reset all the stateful \ + varaibles. Otherwise, it will reset the stateful variables according to the ``data_id``. For example, \ + different environments/episodes in evaluation in ``data_id`` will have different hidden state in RNN. + Arguments: + - data_id (:obj:`Optional[List[int]]`): The id of the data, which is used to reset the stateful variables \ + (i.e., RNN hidden_state in R2D2) specified by ``data_id``. + """ self._collect_model.reset(data_id=data_id) - def _process_transition(self, obs: Any, model_output: dict, timestep: namedtuple) -> dict: - r""" + def _process_transition(self, obs: torch.Tensor, policy_output: Dict[str, torch.Tensor], + timestep: namedtuple) -> Dict[str, torch.Tensor]: + """ Overview: - Generate dict type transition data from inputs. + Process and pack one timestep transition data into a dict, which can be directly used for training and \ + saved in replay buffer. For R2D2, it contains obs, action, prev_state, reward, and done. Arguments: - - obs (:obj:`Any`): Env observation - - model_output (:obj:`dict`): Output of collect model, including at least ['action', 'prev_state'] - - timestep (:obj:`namedtuple`): Output after env step, including at least ['reward', 'done'] \ - (here 'obs' indicates obs after env step). + - obs (:obj:`torch.Tensor`): The env observation of current timestep, such as stacked 2D image in Atari. + - policy_output (:obj:`Dict[str, torch.Tensor]`): The output of the policy network given the observation \ + as input. For R2D2, it contains the action and the prev_state of RNN. + - timestep (:obj:`namedtuple`): The execution result namedtuple returned by the environment step method, \ + except all the elements have been transformed into tensor data. Usually, it contains the next obs, \ + reward, done, info, etc. Returns: - - transition (:obj:`dict`): Dict type transition data. + - transition (:obj:`Dict[str, torch.Tensor]`): The processed transition data of the current timestep. """ transition = { 'obs': obs, - 'action': model_output['action'], - 'prev_state': model_output['prev_state'], + 'action': policy_output['action'], + 'prev_state': policy_output['prev_state'], 'reward': timestep.reward, 'done': timestep.done, } return transition - def _get_train_sample(self, data: list) -> Union[None, List[Any]]: - r""" + def _get_train_sample(self, transitions: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ Overview: - Get the trajectory and the n step return data, then sample from the n_step return data + For a given trajectory (transitions, a list of transition) data, process it into a list of sample that \ + can be used for training directly. In R2D2, a train sample is processed transitions with unroll_len \ + length. This method is usually used in collectors to execute necessary \ + RL data preprocessing before training, which can help learner amortize revelant time consumption. \ + In addition, you can also implement this method as an identity function and do the data processing \ + in ``self._forward_learn`` method. Arguments: - - data (:obj:`list`): The trajectory's cache + - transitions (:obj:`List[Dict[str, Any]`): The trajectory data (a list of transition), each element is \ + the same format as the return value of ``self._process_transition`` method. Returns: - - samples (:obj:`dict`): The training samples generated + - samples (:obj:`List[Dict[str, Any]]`): The processed train samples, each sample is a fixed-length \ + trajectory, and each element in a sample is the similar format as input transitions, but may contain \ + more data for training, such as nstep reward and value_gamma factor. """ - data = get_nstep_return_data(data, self._nstep, gamma=self._gamma) - return get_train_sample(data, self._unroll_len) + transitions = get_nstep_return_data(transitions, self._nstep, gamma=self._gamma) + return get_train_sample(transitions, self._unroll_len) def _init_eval(self) -> None: - r""" + """ Overview: - Evaluate mode init method. Called by ``self.__init__``. - Init eval model with argmax strategy. + Initialize the eval mode of policy, including related attributes and modules. For R2D2, it contains the \ + eval model to greedily select action with argmax q_value mechanism and main the hidden state. + This method will be called in ``__init__`` method if ``eval`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_eval`` method, you'd better name them \ + with prefix ``_eval_`` to avoid conflict with other modes, such as ``self._eval_attr1``. """ self._eval_model = model_wrap(self._model, wrapper_name='hidden_state', state_num=self._cfg.eval.env_num) self._eval_model = model_wrap(self._eval_model, wrapper_name='argmax_sample') self._eval_model.reset() - def _forward_eval(self, data: dict) -> dict: - r""" + def _forward_eval(self, data: Dict[int, Any]) -> Dict[int, Any]: + """ Overview: - Forward function of eval mode, similar to ``self._forward_collect``. + Policy forward function of eval mode (evaluation policy performance by interacting with envs). Forward \ + means that the policy gets some necessary data (mainly observation) from the envs and then returns the \ + action to interact with the envs. ``_forward_eval`` often use argmax sample method to get actions that \ + q_value is the highest. Arguments: - - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \ - values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer. + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. Returns: - - output (:obj:`Dict[int, Any]`): The dict of predicting action for the interaction with env. - ReturnsKeys - - necessary: ``action`` + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action. The \ + key of the dict is the same as the input data, i.e. environment id. + + .. note:: + RNN's hidden states are maintained in the policy, so we don't need pass them into data but to reset the \ + hidden states with ``_reset_eval`` method when the episode ends. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for R2D2Policy: ``ding.policy.tests.test_r2d2``. """ data_id = list(data.keys()) data = default_collate(list(data.values())) @@ -477,9 +626,26 @@ def _forward_eval(self, data: dict) -> dict: return {i: d for i, d in zip(data_id, output)} def _reset_eval(self, data_id: Optional[List[int]] = None) -> None: + """ + Overview: + Reset some stateful variables for eval mode when necessary, such as the hidden state of RNN or the \ + memory bank of some special algortihms. If ``data_id`` is None, it means to reset all the stateful \ + varaibles. Otherwise, it will reset the stateful variables according to the ``data_id``. For example, \ + different environments/episodes in evaluation in ``data_id`` will have different hidden state in RNN. + Arguments: + - data_id (:obj:`Optional[List[int]]`): The id of the data, which is used to reset the stateful variables \ + (i.e., RNN hidden_state in R2D2) specified by ``data_id``. + """ self._eval_model.reset(data_id=data_id) def _monitor_vars_learn(self) -> List[str]: + """ + Overview: + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. + Returns: + - necessary_keys (:obj:`List[str]`): The list of the necessary keys to be logged. + """ return super()._monitor_vars_learn() + [ 'total_loss', 'priority', 'q_s_taken-a_t0', 'target_q_s_max-a_t0', 'q_s_a-mean_t0' ] diff --git a/ding/policy/r2d2_collect_traj.py b/ding/policy/r2d2_collect_traj.py index cf1354dc6c..5dbe7c3eb1 100644 --- a/ding/policy/r2d2_collect_traj.py +++ b/ding/policy/r2d2_collect_traj.py @@ -85,8 +85,6 @@ class R2D2CollectTrajPolicy(Policy): # the timestep of burnin operation unroll_len=80, learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, update_per_collect=1, batch_size=64, learning_rate=0.0001, @@ -268,7 +266,7 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: if len(data['burnin_nstep_obs']) != 0: with torch.no_grad(): - inputs = {'obs': data['burnin_nstep_obs'], 'enable_fast_timestep': True} + inputs = {'obs': data['burnin_nstep_obs']} burnin_output = self._learn_model.forward( inputs, saved_state_timesteps=[self._burnin_step, self._burnin_step + self._nstep] ) @@ -277,12 +275,12 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: ) self._learn_model.reset(data_id=None, state=burnin_output['saved_state'][0]) - inputs = {'obs': data['main_obs'], 'enable_fast_timestep': True} + inputs = {'obs': data['main_obs']} q_value = self._learn_model.forward(inputs)['logit'] self._learn_model.reset(data_id=None, state=burnin_output['saved_state'][1]) self._target_model.reset(data_id=None, state=burnin_output_target['saved_state'][1]) - next_inputs = {'obs': data['target_obs'], 'enable_fast_timestep': True} + next_inputs = {'obs': data['target_obs']} with torch.no_grad(): target_q_value = self._target_model.forward(next_inputs)['logit'] # argmax_action double_dqn diff --git a/ding/policy/r2d2_gtrxl.py b/ding/policy/r2d2_gtrxl.py index 13b3d2daeb..9d67289685 100644 --- a/ding/policy/r2d2_gtrxl.py +++ b/ding/policy/r2d2_gtrxl.py @@ -1,5 +1,5 @@ import copy -import sys +import torch from collections import namedtuple from typing import List, Dict, Any, Tuple, Union, Optional @@ -10,9 +10,6 @@ from ding.utils import POLICY_REGISTRY from ding.utils.data import timestep_collate, default_collate, default_decollate from .base_policy import Policy -import torch - -from ding.model.common.head import * @POLICY_REGISTRY.register('r2d2_gtrxl') @@ -58,11 +55,9 @@ class R2D2GTrXLPolicy(Policy): | ``done`` | calculation. | fake termination env 15 ``collect.n_sample`` int [8, 128] | The number of training samples of a | It varies from | call of collector. | different envs - 16 | ``collect.unroll`` int 25 | unroll length of an iteration | unroll_len>1 - | ``_len`` - 17 | ``collect.seq`` int 20 | Training sequence length | unroll_len>=seq_len>1 + 16 | ``collect.seq`` int 20 | Training sequence length | unroll_len>=seq_len>1 | ``_len`` - 18 | ``learn.init_`` str zero | 'zero' or 'old', how to initialize the | + 17 | ``learn.init_`` str zero | 'zero' or 'old', how to initialize the | | ``memory`` | memory before each training iteration. | == ==================== ======== ============== ======================================== ======================= """ @@ -84,15 +79,13 @@ class R2D2GTrXLPolicy(Policy): discount_factor=0.99, # (int) N-step reward for target q_value estimation nstep=5, - # how many steps to use as burnin + # (int) How many steps to use in burnin phase burnin_step=1, # (int) trajectory length unroll_len=25, # (int) training sequence length seq_len=20, learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, update_per_collect=1, batch_size=64, learning_rate=0.0001, @@ -132,7 +125,7 @@ class R2D2GTrXLPolicy(Policy): ) def default_model(self) -> Tuple[str, List[str]]: - return 'gtrxl_discrete', ['ding.model.template.q_learning'] + return 'gtrxldqn', ['ding.model.template.q_learning'] def _init_learn(self) -> None: """ @@ -163,7 +156,7 @@ def _init_learn(self) -> None: self._seq_len = self._cfg.seq_len self._value_rescale = self._cfg.learn.value_rescale self._init_memory = self._cfg.learn.init_memory - assert self._init_memory in ['zero', 'old'] + assert self._init_memory in ['zero', 'old'], self._init_memory self._target_model = copy.deepcopy(self._model) @@ -357,7 +350,6 @@ def _init_collect(self) -> None: Collect mode init method. Called by ``self.__init__``. Init unroll length and sequence len, collect model. """ - assert 'unroll_len' not in self._cfg.collect, "Use default unroll_len" self._nstep = self._cfg.nstep self._gamma = self._cfg.discount_factor self._unroll_len = self._cfg.unroll_len diff --git a/ding/policy/r2d3.py b/ding/policy/r2d3.py index c3de42b011..e4f3e56a6a 100644 --- a/ding/policy/r2d3.py +++ b/ding/policy/r2d3.py @@ -85,8 +85,6 @@ class R2D3Policy(Policy): # the timestep of burnin operation learn_unroll_len=80, learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, update_per_collect=1, batch_size=64, learning_rate=0.0001, @@ -277,7 +275,7 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: if len(data['burnin_nstep_obs']) != 0: with torch.no_grad(): - inputs = {'obs': data['burnin_nstep_obs'], 'enable_fast_timestep': True} + inputs = {'obs': data['burnin_nstep_obs']} burnin_output = self._learn_model.forward( inputs, saved_state_timesteps=[self._burnin_step, self._burnin_step + self._nstep, self._burnin_step + 1] @@ -288,14 +286,14 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: ) self._learn_model.reset(data_id=None, state=burnin_output['saved_state'][0]) - inputs = {'obs': data['main_obs'], 'enable_fast_timestep': True} + inputs = {'obs': data['main_obs']} q_value = self._learn_model.forward(inputs)['logit'] # n-step self._learn_model.reset(data_id=None, state=burnin_output['saved_state'][1]) self._target_model.reset(data_id=None, state=burnin_output_target['saved_state'][1]) - next_inputs = {'obs': data['target_obs'], 'enable_fast_timestep': True} + next_inputs = {'obs': data['target_obs']} with torch.no_grad(): target_q_value = self._target_model.forward(next_inputs)['logit'] # argmax_action double_dqn @@ -305,7 +303,7 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: self._learn_model.reset(data_id=None, state=burnin_output['saved_state'][2]) self._target_model.reset(data_id=None, state=burnin_output_target['saved_state'][2]) - next_inputs_one_step = {'obs': data['target_obs_one_step'], 'enable_fast_timestep': True} + next_inputs_one_step = {'obs': data['target_obs_one_step']} with torch.no_grad(): target_q_value_one_step = self._target_model.forward(next_inputs_one_step)['logit'] # argmax_action double_dqn diff --git a/ding/policy/rainbow.py b/ding/policy/rainbow.py index 0304eef271..1309d90ad0 100644 --- a/ding/policy/rainbow.py +++ b/ding/policy/rainbow.py @@ -8,7 +8,7 @@ from ding.utils import POLICY_REGISTRY from ding.utils.data import default_collate, default_decollate from .dqn import DQNPolicy -from .common_utils import default_preprocess_learn +from .common_utils import default_preprocess_learn, set_noise_mode @POLICY_REGISTRY.register('rainbow') @@ -86,9 +86,9 @@ class RainbowDQNPolicy(DQNPolicy): discount_factor=0.99, # (int) N-step reward for target q_value estimation nstep=3, + # (bool) Whether to use NoisyNet for exploration in both learning and collecting. Default is True. + noisy_net=True, learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, # How many updates(iterations) to train after collector's one collection. # Bigger "update_per_collect" means bigger off-policy. # collect data -> update policy-> collect data -> ... @@ -202,6 +202,11 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: # ==================== self._learn_model.train() self._target_model.train() + + # Set noise mode for NoisyNet for exploration in learning if enabled in config + set_noise_mode(self._learn_model, True) + set_noise_mode(self._target_model, True) + # reset noise of noisenet for both main model and target model self._reset_noise(self._learn_model) self._reset_noise(self._target_model) @@ -263,12 +268,16 @@ def _forward_collect(self, data: dict, eps: float) -> dict: ReturnsKeys - necessary: ``action`` """ + # Set noise mode for NoisyNet for exploration in collecting if enabled in config + # We need to reset set_noise_mode every _forward_xxx because the model is reused across + # different phases (learn/collect/eval). + set_noise_mode(self._collect_model, True) + data_id = list(data.keys()) data = default_collate(list(data.values())) if self._cuda: data = to_device(data, self._device) self._collect_model.eval() - self._reset_noise(self._collect_model) with torch.no_grad(): output = self._collect_model.forward(data, eps=eps) if self._cuda: diff --git a/ding/policy/sac.py b/ding/policy/sac.py index df377f71fd..5bb870569d 100644 --- a/ding/policy/sac.py +++ b/ding/policy/sac.py @@ -16,139 +16,66 @@ from .common_utils import default_preprocess_learn -@POLICY_REGISTRY.register('sac_discrete') -class SACDiscretePolicy(Policy): - r""" - Overview: - Policy class of discrete SAC algorithm. - - Config: - == ==================== ======== ============= ================================= ======================= - ID Symbol Type Default Value Description Other(Shape) - == ==================== ======== ============= ================================= ======================= - 1 ``type`` str td3 | RL policy register name, refer | this arg is optional, - | to registry ``POLICY_REGISTRY`` | a placeholder - 2 ``cuda`` bool True | Whether to use cuda for network | - 3 | ``random_`` int 10000 | Number of randomly collected | Default to 10000 for - | ``collect_size`` | training samples in replay | SAC, 25000 for DDPG/ - | | buffer when training starts. | TD3. - 4 | ``model.policy_`` int 256 | Linear layer size for policy | - | ``embedding_size`` | network. | - 5 | ``model.soft_q_`` int 256 | Linear layer size for soft q | - | ``embedding_size`` | network. | - 6 | ``model.value_`` int 256 | Linear layer size for value | Defalut to None when - | ``embedding_size`` | network. | model.value_network - | | | is False. - 7 | ``learn.learning`` float 3e-4 | Learning rate for soft q | Defalut to 1e-3, when - | ``_rate_q`` | network. | model.value_network - | | | is True. - 8 | ``learn.learning`` float 3e-4 | Learning rate for policy | Defalut to 1e-3, when - | ``_rate_policy`` | network. | model.value_network - | | | is True. - 9 | ``learn.learning`` float 3e-4 | Learning rate for policy | Defalut to None when - | ``_rate_value`` | network. | model.value_network - | | | is False. - 10 | ``learn.alpha`` float 0.2 | Entropy regularization | alpha is initiali- - | | coefficient. | zation for auto - | | | `\alpha`, when - | | | auto_alpha is True - 11 | ``learn.repara_`` bool True | Determine whether to use | - | ``meterization`` | reparameterization trick. | - 12 | ``learn.`` bool False | Determine whether to use | Temperature parameter - | ``auto_alpha`` | auto temperature parameter | determines the - | | `\alpha`. | relative importance - | | | of the entropy term - | | | against the reward. - 13 | ``learn.-`` bool False | Determine whether to ignore | Use ignore_done only - | ``ignore_done`` | done flag. | in halfcheetah env. - 14 | ``learn.-`` float 0.005 | Used for soft update of the | aka. Interpolation - | ``target_theta`` | target network. | factor in polyak aver - | | | aging for target - | | | networks. - == ==================== ======== ============= ================================= ======================= - """ +@POLICY_REGISTRY.register('discrete_sac') +class DiscreteSACPolicy(Policy): + """ + Overview: + Policy class of discrete SAC algorithm. Paper link: https://arxiv.org/abs/1910.07207. + """ config = dict( # (str) RL policy register name (refer to function "POLICY_REGISTRY"). - type='sac_discrete', - # (bool) Whether to use cuda for network. + type='discrete_sac', + # (bool) Whether to use cuda for network and loss computation. cuda=False, - # (bool type) on_policy: Determine whether on-policy or off-policy. - # on-policy setting influences the behaviour of buffer. - # Default False in SAC. + # (bool) Whether to belong to on-policy or off-policy algorithm, DiscreteSAC is an off-policy algorithm. on_policy=False, - # (bool type) priority: Determine whether to use priority in buffer sample. - # Default False in SAC. + # (bool) Whether to use priority sampling in buffer. Default to False in DiscreteSAC. priority=False, - # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. + # (bool) Whether use Importance Sampling weight to correct biased update. If True, priority must be True. priority_IS_weight=False, - # (int) Number of training samples(randomly collected) in replay buffer when training starts. - # Default 10000 in SAC. + # (int) Number of training samples (randomly collected) in replay buffer when training starts. random_collect_size=10000, - # (bool) Whether to need policy data in process transition + # (bool) Whether to need policy-specific data in process transition. transition_with_policy_data=True, - # (bool) Whether to enable multi-agent training setting - multi_agent=True, + # (bool) Whether to enable multi-agent training setting. + multi_agent=False, model=dict( - # (bool type) twin_critic: Determine whether to use double-soft-q-net for target q computation. - # Please refer to TD3 about Clipped Double-Q Learning trick, which learns two Q-functions instead of one . - # Default to True. + # (bool) Whether to use double-soft-q-net for target q computation. + # For more details, please refer to TD3 about Clipped Double-Q Learning trick. twin_critic=True, - - # (bool type) value_network: Determine whether to use value network as the - # original SAC paper (arXiv 1801.01290). - # using value_network needs to set learning_rate_value, learning_rate_q, - # and learning_rate_policy in `cfg.policy.learn`. - # Default to False. - # value_network=False, ), + # learn_mode config learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, - # How many updates(iterations) to train after collector's one collection. + # (int) How many updates (iterations) to train after collector's one collection. # Bigger "update_per_collect" means bigger off-policy. - # collect data -> update policy-> collect data -> ... update_per_collect=1, - # (int) Minibatch size for gradient descent. + # (int) Minibatch size for one gradient descent. batch_size=256, - - # (float type) learning_rate_q: Learning rate for soft q network. - # Default to 3e-4. - # Please set to 1e-3, when model.value_network is True. + # (float) Learning rate for soft q network. learning_rate_q=3e-4, - # (float type) learning_rate_policy: Learning rate for policy network. - # Default to 3e-4. - # Please set to 1e-3, when model.value_network is True. + # (float) Learning rate for policy network. learning_rate_policy=3e-4, - # (float type) learning_rate_value: Learning rate for value network. - # `learning_rate_value` should be initialized, when model.value_network is True. - # Please set to 3e-4, when model.value_network is True. - learning_rate_value=3e-4, - - # (float type) learning_rate_alpha: Learning rate for auto temperature parameter `\alpha`. - # Default to 3e-4. + # (float) Learning rate for auto temperature parameter `\alpha`. learning_rate_alpha=3e-4, - # (float type) target_theta: Used for soft update of the target network, - # aka. Interpolation factor in polyak averaging for target networks. - # Default to 0.005. + # (float) Used for soft update of the target network, + # aka. Interpolation factor in EMA update for target network. target_theta=0.005, - # (float) discount factor for the discounted sum of rewards, aka. gamma. + # (float) Discount factor for the discounted sum of rewards, aka. gamma. discount_factor=0.99, - - # (float type) alpha: Entropy regularization coefficient. + # (float) Entropy regularization coefficient in SAC. # Please check out the original SAC paper (arXiv 1801.01290): Eq 1 for more details. # If auto_alpha is set to `True`, alpha is initialization for auto `\alpha`. - # Default to 0.2. alpha=0.2, - - # (bool type) auto_alpha: Determine whether to use auto temperature parameter `\alpha` . - # Temperature parameter determines the relative importance of the entropy term against the reward. + # (bool) Whether to use auto temperature parameter `\alpha` . + # Temperature parameter `\alpha` determines the relative importance of the entropy term against the reward. # Please check out the original SAC paper (arXiv 1801.01290): Eq 1 for more details. - # Default to False. - # Note that: Using auto alpha needs to set learning_rate_alpha in `cfg.policy.learn`. + # Note that: Using auto alpha needs to set the above `learning_rate_alpha`. auto_alpha=True, - # (bool type) log_space: Determine whether to use auto `\alpha` in log space. + # (bool) Whether to use auto `\alpha` in log space. log_space=True, + # (float) Target policy entropy value for auto temperature (alpha) adjustment. + target_entropy=None, # (bool) Whether ignore done(usually for max step termination env. e.g. pendulum) # Note: Gym wraps the MuJoCo envs by default with TimeLimit environment wrappers. # These limit HalfCheetah, and several other MuJoCo envs, to max length of 1000. @@ -157,50 +84,64 @@ class SACDiscretePolicy(Policy): # TD-error accurate computation(``gamma * (1 - done) * next_v + reward``), # when the episode step is greater than max episode step. ignore_done=False, - # (float) Weight uniform initialization range in the last output layer + # (float) Weight uniform initialization max range in the last output layer init_w=3e-3, ), + # collect_mode config collect=dict( - # You can use either "n_sample" or "n_episode" in actor.collect. - # Get "n_sample" samples per collect. - # Default n_sample to 1. - # n_sample=1, - # (int) Cut trajectories into pieces with length "unroll_len". + # (int) How many training samples collected in one collection procedure. + # Only one of [n_sample, n_episode] shoule be set. + n_sample=1, + # (int) Split episodes or trajectories into pieces with length `unroll_len`. unroll_len=1, + # (bool) Whether to collect logit in `process_transition`. + # In some algorithm like guided cost learning, we need to use logit to train the reward model. + collector_logit=False, ), - eval=dict(), + eval=dict(), # for compability other=dict( replay_buffer=dict( - # (int type) replay_buffer_size: Max size of replay buffer. + # (int) Maximum size of replay buffer. Usually, larger buffer size is good + # for SAC but cost more storage. replay_buffer_size=1000000, - # (int type) max_use: Max use times of one data in the buffer. - # Data will be removed once used for too many times. - # Default to infinite. - # max_use=256, ), ), ) - r""" - Overview: - Policy class of SAC algorithm. - """ def default_model(self) -> Tuple[str, List[str]]: + """ + Overview: + Return this algorithm default neural network model setting for demonstration. ``__init__`` method will \ + automatically call this method to get the default model setting and create model. + Returns: + - model_info (:obj:`Tuple[str, List[str]]`): The registered model name and model's import_names. + """ if self._cfg.multi_agent: - return 'maqac', ['ding.model.template.maqac'] + return 'discrete_maqac', ['ding.model.template.maqac'] else: return 'discrete_qac', ['ding.model.template.qac'] def _init_learn(self) -> None: - r""" + """ Overview: - Learn mode init method. Called by ``self.__init__``. - Init q, value and policy's optimizers, algorithm config, main and target models. + Initialize the learn mode of policy, including related attributes and modules. For DiscreteSAC, it mainly \ + contains three optimizers, algorithm-specific arguments such as gamma and twin_critic, main and target \ + model. Especially, the ``auto_alpha`` mechanism for balancing max entropy target is also initialized here. + This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. """ - # Init self._priority = self._cfg.priority self._priority_IS_weight = self._cfg.priority_IS_weight - # self._value_network = False # TODO self._cfg.model.value_network self._twin_critic = self._cfg.model.twin_critic self._optimizer_q = Adam( @@ -212,12 +153,14 @@ def _init_learn(self) -> None: lr=self._cfg.learn.learning_rate_policy, ) - # Algorithm config + # Algorithm-Specific Config self._gamma = self._cfg.learn.discount_factor - self._multi_agent = self._cfg.multi_agent - # Init auto alpha if self._cfg.learn.auto_alpha: - self._target_entropy = self._cfg.learn.get('target_entropy', -np.prod(self._cfg.model.action_shape)) + if self._cfg.learn.target_entropy is None: + assert 'action_shape' in self._cfg.model, "DiscreteSAC need network model with action_shape variable" + self._target_entropy = -np.prod(self._cfg.model.action_shape) + else: + self._target_entropy = self._cfg.learn.target_entropy if self._cfg.learn.log_space: self._log_alpha = torch.log(torch.FloatTensor([self._cfg.learn.alpha])) self._log_alpha = self._log_alpha.to(self._device).requires_grad_() @@ -249,16 +192,34 @@ def _init_learn(self) -> None: self._learn_model.reset() self._target_model.reset() - self._forward_learn_cnt = 0 - - def _forward_learn(self, data: dict) -> Dict[str, Any]: - r""" + def _forward_learn(self, data: List[Dict[str, Any]]) -> Dict[str, Any]: + """ Overview: - Forward and backward function of learn mode. + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data from the replay buffer and then returns the output \ + result, including various training information such as loss, action, priority. Arguments: - - data (:obj:`dict`): Dict type data, including at least ['obs', 'action', 'reward', 'next_obs'] + - data (:obj:`List[Dict[int, Any]]`): The input data used for policy forward, including a batch of \ + training samples. For each element in list, the key of the dict is the name of data items and the \ + value is the corresponding data. Usually, the value is torch.Tensor or np.ndarray or there dict/list \ + combinations. In the ``_forward_learn`` method, data often need to first be stacked in the batch \ + dimension by some utility functions such as ``default_preprocess_learn``. \ + For SAC, each element in list is a dict containing at least the following keys: ``obs``, ``action``, \ + ``logit``, ``reward``, ``next_obs``, ``done``. Sometimes, it also contains other keys like ``weight``. Returns: - - info_dict (:obj:`Dict[str, Any]`): Including current lr and loss. + - info_dict (:obj:`Dict[str, Any]`): The information dict that indicated training result, which will be \ + recorded in text log and tensorboard, values must be python scalar or a list of scalars. For the \ + detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for DiscreteSACPolicy: \ + ``ding.policy.tests.test_discrete_sac``. """ loss_dict = {} data = default_preprocess_learn( @@ -281,21 +242,21 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: action = data['action'] # 1. predict q value - q_value = self._learn_model.forward({'obs': obs}, mode='compute_critic')['q_value'] + q_value = self._learn_model.forward(obs, mode='compute_critic')['q_value'] dist = torch.distributions.categorical.Categorical(logits=logit) dist_entropy = dist.entropy() entropy = dist_entropy.mean() - # 2. predict target value depend self._value_network. + # 2. predict target value # target q value. SARSA: first predict next action, then calculate next q value with torch.no_grad(): - policy_output_next = self._learn_model.forward({'obs': next_obs}, mode='compute_actor') - if self._multi_agent: + policy_output_next = self._learn_model.forward(next_obs, mode='compute_actor') + if self._cfg.multi_agent: policy_output_next['logit'][policy_output_next['action_mask'] == 0.0] = -1e8 prob = F.softmax(policy_output_next['logit'], dim=-1) log_prob = torch.log(prob + 1e-8) - target_q_value = self._target_model.forward({'obs': next_obs}, mode='compute_critic')['q_value'] + target_q_value = self._target_model.forward(next_obs, mode='compute_critic')['q_value'] # the value of a policy according to the maximum entropy objective if self._twin_critic: # find min one as target q value @@ -324,15 +285,16 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: self._optimizer_q.step() # 5. evaluate to get action distribution - policy_output = self._learn_model.forward({'obs': data['obs']}, mode='compute_actor') - if self._multi_agent: + policy_output = self._learn_model.forward(obs, mode='compute_actor') + # 6. apply discrete action mask in multi_agent setting + if self._cfg.multi_agent: policy_output['logit'][policy_output['action_mask'] == 0.0] = -1e8 logit = policy_output['logit'] prob = F.softmax(logit, dim=-1) log_prob = F.log_softmax(logit, dim=-1) with torch.no_grad(): - new_q_value = self._learn_model.forward({'obs': data['obs']}, mode='compute_critic')['q_value'] + new_q_value = self._learn_model.forward(obs, mode='compute_critic')['q_value'] if self._twin_critic: new_q_value = torch.min(new_q_value[0], new_q_value[1]) # 7. compute policy loss @@ -367,15 +329,12 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: torch.zeros_like(self._alpha)).requires_grad_() loss_dict['total_loss'] = sum(loss_dict.values()) - info_dict = {} - - # ============= - # after update - # ============= - self._forward_learn_cnt += 1 # target update self._target_model.update(self._learn_model.state_dict()) return { + 'total_loss': loss_dict['total_loss'].item(), + 'policy_loss': loss_dict['policy_loss'].item(), + 'critic_loss': loss_dict['critic_loss'].item(), 'cur_lr_q': self._optimizer_q.defaults['lr'], 'cur_lr_p': self._optimizer_policy.defaults['lr'], 'priority': td_error_per_sample.abs().tolist(), @@ -385,13 +344,18 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: 'q_value_2': target_q_value[1].detach().mean().item(), 'target_value': target_value.detach().mean().item(), 'entropy': entropy.item(), - **info_dict, - **loss_dict } def _state_dict_learn(self) -> Dict[str, Any]: + """ + Overview: + Return the state_dict of learn mode, usually including model, target_model and optimizers. + Returns: + - state_dict (:obj:`Dict[str, Any]`): The dict of current policy learn state, for saving and restoring. + """ ret = { 'model': self._learn_model.state_dict(), + 'target_model': self._target_model.state_dict(), 'optimizer_q': self._optimizer_q.state_dict(), 'optimizer_policy': self._optimizer_policy.state_dict(), } @@ -400,37 +364,68 @@ def _state_dict_learn(self) -> Dict[str, Any]: return ret def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: + """ + Overview: + Load the state_dict variable into policy learn mode. + Arguments: + - state_dict (:obj:`Dict[str, Any]`): The dict of policy learn state saved before. + + .. tip:: + If you want to only load some parts of model, you can simply set the ``strict`` argument in \ + load_state_dict to ``False``, or refer to ``ding.torch_utils.checkpoint_helper`` for more \ + complicated operation. + """ self._learn_model.load_state_dict(state_dict['model']) + self._target_model.load_state_dict(state_dict['target_model']) self._optimizer_q.load_state_dict(state_dict['optimizer_q']) - # if self._value_network: - # self._optimizer_value.load_state_dict(state_dict['optimizer_value']) self._optimizer_policy.load_state_dict(state_dict['optimizer_policy']) if self._auto_alpha: self._alpha_optim.load_state_dict(state_dict['optimizer_alpha']) def _init_collect(self) -> None: - r""" + """ Overview: - Collect mode init method. Called by ``self.__init__``. - Init traj and unroll length, collect model. - Use action noise for exploration. + Initialize the collect mode of policy, including related attributes and modules. For SAC, it contains the \ + collect_model to balance the exploration and exploitation with the epsilon and multinomial sample \ + mechanism, and other algorithm-specific arguments such as unroll_len. \ + This method will be called in ``__init__`` method if ``collect`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_collect`` method, you'd better name them \ + with prefix ``_collect_`` to avoid conflict with other modes, such as ``self._collect_attr1``. """ self._unroll_len = self._cfg.collect.unroll_len - self._multi_agent = self._cfg.multi_agent # Empirically, we found that eps_greedy_multinomial_sample works better than multinomial_sample # and eps_greedy_sample, and we don't divide logit by alpha, # for the details please refer to ding/model/wrapper/model_wrappers self._collect_model = model_wrap(self._model, wrapper_name='eps_greedy_multinomial_sample') self._collect_model.reset() - def _forward_collect(self, data: dict, eps: float) -> dict: - r""" + def _forward_collect(self, data: Dict[int, Any], eps: float) -> Dict[int, Any]: + """ Overview: - Forward function of collect mode. + Policy forward function of collect mode (collecting training data by interacting with envs). Forward means \ + that the policy gets some necessary data (mainly observation) from the envs and then returns the output \ + data, such as the action to interact with the envs. Besides, this policy also needs ``eps`` argument for \ + exploration, i.e., classic epsilon-greedy exploration strategy. Arguments: - - data (:obj:`dict`): Dict type data, including at least ['obs']. + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. + - eps (:obj:`float`): The epsilon value for exploration. Returns: - - output (:obj:`dict`): Dict type data, including at least inferred action according to input obs. + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action and \ + other necessary data for learn mode defined in ``self._process_transition`` method. The key of the \ + dict is the same as the input data, i.e. environment id. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for DiscreteSACPolicy: \ + ``ding.policy.tests.test_discrete_sac``. """ data_id = list(data.keys()) data = default_collate(list(data.values())) @@ -438,54 +433,88 @@ def _forward_collect(self, data: dict, eps: float) -> dict: data = to_device(data, self._device) self._collect_model.eval() with torch.no_grad(): - output = self._collect_model.forward({'obs': data}, mode='compute_actor', eps=eps) + output = self._collect_model.forward(data, mode='compute_actor', eps=eps) if self._cuda: output = to_device(output, 'cpu') output = default_decollate(output) return {i: d for i, d in zip(data_id, output)} - def _process_transition(self, obs: Any, model_output: dict, timestep: namedtuple) -> dict: - r""" + def _process_transition(self, obs: torch.Tensor, policy_output: Dict[str, torch.Tensor], + timestep: namedtuple) -> Dict[str, torch.Tensor]: + """ Overview: - Generate dict type transition data from inputs. + Process and pack one timestep transition data into a dict, which can be directly used for training and \ + saved in replay buffer. For discrete SAC, it contains obs, next_obs, logit, action, reward, done. Arguments: - - obs (:obj:`Any`): Env observation - - model_output (:obj:`dict`): Output of collect model, including at least ['action'] - - timestep (:obj:`namedtuple`): Output after env step, including at least ['obs', 'reward', 'done'] \ - (here 'obs' indicates obs after env step, i.e. next_obs). - Return: - - transition (:obj:`Dict[str, Any]`): Dict type transition data. + - obs (:obj:`torch.Tensor`): The env observation of current timestep, such as stacked 2D image in Atari. + - policy_output (:obj:`Dict[str, torch.Tensor]`): The output of the policy network with the observation \ + as input. For discrete SAC, it contains the action and the logit of the action. + - timestep (:obj:`namedtuple`): The execution result namedtuple returned by the environment step method, \ + except all the elements have been transformed into tensor data. Usually, it contains the next obs, \ + reward, done, info, etc. + Returns: + - transition (:obj:`Dict[str, torch.Tensor]`): The processed transition data of the current timestep. """ transition = { 'obs': obs, 'next_obs': timestep.obs, - 'action': model_output['action'], - 'logit': model_output['logit'], + 'action': policy_output['action'], + 'logit': policy_output['logit'], 'reward': timestep.reward, 'done': timestep.done, } return transition - def _get_train_sample(self, data: list) -> Union[None, List[Any]]: - return get_train_sample(data, self._unroll_len) + def _get_train_sample(self, transitions: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Overview: + For a given trajectory (transitions, a list of transition) data, process it into a list of sample that \ + can be used for training directly. In discrete SAC, a train sample is a processed transition (unroll_len=1). + Arguments: + - transitions (:obj:`List[Dict[str, Any]`): The trajectory data (a list of transition), each element is \ + the same format as the return value of ``self._process_transition`` method. + Returns: + - samples (:obj:`List[Dict[str, Any]]`): The processed train samples, each element is the similar format \ + as input transitions, but may contain more data for training. + """ + return get_train_sample(transitions, self._unroll_len) def _init_eval(self) -> None: - r""" + """ Overview: - Evaluate mode init method. Called by ``self.__init__``. - Init eval model. Unlike learn and collect model, eval model does not need noise. + Initialize the eval mode of policy, including related attributes and modules. For DiscreteSAC, it contains \ + the eval model to greedily select action type with argmax q_value mechanism. + This method will be called in ``__init__`` method if ``eval`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_eval`` method, you'd better name them \ + with prefix ``_eval_`` to avoid conflict with other modes, such as ``self._eval_attr1``. """ self._eval_model = model_wrap(self._model, wrapper_name='argmax_sample') self._eval_model.reset() - def _forward_eval(self, data: dict) -> dict: - r""" + def _forward_eval(self, data: Dict[int, Any]) -> Dict[int, Any]: + """ Overview: - Forward function for eval mode, similar to ``self._forward_collect``. + Policy forward function of eval mode (evaluation policy performance by interacting with envs). Forward \ + means that the policy gets some necessary data (mainly observation) from the envs and then returns the \ + action to interact with the envs. Arguments: - - data (:obj:`dict`): Dict type data, including at least ['obs']. + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. Returns: - - output (:obj:`dict`): Dict type data, including at least inferred action according to input obs. + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action. The \ + key of the dict is the same as the input data, i.e. environment id. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for DiscreteSACPolicy: \ + ``ding.policy.tests.test_discrete_sac``. """ data_id = list(data.keys()) data = default_collate(list(data.values())) @@ -493,18 +522,19 @@ def _forward_eval(self, data: dict) -> dict: data = to_device(data, self._device) self._eval_model.eval() with torch.no_grad(): - output = self._eval_model.forward({'obs': data}, mode='compute_actor') + output = self._eval_model.forward(data, mode='compute_actor') if self._cuda: output = to_device(output, 'cpu') output = default_decollate(output) return {i: d for i, d in zip(data_id, output)} def _monitor_vars_learn(self) -> List[str]: - r""" + """ Overview: - Return variables' name if variables are to used in monitor. + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. Returns: - - vars (:obj:`List[str]`): Variables' name list. + - necessary_keys (:obj:`List[str]`): The list of the necessary keys to be logged. """ twin_critic = ['twin_critic_loss'] if self._twin_critic else [] if self._auto_alpha: @@ -521,142 +551,103 @@ def _monitor_vars_learn(self) -> List[str]: @POLICY_REGISTRY.register('sac') class SACPolicy(Policy): - r""" - Overview: - Policy class of continuous SAC algorithm. - - https://arxiv.org/pdf/1801.01290.pdf + """ + Overview: + Policy class of continuous SAC algorithm. Paper link: https://arxiv.org/pdf/1801.01290.pdf - Config: + Config: == ==================== ======== ============= ================================= ======================= - ID Symbol Type Default Value Description Other(Shape) + ID Symbol Type Default Value Description Other == ==================== ======== ============= ================================= ======================= - 1 ``type`` str td3 | RL policy register name, refer | this arg is optional, + 1 ``type`` str sac | RL policy register name, refer | this arg is optional, | to registry ``POLICY_REGISTRY`` | a placeholder 2 ``cuda`` bool True | Whether to use cuda for network | - 3 | ``random_`` int 10000 | Number of randomly collected | Default to 10000 for + 3 ``on_policy`` bool False | SAC is an off-policy | + | algorithm. | + 4 ``priority`` bool False | Whether to use priority | + | sampling in buffer. | + 5 | ``priority_IS_`` bool False | Whether use Importance Sampling | + | ``weight`` | weight to correct biased update | + 6 | ``random_`` int 10000 | Number of randomly collected | Default to 10000 for | ``collect_size`` | training samples in replay | SAC, 25000 for DDPG/ | | buffer when training starts. | TD3. - 4 | ``model.policy_`` int 256 | Linear layer size for policy | - | ``embedding_size`` | network. | - 5 | ``model.soft_q_`` int 256 | Linear layer size for soft q | - | ``embedding_size`` | network. | - 6 | ``model.value_`` int 256 | Linear layer size for value | Defalut to None when - | ``embedding_size`` | network. | model.value_network - | | | is False. - 7 | ``learn.learning`` float 3e-4 | Learning rate for soft q | Defalut to 1e-3, when - | ``_rate_q`` | network. | model.value_network - | | | is True. - 8 | ``learn.learning`` float 3e-4 | Learning rate for policy | Defalut to 1e-3, when - | ``_rate_policy`` | network. | model.value_network - | | | is True. - 9 | ``learn.learning`` float 3e-4 | Learning rate for policy | Defalut to None when - | ``_rate_value`` | network. | model.value_network - | | | is False. - 10 | ``learn.alpha`` float 0.2 | Entropy regularization | alpha is initiali- + 7 | ``learn.learning`` float 3e-4 | Learning rate for soft q | Defalut to 1e-3 + | ``_rate_q`` | network. | + 8 | ``learn.learning`` float 3e-4 | Learning rate for policy | Defalut to 1e-3 + | ``_rate_policy`` | network. | + 9 | ``learn.alpha`` float 0.2 | Entropy regularization | alpha is initiali- | | coefficient. | zation for auto - | | | `\alpha`, when + | | | alpha, when | | | auto_alpha is True - 11 | ``learn.repara_`` bool True | Determine whether to use | - | ``meterization`` | reparameterization trick. | - 12 | ``learn.`` bool False | Determine whether to use | Temperature parameter + 10 | ``learn.`` bool False | Determine whether to use | Temperature parameter | ``auto_alpha`` | auto temperature parameter | determines the - | | `\alpha`. | relative importance + | | alpha. | relative importance | | | of the entropy term | | | against the reward. - 13 | ``learn.-`` bool False | Determine whether to ignore | Use ignore_done only - | ``ignore_done`` | done flag. | in halfcheetah env. - 14 | ``learn.-`` float 0.005 | Used for soft update of the | aka. Interpolation + 11 | ``learn.-`` bool False | Determine whether to ignore | Use ignore_done only + | ``ignore_done`` | done flag. | in env like Pendulum + 12 | ``learn.-`` float 0.005 | Used for soft update of the | aka. Interpolation | ``target_theta`` | target network. | factor in polyak aver | | | aging for target | | | networks. == ==================== ======== ============= ================================= ======================= - """ + """ config = dict( # (str) RL policy register name (refer to function "POLICY_REGISTRY"). type='sac', - # (bool) Whether to use cuda for network. + # (bool) Whether to use cuda for network and loss computation. cuda=False, - # (bool type) on_policy: Determine whether on-policy or off-policy. - # on-policy setting influences the behaviour of buffer. - # Default False in SAC. + # (bool) Whether to belong to on-policy or off-policy algorithm, SAC is an off-policy algorithm. on_policy=False, - # (bool type) priority: Determine whether to use priority in buffer sample. - # Default False in SAC. + # (bool) Whether to use priority sampling in buffer. Default to False in SAC. priority=False, - # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. + # (bool) Whether use Importance Sampling weight to correct biased update. If True, priority must be True. priority_IS_weight=False, - # (int) Number of training samples(randomly collected) in replay buffer when training starts. - # Default 10000 in SAC. + # (int) Number of training samples (randomly collected) in replay buffer when training starts. random_collect_size=10000, - # (bool) Whether to need policy data in process transition + # (bool) Whether to need policy-specific data in process transition. transition_with_policy_data=True, - # (bool) Whether to enable multi-agent training setting + # (bool) Whether to enable multi-agent training setting. multi_agent=False, model=dict( - # (bool type) twin_critic: Determine whether to use double-soft-q-net for target q computation. - # Please refer to TD3 about Clipped Double-Q Learning trick, which learns two Q-functions instead of one . - # Default to True. + # (bool) Whether to use double-soft-q-net for target q computation. + # For more details, please refer to TD3 about Clipped Double-Q Learning trick. twin_critic=True, - - # (bool type) value_network: Determine whether to use value network as the - # original SAC paper (arXiv 1801.01290). - # using value_network needs to set learning_rate_value, learning_rate_q, - # and learning_rate_policy in `cfg.policy.learn`. - # Default to False. - # value_network=False, - - # (str type) action_space: Use reparameterization trick for continous action + # (str) Use reparameterization trick for continous action. action_space='reparameterization', ), + # learn_mode config learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, - # How many updates(iterations) to train after collector's one collection. + # (int) How many updates (iterations) to train after collector's one collection. # Bigger "update_per_collect" means bigger off-policy. - # collect data -> update policy-> collect data -> ... update_per_collect=1, - # (int) Minibatch size for gradient descent. + # (int) Minibatch size for one gradient descent. batch_size=256, - - # (float type) learning_rate_q: Learning rate for soft q network. - # Default to 3e-4. - # Please set to 1e-3, when model.value_network is True. + # (float) Learning rate for soft q network. learning_rate_q=3e-4, - # (float type) learning_rate_policy: Learning rate for policy network. - # Default to 3e-4. - # Please set to 1e-3, when model.value_network is True. + # (float) Learning rate for policy network. learning_rate_policy=3e-4, - # (float type) learning_rate_value: Learning rate for value network. - # `learning_rate_value` should be initialized, when model.value_network is True. - # Please set to 3e-4, when model.value_network is True. - learning_rate_value=3e-4, - - # (float type) learning_rate_alpha: Learning rate for auto temperature parameter `\alpha`. - # Default to 3e-4. + # (float) Learning rate for auto temperature parameter `\alpha`. learning_rate_alpha=3e-4, - # (float type) target_theta: Used for soft update of the target network, - # aka. Interpolation factor in polyak averaging for target networks. - # Default to 0.005. + # (float) Used for soft update of the target network, + # aka. Interpolation factor in EMA update for target network. target_theta=0.005, # (float) discount factor for the discounted sum of rewards, aka. gamma. discount_factor=0.99, - - # (float type) alpha: Entropy regularization coefficient. + # (float) Entropy regularization coefficient in SAC. # Please check out the original SAC paper (arXiv 1801.01290): Eq 1 for more details. # If auto_alpha is set to `True`, alpha is initialization for auto `\alpha`. - # Default to 0.2. alpha=0.2, - - # (bool type) auto_alpha: Determine whether to use auto temperature parameter `\alpha` . - # Temperature parameter determines the relative importance of the entropy term against the reward. + # (bool) Whether to use auto temperature parameter `\alpha` . + # Temperature parameter `\alpha` determines the relative importance of the entropy term against the reward. # Please check out the original SAC paper (arXiv 1801.01290): Eq 1 for more details. - # Default to False. - # Note that: Using auto alpha needs to set learning_rate_alpha in `cfg.policy.learn`. + # Note that: Using auto alpha needs to set the above `learning_rate_alpha`. auto_alpha=True, - # (bool type) log_space: Determine whether to use auto `\alpha` in log space. + # (bool) Whether to use auto `\alpha` in log space. log_space=True, + # (float) Target policy entropy value for auto temperature (alpha) adjustment. + target_entropy=None, # (bool) Whether ignore done(usually for max step termination env. e.g. pendulum) # Note: Gym wraps the MuJoCo envs by default with TimeLimit environment wrappers. # These limit HalfCheetah, and several other MuJoCo envs, to max length of 1000. @@ -665,71 +656,73 @@ class SACPolicy(Policy): # TD-error accurate computation(``gamma * (1 - done) * next_v + reward``), # when the episode step is greater than max episode step. ignore_done=False, - # (float) Weight uniform initialization range in the last output layer + # (float) Weight uniform initialization max range in the last output layer. init_w=3e-3, ), + # collect_mode config collect=dict( - # If you need the data collected by the collector to contain logit key which reflect the probability of - # the action, you can change the key to be True. - # In Guided cost Learning, we need to use logit to train the reward model, we change the key to be True. - # Default collector_logit to False. - collector_logit=False, - # You can use either "n_sample" or "n_episode" in actor.collect. - # Get "n_sample" samples per collect. - # Default n_sample to 1. + # (int) How many training samples collected in one collection procedure. n_sample=1, - # (int) Cut trajectories into pieces with length "unroll_len". + # (int) Split episodes or trajectories into pieces with length `unroll_len`. unroll_len=1, + # (bool) Whether to collect logit in `process_transition`. + # In some algorithm like guided cost learning, we need to use logit to train the reward model. + collector_logit=False, ), - eval=dict( - evaluator=dict( - # (int) Evaluate every "eval_freq" training iterations. - eval_freq=5000, - ), - ), + eval=dict(), # for compability other=dict( replay_buffer=dict( - # (int type) replay_buffer_size: Max size of replay buffer. + # (int) Maximum size of replay buffer. Usually, larger buffer size is good + # for SAC but cost more storage. replay_buffer_size=1000000, - # (int type) max_use: Max use times of one data in the buffer. - # Data will be removed once used for too many times. - # Default to infinite. - # max_use=256, ), ), ) def default_model(self) -> Tuple[str, List[str]]: + """ + Overview: + Return this algorithm default neural network model setting for demonstration. ``__init__`` method will \ + automatically call this method to get the default model setting and create model. + Returns: + - model_info (:obj:`Tuple[str, List[str]]`): The registered model name and model's import_names. + """ if self._cfg.multi_agent: - return 'maqac_continuous', ['ding.model.template.maqac'] + return 'continuous_maqac', ['ding.model.template.maqac'] else: - return 'qac', ['ding.model.template.qac'] + return 'continuous_qac', ['ding.model.template.qac'] def _init_learn(self) -> None: - r""" + """ Overview: - Learn mode init method. Called by ``self.__init__``. - Init q, value and policy's optimizers, algorithm config, main and target models. + Initialize the learn mode of policy, including related attributes and modules. For SAC, it mainly \ + contains three optimizers, algorithm-specific arguments such as gamma and twin_critic, main and target \ + model. Especially, the ``auto_alpha`` mechanism for balancing max entropy target is also initialized here. + This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. + + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. + + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. + + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. """ - # Init self._priority = self._cfg.priority self._priority_IS_weight = self._cfg.priority_IS_weight - self._value_network = False # TODO self._cfg.model.value_network self._twin_critic = self._cfg.model.twin_critic # Weight Init for the last output layer - init_w = self._cfg.learn.init_w - self._model.actor[2].mu.weight.data.uniform_(-init_w, init_w) - self._model.actor[2].mu.bias.data.uniform_(-init_w, init_w) - self._model.actor[2].log_sigma_layer.weight.data.uniform_(-init_w, init_w) - self._model.actor[2].log_sigma_layer.bias.data.uniform_(-init_w, init_w) - - # Optimizers - if self._value_network: - self._optimizer_value = Adam( - self._model.value_critic.parameters(), - lr=self._cfg.learn.learning_rate_value, - ) + if hasattr(self._model, 'actor_head'): # keep compatibility + init_w = self._cfg.learn.init_w + self._model.actor_head[-1].mu.weight.data.uniform_(-init_w, init_w) + self._model.actor_head[-1].mu.bias.data.uniform_(-init_w, init_w) + self._model.actor_head[-1].log_sigma_layer.weight.data.uniform_(-init_w, init_w) + self._model.actor_head[-1].log_sigma_layer.bias.data.uniform_(-init_w, init_w) + self._optimizer_q = Adam( self._model.critic.parameters(), lr=self._cfg.learn.learning_rate_q, @@ -739,11 +732,14 @@ def _init_learn(self) -> None: lr=self._cfg.learn.learning_rate_policy, ) - # Algorithm config + # Algorithm-Specific Config self._gamma = self._cfg.learn.discount_factor - # Init auto alpha if self._cfg.learn.auto_alpha: - self._target_entropy = self._cfg.learn.get('target_entropy', -np.prod(self._cfg.model.action_shape)) + if self._cfg.learn.target_entropy is None: + assert 'action_shape' in self._cfg.model, "SAC need network model with action_shape variable" + self._target_entropy = -np.prod(self._cfg.model.action_shape) + else: + self._target_entropy = self._cfg.learn.target_entropy if self._cfg.learn.log_space: self._log_alpha = torch.log(torch.FloatTensor([self._cfg.learn.alpha])) self._log_alpha = self._log_alpha.to(self._device).requires_grad_() @@ -775,17 +771,33 @@ def _init_learn(self) -> None: self._learn_model.reset() self._target_model.reset() - self._forward_learn_cnt = 0 - - def _forward_learn(self, data: dict) -> Dict[str, Any]: + def _forward_learn(self, data: List[Dict[str, Any]]) -> Dict[str, Any]: """ Overview: - Forward and backward function of learn mode. + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data from the replay buffer and then returns the output \ + result, including various training information such as loss, action, priority. Arguments: - - data (:obj:`dict`): Dict type data, including at least ['obs', 'action', 'reward', 'next_obs'] + - data (:obj:`List[Dict[int, Any]]`): The input data used for policy forward, including a batch of \ + training samples. For each element in list, the key of the dict is the name of data items and the \ + value is the corresponding data. Usually, the value is torch.Tensor or np.ndarray or there dict/list \ + combinations. In the ``_forward_learn`` method, data often need to first be stacked in the batch \ + dimension by some utility functions such as ``default_preprocess_learn``. \ + For SAC, each element in list is a dict containing at least the following keys: ``obs``, ``action``, \ + ``reward``, ``next_obs``, ``done``. Sometimes, it also contains other keys such as ``weight``. Returns: - - info_dict (:obj:`Dict[str, Any]`): Including current lr, loss, target_q_value and other \ - running information. + - info_dict (:obj:`Dict[str, Any]`): The information dict that indicated training result, which will be \ + recorded in text log and tensorboard, values must be python scalar or a list of scalars. For the \ + detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for SACPolicy: ``ding.policy.tests.test_sac``. """ loss_dict = {} data = default_preprocess_learn( @@ -808,34 +820,27 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: # 1. predict q value q_value = self._learn_model.forward(data, mode='compute_critic')['q_value'] - # 2. predict target value depend self._value_network. - if self._value_network: - v_value = self._learn_model.forward(obs, mode='compute_value_critic')['v_value'] - with torch.no_grad(): - next_v_value = self._target_model.forward(next_obs, mode='compute_value_critic')['v_value'] - target_q_value = next_v_value - else: - # target q value. - with torch.no_grad(): - (mu, sigma) = self._learn_model.forward(next_obs, mode='compute_actor')['logit'] - - dist = Independent(Normal(mu, sigma), 1) - pred = dist.rsample() - next_action = torch.tanh(pred) - y = 1 - next_action.pow(2) + 1e-6 - # keep dimension for loss computation (usually for action space is 1 env. e.g. pendulum) - next_log_prob = dist.log_prob(pred).unsqueeze(-1) - next_log_prob = next_log_prob - torch.log(y).sum(-1, keepdim=True) - - next_data = {'obs': next_obs, 'action': next_action} - target_q_value = self._target_model.forward(next_data, mode='compute_critic')['q_value'] - # the value of a policy according to the maximum entropy objective - if self._twin_critic: - # find min one as target q value - target_q_value = torch.min(target_q_value[0], - target_q_value[1]) - self._alpha * next_log_prob.squeeze(-1) - else: - target_q_value = target_q_value - self._alpha * next_log_prob.squeeze(-1) + # 2. predict target value + with torch.no_grad(): + (mu, sigma) = self._learn_model.forward(next_obs, mode='compute_actor')['logit'] + + dist = Independent(Normal(mu, sigma), 1) + pred = dist.rsample() + next_action = torch.tanh(pred) + y = 1 - next_action.pow(2) + 1e-6 + # keep dimension for loss computation (usually for action space is 1 env. e.g. pendulum) + next_log_prob = dist.log_prob(pred).unsqueeze(-1) + next_log_prob = next_log_prob - torch.log(y).sum(-1, keepdim=True) + + next_data = {'obs': next_obs, 'action': next_action} + target_q_value = self._target_model.forward(next_data, mode='compute_critic')['q_value'] + # the value of a policy according to the maximum entropy objective + if self._twin_critic: + # find min one as target q value + target_q_value = torch.min(target_q_value[0], + target_q_value[1]) - self._alpha * next_log_prob.squeeze(-1) + else: + target_q_value = target_q_value - self._alpha * next_log_prob.squeeze(-1) # 3. compute q loss if self._twin_critic: @@ -850,9 +855,10 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: # 4. update q network self._optimizer_q.zero_grad() - loss_dict['critic_loss'].backward() if self._twin_critic: - loss_dict['twin_critic_loss'].backward() + (loss_dict['critic_loss'] + loss_dict['twin_critic_loss']).backward() + else: + loss_dict['critic_loss'].backward() self._optimizer_q.step() # 5. evaluate to get action distribution @@ -870,28 +876,17 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: if self._twin_critic: new_q_value = torch.min(new_q_value[0], new_q_value[1]) - # 6. (optional) compute value loss and update value network - if self._value_network: - # new_q_value: (bs, ), log_prob: (bs, act_shape) -> target_v_value: (bs, ) - target_v_value = (new_q_value.unsqueeze(-1) - self._alpha * log_prob).mean(dim=-1) - loss_dict['value_loss'] = F.mse_loss(v_value, target_v_value.detach()) - - # update value network - self._optimizer_value.zero_grad() - loss_dict['value_loss'].backward() - self._optimizer_value.step() - - # 7. compute policy loss + # 6. compute policy loss policy_loss = (self._alpha * log_prob - new_q_value.unsqueeze(-1)).mean() loss_dict['policy_loss'] = policy_loss - # 8. update policy network + # 7. update policy network self._optimizer_policy.zero_grad() loss_dict['policy_loss'].backward() self._optimizer_policy.step() - # 9. compute alpha loss + # 8. compute alpha loss if self._auto_alpha: if self._log_space: log_prob = log_prob + self._target_entropy @@ -912,10 +907,6 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: loss_dict['total_loss'] = sum(loss_dict.values()) - # ============= - # after update - # ============= - self._forward_learn_cnt += 1 # target update self._target_model.update(self._learn_model.state_dict()) return { @@ -925,55 +916,86 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: 'td_error': td_error_per_sample.detach().mean().item(), 'alpha': self._alpha.item(), 'target_q_value': target_q_value.detach().mean().item(), + 'transformed_log_prob': log_prob.mean().item(), **loss_dict } def _state_dict_learn(self) -> Dict[str, Any]: + """ + Overview: + Return the state_dict of learn mode, usually including model, target_model and optimizers. + Returns: + - state_dict (:obj:`Dict[str, Any]`): The dict of current policy learn state, for saving and restoring. + """ ret = { 'model': self._learn_model.state_dict(), 'target_model': self._target_model.state_dict(), 'optimizer_q': self._optimizer_q.state_dict(), 'optimizer_policy': self._optimizer_policy.state_dict(), } - if self._value_network: - ret.update({'optimizer_value': self._optimizer_value.state_dict()}) if self._auto_alpha: ret.update({'optimizer_alpha': self._alpha_optim.state_dict()}) return ret def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: + """ + Overview: + Load the state_dict variable into policy learn mode. + Arguments: + - state_dict (:obj:`Dict[str, Any]`): The dict of policy learn state saved before. + + .. tip:: + If you want to only load some parts of model, you can simply set the ``strict`` argument in \ + load_state_dict to ``False``, or refer to ``ding.torch_utils.checkpoint_helper`` for more \ + complicated operation. + """ self._learn_model.load_state_dict(state_dict['model']) self._target_model.load_state_dict(state_dict['target_model']) self._optimizer_q.load_state_dict(state_dict['optimizer_q']) - if self._value_network: - self._optimizer_value.load_state_dict(state_dict['optimizer_value']) self._optimizer_policy.load_state_dict(state_dict['optimizer_policy']) if self._auto_alpha: self._alpha_optim.load_state_dict(state_dict['optimizer_alpha']) def _init_collect(self) -> None: - r""" + """ Overview: - Collect mode init method. Called by ``self.__init__``. - Init traj and unroll length, collect model. - Use action noise for exploration. + Initialize the collect mode of policy, including related attributes and modules. For SAC, it contains the \ + collect_model other algorithm-specific arguments such as unroll_len. \ + This method will be called in ``__init__`` method if ``collect`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_collect`` method, you'd better name them \ + with prefix ``_collect_`` to avoid conflict with other modes, such as ``self._collect_attr1``. """ self._unroll_len = self._cfg.collect.unroll_len self._collect_model = model_wrap(self._model, wrapper_name='base') self._collect_model.reset() - def _forward_collect(self, data: dict) -> dict: - r""" + def _forward_collect(self, data: Dict[int, Any], **kwargs) -> Dict[int, Any]: + """ Overview: - Forward function of collect mode. + Policy forward function of collect mode (collecting training data by interacting with envs). Forward means \ + that the policy gets some necessary data (mainly observation) from the envs and then returns the output \ + data, such as the action to interact with the envs. Arguments: - - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \ - values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer. + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. Returns: - - output (:obj:`Dict[int, Any]`): Dict type data, including at least inferred action according to input obs. - ReturnsKeys - - necessary: ``action`` - - optional: ``logit`` + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action and \ + other necessary data for learn mode defined in ``self._process_transition`` method. The key of the \ + dict is the same as the input data, i.e. environment id. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + ``logit`` in SAC means the mu and sigma of Gaussioan distribution. Here we use this name for consistency. + + .. note:: + For more detailed examples, please refer to our unittest for SACPolicy: ``ding.policy.tests.test_sac``. """ data_id = list(data.keys()) data = default_collate(list(data.values())) @@ -990,17 +1012,22 @@ def _forward_collect(self, data: dict) -> dict: output = default_decollate(output) return {i: d for i, d in zip(data_id, output)} - def _process_transition(self, obs: Any, policy_output: dict, timestep: namedtuple) -> dict: - r""" + def _process_transition(self, obs: torch.Tensor, policy_output: Dict[str, torch.Tensor], + timestep: namedtuple) -> Dict[str, torch.Tensor]: + """ Overview: - Generate dict type transition data from inputs. + Process and pack one timestep transition data into a dict, which can be directly used for training and \ + saved in replay buffer. For continuous SAC, it contains obs, next_obs, action, reward, done. The logit \ + will be also added when ``collector_logit`` is True. Arguments: - - obs (:obj:`Any`): Env observation - - policy_output (:obj:`dict`): Output of policy collect model, including at least ['action'] - - timestep (:obj:`namedtuple`): Output after env step, including at least ['obs', 'reward', 'done'] \ - (here 'obs' indicates obs after env step, i.e. next_obs). - Return: - - transition (:obj:`Dict[str, Any]`): Dict type transition data. + - obs (:obj:`torch.Tensor`): The env observation of current timestep, such as stacked 2D image in Atari. + - policy_output (:obj:`Dict[str, torch.Tensor]`): The output of the policy network with the observation \ + as input. For continuous SAC, it contains the action and the logit (mu and sigma) of the action. + - timestep (:obj:`namedtuple`): The execution result namedtuple returned by the environment step method, \ + except all the elements have been transformed into tensor data. Usually, it contains the next obs, \ + reward, done, info, etc. + Returns: + - transition (:obj:`Dict[str, torch.Tensor]`): The processed transition data of the current timestep. """ if self._cfg.collect.collector_logit: transition = { @@ -1021,30 +1048,59 @@ def _process_transition(self, obs: Any, policy_output: dict, timestep: namedtupl } return transition - def _get_train_sample(self, data: list) -> Union[None, List[Any]]: - return get_train_sample(data, self._unroll_len) + def _get_train_sample(self, transitions: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Overview: + For a given trajectory (transitions, a list of transition) data, process it into a list of sample that \ + can be used for training directly. In continuous SAC, a train sample is a processed transition \ + (unroll_len=1). + Arguments: + - transitions (:obj:`List[Dict[str, Any]`): The trajectory data (a list of transition), each element is \ + the same format as the return value of ``self._process_transition`` method. + Returns: + - samples (:obj:`List[Dict[str, Any]]`): The processed train samples, each element is the similar format \ + as input transitions, but may contain more data for training. + """ + return get_train_sample(transitions, self._unroll_len) def _init_eval(self) -> None: - r""" + """ Overview: - Evaluate mode init method. Called by ``self.__init__``. - Init eval model. Unlike learn and collect model, eval model does not need noise. + Initialize the eval mode of policy, including related attributes and modules. For SAC, it contains the \ + eval model, which is equipped with ``base`` model wrapper to ensure compability. + This method will be called in ``__init__`` method if ``eval`` field is in ``enable_field``. + + .. note:: + If you want to set some spacial member variables in ``_init_eval`` method, you'd better name them \ + with prefix ``_eval_`` to avoid conflict with other modes, such as ``self._eval_attr1``. """ self._eval_model = model_wrap(self._model, wrapper_name='base') self._eval_model.reset() - def _forward_eval(self, data: dict) -> dict: - r""" + def _forward_eval(self, data: Dict[int, Any]) -> Dict[int, Any]: + """ Overview: - Forward function of eval mode, similar to ``self._forward_collect``. + Policy forward function of eval mode (evaluation policy performance by interacting with envs). Forward \ + means that the policy gets some necessary data (mainly observation) from the envs and then returns the \ + action to interact with the envs. Arguments: - - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \ - values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer. + - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ + key of the dict is environment id and the value is the corresponding data of the env. Returns: - - output (:obj:`Dict[int, Any]`): The dict of predicting action for the interaction with env. - ReturnsKeys - - necessary: ``action`` - - optional: ``logit`` + - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action. The \ + key of the dict is the same as the input data, i.e. environment id. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + ``logit`` in SAC means the mu and sigma of Gaussioan distribution. Here we use this name for consistency. + + .. note:: + For more detailed examples, please refer to our unittest for SACPolicy: ``ding.policy.tests.test_sac``. """ data_id = list(data.keys()) data = default_collate(list(data.values())) @@ -1061,16 +1117,17 @@ def _forward_eval(self, data: dict) -> dict: return {i: d for i, d in zip(data_id, output)} def _monitor_vars_learn(self) -> List[str]: - r""" + """ Overview: - Return variables' name if variables are to used in monitor. + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. Returns: - - vars (:obj:`List[str]`): Variables' name list. + - necessary_keys (:obj:`List[str]`): The list of the necessary keys to be logged. """ twin_critic = ['twin_critic_loss'] if self._twin_critic else [] alpha_loss = ['alpha_loss'] if self._auto_alpha else [] - value_loss = ['value_loss'] if self._value_network else [] return [ + 'value_loss' 'alpha_loss', 'policy_loss', 'critic_loss', @@ -1079,214 +1136,49 @@ def _monitor_vars_learn(self) -> List[str]: 'target_q_value', 'alpha', 'td_error', - ] + twin_critic + alpha_loss + value_loss + 'transformed_log_prob', + ] + twin_critic + alpha_loss @POLICY_REGISTRY.register('sqil_sac') class SQILSACPolicy(SACPolicy): - r""" - Overview: - Policy class of continuous SAC algorithm. - - https://arxiv.org/pdf/1801.01290.pdf - - Config: - == ==================== ======== ============= ================================= ======================= - ID Symbol Type Default Value Description Other(Shape) - == ==================== ======== ============= ================================= ======================= - 1 ``type`` str td3 | RL policy register name, refer | this arg is optional, - | to registry ``POLICY_REGISTRY`` | a placeholder - 2 ``cuda`` bool True | Whether to use cuda for network | - 3 | ``random_`` int 10000 | Number of randomly collected | Default to 10000 for - | ``collect_size`` | training samples in replay | SAC, 25000 for DDPG/ - | | buffer when training starts. | TD3. - 4 | ``model.policy_`` int 256 | Linear layer size for policy | - | ``embedding_size`` | network. | - 5 | ``model.soft_q_`` int 256 | Linear layer size for soft q | - | ``embedding_size`` | network. | - 6 | ``model.value_`` int 256 | Linear layer size for value | Defalut to None when - | ``embedding_size`` | network. | model.value_network - | | | is False. - 7 | ``learn.learning`` float 3e-4 | Learning rate for soft q | Defalut to 1e-3, when - | ``_rate_q`` | network. | model.value_network - | | | is True. - 8 | ``learn.learning`` float 3e-4 | Learning rate for policy | Defalut to 1e-3, when - | ``_rate_policy`` | network. | model.value_network - | | | is True. - 9 | ``learn.learning`` float 3e-4 | Learning rate for policy | Defalut to None when - | ``_rate_value`` | network. | model.value_network - | | | is False. - 10 | ``learn.alpha`` float 0.2 | Entropy regularization | alpha is initiali- - | | coefficient. | zation for auto - | | | `\alpha`, when - | | | auto_alpha is True - 11 | ``learn.repara_`` bool True | Determine whether to use | - | ``meterization`` | reparameterization trick. | - 12 | ``learn.`` bool False | Determine whether to use | Temperature parameter - | ``auto_alpha`` | auto temperature parameter | determines the - | | `\alpha`. | relative importance - | | | of the entropy term - | | | against the reward. - 13 | ``learn.-`` bool False | Determine whether to ignore | Use ignore_done only - | ``ignore_done`` | done flag. | in halfcheetah env. - 14 | ``learn.-`` float 0.005 | Used for soft update of the | aka. Interpolation - | ``target_theta`` | target network. | factor in polyak aver - | | | aging for target - | | | networks. - == ==================== ======== ============= ================================= ======================= - """ - - config = dict( - # (str) RL policy register name (refer to function "POLICY_REGISTRY"). - type='sqil_sac', - # (bool) Whether to use cuda for network. - cuda=False, - # (bool type) on_policy: Determine whether on-policy or off-policy. - # on-policy setting influences the behaviour of buffer. - # Default False in SAC. - on_policy=False, - # (bool type) priority: Determine whether to use priority in buffer sample. - # Default False in SAC. - priority=False, - # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. - priority_IS_weight=False, - # (int) Number of training samples(randomly collected) in replay buffer when training starts. - # Default 10000 in SAC. - random_collect_size=10000, - # (bool) Whether to need policy data in process transition - transition_with_policy_data=True, - # (bool) Whether to enable multi-agent training setting - multi_agent=False, - model=dict( - # (bool type) twin_critic: Determine whether to use double-soft-q-net for target q computation. - # Please refer to TD3 about Clipped Double-Q Learning trick, which learns two Q-functions instead of one . - # Default to True. - twin_critic=True, - - # (bool type) value_network: Determine whether to use value network as the - # original SAC paper (arXiv 1801.01290). - # using value_network needs to set learning_rate_value, learning_rate_q, - # and learning_rate_policy in `cfg.policy.learn`. - # Default to False. - # value_network=False, - - # (str type) action_space: Use reparameterization trick for continous action - action_space='reparameterization', - ), - learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, - # How many updates(iterations) to train after collector's one collection. - # Bigger "update_per_collect" means bigger off-policy. - # collect data -> update policy-> collect data -> ... - update_per_collect=1, - # (int) Minibatch size for gradient descent. - batch_size=256, - - # (float type) learning_rate_q: Learning rate for soft q network. - # Default to 3e-4. - # Please set to 1e-3, when model.value_network is True. - learning_rate_q=3e-4, - # (float type) learning_rate_policy: Learning rate for policy network. - # Default to 3e-4. - # Please set to 1e-3, when model.value_network is True. - learning_rate_policy=3e-4, - # (float type) learning_rate_value: Learning rate for value network. - # `learning_rate_value` should be initialized, when model.value_network is True. - # Please set to 3e-4, when model.value_network is True. - learning_rate_value=3e-4, + """ + Overview: + Policy class of continuous SAC algorithm with SQIL extension. + SAC paper link: https://arxiv.org/pdf/1801.01290.pdf + SQIL paper link: https://arxiv.org/abs/1905.11108 + """ - # (float type) learning_rate_alpha: Learning rate for auto temperature parameter `\alpha`. - # Default to 3e-4. - learning_rate_alpha=3e-4, - # (float type) target_theta: Used for soft update of the target network, - # aka. Interpolation factor in polyak averaging for target networks. - # Default to 0.005. - target_theta=0.005, - # (float) discount factor for the discounted sum of rewards, aka. gamma. - discount_factor=0.99, + def _init_learn(self) -> None: + """ + Overview: + Initialize the learn mode of policy, including related attributes and modules. For SAC, it mainly \ + contains three optimizers, algorithm-specific arguments such as gamma and twin_critic, main and target \ + model. Especially, the ``auto_alpha`` mechanism for balancing max entropy target is also initialized here. + This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. - # (float type) alpha: Entropy regularization coefficient. - # Please check out the original SAC paper (arXiv 1801.01290): Eq 1 for more details. - # If auto_alpha is set to `True`, alpha is initialization for auto `\alpha`. - # Default to 0.2. - alpha=0.2, + .. note:: + For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ + and ``_load_state_dict_learn`` methods. - # (bool type) auto_alpha: Determine whether to use auto temperature parameter `\alpha` . - # Temperature parameter determines the relative importance of the entropy term against the reward. - # Please check out the original SAC paper (arXiv 1801.01290): Eq 1 for more details. - # Default to False. - # Note that: Using auto alpha needs to set learning_rate_alpha in `cfg.policy.learn`. - auto_alpha=True, - # (bool type) log_space: Determine whether to use auto `\alpha` in log space. - log_space=True, - # (bool) Whether ignore done(usually for max step termination env. e.g. pendulum) - # Note: Gym wraps the MuJoCo envs by default with TimeLimit environment wrappers. - # These limit HalfCheetah, and several other MuJoCo envs, to max length of 1000. - # However, interaction with HalfCheetah always gets done with False, - # Since we inplace done==True with done==False to keep - # TD-error accurate computation(``gamma * (1 - done) * next_v + reward``), - # when the episode step is greater than max episode step. - ignore_done=False, - # (float) Weight uniform initialization range in the last output layer - init_w=3e-3, - ), - collect=dict( - # If you need the data collected by the collector to contain logit key which reflect the probability of - # the action, you can change the key to be True. - # In Guided cost Learning, we need to use logit to train the reward model, we change the key to be True. - # Default collector_logit to False. - collector_logit=False, - # You can use either "n_sample" or "n_episode" in actor.collect. - # Get "n_sample" samples per collect. - # Default n_sample to 1. - n_sample=1, - # (int) Cut trajectories into pieces with length "unroll_len". - unroll_len=1, - ), - eval=dict( - evaluator=dict( - # (int) Evaluate every "eval_freq" training iterations. - eval_freq=5000, - ), - ), - other=dict( - replay_buffer=dict( - # (int type) replay_buffer_size: Max size of replay buffer. - replay_buffer_size=1000000, - # (int type) max_use: Max use times of one data in the buffer. - # Data will be removed once used for too many times. - # Default to infinite. - # max_use=256, - ), - ), - ) + .. note:: + For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. - def _init_learn(self) -> None: - r""" - Overview: - Learn mode init method. Called by ``self.__init__``. - Init q, value and policy's optimizers, algorithm config, main and target models. + .. note:: + If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ + with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. """ - # Init self._priority = self._cfg.priority self._priority_IS_weight = self._cfg.priority_IS_weight - self._value_network = False # TODO self._cfg.model.value_network self._twin_critic = self._cfg.model.twin_critic # Weight Init for the last output layer init_w = self._cfg.learn.init_w - self._model.actor[2].mu.weight.data.uniform_(-init_w, init_w) - self._model.actor[2].mu.bias.data.uniform_(-init_w, init_w) - self._model.actor[2].log_sigma_layer.weight.data.uniform_(-init_w, init_w) - self._model.actor[2].log_sigma_layer.bias.data.uniform_(-init_w, init_w) - - # Optimizers - if self._value_network: - self._optimizer_value = Adam( - self._model.value_critic.parameters(), - lr=self._cfg.learn.learning_rate_value, - ) + self._model.actor_head[-1].mu.weight.data.uniform_(-init_w, init_w) + self._model.actor_head[-1].mu.bias.data.uniform_(-init_w, init_w) + self._model.actor_head[-1].log_sigma_layer.weight.data.uniform_(-init_w, init_w) + self._model.actor_head[-1].log_sigma_layer.bias.data.uniform_(-init_w, init_w) + self._optimizer_q = Adam( self._model.critic.parameters(), lr=self._cfg.learn.learning_rate_q, @@ -1296,11 +1188,14 @@ def _init_learn(self) -> None: lr=self._cfg.learn.learning_rate_policy, ) - # Algorithm config + # Algorithm-Specific Config self._gamma = self._cfg.learn.discount_factor - # Init auto alpha if self._cfg.learn.auto_alpha: - self._target_entropy = self._cfg.learn.get('target_entropy', -np.prod(self._cfg.model.action_shape)) + if self._cfg.learn.target_entropy is None: + assert 'action_shape' in self._cfg.model, "SQILSACPolicy need network model with action_shape variable" + self._target_entropy = -np.prod(self._cfg.model.action_shape) + else: + self._target_entropy = self._cfg.learn.target_entropy if self._cfg.learn.log_space: self._log_alpha = torch.log(torch.FloatTensor([self._cfg.learn.alpha])) self._log_alpha = self._log_alpha.to(self._device).requires_grad_() @@ -1332,20 +1227,41 @@ def _init_learn(self) -> None: self._learn_model.reset() self._target_model.reset() - self._forward_learn_cnt = 0 # monitor cossimilarity and entropy switch self._monitor_cos = True self._monitor_entropy = True - def _forward_learn(self, data: dict) -> Dict[str, Any]: + def _forward_learn(self, data: List[Dict[str, Any]]) -> Dict[str, Any]: """ Overview: - Forward and backward function of learn mode. + Policy forward function of learn mode (training policy and updating parameters). Forward means \ + that the policy inputs some training batch data from the replay buffer and then returns the output \ + result, including various training information such as loss, action, priority. Arguments: - - data (:obj:`dict`): Dict type data, including at least ['obs', 'action', 'reward', 'next_obs'] + - data (:obj:`List[Dict[int, Any]]`): The input data used for policy forward, including a batch of \ + training samples. For each element in list, the key of the dict is the name of data items and the \ + value is the corresponding data. Usually, the value is torch.Tensor or np.ndarray or there dict/list \ + combinations. In the ``_forward_learn`` method, data often need to first be stacked in the batch \ + dimension by some utility functions such as ``default_preprocess_learn``. \ + For SAC, each element in list is a dict containing at least the following keys: ``obs``, ``action``, \ + ``reward``, ``next_obs``, ``done``. Sometimes, it also contains other keys such as ``weight``. Returns: - - info_dict (:obj:`Dict[str, Any]`): Including current lr, loss, target_q_value and other \ - running information. + - info_dict (:obj:`Dict[str, Any]`): The information dict that indicated training result, which will be \ + recorded in text log and tensorboard, values must be python scalar or a list of scalars. For the \ + detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. + + .. note:: + For SQIL + SAC, input data is composed of two parts with the same size: agent data and expert data. \ + Both of them are relabelled with new reward according to SQIL algorithm. + + .. note:: + The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ + For the data type that not supported, the main reason is that the corresponding model does not support it. \ + You can implement you own model rather than use the default model. For more information, please raise an \ + issue in GitHub repo and we will continue to follow up. + + .. note:: + For more detailed examples, please refer to our unittest for SACPolicy: ``ding.policy.tests.test_sac``. """ loss_dict = {} if self._monitor_cos: @@ -1388,33 +1304,26 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: # 1. predict q value q_value = self._learn_model.forward(data, mode='compute_critic')['q_value'] - # 2. predict target value depend self._value_network. - if self._value_network: - v_value = self._learn_model.forward(obs, mode='compute_value_critic')['v_value'] - with torch.no_grad(): - next_v_value = self._target_model.forward(next_obs, mode='compute_value_critic')['v_value'] - target_q_value = next_v_value - else: - # target q value. - with torch.no_grad(): - (mu, sigma) = self._learn_model.forward(next_obs, mode='compute_actor')['logit'] - dist = Independent(Normal(mu, sigma), 1) - pred = dist.rsample() - next_action = torch.tanh(pred) - y = 1 - next_action.pow(2) + 1e-6 - # keep dimension for loss computation (usually for action space is 1 env. e.g. pendulum) - next_log_prob = dist.log_prob(pred).unsqueeze(-1) - next_log_prob = next_log_prob - torch.log(y).sum(-1, keepdim=True) - - next_data = {'obs': next_obs, 'action': next_action} - target_q_value = self._target_model.forward(next_data, mode='compute_critic')['q_value'] - # the value of a policy according to the maximum entropy objective - if self._twin_critic: - # find min one as target q value - target_q_value = torch.min(target_q_value[0], - target_q_value[1]) - self._alpha * next_log_prob.squeeze(-1) - else: - target_q_value = target_q_value - self._alpha * next_log_prob.squeeze(-1) + # 2. predict target value + with torch.no_grad(): + (mu, sigma) = self._learn_model.forward(next_obs, mode='compute_actor')['logit'] + dist = Independent(Normal(mu, sigma), 1) + pred = dist.rsample() + next_action = torch.tanh(pred) + y = 1 - next_action.pow(2) + 1e-6 + # keep dimension for loss computation (usually for action space is 1 env. e.g. pendulum) + next_log_prob = dist.log_prob(pred).unsqueeze(-1) + next_log_prob = next_log_prob - torch.log(y).sum(-1, keepdim=True) + + next_data = {'obs': next_obs, 'action': next_action} + target_q_value = self._target_model.forward(next_data, mode='compute_critic')['q_value'] + # the value of a policy according to the maximum entropy objective + if self._twin_critic: + # find min one as target q value + target_q_value = torch.min(target_q_value[0], + target_q_value[1]) - self._alpha * next_log_prob.squeeze(-1) + else: + target_q_value = target_q_value - self._alpha * next_log_prob.squeeze(-1) # 3. compute q loss if self._twin_critic: @@ -1429,9 +1338,10 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: # 4. update q network self._optimizer_q.zero_grad() - loss_dict['critic_loss'].backward() if self._twin_critic: - loss_dict['twin_critic_loss'].backward() + (loss_dict['critic_loss'] + loss_dict['twin_critic_loss']).backward() + else: + loss_dict['critic_loss'].backward() self._optimizer_q.step() # 5. evaluate to get action distribution @@ -1465,7 +1375,6 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: if self._twin_critic: expert_new_q_value = torch.min(expert_new_q_value[0], expert_new_q_value[1]) - # all data (mu, sigma) = self._learn_model.forward(data['obs'], mode='compute_actor')['logit'] dist = Independent(Normal(mu, sigma), 1) # for monitor the entropy of policy @@ -1485,22 +1394,11 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: if self._twin_critic: new_q_value = torch.min(new_q_value[0], new_q_value[1]) - # 6. (optional) compute value loss and update value network - if self._value_network: - # new_q_value: (bs, ), log_prob: (bs, act_shape) -> target_v_value: (bs, ) - target_v_value = (new_q_value.unsqueeze(-1) - self._alpha * log_prob).mean(dim=-1) - loss_dict['value_loss'] = F.mse_loss(v_value, target_v_value.detach()) - - # update value network - self._optimizer_value.zero_grad() - loss_dict['value_loss'].backward() - self._optimizer_value.step() - - # 7. compute policy loss + # 6. compute policy loss policy_loss = (self._alpha * log_prob - new_q_value.unsqueeze(-1)).mean() loss_dict['policy_loss'] = policy_loss - # 8. update policy network + # 7. update policy network if self._monitor_cos: agent_policy_loss = (self._alpha * agent_log_prob - agent_new_q_value.unsqueeze(-1)).mean() expert_policy_loss = (self._alpha * expert_log_prob - expert_new_q_value.unsqueeze(-1)).mean() @@ -1518,7 +1416,7 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: loss_dict['policy_loss'].backward() self._optimizer_policy.step() - # 9. compute alpha loss + # 8. compute alpha loss if self._auto_alpha: if self._log_space: log_prob = log_prob + self._target_entropy @@ -1539,10 +1437,6 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: loss_dict['total_loss'] = sum(loss_dict.values()) - # ============= - # after update - # ============= - self._forward_learn_cnt += 1 # target update self._target_model.update(self._learn_model.state_dict()) var_monitor = { @@ -1567,18 +1461,19 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: return var_monitor def _monitor_vars_learn(self) -> List[str]: - r""" + """ Overview: - Return variables' name if variables are to used in monitor. + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. Returns: - - vars (:obj:`List[str]`): Variables' name list. + - necessary_keys (:obj:`List[str]`): The list of the necessary keys to be logged. """ twin_critic = ['twin_critic_loss'] if self._twin_critic else [] alpha_loss = ['alpha_loss'] if self._auto_alpha else [] - value_loss = ['value_loss'] if self._value_network else [] cos_similarity = ['cos_similarity'] if self._monitor_cos else [] entropy = ['entropy'] if self._monitor_entropy else [] return [ + 'value_loss' 'alpha_loss', 'policy_loss', 'critic_loss', @@ -1593,4 +1488,4 @@ def _monitor_vars_learn(self) -> List[str]: 'sigma', 'q_value0', 'q_value1', - ] + twin_critic + alpha_loss + value_loss + cos_similarity + entropy + ] + twin_critic + alpha_loss + cos_similarity + entropy diff --git a/ding/policy/sql.py b/ding/policy/sql.py index d8eb76ef5e..dc6170dfb7 100644 --- a/ding/policy/sql.py +++ b/ding/policy/sql.py @@ -35,8 +35,7 @@ class SQLPolicy(Policy): # (int) N-step reward for target q_value estimation nstep=1, learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, + # How many updates(iterations) to train after collector's one collection. # Bigger "update_per_collect" means bigger off-policy. # collect data -> update policy-> collect data -> ... @@ -152,7 +151,7 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: # ==================== self._optimizer.zero_grad() loss.backward() - if self._cfg.learn.multi_gpu: + if self._cfg.multi_gpu: self.sync_gradients(self._learn_model) self._optimizer.step() diff --git a/ding/policy/sqn.py b/ding/policy/sqn.py index 23a91d3f98..5241ee993e 100644 --- a/ding/policy/sqn.py +++ b/ding/policy/sqn.py @@ -31,7 +31,6 @@ class SQNPolicy(Policy): # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. priority_IS_weight=False, learn=dict( - multi_gpu=False, update_per_collect=16, batch_size=64, learning_rate_q=0.001, diff --git a/ding/policy/td3.py b/ding/policy/td3.py index 30cfe82d89..7359190282 100644 --- a/ding/policy/td3.py +++ b/ding/policy/td3.py @@ -1,20 +1,15 @@ +from typing import List from ding.utils import POLICY_REGISTRY from .ddpg import DDPGPolicy @POLICY_REGISTRY.register('td3') class TD3Policy(DDPGPolicy): - r""" + """ Overview: - Policy class of TD3 algorithm. - - Since DDPG and TD3 share many common things, we can easily derive this TD3 + Policy class of TD3 algorithm. Since DDPG and TD3 share many common things, we can easily derive this TD3 \ class from DDPG class by changing ``_actor_update_freq``, ``_twin_critic`` and noise in model wrapper. - - https://arxiv.org/pdf/1802.09477.pdf - - Property: - learn_mode, collect_mode, eval_mode + Paper link: https://arxiv.org/pdf/1802.09477.pdf Config: @@ -67,9 +62,7 @@ class from DDPG class by changing ``_actor_update_freq``, ``_twin_critic`` and n type='td3', # (bool) Whether to use cuda for network. cuda=False, - # (bool type) on_policy: Determine whether on-policy or off-policy. - # on-policy setting influences the behaviour of buffer. - # Default False in TD3. + # (bool) on_policy: Determine whether on-policy or off-policy. Default False in TD3. on_policy=False, # (bool) Whether use priority(priority sample, IS weight, update priority) # Default False in TD3. @@ -79,19 +72,23 @@ class from DDPG class by changing ``_actor_update_freq``, ``_twin_critic`` and n # (int) Number of training samples(randomly collected) in replay buffer when training starts. # Default 25000 in DDPG/TD3. random_collect_size=25000, + # (bool) Whether to need policy data in process transition. + transition_with_policy_data=False, # (str) Action space type action_space='continuous', # ['continuous', 'hybrid'] # (bool) Whether use batch normalization for reward reward_batch_norm=False, + # (bool) Whether to enable multi-agent training setting + multi_agent=False, model=dict( # (bool) Whether to use two critic networks or only one. # Clipped Double Q-Learning for Actor-Critic in original TD3 paper(https://arxiv.org/pdf/1802.09477.pdf). # Default True for TD3, False for DDPG. twin_critic=True, ), + # learn_mode config learn=dict( - multi_gpu=False, - # How many updates(iterations) to train after collector's one collection. + # (int) How many updates(iterations) to train after collector's one collection. # Bigger "update_per_collect" means bigger off-policy. # collect data -> update policy-> collect data -> ... update_per_collect=1, @@ -109,7 +106,7 @@ class from DDPG class by changing ``_actor_update_freq``, ``_twin_critic`` and n # TD-error accurate computation(``gamma * (1 - done) * next_v + reward``), # when the episode step is greater than max episode step. ignore_done=False, - # (float type) target_theta: Used for soft update of the target network, + # (float) target_theta: Used for soft update of the target network, # aka. Interpolation factor in polyak averaging for target networks. # Default to 0.005. target_theta=0.005, @@ -127,27 +124,37 @@ class from DDPG class by changing ``_actor_update_freq``, ``_twin_critic`` and n noise_sigma=0.2, # (dict) Limit for range of target policy smoothing noise, aka. noise_clip. noise_range=dict( + # (int) min value of noise min=-0.5, + # (int) max value of noise max=0.5, ), ), + # collect_mode config collect=dict( + # (int) How many training samples collected in one collection procedure. + # Only one of [n_sample, n_episode] shoule be set. # n_sample=1, # (int) Cut trajectories into pieces with length "unroll_len". unroll_len=1, # (float) It is a must to add noise during collection. So here omits "noise" and only set "noise_sigma". noise_sigma=0.1, ), - eval=dict( - evaluator=dict( - # (int) Evaluate every "eval_freq" training iterations. - eval_freq=5000, - ), - ), + eval=dict(), # for compability other=dict( replay_buffer=dict( - # (int) Maximum size of replay buffer. + # (int) Maximum size of replay buffer. Usually, larger buffer size is better. replay_buffer_size=100000, ), ), ) + + def _monitor_vars_learn(self) -> List[str]: + """ + Overview: + Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ + as text logger, tensorboard logger, will use these keys to save the corresponding data. + Returns: + - necessary_keys (:obj:`List[str]`): The list of the necessary keys to be logged. + """ + return ["q_value", "loss", "lr", "entropy", "target_q_value", "td_error"] diff --git a/ding/policy/td3_bc.py b/ding/policy/td3_bc.py index e8e30d8460..e30b6bfc07 100644 --- a/ding/policy/td3_bc.py +++ b/ding/policy/td3_bc.py @@ -111,7 +111,7 @@ class from DDPG class by changing ``_actor_update_freq``, ``_twin_critic`` and n critic_head_hidden_size=256, ), learn=dict( - multi_gpu=False, + # How many updates(iterations) to train after collector's one collection. # Bigger "update_per_collect" means bigger off-policy. # collect data -> update policy-> collect data -> ... @@ -173,11 +173,13 @@ class from DDPG class by changing ``_actor_update_freq``, ``_twin_critic`` and n ), ) + def default_model(self) -> Tuple[str, List[str]]: + return 'continuous_qac', ['ding.model.template.qac'] + def _init_learn(self) -> None: - r""" + """ Overview: - Learn mode init method. Called by ``self.__init__``. - Init actor and critic optimizers, algorithm config. + Learn mode init method. Called by ``self.__init__``. Init actor and critic optimizers, algorithm config. """ super(TD3BCPolicy, self)._init_learn() self._alpha = self._cfg.learn.alpha @@ -195,6 +197,9 @@ def _init_learn(self) -> None: clip_value=1.0, ) + self.noise_sigma = self._cfg.learn.noise_sigma + self.noise_range = self._cfg.learn.noise_range + def _forward_learn(self, data: dict) -> Dict[str, Any]: r""" Overview: @@ -234,6 +239,9 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: # target q value. with torch.no_grad(): next_action = self._target_model.forward(next_obs, mode='compute_actor')['action'] + noise = (torch.randn_like(next_action) * + self.noise_sigma).clamp(self.noise_range['min'], self.noise_range['max']) + next_action = (next_action + noise).clamp(-1, 1) next_data = {'obs': next_obs, 'action': next_action} target_q_value = self._target_model.forward(next_data, mode='compute_critic')['q_value'] if self._twin_critic: diff --git a/ding/policy/td3_vae.py b/ding/policy/td3_vae.py index da3d01ee82..7d029c0a91 100644 --- a/ding/policy/td3_vae.py +++ b/ding/policy/td3_vae.py @@ -105,7 +105,7 @@ class from DDPG class by changing ``_actor_update_freq``, ``_twin_critic`` and n twin_critic=True, ), learn=dict( - multi_gpu=False, + # How many updates(iterations) to train after collector's one collection. # Bigger "update_per_collect" means bigger off-policy. # collect data -> update policy-> collect data -> ... @@ -168,7 +168,7 @@ class from DDPG class by changing ``_actor_update_freq``, ``_twin_critic`` and n ) def default_model(self) -> Tuple[str, List[str]]: - return 'qac', ['ding.model.template.qac'] + return 'continuous_qac', ['ding.model.template.qac'] def _init_learn(self) -> None: r""" diff --git a/ding/policy/tests/test_common_utils.py b/ding/policy/tests/test_common_utils.py new file mode 100644 index 0000000000..38bf67ed98 --- /dev/null +++ b/ding/policy/tests/test_common_utils.py @@ -0,0 +1,179 @@ +import unittest +import pytest +import numpy as np +import torch +import treetensor.torch as ttorch + +from ding.policy.common_utils import default_preprocess_learn + +shape_test = [ + [2], + [1], +] + +dtype_test = [ + "int64", + "float32", +] + +data_type_test = [ + "numpy", + "torch", + "treetensor", +] + + +def get_action(shape, dtype, class_type): + if class_type == "numpy": + if dtype == "int64": + dtype = np.int64 + elif dtype == "float32": + dtype = np.float32 + return np.random.randn(*shape).astype(dtype) + else: + if dtype == "int64": + dtype = torch.int64 + elif dtype == "float32": + dtype = torch.float32 + + if class_type == "torch": + return torch.randn(*shape).type(dtype) + elif class_type == "treetensor": + return ttorch.randn(*shape).type(dtype) + + +@pytest.mark.unittest +def test_default_preprocess_learn_action(): + + for shape in shape_test: + for dtype in dtype_test: + for data_type in data_type_test: + + data = [ + { + 'obs': np.random.randn(4, 84, 84), + 'action': get_action(shape, dtype, data_type), + 'reward': 1.0, + 'next_obs': np.random.randn(4, 84, 84), + 'done': False, + 'weight': 1.0, + 'value': 1.0, + 'adv': 1.0, + } for _ in range(10) + ] + use_priority_IS_weight = False + use_priority = False + use_nstep = False + ignore_done = False + data = default_preprocess_learn(data, use_priority_IS_weight, use_priority, use_nstep, ignore_done) + + assert data['obs'].shape == torch.Size([10, 4, 84, 84]) + if dtype in ["int64"] and shape[0] == 1: + assert data['action'].shape == torch.Size([10]) + else: + assert data['action'].shape == torch.Size([10, *shape]) + assert data['reward'].shape == torch.Size([10]) + assert data['next_obs'].shape == torch.Size([10, 4, 84, 84]) + assert data['done'].shape == torch.Size([10]) + assert data['weight'].shape == torch.Size([10]) + assert data['value'].shape == torch.Size([10]) + assert data['adv'].shape == torch.Size([10]) + + +@pytest.mark.unittest +def test_default_preprocess_learn_reward_done_adv_1d(): + + data = [ + { + 'obs': np.random.randn(4, 84, 84), + 'action': np.random.randn(2), + 'reward': np.array([1.0]), + 'next_obs': np.random.randn(4, 84, 84), + 'done': False, + 'value': np.array([1.0]), + 'adv': np.array([1.0]), + } for _ in range(10) + ] + use_priority_IS_weight = False + use_priority = False + use_nstep = False + ignore_done = False + data = default_preprocess_learn(data, use_priority_IS_weight, use_priority, use_nstep, ignore_done) + + assert data['reward'].shape == torch.Size([10]) + assert data['done'].shape == torch.Size([10]) + assert data['weight'] is None + assert data['value'].shape == torch.Size([10]) + assert data['adv'].shape == torch.Size([10]) + + +@pytest.mark.unittest +def test_default_preprocess_learn_ignore_done(): + data = [ + { + 'obs': np.random.randn(4, 84, 84), + 'action': np.random.randn(2), + 'reward': np.array([1.0]), + 'next_obs': np.random.randn(4, 84, 84), + 'done': True, + 'value': np.array([1.0]), + 'adv': np.array([1.0]), + } for _ in range(10) + ] + use_priority_IS_weight = False + use_priority = False + use_nstep = False + ignore_done = True + data = default_preprocess_learn(data, use_priority_IS_weight, use_priority, use_nstep, ignore_done) + + assert data['done'].dtype == torch.float32 + assert torch.sum(data['done']) == 0 + + +@pytest.mark.unittest +def test_default_preprocess_learn_use_priority_IS_weight(): + data = [ + { + 'obs': np.random.randn(4, 84, 84), + 'action': np.random.randn(2), + 'reward': 1.0, + 'next_obs': np.random.randn(4, 84, 84), + 'done': False, + 'priority_IS': 1.0, + 'value': 1.0, + 'adv': 1.0, + } for _ in range(10) + ] + use_priority_IS_weight = True + use_priority = True + use_nstep = False + ignore_done = False + data = default_preprocess_learn(data, use_priority_IS_weight, use_priority, use_nstep, ignore_done) + + assert data['weight'].shape == torch.Size([10]) + assert torch.sum(data['weight']) == torch.tensor(10.0) + + +@pytest.mark.unittest +def test_default_preprocess_learn_nstep(): + data = [ + { + 'obs': np.random.randn(4, 84, 84), + 'action': np.random.randn(2), + 'reward': np.array([1.0, 2.0, 0.0]), + 'next_obs': np.random.randn(4, 84, 84), + 'done': False, + 'value': 1.0, + 'adv': 1.0, + } for _ in range(10) + ] + use_priority_IS_weight = False + use_priority = False + use_nstep = True + ignore_done = False + data = default_preprocess_learn(data, use_priority_IS_weight, use_priority, use_nstep, ignore_done) + + assert data['reward'].shape == torch.Size([3, 10]) + assert data['reward'][0][0] == torch.tensor(1.0) + assert data['reward'][1][0] == torch.tensor(2.0) + assert data['reward'][2][0] == torch.tensor(0.0) diff --git a/ding/policy/tests/test_cql.py b/ding/policy/tests/test_cql.py index cfac38cca3..248653da6a 100644 --- a/ding/policy/tests/test_cql.py +++ b/ding/policy/tests/test_cql.py @@ -3,7 +3,7 @@ import pytest import torch from easydict import EasyDict -from ding.policy.cql import CQLPolicy, CQLDiscretePolicy +from ding.policy.cql import CQLPolicy, DiscreteCQLPolicy from ding.utils.data import offline_data_save_type from tensorboardX import SummaryWriter from ding.model.wrapper.model_wrappers import ArgmaxSampleWrapper, EpsGreedySampleWrapper, TargetNetworkWrapper @@ -15,7 +15,7 @@ obs_space = 5 action_space = 3 -cfg1 = EasyDict(CQLPolicy.config) +cfg1 = EasyDict(CQLPolicy.default_config()) cfg1.model.obs_shape = obs_space cfg1.model.action_shape = action_space @@ -23,7 +23,7 @@ cfg2.learn.auto_alpha = False cfg2.learn.log_space = False -cfg3 = EasyDict(CQLDiscretePolicy.config) +cfg3 = EasyDict(DiscreteCQLPolicy.default_config()) cfg3.model = {} cfg3.model.obs_shape = obs_space cfg3.model.action_shape = action_space @@ -89,7 +89,7 @@ def get_transition_discrete(size=20): @pytest.mark.parametrize('cfg', [cfg3, cfg4]) @pytest.mark.unittest def test_cql_discrete(cfg): - policy = CQLDiscretePolicy(cfg, enable_field=['collect', 'eval', 'learn']) + policy = DiscreteCQLPolicy(cfg, enable_field=['collect', 'eval', 'learn']) assert type(policy._learn_model) == ArgmaxSampleWrapper assert type(policy._target_model) == TargetNetworkWrapper assert type(policy._collect_model) == EpsGreedySampleWrapper diff --git a/ding/policy/tests/test_stdim.py b/ding/policy/tests/test_stdim.py index 546a67810d..1ff459c922 100644 --- a/ding/policy/tests/test_stdim.py +++ b/ding/policy/tests/test_stdim.py @@ -8,9 +8,7 @@ obs_shape = 4 action_shape = 2 -cfg1 = PPOSTDIMPolicy.config -cfg1["model"] = {} -cfg1 = EasyDict(PPOSTDIMPolicy.config) +cfg1 = EasyDict(PPOSTDIMPolicy.default_config()) cfg1.model.obs_shape = obs_shape cfg1.model.action_shape = action_shape diff --git a/ding/policy/wqmix.py b/ding/policy/wqmix.py index 634238d275..6a27cdce94 100644 --- a/ding/policy/wqmix.py +++ b/ding/policy/wqmix.py @@ -57,8 +57,6 @@ class WQMIXPolicy(QMIXPolicy): # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. priority_IS_weight=False, learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, update_per_collect=20, batch_size=32, learning_rate=0.0005, diff --git a/ding/reward_model/base_reward_model.py b/ding/reward_model/base_reward_model.py index 9dfe226f80..963bacf1d7 100644 --- a/ding/reward_model/base_reward_model.py +++ b/ding/reward_model/base_reward_model.py @@ -1,8 +1,11 @@ from abc import ABC, abstractmethod +from typing import Dict from easydict import EasyDict +from ditk import logging +import os import copy from typing import Any -from ding.utils import REWARD_MODEL_REGISTRY, import_module +from ding.utils import REWARD_MODEL_REGISTRY, import_module, save_file class BaseRewardModel(ABC): @@ -90,6 +93,28 @@ def reward_deepcopy(self, train_data) -> Any: ] return train_data_reward_deepcopy + def state_dict(self) -> Dict: + # this method should be overrided by subclass. + return {} + + def load_state_dict(self, _state_dict) -> None: + # this method should be overrided by subclass. + pass + + def save(self, path: str = None, name: str = 'best'): + if path is None: + path = self.cfg.exp_name + path = os.path.join(path, 'reward_model', 'ckpt') + if not os.path.exists(path): + try: + os.makedirs(path) + except FileExistsError: + pass + path = os.path.join(path, 'ckpt_{}.pth.tar'.format(name)) + state_dict = self.state_dict() + save_file(path, state_dict) + logging.info('Saved reward model ckpt in {}'.format(path)) + def create_reward_model(cfg: dict, device: str, tb_logger: 'SummaryWriter') -> BaseRewardModel: # noqa """ diff --git a/ding/reward_model/drex_reward_model.py b/ding/reward_model/drex_reward_model.py index a58b302e60..645b469088 100644 --- a/ding/reward_model/drex_reward_model.py +++ b/ding/reward_model/drex_reward_model.py @@ -1,11 +1,7 @@ import copy from easydict import EasyDict -import numpy as np import pickle -import torch -import torch.nn as nn - from ding.utils import REWARD_MODEL_REGISTRY from .trex_reward_model import TrexRewardModel @@ -13,15 +9,46 @@ @REWARD_MODEL_REGISTRY.register('drex') class DrexRewardModel(TrexRewardModel): + """ + Overview: + The Drex reward model class (https://arxiv.org/pdf/1907.03976.pdf) + Interface: + ``estimate``, ``train``, ``load_expert_data``, ``collect_data``, ``clear_date``, \ + ``__init__``, ``_train``, + Config: + == ==================== ====== ============= ======================================= =============== + ID Symbol Type Default Value Description Other(Shape) + == ==================== ====== ============= ======================================= =============== + 1 ``type`` str drex | Reward model register name, refer | + | to registry ``REWARD_MODEL_REGISTRY`` | + 3 | ``learning_rate`` float 0.00001 | learning rate for optimizer | + 4 | ``update_per_`` int 100 | Number of updates per collect | + | ``collect`` | | + 5 | ``batch_size`` int 64 | How many samples in a training batch | + 6 | ``hidden_size`` int 128 | Linear model hidden size | + 7 | ``num_trajs`` int 0 | Number of downsampled full | + | trajectories | + 8 | ``num_snippets`` int 6000 | Number of short subtrajectories | + | to sample | + == ==================== ====== ============= ======================================= ================ + """ config = dict( + # (str) Reward model register name, refer to registry ``REWARD_MODEL_REGISTRY``. type='drex', + # (float) The step size of gradient descent. learning_rate=1e-5, + # (int) How many updates(iterations) to train after collector's one collection. + # Bigger "update_per_collect" means bigger off-policy. + # collect data -> update policy-> collect data -> ... update_per_collect=100, + # (int) How many samples in a training batch. batch_size=64, - target_new_data_count=64, + # (int) Linear model hidden size hidden_size=128, - num_trajs=0, # number of downsampled full trajectories - num_snippets=6000, # number of short subtrajectories to sample + # (int) Number of downsampled full trajectories. + num_trajs=0, + # (int) Number of short subtrajectories to sample. + num_snippets=6000, ) bc_cfg = None diff --git a/ding/reward_model/gail_irl_model.py b/ding/reward_model/gail_irl_model.py index c54dfdb9c4..6533e114dd 100644 --- a/ding/reward_model/gail_irl_model.py +++ b/ding/reward_model/gail_irl_model.py @@ -3,7 +3,6 @@ import random from collections.abc import Iterable from easydict import EasyDict -import numpy as np import torch import torch.nn as nn @@ -113,35 +112,51 @@ class GailRewardModel(BaseRewardModel): ``estimate``, ``train``, ``load_expert_data``, ``collect_data``, ``clear_date``, \ ``__init__``, ``state_dict``, ``load_state_dict``, ``learn`` Config: - == ==================== ======== ============= ================================= ======================= - ID Symbol Type Default Value Description Other(Shape) - == ==================== ======== ============= ================================= ======================= - 1 ``type`` str gail | RL policy register name, refer | this arg is optional, - | to registry ``POLICY_REGISTRY`` | a placeholder - 2 | ``expert_data_`` str expert_data. | Path to the expert dataset | Should be a '.pkl' - | ``path`` .pkl | | file - 3 | ``update_per_`` int 100 | Number of updates per collect | - | ``collect`` | | - 4 | ``batch_size`` int 64 | Training batch size | - 5 | ``input_size`` int | Size of the input: | - | | obs_dim + act_dim | - 6 | ``target_new_`` int 64 | Collect steps per iteration | - | ``data_count`` | | - 7 | ``hidden_size`` int 128 | Linear model hidden size | - 8 | ``collect_count`` int 100000 | Expert dataset size | One entry is a (s,a) - | | | tuple - == ==================== ======== ============= ================================= ======================= - - """ + == ==================== ======== ============= =================================== ======================= + ID Symbol Type Default Value Description Other(Shape) + == ==================== ======== ============= =================================== ======================= + 1 ``type`` str gail | RL policy register name, refer | this arg is optional, + | to registry ``POLICY_REGISTRY`` | a placeholder + 2 | ``expert_data_`` str expert_data. | Path to the expert dataset | Should be a '.pkl' + | ``path`` .pkl | | file + 3 | ``learning_rate`` float 0.001 | The step size of gradient descent | + 4 | ``update_per_`` int 100 | Number of updates per collect | + | ``collect`` | | + 5 | ``batch_size`` int 64 | Training batch size | + 6 | ``input_size`` int | Size of the input: | + | | obs_dim + act_dim | + 7 | ``target_new_`` int 64 | Collect steps per iteration | + | ``data_count`` | | + 8 | ``hidden_size`` int 128 | Linear model hidden size | + 9 | ``collect_count`` int 100000 | Expert dataset size | One entry is a (s,a) + | | | tuple + 10 | ``clear_buffer_`` int 1 | clear buffer per fixed iters | make sure replay + | ``per_iters`` | buffer's data count + | | isn't too few. + | | (code work in entry) + == ==================== ======== ============= =================================== ======================= + """ config = dict( + # (str) RL policy register name, refer to registry ``POLICY_REGISTRY``. type='gail', + # (float) The step size of gradient descent. learning_rate=1e-3, + # (int) How many updates(iterations) to train after collector's one collection. + # Bigger "update_per_collect" means bigger off-policy. + # collect data -> update policy-> collect data -> ... update_per_collect=100, + # (int) How many samples in a training batch. batch_size=64, + # (int) Size of the input: obs_dim + act_dim. input_size=4, + # (int) Collect steps per iteration. target_new_data_count=64, + # (int) Linear model hidden size. hidden_size=128, + # (int) Expert dataset size. collect_count=100000, + # (int) Clear buffer per fixed iters. + clear_buffer_per_iters=1, ) def __init__(self, config: EasyDict, device: str, tb_logger: 'SummaryWriter') -> None: # noqa diff --git a/ding/reward_model/guided_cost_reward_model.py b/ding/reward_model/guided_cost_reward_model.py index e0213f9966..437e198f53 100644 --- a/ding/reward_model/guided_cost_reward_model.py +++ b/ding/reward_model/guided_cost_reward_model.py @@ -1,17 +1,14 @@ -from typing import List, Dict, Any, Tuple, Union, Optional +from typing import List, Dict, Any from easydict import EasyDict -import random import torch import torch.nn as nn import torch.optim as optim import torch.nn.functional as F from torch.distributions import Independent, Normal -import copy -from ding.utils import SequenceType, REWARD_MODEL_REGISTRY -from ding.utils.data import default_collate, default_decollate -from ding.model import FCEncoder, ConvEncoder +from ding.utils import REWARD_MODEL_REGISTRY +from ding.utils.data import default_collate from .base_reward_model import BaseRewardModel @@ -38,23 +35,54 @@ def forward(self, x): @REWARD_MODEL_REGISTRY.register('guided_cost') class GuidedCostRewardModel(BaseRewardModel): - r""" + """ Overview: - Policy class of Guided cost algorithm. - - https://arxiv.org/pdf/1603.00448.pdf + Policy class of Guided cost algorithm. (https://arxiv.org/pdf/1603.00448.pdf) + Interface: + ``estimate``, ``train``, ``collect_data``, ``clear_date``, \ + ``__init__``, ``state_dict``, ``load_state_dict``, ``learn``\ + ``state_dict_reward_model``, ``load_state_dict_reward_model`` + Config: + == ==================== ======== ============= ======================================== ================ + ID Symbol Type Default Value Description Other(Shape) + == ==================== ======== ============= ======================================== ================ + 1 ``type`` str guided_cost | Reward model register name, refer | + | to registry ``REWARD_MODEL_REGISTRY`` | + 2 | ``continuous`` bool True | Whether action is continuous | + 3 | ``learning_rate`` float 0.001 | learning rate for optimizer | + 4 | ``update_per_`` int 100 | Number of updates per collect | + | ``collect`` | | + 5 | ``batch_size`` int 64 | Training batch size | + 6 | ``hidden_size`` int 128 | Linear model hidden size | + 7 | ``action_shape`` int 1 | Action space shape | + 8 | ``log_every_n`` int 50 | add loss to log every n iteration | + | ``_train`` | | + 9 | ``store_model_`` int 100 | save model every n iteration | + | ``every_n_train`` | + == ==================== ======== ============= ======================================== ================ """ config = dict( + # (str) Reward model register name, refer to registry ``REWARD_MODEL_REGISTRY``. type='guided_cost', + # (float) The step size of gradient descent. learning_rate=1e-3, + # (int) Action space shape, such as 1. action_shape=1, + # (bool) Whether action is continuous. continuous=True, + # (int) How many samples in a training batch. batch_size=64, + # (int) Linear model hidden size. hidden_size=128, + # (int) How many updates(iterations) to train after collector's one collection. + # Bigger "update_per_collect" means bigger off-policy. + # collect data -> update policy-> collect data -> ... update_per_collect=100, + # (int) Add loss to log every n iteration. log_every_n_train=50, + # (int) Save model every n iteration. store_model_every_n_train=100, ) diff --git a/ding/reward_model/icm_reward_model.py b/ding/reward_model/icm_reward_model.py index e45dc0ae03..9cc6e23e9b 100644 --- a/ding/reward_model/icm_reward_model.py +++ b/ding/reward_model/icm_reward_model.py @@ -5,7 +5,6 @@ import torch import torch.nn as nn import torch.optim as optim -import torch.nn.functional as F from ding.utils import SequenceType, REWARD_MODEL_REGISTRY from ding.model import FCEncoder, ConvEncoder @@ -28,7 +27,7 @@ def collect_states(iterator: list) -> Tuple[list, list, list]: class ICMNetwork(nn.Module): - r""" + """ Intrinsic Curiosity Model (ICM Module) Implementation of: [1] Curiosity-driven Exploration by Self-supervised Prediction @@ -130,27 +129,71 @@ class ICMRewardModel(BaseRewardModel): The ICM reward model class (https://arxiv.org/pdf/1705.05363.pdf) Interface: ``estimate``, ``train``, ``collect_data``, ``clear_data``, \ - ``__init__``, ``_train``, + ``__init__``, ``_train``, ``load_state_dict``, ``state_dict`` + Config: + == ==================== ======== ============= ==================================== ======================= + ID Symbol Type Default Value Description Other(Shape) + == ==================== ======== ============= ==================================== ======================= + 1 ``type`` str icm | Reward model register name, | + | refer to registry | + | ``REWARD_MODEL_REGISTRY`` | + 2 | ``intrinsic_`` str add | the intrinsic reward type | including add, new + | ``reward_type`` | | , or assign + 3 | ``learning_rate`` float 0.001 | The step size of gradient descent | + 4 | ``obs_shape`` Tuple( 6 | the observation shape | + [int, + list]) + 5 | ``action_shape`` int 7 | the action space shape | + 6 | ``batch_size`` int 64 | Training batch size | + 7 | ``hidden`` list [64, 64, | the MLP layer shape | + | ``_size_list`` (int) 128] | | + 8 | ``update_per_`` int 100 | Number of updates per collect | + | ``collect`` | | + 9 | ``reverse_scale`` float 1 | the importance weight of the | + | forward and reverse loss | + 10 | ``intrinsic_`` float 0.003 | the weight of intrinsic reward | r = w*r_i + r_e + ``reward_weight`` + 11 | ``extrinsic_`` bool True | Whether to normlize + ``reward_norm`` | extrinsic reward + 12 | ``extrinsic_`` int 1 | the upper bound of the reward + ``reward_norm_max`` | normalization + 13 | ``clear_buffer`` int 1 | clear buffer per fixed iters | make sure replay + ``_per_iters`` | buffer's data count + | isn't too few. + | (code work in entry) + == ==================== ======== ============= ==================================== ======================= """ config = dict( - # (str) the type of the exploration method + # (str) Reward model register name, refer to registry ``REWARD_MODEL_REGISTRY``. type='icm', - # (str) the intrinsic reward type, including add, new, or assign + # (str) The intrinsic reward type, including add, new, or assign. intrinsic_reward_type='add', - # (float) learning rate of the optimizer + # (float) The step size of gradient descent. learning_rate=1e-3, - # (Tuple[int, list]), the observation shape, + # (Tuple[int, list]), The observation shape. obs_shape=6, - # (int) the action shape, support discrete action only in this version + # (int) The action shape, support discrete action only in this version. action_shape=7, - # (float) batch size + # (float) Batch size. batch_size=64, - # (list) the MLP layer shape + # (list) The MLP layer shape. hidden_size_list=[64, 64, 128], - # (int) update how many times after each collect + # (int) How many updates(iterations) to train after collector's one collection. + # Bigger "update_per_collect" means bigger off-policy. + # collect data -> update policy-> collect data -> ... update_per_collect=100, - # (float) the importance weight of the forward and reverse loss + # (float) The importance weight of the forward and reverse loss. reverse_scale=1, + # (float) The weight of intrinsic reward. + # r = intrinsic_reward_weight * r_i + r_e. + intrinsic_reward_weight=0.003, # 1/300 + # (bool) Whether to normlize extrinsic reward. + # Normalize the reward to [0, extrinsic_reward_norm_max]. + extrinsic_reward_norm=True, + # (int) The upper bound of the reward normalization. + extrinsic_reward_norm_max=1, + # (int) Clear buffer per fixed iters. + clear_buffer_per_iters=100, ) def __init__(self, config: EasyDict, device: str, tb_logger: 'SummaryWriter') -> None: # noqa @@ -171,8 +214,12 @@ def __init__(self, config: EasyDict, device: str, tb_logger: 'SummaryWriter') -> self.ce = nn.CrossEntropyLoss(reduction="mean") self.forward_mse = nn.MSELoss(reduction='none') self.reverse_scale = config.reverse_scale + self.res = nn.Softmax(dim=-1) + self.estimate_cnt_icm = 0 + self.train_cnt_icm = 0 def _train(self) -> None: + self.train_cnt_icm += 1 train_data_list = [i for i in range(0, len(self.train_states))] train_data_index = random.sample(train_data_list, self.cfg.batch_size) data_states: list = [self.train_states[i] for i in train_data_index] @@ -187,6 +234,13 @@ def _train(self) -> None: ) inverse_loss = self.ce(pred_action_logit, data_actions.long()) forward_loss = self.forward_mse(pred_next_state_feature, real_next_state_feature.detach()).mean() + self.tb_logger.add_scalar('icm_reward/forward_loss', forward_loss, self.train_cnt_icm) + self.tb_logger.add_scalar('icm_reward/inverse_loss', inverse_loss, self.train_cnt_icm) + action = torch.argmax(self.res(pred_action_logit), -1) + accuracy = torch.sum(action == data_actions.squeeze(-1)).item() / data_actions.shape[0] + self.tb_logger.add_scalar('icm_reward/action_accuracy', accuracy, self.train_cnt_icm) + loss = self.reverse_scale * inverse_loss + forward_loss + self.tb_logger.add_scalar('icm_reward/total_loss', loss, self.train_cnt_icm) loss = self.reverse_scale * inverse_loss + forward_loss self.opt.zero_grad() loss.backward() @@ -195,7 +249,6 @@ def _train(self) -> None: def train(self) -> None: for _ in range(self.cfg.update_per_collect): self._train() - self.clear_data() def estimate(self, data: list) -> List[Dict]: # NOTE: deepcopy reward part of data is very important, @@ -207,17 +260,32 @@ def estimate(self, data: list) -> List[Dict]: actions = torch.cat(actions).to(self.device) with torch.no_grad(): real_next_state_feature, pred_next_state_feature, _ = self.reward_model(states, next_states, actions) - reward = self.forward_mse(real_next_state_feature, pred_next_state_feature).mean(dim=1) - reward = (reward - reward.min()) / (reward.max() - reward.min() + 1e-8) - reward = reward.to(train_data_augmented[0]['reward'].device) - reward = torch.chunk(reward, reward.shape[0], dim=0) - for item, rew in zip(train_data_augmented, reward): + raw_icm_reward = self.forward_mse(real_next_state_feature, pred_next_state_feature).mean(dim=1) + self.estimate_cnt_icm += 1 + self.tb_logger.add_scalar('icm_reward/raw_icm_reward_max', raw_icm_reward.max(), self.estimate_cnt_icm) + self.tb_logger.add_scalar('icm_reward/raw_icm_reward_mean', raw_icm_reward.mean(), self.estimate_cnt_icm) + self.tb_logger.add_scalar('icm_reward/raw_icm_reward_min', raw_icm_reward.min(), self.estimate_cnt_icm) + self.tb_logger.add_scalar('icm_reward/raw_icm_reward_std', raw_icm_reward.std(), self.estimate_cnt_icm) + icm_reward = (raw_icm_reward - raw_icm_reward.min()) / (raw_icm_reward.max() - raw_icm_reward.min() + 1e-8) + self.tb_logger.add_scalar('icm_reward/icm_reward_max', icm_reward.max(), self.estimate_cnt_icm) + self.tb_logger.add_scalar('icm_reward/icm_reward_mean', icm_reward.mean(), self.estimate_cnt_icm) + self.tb_logger.add_scalar('icm_reward/icm_reward_min', icm_reward.min(), self.estimate_cnt_icm) + self.tb_logger.add_scalar('icm_reward/icm_reward_std', icm_reward.std(), self.estimate_cnt_icm) + icm_reward = (raw_icm_reward - raw_icm_reward.min()) / (raw_icm_reward.max() - raw_icm_reward.min() + 1e-8) + icm_reward = icm_reward.to(self.device) + for item, icm_rew in zip(train_data_augmented, icm_reward): if self.intrinsic_reward_type == 'add': - item['reward'] += rew + if self.cfg.extrinsic_reward_norm: + item['reward'] = item[ + 'reward'] / self.cfg.extrinsic_reward_norm_max + icm_rew * self.cfg.intrinsic_reward_weight + else: + item['reward'] = item['reward'] + icm_rew * self.cfg.intrinsic_reward_weight elif self.intrinsic_reward_type == 'new': - item['intrinsic_reward'] = rew + item['intrinsic_reward'] = icm_rew + if self.cfg.extrinsic_reward_norm: + item['reward'] = item['reward'] / self.cfg.extrinsic_reward_norm_max elif self.intrinsic_reward_type == 'assign': - item['reward'] = rew + item['reward'] = icm_rew return train_data_augmented @@ -233,3 +301,9 @@ def clear_data(self) -> None: self.train_states.clear() self.train_next_states.clear() self.train_actions.clear() + + def state_dict(self) -> Dict: + return self.reward_model.state_dict() + + def load_state_dict(self, _state_dict: Dict) -> None: + self.reward_model.load_state_dict(_state_dict) diff --git a/ding/reward_model/ngu_reward_model.py b/ding/reward_model/ngu_reward_model.py index 0ebe6bc87e..5a8758bdb7 100644 --- a/ding/reward_model/ngu_reward_model.py +++ b/ding/reward_model/ngu_reward_model.py @@ -1,6 +1,6 @@ import copy import random -from typing import Union, Tuple, Any, Dict, List +from typing import Union, Tuple, Dict, List import numpy as np import torch diff --git a/ding/reward_model/pdeil_irl_model.py b/ding/reward_model/pdeil_irl_model.py index ec264e27cc..b09416f5c2 100644 --- a/ding/reward_model/pdeil_irl_model.py +++ b/ding/reward_model/pdeil_irl_model.py @@ -1,8 +1,8 @@ +from typing import List, Dict +from ditk import logging import numpy as np import torch import pickle -from typing import List, Dict -import scipy.stats as stats try: from sklearn.svm import SVC except ImportError: @@ -16,16 +16,42 @@ class PdeilRewardModel(BaseRewardModel): """ Overview: - The Pdeil reward model class + The Pdeil reward model class (https://arxiv.org/abs/2112.06746) Interface: ``estimate``, ``train``, ``load_expert_data``, ``collect_data``, ``clear_date``, \ ``__init__``, ``_train``, ``_batch_mn_pdf`` + Config: + == ==================== ===== ============= ======================================= ======================= + ID Symbol Type Default Value Description Other(Shape) + == ==================== ===== ============= ======================================= ======================= + 1 ``type`` str pdeil | Reward model register name, refer | + | to registry ``REWARD_MODEL_REGISTRY`` | + 2 | ``expert_data_`` str expert_data. | Path to the expert dataset | Should be a '.pkl' + | ``path`` .pkl | | file + 3 | ``discrete_`` bool False | Whether the action is discrete | + | ``action`` | | + 4 | ``alpha`` float 0.5 | coefficient for Probability | + | | Density Estimator | + 5 | ``clear_buffer`` int 1 | clear buffer per fixed iters | make sure replay + ``_per_iters`` | buffer's data count + | isn't too few. + | (code work in entry) + == ==================== ===== ============= ======================================= ======================= """ config = dict( + # (str) Reward model register name, refer to registry ``REWARD_MODEL_REGISTRY``. type='pdeil', + # (str) Path to the expert dataset. # expert_data_path='expert_data.pkl', + # (bool) Whether the action is discrete. discrete_action=False, + # (float) Coefficient for Probability Density Estimator. + # alpha + beta = 1, alpha is in [0,1] + # when alpha is close to 0, the estimator has high variance and low bias; + # when alpha is close to 1, the estimator has high bias and low variance. alpha=0.5, + # (int) Clear buffer per fixed iters. + clear_buffer_per_iters=1, ) def __init__(self, cfg: dict, device, tb_logger: 'SummaryWriter') -> None: # noqa @@ -45,6 +71,13 @@ def __init__(self, cfg: dict, device, tb_logger: 'SummaryWriter') -> None: # no - tb_logger (:obj:`str`): Logger, defaultly set as 'SummaryWriter' for model summary """ super(PdeilRewardModel, self).__init__() + try: + import scipy.stats as stats + self.stats = stats + except ImportError: + import sys + logging.warning("Please install scipy first, such as `pip3 install scipy`.") + sys.exit(1) self.cfg: dict = cfg self.e_u_s = None self.e_sigma_s = None @@ -119,7 +152,9 @@ def _batch_mn_pdf(self, x: np.ndarray, mean: np.ndarray, cov: np.ndarray) -> np. Overview: Get multivariate normal pdf of given np array. """ - return np.asarray(stats.multivariate_normal.pdf(x, mean=mean, cov=cov, allow_singular=False), dtype=np.float32) + return np.asarray( + self.stats.multivariate_normal.pdf(x, mean=mean, cov=cov, allow_singular=False), dtype=np.float32 + ) def estimate(self, data: list) -> List[Dict]: """ diff --git a/ding/reward_model/pwil_irl_model.py b/ding/reward_model/pwil_irl_model.py index 2c513462a8..8738ee2d81 100644 --- a/ding/reward_model/pwil_irl_model.py +++ b/ding/reward_model/pwil_irl_model.py @@ -35,17 +35,48 @@ class PwilRewardModel(BaseRewardModel): Interface: ``estimate``, ``train``, ``load_expert_data``, ``collect_data``, ``clear_date``, \ ``__init__``, ``_train``, ``_get_state_distance``, ``_get_action_distance`` + Config: + == ================== ===== ============= ======================================= ======================= + ID Symbol Type Default Value Description Other(Shape) + == ================== ===== ============= ======================================= ======================= + 1 ``type`` str pwil | Reward model register name, refer | + | to registry ``REWARD_MODEL_REGISTRY`` | + 2 | ``expert_data_`` str expert_data. | Path to the expert dataset | Should be a '.pkl' + | ``path`` .pkl | | file + 3 | ``sample_size`` int 1000 | sample data from expert dataset | + | with fixed size | + 4 | ``alpha`` int 5 | factor alpha | + 5 | ``beta`` int 5 | factor beta | + 6 | ``s_size`` int 4 | state size | + 7 | ``a_size`` int 2 | action size | + 8 | ``clear_buffer`` int 1 | clear buffer per fixed iters | make sure replay + ``_per_iters`` | buffer's data count + | isn't too few. + | (code work in entry) + == ================== ===== ============= ======================================= ======================= Properties: - reward_table (:obj: `Dict`): In this algorithm, reward model is a dictionary. """ config = dict( + # (str) Reward model register name, refer to registry ``REWARD_MODEL_REGISTRY``. type='pwil', + # (str) Path to the expert dataset. # expert_data_path='expert_data.pkl', + # (int) Sample data from expert dataset with fixed size. sample_size=1000, + # r = alpha * exp((-beta*T/sqrt(|s_size|+ |a_size|))*c_i) + # key idea for this reward is to minimize. + # the Wasserstein distance between the state-action distribution. + # (int) Factor alpha. alpha=5, + # (int) Factor beta. beta=5, + #（int）State size. # s_size=4, + # (int) Action size. # a_size=2, + # (int) Clear buffer per fixed iters. + clear_buffer_per_iters=1, ) def __init__(self, config: Dict, device: str, tb_logger: 'SummaryWriter') -> None: # noqa diff --git a/ding/reward_model/red_irl_model.py b/ding/reward_model/red_irl_model.py index fdfea10b8d..a7daeeceec 100644 --- a/ding/reward_model/red_irl_model.py +++ b/ding/reward_model/red_irl_model.py @@ -34,19 +34,54 @@ class RedRewardModel(BaseRewardModel): Interface: ``estimate``, ``train``, ``load_expert_data``, ``collect_data``, ``clear_date``, \ ``__init__``, ``_train`` + Config: + == ================== ===== ============= ======================================= ======================= + ID Symbol Type Default Value Description Other(Shape) + == ================== ===== ============= ======================================= ======================= + 1 ``type`` str red | Reward model register name, refer | + | to registry ``REWARD_MODEL_REGISTRY`` | + 2 | ``expert_data_`` str expert_data | Path to the expert dataset | Should be a '.pkl' + | ``path`` .pkl | | file + 3 | ``sample_size`` int 1000 | sample data from expert dataset | + | with fixed size | + 4 | ``sigma`` int 5 | hyperparameter of r(s,a) | r(s,a) = exp( + | -sigma* L(s,a)) + 5 | ``batch_size`` int 64 | Training batch size | + 6 | ``hidden_size`` int 128 | Linear model hidden size | + 7 | ``update_per_`` int 100 | Number of updates per collect | + | ``collect`` | | + 8 | ``clear_buffer`` int 1 | clear buffer per fixed iters | make sure replay + ``_per_iters`` | buffer's data count + | isn't too few. + | (code work in entry) + == ================== ===== ============= ======================================= ======================= Properties: - online_net (:obj: `SENet`): The reward model, in default initialized once as the training begins. """ config = dict( + # (str) Reward model register name, refer to registry ``REWARD_MODEL_REGISTRY``. type='red', + # (int) Linear model input size. # input_size=4, + # (int) Sample data from expert dataset with fixed size. sample_size=1000, + # (int) Linear model hidden size. hidden_size=128, + # (float) The step size of gradient descent. learning_rate=1e-3, + # (int) How many updates(iterations) to train after collector's one collection. + # Bigger "update_per_collect" means bigger off-policy. + # collect data -> update policy-> collect data -> ... update_per_collect=100, + # (str) Path to the expert dataset # expert_data_path='expert_data.pkl', + # (int) How many samples in a training batch. batch_size=64, + # (float) Hyperparameter at estimated score of r(s,a). + # r(s,a) = exp(-sigma* L(s,a)) sigma=0.5, + # (int) Clear buffer per fixed iters. + clear_buffer_per_iters=1, ) def __init__(self, config: Dict, device: str, tb_logger: 'SummaryWriter') -> None: # noqa diff --git a/ding/reward_model/rnd_reward_model.py b/ding/reward_model/rnd_reward_model.py index 9436c06d97..00bb1542fd 100644 --- a/ding/reward_model/rnd_reward_model.py +++ b/ding/reward_model/rnd_reward_model.py @@ -1,4 +1,4 @@ -from typing import Union, Tuple, List, Dict, Any +from typing import Union, Tuple, List, Dict from easydict import EasyDict import random @@ -12,7 +12,7 @@ from .base_reward_model import BaseRewardModel from ding.utils import RunningMeanStd from ding.torch_utils.data_helper import to_tensor -import copy +import numpy as np def collect_states(iterator): @@ -50,29 +50,76 @@ def forward(self, obs: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: @REWARD_MODEL_REGISTRY.register('rnd') class RndRewardModel(BaseRewardModel): + """ + Overview: + The RND reward model class (https://arxiv.org/abs/1810.12894v1) + Interface: + ``estimate``, ``train``, ``collect_data``, ``clear_data``, \ + ``__init__``, ``_train``, ``load_state_dict``, ``state_dict`` + Config: + == ==================== ===== ============= ======================================= ======================= + ID Symbol Type Default Value Description Other(Shape) + == ==================== ===== ============= ======================================= ======================= + 1 ``type`` str rnd | Reward model register name, refer | + | to registry ``REWARD_MODEL_REGISTRY`` | + 2 | ``intrinsic_`` str add | the intrinsic reward type | including add, new + | ``reward_type`` | | , or assign + 3 | ``learning_rate`` float 0.001 | The step size of gradient descent | + 4 | ``batch_size`` int 64 | Training batch size | + 5 | ``hidden`` list [64, 64, | the MLP layer shape | + | ``_size_list`` (int) 128] | | + 6 | ``update_per_`` int 100 | Number of updates per collect | + | ``collect`` | | + 7 | ``obs_norm`` bool True | Observation normalization | + 8 | ``obs_norm_`` int 0 | min clip value for obs normalization | + | ``clamp_min`` + 9 | ``obs_norm_`` int 1 | max clip value for obs normalization | + | ``clamp_max`` + 10 | ``intrinsic_`` float 0.01 | the weight of intrinsic reward | r = w*r_i + r_e + ``reward_weight`` + 11 | ``extrinsic_`` bool True | Whether to normlize extrinsic reward + ``reward_norm`` + 12 | ``extrinsic_`` int 1 | the upper bound of the reward + ``reward_norm_max`` | normalization + == ==================== ===== ============= ======================================= ======================= + """ config = dict( + # (str) Reward model register name, refer to registry ``REWARD_MODEL_REGISTRY``. type='rnd', + # (str) The intrinsic reward type, including add, new, or assign. intrinsic_reward_type='add', + # (float) The step size of gradient descent. learning_rate=1e-3, + # (float) Batch size. batch_size=64, + # (list(int)) Sequence of ``hidden_size`` of reward network. + # If obs.shape == 1, use MLP layers. + # If obs.shape == 3, use conv layer and final dense layer. hidden_size_list=[64, 64, 128], + # (int) How many updates(iterations) to train after collector's one collection. + # Bigger "update_per_collect" means bigger off-policy. + # collect data -> update policy-> collect data -> ... update_per_collect=100, + # (bool) Observation normalization: transform obs to mean 0, std 1. obs_norm=True, + # (int) Min clip value for observation normalization. obs_norm_clamp_min=-1, + # (int) Max clip value for observation normalization. obs_norm_clamp_max=1, - intrinsic_reward_weight=None, - # means the relative weight of RND intrinsic_reward. - # If intrinsic_reward_weight=None, we will automatically set it based on - # the absolute value of the difference between max and min extrinsic reward in the sampled mini-batch - # please refer to estimate() method for details. - intrinsic_reward_rescale=0.01, - # means the rescale value of RND intrinsic_reward only used when intrinsic_reward_weight is None + # Means the relative weight of RND intrinsic_reward. + # (float) The weight of intrinsic reward + # r = intrinsic_reward_weight * r_i + r_e. + intrinsic_reward_weight=0.01, + # (bool) Whether to normlize extrinsic reward. + # Normalize the reward to [0, extrinsic_reward_norm_max]. + extrinsic_reward_norm=True, + # (int) The upper bound of the reward normalization. + extrinsic_reward_norm_max=1, ) def __init__(self, config: EasyDict, device: str = 'cpu', tb_logger: 'SummaryWriter' = None) -> None: # noqa super(RndRewardModel, self).__init__() self.cfg = config - self.intrinsic_reward_rescale = self.cfg.intrinsic_reward_rescale assert device == "cpu" or device.startswith("cuda") self.device = device if tb_logger is None: # TODO @@ -87,6 +134,7 @@ def __init__(self, config: EasyDict, device: str = 'cpu', tb_logger: 'SummaryWri self.opt = optim.Adam(self.reward_model.predictor.parameters(), config.learning_rate) self._running_mean_std_rnd_reward = RunningMeanStd(epsilon=1e-4) self.estimate_cnt_rnd = 0 + self.train_cnt_icm = 0 self._running_mean_std_rnd_obs = RunningMeanStd(epsilon=1e-4) def _train(self) -> None: @@ -102,6 +150,7 @@ def _train(self) -> None: predict_feature, target_feature = self.reward_model(train_data) loss = F.mse_loss(predict_feature, target_feature.detach()) + self.tb_logger.add_scalar('rnd_reward/loss', loss, self.train_cnt_icm) self.opt.zero_grad() loss.backward() self.opt.step() @@ -109,6 +158,7 @@ def _train(self) -> None: def train(self) -> None: for _ in range(self.cfg.update_per_collect): self._train() + self.train_cnt_icm += 1 def estimate(self, data: list) -> List[Dict]: """ @@ -132,14 +182,16 @@ def estimate(self, data: list) -> List[Dict]: self._running_mean_std_rnd_reward.update(mse.cpu().numpy()) # Note: according to the min-max normalization, transform rnd reward to [0,1] - rnd_reward = (mse - mse.min()) / (mse.max() - mse.min() + 1e-11) + rnd_reward = (mse - mse.min()) / (mse.max() - mse.min() + 1e-8) + # save the rnd_reward statistics into tb_logger self.estimate_cnt_rnd += 1 self.tb_logger.add_scalar('rnd_reward/rnd_reward_max', rnd_reward.max(), self.estimate_cnt_rnd) self.tb_logger.add_scalar('rnd_reward/rnd_reward_mean', rnd_reward.mean(), self.estimate_cnt_rnd) self.tb_logger.add_scalar('rnd_reward/rnd_reward_min', rnd_reward.min(), self.estimate_cnt_rnd) + self.tb_logger.add_scalar('rnd_reward/rnd_reward_std', rnd_reward.std(), self.estimate_cnt_rnd) - rnd_reward = rnd_reward.to(train_data_augmented[0]['reward'].device) + rnd_reward = rnd_reward.to(self.device) rnd_reward = torch.chunk(rnd_reward, rnd_reward.shape[0], dim=0) """ NOTE: Following normalization approach to extrinsic reward seems be not reasonable, @@ -148,30 +200,26 @@ def estimate(self, data: list) -> List[Dict]: # rewards = torch.stack([data[i]['reward'] for i in range(len(data))]) # rewards = (rewards - torch.min(rewards)) / (torch.max(rewards) - torch.min(rewards)) - # TODO(pu): how to set intrinsic_reward_rescale automatically? - if self.cfg.intrinsic_reward_weight is None: - """ - NOTE: the following way of setting self.cfg.intrinsic_reward_weight is only suitable for the dense - reward env like lunarlander, not suitable for the dense reward env. - In sparse reward env, e.g. minigrid, if the agent reaches the goal, it obtain reward ~1, otherwise 0. - Thus, in sparse reward env, it's reasonable to set the intrinsic_reward_weight approximately equal to - the inverse of max_episode_steps. - """ - self.cfg.intrinsic_reward_weight = self.intrinsic_reward_rescale * max( - 1, - abs( - max([train_data_augmented[i]['reward'] for i in range(len(train_data_augmented))]) - - min([train_data_augmented[i]['reward'] for i in range(len(train_data_augmented))]) - ) - ) for item, rnd_rew in zip(train_data_augmented, rnd_reward): if self.intrinsic_reward_type == 'add': - item['reward'] = item['reward'] + rnd_rew * self.cfg.intrinsic_reward_weight + if self.cfg.extrinsic_reward_norm: + item['reward'] = item[ + 'reward'] / self.cfg.extrinsic_reward_norm_max + rnd_rew * self.cfg.intrinsic_reward_weight + else: + item['reward'] = item['reward'] + rnd_rew * self.cfg.intrinsic_reward_weight elif self.intrinsic_reward_type == 'new': item['intrinsic_reward'] = rnd_rew + if self.cfg.extrinsic_reward_norm: + item['reward'] = item['reward'] / self.cfg.extrinsic_reward_norm_max elif self.intrinsic_reward_type == 'assign': item['reward'] = rnd_rew + # save the augmented_reward statistics into tb_logger + rew = [item['reward'].cpu().numpy() for item in train_data_augmented] + self.tb_logger.add_scalar('augmented_reward/reward_max', np.max(rew), self.estimate_cnt_rnd) + self.tb_logger.add_scalar('augmented_reward/reward_mean', np.mean(rew), self.estimate_cnt_rnd) + self.tb_logger.add_scalar('augmented_reward/reward_min', np.min(rew), self.estimate_cnt_rnd) + self.tb_logger.add_scalar('augmented_reward/reward_std', np.std(rew), self.estimate_cnt_rnd) return train_data_augmented def collect_data(self, data: list) -> None: @@ -179,3 +227,9 @@ def collect_data(self, data: list) -> None: def clear_data(self) -> None: self.train_obs.clear() + + def state_dict(self) -> Dict: + return self.reward_model.state_dict() + + def load_state_dict(self, _state_dict: Dict) -> None: + self.reward_model.load_state_dict(_state_dict) diff --git a/ding/reward_model/trex_reward_model.py b/ding/reward_model/trex_reward_model.py index 9fc069626c..635dc5e75e 100644 --- a/ding/reward_model/trex_reward_model.py +++ b/ding/reward_model/trex_reward_model.py @@ -1,24 +1,17 @@ -from collections.abc import Iterable -from easydict import EasyDict -import numpy as np -import pickle from copy import deepcopy from typing import Tuple, Optional, List, Dict +from easydict import EasyDict +import pickle +import os +import numpy as np import torch import torch.nn as nn -import torch.nn.functional as F import torch.optim as optim -from torch.distributions import Normal, Independent -from torch.distributions.categorical import Categorical from ding.utils import REWARD_MODEL_REGISTRY -from ding.model.template.q_learning import DQN -from ding.model.template.vac import VAC -from ding.model.template.qac import QAC from ding.utils import SequenceType from ding.model.common import FCEncoder -from ding.utils.data import offline_data_save_type from ding.utils import build_logger from ding.utils.data import default_collate @@ -35,11 +28,10 @@ class TrexConvEncoder(nn.Module): """ def __init__( - self, - obs_shape: SequenceType, - hidden_size_list: SequenceType = [16, 16, 16, 16, 64, 1], - activation: Optional[nn.Module] = nn.LeakyReLU(), - norm_type: Optional[str] = None + self, + obs_shape: SequenceType, + hidden_size_list: SequenceType = [16, 16, 16, 16, 64, 1], + activation: Optional[nn.Module] = nn.LeakyReLU() ) -> None: r""" Overview: @@ -51,9 +43,7 @@ def __init__( - hidden_size_list (:obj:`SequenceType`): The collection of ``hidden_size`` - activation (:obj:`nn.Module`): The type of activation to use in the conv ``layers``, - if ``None`` then default set to ``nn.ReLU()`` - - norm_type (:obj:`str`): - The type of normalization to use, see ``ding.torch_utils.ResBlock`` for more details + if ``None`` then default set to ``nn.LeakyReLU()`` """ super(TrexConvEncoder, self).__init__() self.obs_shape = obs_shape @@ -147,13 +137,32 @@ class TrexRewardModel(BaseRewardModel): Interface: ``estimate``, ``train``, ``load_expert_data``, ``collect_data``, ``clear_date``, \ ``__init__``, ``_train``, + Config: + == ==================== ====== ============= ============================================ ============= + ID Symbol Type Default Value Description Other(Shape) + == ==================== ====== ============= ============================================ ============= + 1 ``type`` str trex | Reward model register name, refer | + | to registry ``REWARD_MODEL_REGISTRY`` | + 3 | ``learning_rate`` float 0.00001 | learning rate for optimizer | + 4 | ``update_per_`` int 100 | Number of updates per collect | + | ``collect`` | | + 5 | ``num_trajs`` int 0 | Number of downsampled full trajectories | + 6 | ``num_snippets`` int 6000 | Number of short subtrajectories to sample | + == ==================== ====== ============= ============================================ ============= """ config = dict( + # (str) Reward model register name, refer to registry ``REWARD_MODEL_REGISTRY``. type='trex', + # (float) The step size of gradient descent. learning_rate=1e-5, + # (int) How many updates(iterations) to train after collector's one collection. + # Bigger "update_per_collect" means bigger off-policy. + # collect data -> update policy-> collect data -> ... update_per_collect=100, - num_trajs=0, # number of downsampled full trajectories - num_snippets=6000, # number of short subtrajectories to sample + # (int) Number of downsampled full trajectories. + num_trajs=0, + # (int) Number of short subtrajectories to sample. + num_snippets=6000, ) def __init__(self, config: EasyDict, device: str, tb_logger: 'SummaryWriter') -> None: # noqa @@ -196,14 +205,14 @@ def __init__(self, config: EasyDict, device: str, tb_logger: 'SummaryWriter') -> def load_expert_data(self) -> None: """ Overview: - Getting the expert data from ``config.data_path`` attribute in self + Getting the expert data. Effects: This is a side effect function which updates the expert data attribute \ (i.e. ``self.expert_data``) with ``fn:concat_state_action_pairs`` """ - with open(self.cfg.reward_model.data_path + '/episodes_data.pkl', 'rb') as f: + with open(os.path.join(self.cfg.exp_name, 'episodes_data.pkl'), 'rb') as f: self.pre_expert_data = pickle.load(f) - with open(self.cfg.reward_model.data_path + '/learning_returns.pkl', 'rb') as f: + with open(os.path.join(self.cfg.exp_name, 'learning_returns.pkl'), 'rb') as f: self.learning_returns = pickle.load(f) self.create_training_data() @@ -316,12 +325,13 @@ def _train(self): item_loss = loss.item() cum_loss += item_loss if i % 100 == 99: - self._logger.info("epoch {}:{} loss {}".format(epoch, i, cum_loss)) + self._logger.info("[epoch {}:{}] loss {}".format(epoch, i, cum_loss)) self._logger.info("abs_returns: {}".format(abs_rewards)) cum_loss = 0.0 self._logger.info("check pointing") - torch.save(self.reward_model.state_dict(), self.cfg.reward_model.reward_model_path) - torch.save(self.reward_model.state_dict(), self.cfg.reward_model.reward_model_path) + if not os.path.exists(os.path.join(self.cfg.exp_name, 'ckpt_reward_model')): + os.makedirs(os.path.join(self.cfg.exp_name, 'ckpt_reward_model')) + torch.save(self.reward_model.state_dict(), os.path.join(self.cfg.exp_name, 'ckpt_reward_model/latest.pth.tar')) self._logger.info("finished training") def train(self): diff --git a/ding/rl_utils/README.md b/ding/rl_utils/README.md new file mode 100644 index 0000000000..56c4189993 --- /dev/null +++ b/ding/rl_utils/README.md @@ -0,0 +1,103 @@ +# Testing Log Probability Methods with GPU + +The script `test_log_prob_fn` benchmarks different methods for calculating log probabilities (`naive_method`, `efficient_method`, `less_efficient_method`) using both `float32` and `bfloat16` precision formats on a GPU. + +## Overview + +The script performs benchmarks on three different methods for calculating log probabilities and reports the time and peak GPU memory usage for each method: + +- **Naive Method** +- **Efficient Method** +- **Less Efficient Method** + +It runs the tests for two types of precision: + +- `float32` +- `bfloat16` + +## Test Functions + +There are two main test functions in the script: + +1. **`test_log_prob_methods_float32`**: This function benchmarks the three methods using `float32` precision. +2. **`test_log_prob_methods_bfloat16`**: This function benchmarks the three methods using `bfloat16` precision. + +### Workflow + +1. **Parameters Setup**: The tests are executed using a batch size of `16`, a sequence length of `1024`, and a dictionary size of `32768`. The data is randomly generated for benchmarking. +2. **GPU Memory Tracking**: The GPU memory is tracked using `torch.cuda.max_memory_allocated()` to measure the peak memory usage during the benchmark. +3. **Method Execution**: Each method is run multiple times (10 iterations) to measure the execution time and to ensure stability. +4. **Results Validation**: The results from each method are compared with the `Naive` method to check for correctness, with a tolerance value applied for `bfloat16` precision. + +### Benchmarked Methods: + +- **Naive Method**: The basic, unoptimized method for calculating log probabilities. +- **Efficient Method**: An optimized version of the naive method to reduce memory usage. +- **Less Efficient Method**: A method with a higher memory consumption compared to the efficient method. + +### GPU Memory Usage: + +The function `get_gpu_memory()` is used to fetch the current peak GPU memory usage during the execution of each method. + +## Output Example + +### Testing with `float32` Precision + +``` +================================================== +Testing with float32 precision +================================================== + +Naive: +Time: 5.07 ± 0.83 ms +Peak GPU Memory: 4096.31 MB + +Efficient: +Time: 15.76 ± 21.19 ms +Peak GPU Memory: 2176.44 MB + +Less_Efficient: +Time: 14.63 ± 5.06 ms +Peak GPU Memory: 4608.39 MB +PASSED [100%] +``` + +### Testing with `bfloat16` Precision + +``` +================================================== +Testing with bfloat16 precision +================================================== + +Naive: +Time: 1.42 ± 0.00 ms +Peak GPU Memory: 2048.22 MB + +Efficient: +Time: 1.83 ± 0.01 ms +Peak GPU Memory: 1152.25 MB + +Less_Efficient: +Time: 8.67 ± 0.07 ms +Peak GPU Memory: 2560.27 MB +``` + +## Results Analysis + +- Execution Time + - The Naive method is the fastest in both precisions but sacrifices memory efficiency. + - The Efficient method balances memory usage and execution time, though it is slower than the Naive method. + - The Less Efficient method is slower than both the Naive and Efficient methods and consumes the most memory, making it the least desirable for both speed and memory usage. +- GPU Memory + - The Efficient method consistently uses the least memory, especially in the `bfloat16` precision where it achieves the lowest memory consumption. + - The Naive method uses more memory than the Efficient method but has lower execution times. + - The Less Efficient method consumes the most memory in both precision formats. + +## How to Run the Tests + +To run the tests: + +```bash +pytest -v -s tests/test_log_prob_utils.py +``` + diff --git a/ding/rl_utils/__init__.py b/ding/rl_utils/__init__.py index da375d1025..e86f6c1786 100644 --- a/ding/rl_utils/__init__.py +++ b/ding/rl_utils/__init__.py @@ -1,22 +1,27 @@ from .exploration import get_epsilon_greedy_fn, create_noise_generator -from .ppo import ppo_data, ppo_loss, ppo_info, ppo_policy_data, ppo_policy_error, ppo_value_data, ppo_value_error,\ - ppo_error, ppo_error_continuous, ppo_policy_error_continuous +from .ppo import ppo_data, ppo_loss, ppo_info, ppo_policy_data, ppo_policy_error, ppo_value_data, ppo_value_error, \ + ppo_error, ppo_error_continuous, ppo_policy_error_continuous, ppo_data_continuous, ppo_policy_data_continuous +from .happo import happo_data, happo_policy_data, happo_value_data, happo_loss, happo_policy_loss, happo_info, \ + happo_error, happo_policy_error, happo_value_error, happo_error_continuous, happo_policy_error_continuous from .ppg import ppg_data, ppg_joint_loss, ppg_joint_error from .gae import gae_data, gae -from .a2c import a2c_data, a2c_error +from .a2c import a2c_data, a2c_error, a2c_error_continuous from .coma import coma_data, coma_error -from .td import q_nstep_td_data, q_nstep_td_error, q_1step_td_data, q_1step_td_error, td_lambda_data, td_lambda_error,\ +from .td import q_nstep_td_data, q_nstep_td_error, q_1step_td_data, \ + q_1step_td_error, m_q_1step_td_data, m_q_1step_td_error, td_lambda_data, td_lambda_error, \ q_nstep_td_error_with_rescale, v_1step_td_data, v_1step_td_error, v_nstep_td_data, v_nstep_td_error, \ generalized_lambda_returns, dist_1step_td_data, dist_1step_td_error, dist_nstep_td_error, dist_nstep_td_data, \ - nstep_return_data, nstep_return, iqn_nstep_td_data, iqn_nstep_td_error, qrdqn_nstep_td_data, qrdqn_nstep_td_error,\ + nstep_return_data, nstep_return, iqn_nstep_td_data, iqn_nstep_td_error, qrdqn_nstep_td_data, qrdqn_nstep_td_error, \ fqf_nstep_td_data, fqf_nstep_td_error, fqf_calculate_fraction_loss, evaluate_quantile_at_action, \ - q_nstep_sql_td_error, dqfd_nstep_td_error, dqfd_nstep_td_data, q_v_1step_td_error, q_v_1step_td_data,\ - dqfd_nstep_td_error_with_rescale, discount_cumsum + q_nstep_sql_td_error, dqfd_nstep_td_error, dqfd_nstep_td_data, q_v_1step_td_error, q_v_1step_td_data, \ + dqfd_nstep_td_error_with_rescale, discount_cumsum, bdq_nstep_td_error from .vtrace import vtrace_loss, compute_importance_weights from .upgo import upgo_loss from .adder import get_gae, get_gae_with_default_last_value, get_nstep_return_data, get_train_sample -from .value_rescale import value_transform, value_inv_transform -from .vtrace import vtrace_data, vtrace_error +from .value_rescale import value_transform, value_inv_transform, symlog, inv_symlog +from .vtrace import vtrace_data, vtrace_error_discrete_action, vtrace_error_continuous_action from .beta_function import beta_function_map from .retrace import compute_q_retraces from .acer import acer_policy_error, acer_value_error, acer_trust_region_update +from .sampler import ArgmaxSampler, MultinomialSampler, MuSampler, ReparameterizationSampler, HybridStochasticSampler, \ + HybridDeterminsticSampler diff --git a/ding/rl_utils/a2c.py b/ding/rl_utils/a2c.py index 0ed1deb29d..09cb199553 100644 --- a/ding/rl_utils/a2c.py +++ b/ding/rl_utils/a2c.py @@ -1,6 +1,7 @@ from collections import namedtuple import torch import torch.nn.functional as F +from torch.distributions import Independent, Normal a2c_data = namedtuple('a2c_data', ['logit', 'action', 'value', 'adv', 'return_', 'weight']) a2c_loss = namedtuple('a2c_loss', ['policy_loss', 'value_loss', 'entropy_loss']) @@ -9,7 +10,7 @@ def a2c_error(data: namedtuple) -> namedtuple: """ Overview: - Implementation of A2C(Advantage Actor-Critic) (arXiv:1602.01783) + Implementation of A2C(Advantage Actor-Critic) (arXiv:1602.01783) for discrete action space Arguments: - data (:obj:`namedtuple`): a2c input data with fieids shown in ``a2c_data`` Returns: @@ -24,6 +25,16 @@ def a2c_error(data: namedtuple) -> namedtuple: - policy_loss (:obj:`torch.FloatTensor`): :math:`()`, 0-dim tensor - value_loss (:obj:`torch.FloatTensor`): :math:`()` - entropy_loss (:obj:`torch.FloatTensor`): :math:`()` + Examples: + >>> data = a2c_data( + >>> logit=torch.randn(2, 3), + >>> action=torch.randint(0, 3, (2, )), + >>> value=torch.randn(2, ), + >>> adv=torch.randn(2, ), + >>> return_=torch.randn(2, ), + >>> weight=torch.ones(2, ), + >>> ) + >>> loss = a2c_error(data) """ logit, action, value, adv, return_, weight = data if weight is None: @@ -34,3 +45,44 @@ def a2c_error(data: namedtuple) -> namedtuple: policy_loss = -(logp * adv * weight).mean() value_loss = (F.mse_loss(return_, value, reduction='none') * weight).mean() return a2c_loss(policy_loss, value_loss, entropy_loss) + + +def a2c_error_continuous(data: namedtuple) -> namedtuple: + """ + Overview: + Implementation of A2C(Advantage Actor-Critic) (arXiv:1602.01783) for continuous action space + Arguments: + - data (:obj:`namedtuple`): a2c input data with fieids shown in ``a2c_data`` + Returns: + - a2c_loss (:obj:`namedtuple`): the a2c loss item, all of them are the differentiable 0-dim tensor + Shapes: + - logit (:obj:`torch.FloatTensor`): :math:`(B, N)`, where B is batch size and N is action dim + - action (:obj:`torch.LongTensor`): :math:`(B, N)` + - value (:obj:`torch.FloatTensor`): :math:`(B, )` + - adv (:obj:`torch.FloatTensor`): :math:`(B, )` + - return (:obj:`torch.FloatTensor`): :math:`(B, )` + - weight (:obj:`torch.FloatTensor` or :obj:`None`): :math:`(B, )` + - policy_loss (:obj:`torch.FloatTensor`): :math:`()`, 0-dim tensor + - value_loss (:obj:`torch.FloatTensor`): :math:`()` + - entropy_loss (:obj:`torch.FloatTensor`): :math:`()` + Examples: + >>> data = a2c_data( + >>> logit={'mu': torch.randn(2, 3), 'sigma': torch.sqrt(torch.randn(2, 3)**2)}, + >>> action=torch.randn(2, 3), + >>> value=torch.randn(2, ), + >>> adv=torch.randn(2, ), + >>> return_=torch.randn(2, ), + >>> weight=torch.ones(2, ), + >>> ) + >>> loss = a2c_error_continuous(data) + """ + logit, action, value, adv, return_, weight = data + if weight is None: + weight = torch.ones_like(value) + + dist = Independent(Normal(logit['mu'], logit['sigma']), 1) + logp = dist.log_prob(action) + entropy_loss = (dist.entropy() * weight).mean() + policy_loss = -(logp * adv * weight).mean() + value_loss = (F.mse_loss(return_, value, reduction='none') * weight).mean() + return a2c_loss(policy_loss, value_loss, entropy_loss) diff --git a/ding/rl_utils/acer.py b/ding/rl_utils/acer.py index 029ea9a448..ba83fc9393 100644 --- a/ding/rl_utils/acer.py +++ b/ding/rl_utils/acer.py @@ -16,7 +16,7 @@ def acer_policy_error( ) -> Tuple[torch.Tensor, torch.Tensor]: """ Overview: - Get ACER policy loss + Get ACER policy loss. Arguments: - q_values (:obj:`torch.Tensor`): Q values - q_retraces (:obj:`torch.Tensor`): Q values (be calculated by retrace method) @@ -37,6 +37,14 @@ def acer_policy_error( - ratio (:obj:`torch.FloatTensor`): :math:`(T, B, N)` - actor_loss (:obj:`torch.FloatTensor`): :math:`(T, B, 1)` - bc_loss (:obj:`torch.FloatTensor`): :math:`(T, B, 1)` + Examples: + >>> q_values=torch.randn(2, 3, 4), + >>> q_retraces=torch.randn(2, 3, 1), + >>> v_pred=torch.randn(2, 3, 1), + >>> target_pi=torch.randn(2, 3, 4), + >>> actions=torch.randint(0, 4, (2, 3)), + >>> ratio=torch.randn(2, 3, 4), + >>> loss = acer_policy_error(q_values, q_retraces, v_pred, target_pi, actions, ratio) """ actions = actions.unsqueeze(-1) with torch.no_grad(): @@ -56,7 +64,7 @@ def acer_policy_error( def acer_value_error(q_values, q_retraces, actions): """ Overview: - Get ACER critic loss + Get ACER critic loss. Arguments: - q_values (:obj:`torch.Tensor`): Q values - q_retraces (:obj:`torch.Tensor`): Q values (be calculated by retrace method) @@ -69,6 +77,11 @@ def acer_value_error(q_values, q_retraces, actions): - q_retraces (:obj:`torch.FloatTensor`): :math:`(T, B, 1)` - actions (:obj:`torch.LongTensor`): :math:`(T, B)` - critic_loss (:obj:`torch.FloatTensor`): :math:`(T, B, 1)` + Examples: + >>> q_values=torch.randn(2, 3, 4) + >>> q_retraces=torch.randn(2, 3, 1) + >>> actions=torch.randint(0, 4, (2, 3)) + >>> loss = acer_value_error(q_values, q_retraces, actions) """ actions = actions.unsqueeze(-1) critic_loss = 0.5 * (q_retraces - q_values.gather(-1, actions)).pow(2) @@ -92,6 +105,12 @@ def acer_trust_region_update( Shapes: - target_pi (:obj:`torch.FloatTensor`): :math:`(T, B, N)` - avg_pi (:obj:`torch.FloatTensor`): :math:`(T, B, N)` + - update_gradients (:obj:`list(torch.FloatTensor)`): :math:`(T, B, N)` + Examples: + >>> actor_gradients=[torch.randn(2, 3, 4)] + >>> target_pi=torch.randn(2, 3, 4) + >>> avg_pi=torch.randn(2, 3, 4) + >>> loss = acer_trust_region_update(actor_gradients, target_pi, avg_pi, 0.1) """ with torch.no_grad(): KL_gradients = [torch.exp(avg_logit)] diff --git a/ding/rl_utils/adder.py b/ding/rl_utils/adder.py index c82fcee774..eea5f5f49f 100644 --- a/ding/rl_utils/adder.py +++ b/ding/rl_utils/adder.py @@ -4,7 +4,7 @@ import torch from ding.utils import list_split, lists_to_dicts -from .gae import gae, gae_data +from ding.rl_utils.gae import gae, gae_data class Adder(object): @@ -23,13 +23,23 @@ def get_gae(cls, data: List[Dict[str, Any]], last_value: torch.Tensor, gamma: fl Overview: Get GAE advantage for stacked transitions(T timestep, 1 batch). Call ``gae`` for calculation. Arguments: - - data (:obj:`list`): Transitions list, each element is a transition dict with at least ['value', 'reward'] + - data (:obj:`list`): Transitions list, each element is a transition dict with at least \ + ``['value', 'reward']``. - last_value (:obj:`torch.Tensor`): The last value(i.e.: the T+1 timestep) - - gamma (:obj:`float`): The future discount factor - - gae_lambda (:obj:`float`): GAE lambda parameter + - gamma (:obj:`float`): The future discount factor, should be in [0, 1], defaults to 0.99. + - gae_lambda (:obj:`float`): GAE lambda parameter, should be in [0, 1], defaults to 0.97, \ + when lambda -> 0, it induces bias, but when lambda -> 1, it has high variance due to the sum of terms. - cuda (:obj:`bool`): Whether use cuda in GAE computation Returns: - data (:obj:`list`): transitions list like input one, but each element owns extra advantage key 'adv' + Examples: + >>> B, T = 2, 3 # batch_size, timestep + >>> data = [dict(value=torch.randn(B), reward=torch.randn(B)) for _ in range(T)] + >>> last_value = torch.randn(B) + >>> gamma = 0.99 + >>> gae_lambda = 0.95 + >>> cuda = False + >>> data = Adder.get_gae(data, last_value, gamma, gae_lambda, cuda) """ value = torch.stack([d['value'] for d in data]) next_value = torch.stack([d['value'] for d in data][1:] + [last_value]) @@ -54,18 +64,27 @@ def get_gae_with_default_last_value(cls, data: deque, done: bool, gamma: float, Overview: Like ``get_gae`` above to get GAE advantage for stacked transitions. However, this function is designed in case ``last_value`` is not passed. If transition is not done yet, it wouold assign last value in ``data`` - as ``last_value``, discard the last element in ``data``(i.e. len(data) would decrease by 1), and then call + as ``last_value``, discard the last element in ``data`` (i.e. len(data) would decrease by 1), and then call ``get_gae``. Otherwise it would make ``last_value`` equal to 0. Arguments: - data (:obj:`deque`): Transitions list, each element is a transition dict with \ at least['value', 'reward'] - done (:obj:`bool`): Whether the transition reaches the end of an episode(i.e. whether the env is done) - - gamma (:obj:`float`): The future discount factor - - gae_lambda (:obj:`float`): GAE lambda parameter + - gamma (:obj:`float`): The future discount factor, should be in [0, 1], defaults to 0.99. + - gae_lambda (:obj:`float`): GAE lambda parameter, should be in [0, 1], defaults to 0.97, \ + when lambda -> 0, it induces bias, but when lambda -> 1, it has high variance due to the sum of terms. - cuda (:obj:`bool`): Whether use cuda in GAE computation Returns: - data (:obj:`List[Dict[str, Any]]`): transitions list like input one, but each element owns \ extra advantage key 'adv' + Examples: + >>> B, T = 2, 3 # batch_size, timestep + >>> data = [dict(value=torch.randn(B), reward=torch.randn(B)) for _ in range(T)] + >>> done = False + >>> gamma = 0.99 + >>> gae_lambda = 0.95 + >>> cuda = False + >>> data = Adder.get_gae_with_default_last_value(data, done, gamma, gae_lambda, cuda) """ if done: last_value = torch.zeros_like(data[-1]['value']) @@ -85,17 +104,25 @@ def get_nstep_return_data( ) -> deque: """ Overview: - Process raw traj data by updating keys ['next_obs', 'reward', 'done'] in data's dict element. + Process raw traj data by updating keys ``['next_obs', 'reward', 'done']`` in data's dict element. Arguments: - data (:obj:`deque`): Transitions list, each element is a transition dict - nstep (:obj:`int`): Number of steps. If equals to 1, return ``data`` directly; \ Otherwise update with nstep value. Returns: - data (:obj:`deque`): Transitions list like input one, but each element updated with nstep value. + Examples: + >>> data = [dict( + >>> obs=torch.randn(B), + >>> reward=torch.randn(1), + >>> next_obs=torch.randn(B), + >>> done=False) for _ in range(T)] + >>> nstep = 2 + >>> data = Adder.get_nstep_return_data(data, nstep) """ if nstep == 1: return data - fake_reward = torch.zeros(1) + fake_reward = torch.zeros_like(data[0]['reward']) next_obs_flag = 'next_obs' in data[0] for i in range(len(data) - nstep): # update keys ['next_obs', 'reward', 'done'] with their n-step value @@ -104,7 +131,10 @@ def get_nstep_return_data( if cum_reward: data[i]['reward'] = sum([data[i + j]['reward'] * (gamma ** j) for j in range(nstep)]) else: - data[i]['reward'] = torch.cat([data[i + j]['reward'] for j in range(nstep)]) + # data[i]['reward'].shape = (1) or (agent_num, 1) + # single agent env: shape (1) -> (n_step) + # multi-agent env: shape (agent_num, 1) -> (agent_num, n_step) + data[i]['reward'] = torch.cat([data[i + j]['reward'] for j in range(nstep)], dim=-1) data[i]['done'] = data[i + nstep - 1]['done'] if correct_terminate_gamma: data[i]['value_gamma'] = gamma ** nstep @@ -116,7 +146,8 @@ def get_nstep_return_data( else: data[i]['reward'] = torch.cat( [data[i + j]['reward'] - for j in range(len(data) - i)] + [fake_reward for _ in range(nstep - (len(data) - i))] + for j in range(len(data) - i)] + [fake_reward for _ in range(nstep - (len(data) - i))], + dim=-1 ) data[i]['done'] = data[-1]['done'] if correct_terminate_gamma: @@ -133,7 +164,7 @@ def get_train_sample( ) -> List[Dict[str, Any]]: """ Overview: - Process raw traj data by updating keys ['next_obs', 'reward', 'done'] in data's dict element. + Process raw traj data by updating keys ``['next_obs', 'reward', 'done']`` in data's dict element. If ``unroll_len`` equals to 1, which means no process is needed, can directly return ``data``. Otherwise, ``data`` will be splitted according to ``unroll_len``, process residual part according to ``last_fn_type`` and call ``lists_to_dicts`` to form sampled training data. diff --git a/ding/rl_utils/beta_function.py b/ding/rl_utils/beta_function.py index 4096228321..65f2f08fed 100644 --- a/ding/rl_utils/beta_function.py +++ b/ding/rl_utils/beta_function.py @@ -14,6 +14,15 @@ # For CPW, eta = 0.71 most closely match human subjects # this function is locally concave for small values of τ and becomes locally convex for larger values of τ def cpw(x: Union[torch.Tensor, float], eta: float = 0.71) -> Union[torch.Tensor, float]: + """ + Overview: + The implementation of CPW function. + Arguments: + - x (:obj:`Union[torch.Tensor, float]`): The input value. + - eta (:obj:`float`): The hyperparameter of CPW function. + Returns: + - output (:obj:`Union[torch.Tensor, float]`): The output value. + """ return (x ** eta) / ((x ** eta + (1 - x) ** eta) ** (1 / eta)) @@ -22,6 +31,15 @@ def cpw(x: Union[torch.Tensor, float], eta: float = 0.71) -> Union[torch.Tensor, # CVaR is risk-averse def CVaR(x: Union[torch.Tensor, float], eta: float = 0.71) -> Union[torch.Tensor, float]: + """ + Overview: + The implementation of CVaR function, which is a risk-averse function. + Arguments: + - x (:obj:`Union[torch.Tensor, float]`): The input value. + - eta (:obj:`float`): The hyperparameter of CVaR function. + Returns: + - output (:obj:`Union[torch.Tensor, float]`): The output value. + """ assert eta <= 1.0 return x * eta @@ -31,6 +49,15 @@ def CVaR(x: Union[torch.Tensor, float], eta: float = 0.71) -> Union[torch.Tensor # risk-averse (eta < 0) or risk-seeking (eta > 0) def Pow(x: Union[torch.Tensor, float], eta: float = 0.0) -> Union[torch.Tensor, float]: + """ + Overview: + The implementation of Pow function, which is risk-averse when eta < 0 and risk-seeking when eta > 0. + Arguments: + - x (:obj:`Union[torch.Tensor, float]`): The input value. + - eta (:obj:`float`): The hyperparameter of Pow function. + Returns: + - output (:obj:`Union[torch.Tensor, float]`): The output value. + """ if eta >= 0: return x ** (1 / (1 + eta)) else: diff --git a/ding/rl_utils/coma.py b/ding/rl_utils/coma.py index 6a922981c2..0ee1778293 100644 --- a/ding/rl_utils/coma.py +++ b/ding/rl_utils/coma.py @@ -1,7 +1,7 @@ from collections import namedtuple import torch import torch.nn.functional as F -from .td import generalized_lambda_returns +from ding.rl_utils.td import generalized_lambda_returns coma_data = namedtuple('coma_data', ['logit', 'action', 'q_value', 'target_q_value', 'reward', 'weight']) coma_loss = namedtuple('coma_loss', ['policy_loss', 'q_value_loss', 'entropy_loss']) @@ -26,6 +26,18 @@ def coma_error(data: namedtuple, gamma: float, lambda_: float) -> namedtuple: - policy_loss (:obj:`torch.FloatTensor`): :math:`()`, 0-dim tensor - value_loss (:obj:`torch.FloatTensor`): :math:`()` - entropy_loss (:obj:`torch.FloatTensor`): :math:`()` + Examples: + >>> action_dim = 4 + >>> agent_num = 3 + >>> data = coma_data( + >>> logit=torch.randn(2, 3, agent_num, action_dim), + >>> action=torch.randint(0, action_dim, (2, 3, agent_num)), + >>> q_value=torch.randn(2, 3, agent_num, action_dim), + >>> target_q_value=torch.randn(2, 3, agent_num, action_dim), + >>> reward=torch.randn(2, 3), + >>> weight=torch.ones(2, 3, agent_num), + >>> ) + >>> loss = coma_error(data, 0.99, 0.99) """ logit, action, q_value, target_q_value, reward, weight = data if weight is None: diff --git a/ding/rl_utils/efficientzero/game.py b/ding/rl_utils/efficientzero/game.py deleted file mode 100644 index b096da80dc..0000000000 --- a/ding/rl_utils/efficientzero/game.py +++ /dev/null @@ -1,193 +0,0 @@ -""" -The following code is adapted from https://github.com/YeWR/EfficientZero/blob/main/core/game.py -""" - -import copy -import numpy as np -from ding.utils.compression_helper import jpeg_data_decompressor - - -class Game: - - def __init__(self, env, action_space_size: int, config=None): - self.env = env - self.action_space_size = action_space_size - self.config = config - - def legal_actions(self): - raise NotImplementedError - - def step(self, action): - raise NotImplementedError - - def reset(self): - raise NotImplementedError() - - def close(self, *args, **kwargs): - self.env.close(*args, **kwargs) - - def render(self, *args, **kwargs): - self.env.render(*args, **kwargs) - - -class GameHistory: - """ - A block of game history from a full trajectories. - The horizons of Atari games are quite large. Split the whole trajectory into several history blocks. - """ - - def __init__(self, action_space, max_length=200, config=None): - """ - Parameters - ---------- - action_space: int - action space - max_length: int - max transition number of the history block - """ - self.action_space = action_space - self.max_length = max_length - self.config = config - - self.stacked_observations = config.stacked_observations - self.discount = config.discount - self.action_space_size = config.action_space_size - self.zero_obs_shape = (config.obs_shape[-2], config.obs_shape[-1], config.image_channel) - - self.child_visits = [] - self.root_values = [] - - self.actions = [] - self.obs_history = [] - self.rewards = [] - - def init(self, init_observations): - """Initialize a history block, stack the previous stacked_observations frames. - Parameters - ---------- - init_observations: list - list of the stack observations in the previous time steps - """ - self.child_visits = [] - self.root_values = [] - - self.actions = [] - self.obs_history = [] - self.rewards = [] - self.target_values = [] - self.target_rewards = [] - self.target_policies = [] - - assert len(init_observations) == self.stacked_observations - - for observation in init_observations: - self.obs_history.append(copy.deepcopy(observation)) - - def pad_over(self, next_block_observations, next_block_rewards, next_block_root_values, next_block_child_visits): - """To make sure the correction of value targets, we need to add (o_t, r_t, etc) from the next history block - , which is necessary for the bootstrapped values at the end states of this history block. - Eg: len = 100; target value v_100 = r_100 + gamma^1 r_101 + ... + gamma^4 r_104 + gamma^5 v_105, - but r_101, r_102, ... are from the next history block. - Parameters - ---------- - next_block_observations: list - o_t from the next history block - next_block_rewards: list - r_t from the next history block - next_block_root_values: list - root values of MCTS from the next history block - next_block_child_visits: list - root visit count distributions of MCTS from the next history block - """ - assert len(next_block_observations) <= self.config.num_unroll_steps - assert len(next_block_child_visits) <= self.config.num_unroll_steps - assert len(next_block_root_values) <= self.config.num_unroll_steps + self.config.td_steps - assert len(next_block_rewards) <= self.config.num_unroll_steps + self.config.td_steps - 1 - - # notice: next block observation should start from (stacked_observation - 1) in next trajectory - for observation in next_block_observations: - self.obs_history.append(copy.deepcopy(observation)) - - for reward in next_block_rewards: - self.rewards.append(reward) - - for value in next_block_root_values: - self.root_values.append(value) - - for child_visits in next_block_child_visits: - self.child_visits.append(child_visits) - - def is_full(self): - # history block is full - return self.__len__() >= self.max_length - - def legal_actions(self): - return [_ for _ in range(self.action_space.n)] - - def append(self, action, obs, reward): - # append a transition tuple - self.actions.append(action) - self.obs_history.append(obs) - self.rewards.append(reward) - - def obs(self, i, extra_len=0, padding=False): - """To obtain an observation of correct format: o[t, t + stack frames + extra len] - Parameters - ---------- - i: int - time step i - extra_len: int - extra len of the obs frames - padding: bool - True -> padding frames if (t + stack frames) are out of trajectory - """ - # frames = self.obs_history[index:index + self.stacked_observations] - - frames = self.obs_history[i:i + self.stacked_observations + extra_len] - if padding: - pad_len = self.stacked_observations + extra_len - len(frames) - if pad_len > 0: - pad_frames = np.array([frames[-1] for _ in range(pad_len)]) - frames = np.concatenate((frames, pad_frames)) - if self.config.cvt_string: # TODO - frames = [jpeg_data_decompressor(obs, self.config.gray_scale) for obs in frames] - return frames - - def zero_obs(self): - # return a zero frame - return [np.zeros(self.zero_obs_shape, dtype=np.uint8) for _ in range(self.stacked_observations)] - - def step_obs(self): - # return an observation of correct format for model inference - index = len(self.rewards) - frames = self.obs_history[index:index + self.stacked_observations] - if self.config.cvt_string: # TODO - frames = [jpeg_data_decompressor(obs, self.config.gray_scale) for obs in frames] - return frames - - def get_targets(self, i): - # return the value/rewrad/policy targets at step i - return self.target_values[i], self.target_rewards[i], self.target_policies[i] - - def game_over(self): - # post processing the data when a history block is full - # obs_history should be sent into the ray memory. Otherwise, it will cost large amounts of time in copying obs. - self.rewards = np.array(self.rewards) - # self.obs_history = ray.put(np.array(self.obs_history)) - self.obs_history = np.array(self.obs_history) - self.actions = np.array(self.actions) - self.child_visits = np.array(self.child_visits) - self.root_values = np.array(self.root_values) - - def store_search_stats(self, visit_counts, root_value, idx: int = None): - # store the visit count distributions and value of the root node after MCTS - sum_visits = sum(visit_counts) - if idx is None: - self.child_visits.append([visit_count / sum_visits for visit_count in visit_counts]) - self.root_values.append(root_value) - else: - self.child_visits[idx] = [visit_count / sum_visits for visit_count in visit_counts] - self.root_values[idx] = root_value - - def __len__(self): - return len(self.actions) diff --git a/ding/rl_utils/efficientzero/utils.py b/ding/rl_utils/efficientzero/utils.py deleted file mode 100644 index b48f00d9a6..0000000000 --- a/ding/rl_utils/efficientzero/utils.py +++ /dev/null @@ -1,187 +0,0 @@ -""" -The following code is adapted from https://github.com/YeWR/EfficientZero/core/utils.py -""" - -import os -import gym -import torch -import random -import shutil -import logging - -import numpy as np - -from scipy.stats import entropy - - -def get_augmented_data(board_size, play_data): - """ - Overview: - augment the data set by rotation and flipping - Arguments: - play_data: [(state, mcts_prob, winner_z), ..., ...] - """ - extend_data = [] - for data in play_data: - state = data['state'] - mcts_prob = data['mcts_prob'] - winner = data['winner'] - for i in [1, 2, 3, 4]: - # rotate counterclockwise - equi_state = np.array([np.rot90(s, i) for s in state]) - equi_mcts_prob = np.rot90(np.flipud(mcts_prob.reshape(board_size, board_size)), i) - extend_data.append( - { - 'state': equi_state, - 'mcts_prob': np.flipud(equi_mcts_prob).flatten(), - 'winner': winner - } - ) - # flip horizontally - equi_state = np.array([np.fliplr(s) for s in equi_state]) - equi_mcts_prob = np.fliplr(equi_mcts_prob) - extend_data.append( - { - 'state': equi_state, - 'mcts_prob': np.flipud(equi_mcts_prob).flatten(), - 'winner': winner - } - ) - return extend_data - - -class LinearSchedule(object): - - def __init__(self, schedule_timesteps, final_p, initial_p=1.0): - """Linear interpolation between initial_p and final_p over - schedule_timesteps. After this many timesteps pass final_p is - returned. - Parameters - ---------- - schedule_timesteps: int - Number of timesteps for which to linearly anneal initial_p - to final_p - initial_p: float - initial output value - final_p: float - final output value - """ - self.schedule_timesteps = schedule_timesteps - self.final_p = final_p - self.initial_p = initial_p - - def value(self, t): - """See Schedule.value""" - fraction = min(float(t) / self.schedule_timesteps, 1.0) - return self.initial_p + fraction * (self.final_p - self.initial_p) - - -def set_seed(seed): - # set seed - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - torch.cuda.manual_seed(seed) - torch.backends.cudnn.deterministic = True - - -def make_results_dir(exp_path, args): - # make the result directory - os.makedirs(exp_path, exist_ok=True) - if args.opr == 'train' and os.path.exists(exp_path) and os.listdir(exp_path): - if not args.force: - raise FileExistsError('{} is not empty. Please use --force to overwrite it'.format(exp_path)) - else: - print('Warning, path exists! Rewriting...') - shutil.rmtree(exp_path) - os.makedirs(exp_path) - log_path = os.path.join(exp_path, 'logs') - os.makedirs(log_path, exist_ok=True) - os.makedirs(os.path.join(exp_path, 'model'), exist_ok=True) - return exp_path, log_path - - -def init_logger(base_path): - # initialize the logger - formatter = logging.Formatter('[%(asctime)s][%(name)s][%(levelname)s][%(filename)s>%(funcName)s] ==> %(message)s') - for mode in ['train', 'test', 'train_test', 'root']: - file_path = os.path.join(base_path, mode + '.log') - logger = logging.getLogger(mode) - handler = logging.StreamHandler() - handler.setFormatter(formatter) - logger.addHandler(handler) - handler = logging.FileHandler(file_path, mode='a') - handler.setFormatter(formatter) - logger.addHandler(handler) - logger.setLevel(logging.DEBUG) - - -def select_action(visit_counts, temperature=1, deterministic=True): - """select action from the root visit counts. - Parameters - ---------- - temperature: float - the temperature for the distribution - deterministic: bool - True -> select the argmax - False -> sample from the distribution - """ - action_probs = [visit_count_i ** (1 / temperature) for visit_count_i in visit_counts] - total_count = sum(action_probs) - action_probs = [x / total_count for x in action_probs] - if deterministic: - action_pos = np.argmax([v for v in visit_counts]) - else: - action_pos = np.random.choice(len(visit_counts), p=action_probs) - - count_entropy = entropy(action_probs, base=2) - return action_pos, count_entropy - - -def prepare_observation_lst(observation_lst): - """Prepare the observations to satisfy the input fomat of torch - [B, S, W, H, C] -> [B, S x C, W, H] - batch, stack num, width, height, channel - """ - # B, S, W, H, C - observation_lst = np.array(observation_lst, dtype=np.uint8) - observation_lst = np.moveaxis(observation_lst, -1, 2) - - shape = observation_lst.shape - observation_lst = observation_lst.reshape((shape[0], -1, shape[-2], shape[-1])) - - return observation_lst - - -def concat_output_value(output_lst): - # concat the values of the model output list - value_lst = [] - for output in output_lst: - value_lst.append(output.value) - - value_lst = np.concatenate(value_lst) - - return value_lst - - -def concat_output(output_lst): - # concat the model output - value_lst, reward_lst, policy_logits_lst, hidden_state_lst = [], [], [], [] - reward_hidden_c_lst, reward_hidden_h_lst = [], [] - for output in output_lst: - value_lst.append(output.value) - reward_lst.append(output.value_prefix) - policy_logits_lst.append(output.policy_logits) - hidden_state_lst.append(output.hidden_state) - reward_hidden_c_lst.append(output.reward_hidden[0].squeeze(0)) - reward_hidden_h_lst.append(output.reward_hidden[1].squeeze(0)) - - value_lst = np.concatenate(value_lst) - reward_lst = np.concatenate(reward_lst) - policy_logits_lst = np.concatenate(policy_logits_lst) - # hidden_state_lst = torch.cat(hidden_state_lst, 0) - hidden_state_lst = np.concatenate(hidden_state_lst) - reward_hidden_c_lst = np.expand_dims(np.concatenate(reward_hidden_c_lst), axis=0) - reward_hidden_h_lst = np.expand_dims(np.concatenate(reward_hidden_h_lst), axis=0) - - return value_lst, reward_lst, policy_logits_lst, hidden_state_lst, (reward_hidden_c_lst, reward_hidden_h_lst) diff --git a/ding/rl_utils/exploration.py b/ding/rl_utils/exploration.py index fe49a892f9..ee45798c5b 100644 --- a/ding/rl_utils/exploration.py +++ b/ding/rl_utils/exploration.py @@ -12,13 +12,13 @@ def get_epsilon_greedy_fn(start: float, end: float, decay: int, type_: str = 'ex Overview: Generate an epsilon_greedy function with decay, which inputs current timestep and outputs current epsilon. Arguments: - - start (:obj:`float`): Epsilon start value. For 'linear', it should be 1.0. + - start (:obj:`float`): Epsilon start value. For ``linear`` , it should be 1.0. - end (:obj:`float`): Epsilon end value. - decay (:obj:`int`): Controls the speed that epsilon decreases from ``start`` to ``end``. \ We recommend epsilon decays according to env step rather than iteration. - - type (:obj:`str`): How epsilon decays, now supports ['linear', 'exp'(exponential)] + - type (:obj:`str`): How epsilon decays, now supports ``['linear', 'exp'(exponential)]`` . Returns: - - eps_fn (:obj:`function`): The epsilon greedy function with decay + - eps_fn (:obj:`function`): The epsilon greedy function with decay. """ assert type_ in ['linear', 'exp'], type_ if type_ == 'exp': @@ -48,7 +48,7 @@ class BaseNoise(ABC): def __init__(self) -> None: """ Overview: - Initialization method + Initialization method. """ super().__init__() @@ -56,19 +56,19 @@ def __init__(self) -> None: def __call__(self, shape: tuple, device: str) -> torch.Tensor: """ Overview: - Generate noise according to action tensor's shape, device + Generate noise according to action tensor's shape, device. Arguments: - - shape (:obj:`tuple`): size of the action tensor, output noise's size should be the same - - device (:obj:`str`): device of the action tensor, output noise's device should be the same as it + - shape (:obj:`tuple`): size of the action tensor, output noise's size should be the same. + - device (:obj:`str`): device of the action tensor, output noise's device should be the same as it. Returns: - noise (:obj:`torch.Tensor`): generated action noise, \ - have the same shape and device with the input action tensor + have the same shape and device with the input action tensor. """ raise NotImplementedError class GaussianNoise(BaseNoise): - r""" + """ Overview: Derived class for generating gaussian noise, which satisfies :math:`X \sim N(\mu, \sigma^2)` Interface: @@ -78,10 +78,10 @@ class GaussianNoise(BaseNoise): def __init__(self, mu: float = 0.0, sigma: float = 1.0) -> None: """ Overview: - Initialize :math:`\mu` and :math:`\sigma` in Gaussian Distribution + Initialize :math:`\mu` and :math:`\sigma` in Gaussian Distribution. Arguments: - - mu (:obj:`float`): :math:`\mu` , mean value - - sigma (:obj:`float`): :math:`\sigma` , standard deviation, should be positive + - mu (:obj:`float`): :math:`\mu` , mean value. + - sigma (:obj:`float`): :math:`\sigma` , standard deviation, should be positive. """ super(GaussianNoise, self).__init__() self._mu = mu @@ -125,14 +125,15 @@ def __init__( """ Overview: Initialize ``_alpha`` :math:`=\theta * dt\`, - ``beta`` :math:`= \sigma * \sqrt{dt}`, in Ornstein-Uhlenbeck process + ``beta`` :math:`= \sigma * \sqrt{dt}`, in Ornstein-Uhlenbeck process. Arguments: - - mu (:obj:`float`): :math:`\mu` , mean value - - sigma (:obj:`float`): :math:`\sigma` , standard deviation of the perturbation noise - - theta (:obj:`float`): how strongly the noise reacts to perturbations, \ - greater value means stronger reaction - - dt (:obj:`float`): derivative of time t - - x0 (:obj:`float` or :obj:`torch.Tensor`): initial action + - mu (:obj:`float`): :math:`\mu` , mean value. + - sigma (:obj:`float`): :math:`\sigma` , standard deviation of the perturbation noise. + - theta (:obj:`float`): How strongly the noise reacts to perturbations, \ + greater value means stronger reaction. + - dt (:obj:`float`): The derivative of time t. + - x0 (:obj:`Union[float, torch.Tensor]`): The initial state of the noise, \ + should be a scalar or tensor with the same shape as the action tensor. """ super().__init__() self._mu = mu @@ -144,21 +145,21 @@ def __init__( def reset(self) -> None: """ Overview: - Reset ``_x`` to the initial state ``_x0`` + Reset ``_x`` to the initial state ``_x0``. """ self._x = deepcopy(self._x0) def __call__(self, shape: tuple, device: str, mu: Optional[float] = None) -> torch.Tensor: """ Overview: - Generate gaussian noise according to action tensor's shape, device + Generate gaussian noise according to action tensor's shape, device. Arguments: - - shape (:obj:`tuple`): size of the action tensor, output noise's size should be the same - - device (:obj:`str`): device of the action tensor, output noise's device should be the same as it - - mu (:obj:`float`): new mean value :math:`\mu`, you can set it to `None` if don't need it + - shape (:obj:`tuple`): The size of the action tensor, output noise's size should be the same. + - device (:obj:`str`): The device of the action tensor, output noise's device should be the same as it. + - mu (:obj:`float`): The new mean value :math:`\mu`, you can set it to `None` if don't need it. Returns: - noise (:obj:`torch.Tensor`): generated action noise, \ - have the same shape and device with the input action tensor + have the same shape and device with the input action tensor. """ if self._x is None or \ (isinstance(self._x, torch.Tensor) and self._x.shape != shape): @@ -172,13 +173,17 @@ def __call__(self, shape: tuple, device: str, mu: Optional[float] = None) -> tor @property def x0(self) -> Union[float, torch.Tensor]: + """ + Overview: + Get ``self._x0``. + """ return self._x0 @x0.setter def x0(self, _x0: Union[float, torch.Tensor]) -> None: """ Overview: - Set ``self._x0`` and reset ``self.x`` to ``self._x0`` as well + Set ``self._x0`` and reset ``self.x`` to ``self._x0`` as well. """ self._x0 = _x0 self.reset() @@ -194,10 +199,10 @@ def create_noise_generator(noise_type: str, noise_kwargs: dict) -> BaseNoise: or raise an KeyError. In other words, a derived noise generator must first register, then call ``create_noise generator`` to get the instance object. Arguments: - - noise_type (:obj:`str`): the type of noise generator to be created + - noise_type (:obj:`str`): the type of noise generator to be created. Returns: - noise (:obj:`BaseNoise`): the created new noise generator, should be an instance of one of \ - noise_mapping's values + noise_mapping's values. """ if noise_type not in noise_mapping.keys(): raise KeyError("not support noise type: {}".format(noise_type)) diff --git a/ding/rl_utils/gae.py b/ding/rl_utils/gae.py index f9607fb90a..800fcae354 100644 --- a/ding/rl_utils/gae.py +++ b/ding/rl_utils/gae.py @@ -39,6 +39,12 @@ def gae(data: namedtuple, gamma: float = 0.99, lambda_: float = 0.97) -> torch.F - next_value (:obj:`torch.FloatTensor`): :math:`(T, B)` - reward (:obj:`torch.FloatTensor`): :math:`(T, B)` - adv (:obj:`torch.FloatTensor`): :math:`(T, B)` + Examples: + >>> value = torch.randn(2, 3) + >>> next_value = torch.randn(2, 3) + >>> reward = torch.randn(2, 3) + >>> data = gae_data(value, next_value, reward, None, None) + >>> adv = gae(data) """ value, next_value, reward, done, traj_flag = data if done is None: diff --git a/ding/rl_utils/grpo.py b/ding/rl_utils/grpo.py new file mode 100644 index 0000000000..c4bad1a56f --- /dev/null +++ b/ding/rl_utils/grpo.py @@ -0,0 +1,75 @@ +from typing import Tuple +from collections import namedtuple +import torch +from .log_prob_utils import efficient_method, naive_method, less_efficient_method, LogProbFunction + +grpo_policy_data = namedtuple('grpo_policy_data', ['logit_new', 'logit_old', 'logit_ref', 'action', 'adv', 'weight']) +grpo_info = namedtuple('grpo_info', ['approx_kl', 'clipfrac']) + + +def grpo_policy_error( + data: namedtuple, + log_prob_fn: LogProbFunction = efficient_method, # Method to calculate the log probabilities + clip_ratio: float = 0.2, + beta: float = 0.1 # Weight coefficient for KL divergence +) -> Tuple[namedtuple, namedtuple]: + """ + Overview: + Group Relative Policy Optimization (GRPO) algorithm, see https://arxiv.org/abs/2402.03300. + Arguments: + - data (:obj:`namedtuple`): the grpo input data with fields shown in ``grpo_policy_data``. + - clip_ratio (:obj:`float`): the ppo clip ratio for the constraint of policy update, defaults to 0.2. + - beta (:obj:`float`): weight coefficient for KL divergence regularization, defaults to 0.1. + - log_prob_fn (:obj:`LogProbFunction`): The method to calculate the log probabilities, \ + defaults to `efficient_method`. + Returns: + - loss (:obj:`torch.FloatTensor`): the rloo policy loss, a differentiable 0-dim tensor. + - grpo_info (:obj:`namedtuple`): the grpo optim information for monitoring, all of them are Python scalar. + Shapes: + - logit_new (:obj:`torch.FloatTensor`): :math:`(B, S, V)`, where B is batch size, S is sequence length, \ + and V is vocabulary size. + - logit_old (:obj:`torch.FloatTensor`): :math:`(B, S, V)`. + - logit_ref (:obj:`torch.FloatTensor`): :math:`(B, S, V)`. + - action (:obj:`torch.LongTensor`): :math:`(B, S)`. + - adv (:obj:`torch.FloatTensor`): :math:`(B, )`. + - weight (:obj:`torch.FloatTensor` or :obj:`None`): :math:`(B, S)`. + - policy_loss (:obj:`torch.FloatTensor`): :math:`()`, 0-dim tensor. + - mean_kl (:obj:`float`): mean KL divergence between current and reference policy. + - mean_ratio (:obj:`float`): mean probability ratio. + - mean_clipped (:obj:`float`): proportion of clipped probability ratios. + """ + + # Calculate log probabilities for selected token + per_token_logps = log_prob_fn(data.logit_new, data.action) + per_token_ref_logps = log_prob_fn(data.logit_ref, data.action) + per_token_old_logps = log_prob_fn(data.logit_old, data.action) + + # Calculate KL divergence: exp(q-p) - (q-p) - 1, + # where p is current policy and q is reference policy + per_token_kl = (torch.exp(per_token_ref_logps - per_token_logps) - (per_token_ref_logps - per_token_logps) - 1) + + # Calculate policy ratio + ratio = torch.exp(per_token_logps - per_token_old_logps) + ratio_clipped = torch.clamp(ratio, 1 - clip_ratio, 1 + clip_ratio) + + # Calculate loss for each token + advantages = data.adv.unsqueeze(1) # [B, 1] + per_token_loss_unclipped = ratio * advantages + per_token_loss_clipped = ratio_clipped * advantages + per_token_loss = -torch.min(per_token_loss_unclipped, per_token_loss_clipped) + + # Add KL divergence regularization term + per_token_loss = per_token_loss + beta * per_token_kl + + # Calculate average loss using weight mask + weight = data.weight if data.weight is not None \ + else torch.ones_like(per_token_loss) + loss = ((per_token_loss * weight).sum(dim=1) / weight.sum(dim=1)).mean() + + # Calculate additional metrics + with torch.no_grad(): + approx_kl = (per_token_old_logps - per_token_logps).mean().item() + clipped = ratio.gt(1 + clip_ratio) | ratio.lt(1 - clip_ratio) + clipfrac = torch.as_tensor(clipped).float().mean().item() + + return loss, grpo_info(approx_kl=approx_kl, clipfrac=clipfrac) diff --git a/ding/rl_utils/happo.py b/ding/rl_utils/happo.py new file mode 100644 index 0000000000..b37ddc7528 --- /dev/null +++ b/ding/rl_utils/happo.py @@ -0,0 +1,347 @@ +from collections import namedtuple +from typing import Optional, Tuple +import torch +import torch.nn as nn +from torch.distributions import Independent, Normal +from ding.hpc_rl import hpc_wrapper + +happo_value_data = namedtuple('happo_value_data', ['value_new', 'value_old', 'return_', 'weight']) +happo_loss = namedtuple('happo_loss', ['policy_loss', 'value_loss', 'entropy_loss']) +happo_policy_loss = namedtuple('happo_policy_loss', ['policy_loss', 'entropy_loss']) +happo_info = namedtuple('happo_info', ['approx_kl', 'clipfrac']) +happo_data = namedtuple( + 'happo_data', ['logit_new', 'logit_old', 'action', 'value_new', 'value_old', 'adv', 'return_', 'weight', 'factor'] +) +happo_policy_data = namedtuple('happo_policy_data', ['logit_new', 'logit_old', 'action', 'adv', 'weight', 'factor']) + + +def happo_error( + data: namedtuple, + clip_ratio: float = 0.2, + use_value_clip: bool = True, + dual_clip: Optional[float] = None, +) -> Tuple[namedtuple, namedtuple]: + """ + Overview: + Implementation of Proximal Policy Optimization (arXiv:1707.06347) with value_clip and dual_clip + Arguments: + - data (:obj:`namedtuple`): the ppo input data with fieids shown in ``ppo_data`` + - clip_ratio (:obj:`float`): the ppo clip ratio for the constraint of policy update, defaults to 0.2 + - use_value_clip (:obj:`bool`): whether to use clip in value loss with the same ratio as policy + - dual_clip (:obj:`float`): a parameter c mentioned in arXiv:1912.09729 Equ. 5, shoule be in [1, inf),\ + defaults to 5.0, if you don't want to use it, set this parameter to None + Returns: + - happo_loss (:obj:`namedtuple`): the ppo loss item, all of them are the differentiable 0-dim tensor + - happo_info (:obj:`namedtuple`): the ppo optim information for monitoring, all of them are Python scalar + Shapes: + - logit_new (:obj:`torch.FloatTensor`): :math:`(B, N)`, where B is batch size and N is action dim + - logit_old (:obj:`torch.FloatTensor`): :math:`(B, N)` + - action (:obj:`torch.LongTensor`): :math:`(B, )` + - value_new (:obj:`torch.FloatTensor`): :math:`(B, )` + - value_old (:obj:`torch.FloatTensor`): :math:`(B, )` + - adv (:obj:`torch.FloatTensor`): :math:`(B, )` + - return (:obj:`torch.FloatTensor`): :math:`(B, )` + - weight (:obj:`torch.FloatTensor` or :obj:`None`): :math:`(B, )` + - policy_loss (:obj:`torch.FloatTensor`): :math:`()`, 0-dim tensor + - value_loss (:obj:`torch.FloatTensor`): :math:`()` + - entropy_loss (:obj:`torch.FloatTensor`): :math:`()` + Examples: + >>> action_dim = 4 + >>> data = happo_data( + >>> logit_new=torch.randn(3, action_dim), + >>> logit_old=torch.randn(3, action_dim), + >>> action=torch.randint(0, action_dim, (3,)), + >>> value_new=torch.randn(3), + >>> value_old=torch.randn(3), + >>> adv=torch.randn(3), + >>> return_=torch.randn(3), + >>> weight=torch.ones(3), + >>> factor=torch.ones(3, 1), + >>> ) + >>> loss, info = happo_error(data) + + .. note:: + + adv is already normalized value (adv - adv.mean()) / (adv.std() + 1e-8), and there are many + ways to calculate this mean and std, like among data buffer or train batch, so we don't couple + this part into happo_error, you can refer to our examples for different ways. + """ + assert dual_clip is None or dual_clip > 1.0, "dual_clip value must be greater than 1.0, but get value: {}".format( + dual_clip + ) + logit_new, logit_old, action, value_new, value_old, adv, return_, weight, factor = data + policy_data = happo_policy_data(logit_new, logit_old, action, adv, weight, factor) + policy_output, policy_info = happo_policy_error(policy_data, clip_ratio, dual_clip) + value_data = happo_value_data(value_new, value_old, return_, weight) + value_loss = happo_value_error(value_data, clip_ratio, use_value_clip) + + return happo_loss(policy_output.policy_loss, value_loss, policy_output.entropy_loss), policy_info + + +def happo_policy_error( + data: namedtuple, + clip_ratio: float = 0.2, + dual_clip: Optional[float] = None, +) -> Tuple[namedtuple, namedtuple]: + ''' + Overview: + Get PPO policy loss + Arguments: + - data (:obj:`namedtuple`): ppo input data with fieids shown in ``ppo_policy_data`` + - clip_ratio (:obj:`float`): clip value for ratio + - dual_clip (:obj:`float`): a parameter c mentioned in arXiv:1912.09729 Equ. 5, shoule be in [1, inf),\ + defaults to 5.0, if you don't want to use it, set this parameter to None + Returns: + - happo_policy_loss (:obj:`namedtuple`): the ppo policy loss item, all of them are the differentiable \ + 0-dim tensor. + - happo_info (:obj:`namedtuple`): the ppo optim information for monitoring, all of them are Python scalar + Shapes: + - logit_new (:obj:`torch.FloatTensor`): :math:`(B, N)`, where B is batch size and N is action dim + - logit_old (:obj:`torch.FloatTensor`): :math:`(B, N)` + - action (:obj:`torch.LongTensor`): :math:`(B, )` + - adv (:obj:`torch.FloatTensor`): :math:`(B, )` + - weight (:obj:`torch.FloatTensor` or :obj:`None`): :math:`(B, )` + - policy_loss (:obj:`torch.FloatTensor`): :math:`()`, 0-dim tensor + - entropy_loss (:obj:`torch.FloatTensor`): :math:`()` + Examples: + >>> action_dim = 4 + >>> data = ppo_policy_data( + >>> logit_new=torch.randn(3, action_dim), + >>> logit_old=torch.randn(3, action_dim), + >>> action=torch.randint(0, action_dim, (3,)), + >>> adv=torch.randn(3), + >>> weight=torch.ones(3), + >>> factor=torch.ones(3, 1), + >>> ) + >>> loss, info = happo_policy_error(data) + ''' + logit_new, logit_old, action, adv, weight, factor = data + if weight is None: + weight = torch.ones_like(adv) + dist_new = torch.distributions.categorical.Categorical(logits=logit_new) + dist_old = torch.distributions.categorical.Categorical(logits=logit_old) + logp_new = dist_new.log_prob(action) + logp_old = dist_old.log_prob(action) + dist_new_entropy = dist_new.entropy() + if dist_new_entropy.shape != weight.shape: + dist_new_entropy = dist_new.entropy().mean(dim=1) + entropy_loss = (dist_new_entropy * weight).mean() + # policy_loss + ratio = torch.exp(logp_new - logp_old) + if ratio.shape != adv.shape: + ratio = ratio.mean(dim=1) + surr1 = ratio * adv + surr2 = ratio.clamp(1 - clip_ratio, 1 + clip_ratio) * adv + # shape factor: (B,1) surr1: (B,) + clip1 = torch.min(surr1, surr2) * factor.squeeze(1) + if dual_clip is not None: + clip2 = torch.max(clip1, dual_clip * adv) + # only use dual_clip when adv < 0 + policy_loss = -(torch.where(adv < 0, clip2, clip1) * weight).mean() + else: + policy_loss = (-clip1 * weight).mean() + with torch.no_grad(): + approx_kl = (logp_old - logp_new).mean().item() + clipped = ratio.gt(1 + clip_ratio) | ratio.lt(1 - clip_ratio) + clipfrac = torch.as_tensor(clipped).float().mean().item() + return happo_policy_loss(policy_loss, entropy_loss), happo_info(approx_kl, clipfrac) + + +def happo_value_error( + data: namedtuple, + clip_ratio: float = 0.2, + use_value_clip: bool = True, +) -> torch.Tensor: + ''' + Overview: + Get PPO value loss + Arguments: + - data (:obj:`namedtuple`): ppo input data with fieids shown in ``happo_value_data`` + - clip_ratio (:obj:`float`): clip value for ratio + - use_value_clip (:obj:`bool`): whether use value clip + Returns: + - value_loss (:obj:`torch.FloatTensor`): the ppo value loss item, \ + all of them are the differentiable 0-dim tensor + Shapes: + - value_new (:obj:`torch.FloatTensor`): :math:`(B, )`, where B is batch size + - value_old (:obj:`torch.FloatTensor`): :math:`(B, )` + - return (:obj:`torch.FloatTensor`): :math:`(B, )` + - weight (:obj:`torch.FloatTensor` or :obj:`None`): :math:`(B, )` + - value_loss (:obj:`torch.FloatTensor`): :math:`()`, 0-dim tensor + Examples: + >>> action_dim = 4 + >>> data = happo_value_data( + >>> value_new=torch.randn(3), + >>> value_old=torch.randn(3), + >>> return_=torch.randn(3), + >>> weight=torch.ones(3), + >>> ) + >>> loss, info = happo_value_error(data) + ''' + value_new, value_old, return_, weight = data + if weight is None: + weight = torch.ones_like(value_old) + # value_loss + if use_value_clip: + value_clip = value_old + (value_new - value_old).clamp(-clip_ratio, clip_ratio) + v1 = (return_ - value_new).pow(2) + v2 = (return_ - value_clip).pow(2) + value_loss = 0.5 * (torch.max(v1, v2) * weight).mean() + else: + value_loss = 0.5 * ((return_ - value_new).pow(2) * weight).mean() + return value_loss + + +def happo_error_continuous( + data: namedtuple, + clip_ratio: float = 0.2, + use_value_clip: bool = True, + dual_clip: Optional[float] = None, +) -> Tuple[namedtuple, namedtuple]: + """ + Overview: + Implementation of Proximal Policy Optimization (arXiv:1707.06347) with value_clip and dual_clip + Arguments: + - data (:obj:`namedtuple`): the ppo input data with fieids shown in ``ppo_data`` + - clip_ratio (:obj:`float`): the ppo clip ratio for the constraint of policy update, defaults to 0.2 + - use_value_clip (:obj:`bool`): whether to use clip in value loss with the same ratio as policy + - dual_clip (:obj:`float`): a parameter c mentioned in arXiv:1912.09729 Equ. 5, shoule be in [1, inf),\ + defaults to 5.0, if you don't want to use it, set this parameter to None + Returns: + - happo_loss (:obj:`namedtuple`): the ppo loss item, all of them are the differentiable 0-dim tensor + - happo_info (:obj:`namedtuple`): the ppo optim information for monitoring, all of them are Python scalar + Shapes: + - mu_sigma_new (:obj:`tuple`): :math:`((B, N), (B, N))`, where B is batch size and N is action dim + - mu_sigma_old (:obj:`tuple`): :math:`((B, N), (B, N))`, where B is batch size and N is action dim + - action (:obj:`torch.LongTensor`): :math:`(B, )` + - value_new (:obj:`torch.FloatTensor`): :math:`(B, )` + - value_old (:obj:`torch.FloatTensor`): :math:`(B, )` + - adv (:obj:`torch.FloatTensor`): :math:`(B, )` + - return (:obj:`torch.FloatTensor`): :math:`(B, )` + - weight (:obj:`torch.FloatTensor` or :obj:`None`): :math:`(B, )` + - policy_loss (:obj:`torch.FloatTensor`): :math:`()`, 0-dim tensor + - value_loss (:obj:`torch.FloatTensor`): :math:`()` + - entropy_loss (:obj:`torch.FloatTensor`): :math:`()` + Examples: + >>> action_dim = 4 + >>> data = ppo_data_continuous( + >>> mu_sigma_new= dict(mu=torch.randn(3, action_dim), sigma=torch.randn(3, action_dim)**2), + >>> mu_sigma_old= dict(mu=torch.randn(3, action_dim), sigma=torch.randn(3, action_dim)**2), + >>> action=torch.randn(3, action_dim), + >>> value_new=torch.randn(3), + >>> value_old=torch.randn(3), + >>> adv=torch.randn(3), + >>> return_=torch.randn(3), + >>> weight=torch.ones(3), + >>> ) + >>> loss, info = happo_error(data) + + .. note:: + + adv is already normalized value (adv - adv.mean()) / (adv.std() + 1e-8), and there are many + ways to calculate this mean and std, like among data buffer or train batch, so we don't couple + this part into happo_error, you can refer to our examples for different ways. + """ + assert dual_clip is None or dual_clip > 1.0, "dual_clip value must be greater than 1.0, but get value: {}".format( + dual_clip + ) + mu_sigma_new, mu_sigma_old, action, value_new, value_old, adv, return_, weight, factor_batch = data + if weight is None: + weight = torch.ones_like(adv) + + dist_new = Normal(mu_sigma_new['mu'], mu_sigma_new['sigma']) + if len(mu_sigma_old['mu'].shape) == 1: + dist_old = Normal(mu_sigma_old['mu'].unsqueeze(-1), mu_sigma_old['sigma'].unsqueeze(-1)) + else: + dist_old = Normal(mu_sigma_old['mu'], mu_sigma_old['sigma']) + logp_new = dist_new.log_prob(action) + logp_old = dist_old.log_prob(action) + entropy_loss = (dist_new.entropy() * weight.unsqueeze(1)).mean() + + # policy_loss + ratio = torch.exp(logp_new - logp_old) + ratio = torch.prod(ratio, dim=-1) + surr1 = ratio * adv + surr2 = ratio.clamp(1 - clip_ratio, 1 + clip_ratio) * adv + if dual_clip is not None: + # shape factor: (B,1) surr1: (B,) + policy_loss = (-torch.max(factor_batch.squeeze(1) * torch.min(surr1, surr2), dual_clip * adv) * weight).mean() + else: + policy_loss = (-factor_batch.squeeze(1) * torch.min(surr1, surr2) * weight).mean() + with torch.no_grad(): + approx_kl = (logp_old - logp_new).mean().item() + clipped = ratio.gt(1 + clip_ratio) | ratio.lt(1 - clip_ratio) + clipfrac = torch.as_tensor(clipped).float().mean().item() + # value_loss + if use_value_clip: + value_clip = value_old + (value_new - value_old).clamp(-clip_ratio, clip_ratio) + v1 = (return_ - value_new).pow(2) + v2 = (return_ - value_clip).pow(2) + value_loss = 0.5 * (torch.max(v1, v2) * weight).mean() + else: + value_loss = 0.5 * ((return_ - value_new).pow(2) * weight).mean() + + return happo_loss(policy_loss, value_loss, entropy_loss), happo_info(approx_kl, clipfrac) + + +def happo_policy_error_continuous(data: namedtuple, + clip_ratio: float = 0.2, + dual_clip: Optional[float] = None) -> Tuple[namedtuple, namedtuple]: + """ + Overview: + Implementation of Proximal Policy Optimization (arXiv:1707.06347) with dual_clip + Arguments: + - data (:obj:`namedtuple`): the ppo input data with fieids shown in ``ppo_data`` + - clip_ratio (:obj:`float`): the ppo clip ratio for the constraint of policy update, defaults to 0.2 + - dual_clip (:obj:`float`): a parameter c mentioned in arXiv:1912.09729 Equ. 5, shoule be in [1, inf),\ + defaults to 5.0, if you don't want to use it, set this parameter to None + Returns: + - happo_loss (:obj:`namedtuple`): the ppo loss item, all of them are the differentiable 0-dim tensor + - happo_info (:obj:`namedtuple`): the ppo optim information for monitoring, all of them are Python scalar + Shapes: + - mu_sigma_new (:obj:`tuple`): :math:`((B, N), (B, N))`, where B is batch size and N is action dim + - mu_sigma_old (:obj:`tuple`): :math:`((B, N), (B, N))`, where B is batch size and N is action dim + - action (:obj:`torch.LongTensor`): :math:`(B, )` + - adv (:obj:`torch.FloatTensor`): :math:`(B, )` + - weight (:obj:`torch.FloatTensor` or :obj:`None`): :math:`(B, )` + - policy_loss (:obj:`torch.FloatTensor`): :math:`()`, 0-dim tensor + - entropy_loss (:obj:`torch.FloatTensor`): :math:`()` + Examples: + >>> action_dim = 4 + >>> data = ppo_policy_data_continuous( + >>> mu_sigma_new=dict(mu=torch.randn(3, action_dim), sigma=torch.randn(3, action_dim)**2), + >>> mu_sigma_old=dict(mu=torch.randn(3, action_dim), sigma=torch.randn(3, action_dim)**2), + >>> action=torch.randn(3, action_dim), + >>> adv=torch.randn(3), + >>> weight=torch.ones(3), + >>> ) + >>> loss, info = happo_policy_error_continuous(data) + """ + assert dual_clip is None or dual_clip > 1.0, "dual_clip value must be greater than 1.0, but get value: {}".format( + dual_clip + ) + mu_sigma_new, mu_sigma_old, action, adv, weight = data + if weight is None: + weight = torch.ones_like(adv) + + dist_new = Independent(Normal(mu_sigma_new['mu'], mu_sigma_new['sigma']), 1) + if len(mu_sigma_old['mu'].shape) == 1: + dist_old = Independent(Normal(mu_sigma_old['mu'].unsqueeze(-1), mu_sigma_old['sigma'].unsqueeze(-1)), 1) + else: + dist_old = Independent(Normal(mu_sigma_old['mu'], mu_sigma_old['sigma']), 1) + logp_new = dist_new.log_prob(action) + logp_old = dist_old.log_prob(action) + entropy_loss = (dist_new.entropy() * weight).mean() + # policy_loss + ratio = torch.exp(logp_new - logp_old) + surr1 = ratio * adv + surr2 = ratio.clamp(1 - clip_ratio, 1 + clip_ratio) * adv + if dual_clip is not None: + policy_loss = (-torch.max(torch.min(surr1, surr2), dual_clip * adv) * weight).mean() + else: + policy_loss = (-torch.min(surr1, surr2) * weight).mean() + with torch.no_grad(): + approx_kl = (logp_old - logp_new).mean().item() + clipped = ratio.gt(1 + clip_ratio) | ratio.lt(1 - clip_ratio) + clipfrac = torch.as_tensor(clipped).float().mean().item() + return happo_policy_loss(policy_loss, entropy_loss), happo_info(approx_kl, clipfrac) diff --git a/ding/rl_utils/isw.py b/ding/rl_utils/isw.py index 259db6d1d3..e0745f1031 100644 --- a/ding/rl_utils/isw.py +++ b/ding/rl_utils/isw.py @@ -1,33 +1,59 @@ +from typing import Union import torch +from torch.distributions import Categorical, Independent, Normal -def compute_importance_weights(target_output, behaviour_output, action, requires_grad=False): +def compute_importance_weights( + target_output: Union[torch.Tensor, dict], + behaviour_output: Union[torch.Tensor, dict], + action: torch.Tensor, + action_space_type: str = 'discrete', + requires_grad: bool = False +): """ Overview: Computing importance sampling weight with given output and action Arguments: - - target_output (:obj:`torch.Tensor`): the output taking the action by the current policy network,\ - usually this output is network output logit - - behaviour_output (:obj:`torch.Tensor`): the output taking the action by the behaviour policy network,\ - usually this output is network output logit, which is used to produce the trajectory(collector) + - target_output (:obj:`Union[torch.Tensor,dict]`): the output taking the action \ + by the current policy network, \ + usually this output is network output logit if action space is discrete, \ + or is a dict containing parameters of action distribution if action space is continuous. + - behaviour_output (:obj:`Union[torch.Tensor,dict]`): the output taking the action \ + by the behaviour policy network,\ + usually this output is network output logit, if action space is discrete, \ + or is a dict containing parameters of action distribution if action space is continuous. - action (:obj:`torch.Tensor`): the chosen action(index for the discrete action space) in trajectory,\ i.e.: behaviour_action + - action_space_type (:obj:`str`): action space types in ['discrete', 'continuous'] - requires_grad (:obj:`bool`): whether requires grad computation Returns: - rhos (:obj:`torch.Tensor`): Importance sampling weight Shapes: - - target_output (:obj:`torch.FloatTensor`): :math:`(T, B, N)`, where T is timestep, B is batch size and\ - N is action dim - - behaviour_output (:obj:`torch.FloatTensor`): :math:`(T, B, N)` + - target_output (:obj:`Union[torch.FloatTensor,dict]`): :math:`(T, B, N)`, \ + where T is timestep, B is batch size and N is action dim + - behaviour_output (:obj:`Union[torch.FloatTensor,dict]`): :math:`(T, B, N)` - action (:obj:`torch.LongTensor`): :math:`(T, B)` - rhos (:obj:`torch.FloatTensor`): :math:`(T, B)` + Examples: + >>> target_output = torch.randn(2, 3, 4) + >>> behaviour_output = torch.randn(2, 3, 4) + >>> action = torch.randint(0, 4, (2, 3)) + >>> rhos = compute_importance_weights(target_output, behaviour_output, action) """ grad_context = torch.enable_grad() if requires_grad else torch.no_grad() assert isinstance(action, torch.Tensor) + assert action_space_type in ['discrete', 'continuous'] with grad_context: - dist_target = torch.distributions.Categorical(logits=target_output) - dist_behaviour = torch.distributions.Categorical(logits=behaviour_output) - rhos = dist_target.log_prob(action) - dist_behaviour.log_prob(action) - rhos = torch.exp(rhos) - return rhos + if action_space_type == 'continuous': + dist_target = Independent(Normal(loc=target_output['mu'], scale=target_output['sigma']), 1) + dist_behaviour = Independent(Normal(loc=behaviour_output['mu'], scale=behaviour_output['sigma']), 1) + rhos = dist_target.log_prob(action) - dist_behaviour.log_prob(action) + rhos = torch.exp(rhos) + return rhos + elif action_space_type == 'discrete': + dist_target = Categorical(logits=target_output) + dist_behaviour = Categorical(logits=behaviour_output) + rhos = dist_target.log_prob(action) - dist_behaviour.log_prob(action) + rhos = torch.exp(rhos) + return rhos diff --git a/ding/rl_utils/log_prob_utils.py b/ding/rl_utils/log_prob_utils.py new file mode 100644 index 0000000000..f4b35af563 --- /dev/null +++ b/ding/rl_utils/log_prob_utils.py @@ -0,0 +1,87 @@ +from typing import List, Callable, Optional, Any +import torch +from torch import Tensor + +LogitsProcessor = Callable[[Tensor, Tensor], Tensor] + + +def naive_method(logits: Tensor, index: Tensor) -> Tensor: + """Calculate per-token log probabilities using naive method. + + Args: + logits: Token logits of shape [B, S, V] or [S, V] where: + B = batch size + S = sequence length + V = vocabulary size + index: Selected token indices of shape [B, S] or [S] + + Returns: + Tensor: Log probabilities for selected tokens of shape [B, S] or [S] + """ + # Calculate log probabilities for each token + log_prob_new: Tensor = torch.log_softmax(logits, dim=-1) + # Get log probabilities for selected actions + index = index.unsqueeze(-1) # [B, S, 1] or [S, 1] + per_token_logps: Tensor = torch.gather(log_prob_new, -1, index).squeeze(-1) + return per_token_logps + + +def efficient_method(logits: Tensor, index: Tensor) -> Tensor: + """Calculate per-token log probabilities efficiently. + + Args: + logits: Token logits of shape [B, S, V] or [S, V] where: + B = batch size + S = sequence length + V = vocabulary size + index: Selected token indices of shape [B, S] or [S] + + Returns: + Tensor: Log probabilities for selected tokens of shape [B, S] or [S] + """ + if logits.dtype in [torch.float32, torch.float64]: + selected_logits: Tensor = torch.gather(logits, dim=-1, index=index.unsqueeze(-1)).squeeze(-1) + + # Loop to reduce peak mem consumption + logsumexp_values: Tensor = torch.stack([torch.logsumexp(lg, dim=-1) for lg in logits]) + + # log_softmax(x_i) = x_i - logsumexp(x) + per_token_logps: Tensor = selected_logits - logsumexp_values + else: + # logsumexp approach is unstable with bfloat16 + per_token_logps: List[Tensor] = [] + + # Loop to reduce peak mem consumption + for row_logits, row_labels in zip(logits, index): # Iterate over sequence length + row_logps: Tensor = torch.log_softmax(row_logits, dim=-1) + row_per_token_logps: Tensor = row_logps.gather(dim=-1, index=row_labels.unsqueeze(-1)).squeeze(-1) + per_token_logps.append(row_per_token_logps) + + per_token_logps = torch.stack(per_token_logps) + + return per_token_logps + + +def less_efficient_method(logits: Tensor, index: Tensor) -> Tensor: + """Calculate per-token log probabilities using categorical distribution. + + Args: + logits: Token logits of shape [B, S, V] or [S, V] where: + B = batch size + S = sequence length + V = vocabulary size + index: Selected token indices of shape [B, S] or [S] + + Returns: + Tensor: Log probabilities for selected tokens of shape [B, S] or [S] + """ + dist = torch.distributions.categorical.Categorical(logits=logits) + logp: Tensor = dist.log_prob(index) + return logp + + +# 定义一个统一的类型 +LogProbFunction = Callable[[Tensor, Tensor], Tensor] + +# 导出所有方法 +__all__ = ['naive_method', 'efficient_method', 'less_efficient_method', 'LogProbFunction'] diff --git a/ding/rl_utils/ppg.py b/ding/rl_utils/ppg.py index 5bd307584a..286266e57b 100644 --- a/ding/rl_utils/ppg.py +++ b/ding/rl_utils/ppg.py @@ -12,6 +12,38 @@ def ppg_joint_error( clip_ratio: float = 0.2, use_value_clip: bool = True, ) -> Tuple[namedtuple, namedtuple]: + ''' + Overview: + Get PPG joint loss + Arguments: + - data (:obj:`namedtuple`): ppg input data with fieids shown in ``ppg_data`` + - clip_ratio (:obj:`float`): clip value for ratio + - use_value_clip (:obj:`bool`): whether use value clip + Returns: + - ppg_joint_loss (:obj:`namedtuple`): the ppg loss item, all of them are the differentiable 0-dim tensor + Shapes: + - logit_new (:obj:`torch.FloatTensor`): :math:`(B, N)`, where B is batch size and N is action dim + - logit_old (:obj:`torch.FloatTensor`): :math:`(B, N)` + - action (:obj:`torch.LongTensor`): :math:`(B,)` + - value_new (:obj:`torch.FloatTensor`): :math:`(B, 1)` + - value_old (:obj:`torch.FloatTensor`): :math:`(B, 1)` + - return (:obj:`torch.FloatTensor`): :math:`(B, 1)` + - weight (:obj:`torch.FloatTensor`): :math:`(B,)` + - auxiliary_loss (:obj:`torch.FloatTensor`): :math:`()`, 0-dim tensor + - behavioral_cloning_loss (:obj:`torch.FloatTensor`): :math:`()` + Examples: + >>> action_dim = 4 + >>> data = ppg_data( + >>> logit_new=torch.randn(3, action_dim), + >>> logit_old=torch.randn(3, action_dim), + >>> action=torch.randint(0, action_dim, (3,)), + >>> value_new=torch.randn(3, 1), + >>> value_old=torch.randn(3, 1), + >>> return_=torch.randn(3, 1), + >>> weight=torch.ones(3), + >>> ) + >>> loss = ppg_joint_error(data, 0.99, 0.99) + ''' logit_new, logit_old, action, value_new, value_old, return_, weight = data if weight is None: diff --git a/ding/rl_utils/ppo.py b/ding/rl_utils/ppo.py index 5b62d5a489..6fbe287d09 100644 --- a/ding/rl_utils/ppo.py +++ b/ding/rl_utils/ppo.py @@ -1,19 +1,59 @@ from collections import namedtuple from typing import Optional, Tuple import torch +import torch.nn as nn from torch.distributions import Independent, Normal from ding.hpc_rl import hpc_wrapper ppo_data = namedtuple( - 'ppo_data', ['logit_new', 'logit_old', 'action', 'value_new', 'value_old', 'adv', 'return_', 'weight'] + 'ppo_data', + ['logit_new', 'logit_old', 'action', 'value_new', 'value_old', 'adv', 'return_', 'weight', 'logit_pretrained'] +) +ppo_data_continuous = namedtuple( + 'ppo_data_continuous', [ + 'mu_sigma_new', 'mu_sigma_old', 'action', 'value_new', 'value_old', 'adv', 'return_', 'weight', + 'logit_pretrained' + ] +) +ppo_policy_data = namedtuple( + 'ppo_policy_data', ['logit_new', 'logit_old', 'action', 'adv', 'weight', 'logit_pretrained'] +) +ppo_policy_data_continuous = namedtuple( + 'ppo_policy_data_continuous', ['mu_sigma_new', 'mu_sigma_old', 'action', 'adv', 'weight', 'logit_pretrained'] ) -ppo_policy_data = namedtuple('ppo_policy_data', ['logit_new', 'logit_old', 'action', 'adv', 'weight']) ppo_value_data = namedtuple('ppo_value_data', ['value_new', 'value_old', 'return_', 'weight']) -ppo_loss = namedtuple('ppo_loss', ['policy_loss', 'value_loss', 'entropy_loss']) -ppo_policy_loss = namedtuple('ppo_policy_loss', ['policy_loss', 'entropy_loss']) +ppo_loss = namedtuple('ppo_loss', ['policy_loss', 'value_loss', 'entropy_loss', 'kl_div']) +ppo_policy_loss = namedtuple('ppo_policy_loss', ['policy_loss', 'entropy_loss', 'kl_div']) ppo_info = namedtuple('ppo_info', ['approx_kl', 'clipfrac']) +def calculate_kl_div(log_ratio: torch.Tensor, kl_type: str) -> torch.Tensor: + """ + Overview: + Calculate different Monte-Carlo estimators for KL-divergence KL(q, p) = E_q[log(q/p)], + where q is the current policy and p is the pretrained policy. + The implementation is based on John Schulman's blog post "Approximating KL Divergence". + Reference: http://joschu.net/blog/kl-approx.html + Arguments: + - log_ratio (:obj:`torch.Tensor`): The log-ratio of probabilities, which should be + log(q/p) = logp_new - logp_pretrained. + - kl_type (:obj:`str`): The type of KL divergence estimator to use. + - 'k1': The standard, unbiased but high-variance estimator: `E_q[log(q/p)]`. + - 'k2': A biased, low-variance estimator from a second-order approximation: `E_q[1/2 * (log(p/q))^2]`. + - 'k3': An unbiased, low-variance estimator: `E_q[(p/q - 1) - log(p/q)]`. + Returns: + - kl_div (:obj:`torch.Tensor`): The calculated KL divergence estimate. + """ + if kl_type == 'k1': + return log_ratio.mean() + elif kl_type == 'k2': + return (log_ratio ** 2 / 2).mean() + elif kl_type == 'k3': + return (torch.exp(-log_ratio) - 1 + log_ratio).mean() + else: + raise ValueError(f"Unknown kl_type: {kl_type}") + + def shape_fn_ppo(args, kwargs): r""" Overview: @@ -38,7 +78,8 @@ def ppo_error( data: namedtuple, clip_ratio: float = 0.2, use_value_clip: bool = True, - dual_clip: Optional[float] = None + dual_clip: Optional[float] = None, + kl_type: str = 'k1' ) -> Tuple[namedtuple, namedtuple]: """ Overview: @@ -49,6 +90,7 @@ def ppo_error( - use_value_clip (:obj:`bool`): whether to use clip in value loss with the same ratio as policy - dual_clip (:obj:`float`): a parameter c mentioned in arXiv:1912.09729 Equ. 5, shoule be in [1, inf),\ defaults to 5.0, if you don't want to use it, set this parameter to None + - kl_type (:obj:`str`): which kl loss to use, default set to 'k1'. Returns: - ppo_loss (:obj:`namedtuple`): the ppo loss item, all of them are the differentiable 0-dim tensor - ppo_info (:obj:`namedtuple`): the ppo optim information for monitoring, all of them are Python scalar @@ -63,6 +105,20 @@ def ppo_error( - weight (:obj:`torch.FloatTensor` or :obj:`None`): :math:`(B, )` - policy_loss (:obj:`torch.FloatTensor`): :math:`()`, 0-dim tensor - value_loss (:obj:`torch.FloatTensor`): :math:`()` + - entropy_loss (:obj:`torch.FloatTensor`): :math:`()` + Examples: + >>> action_dim = 4 + >>> data = ppo_data( + >>> logit_new=torch.randn(3, action_dim), + >>> logit_old=torch.randn(3, action_dim), + >>> action=torch.randint(0, action_dim, (3,)), + >>> value_new=torch.randn(3), + >>> value_old=torch.randn(3), + >>> adv=torch.randn(3), + >>> return_=torch.randn(3), + >>> weight=torch.ones(3), + >>> ) + >>> loss, info = ppo_error(data) .. note:: @@ -73,29 +129,78 @@ def ppo_error( assert dual_clip is None or dual_clip > 1.0, "dual_clip value must be greater than 1.0, but get value: {}".format( dual_clip ) - logit_new, logit_old, action, value_new, value_old, adv, return_, weight = data - policy_data = ppo_policy_data(logit_new, logit_old, action, adv, weight) - policy_output, policy_info = ppo_policy_error(policy_data, clip_ratio, dual_clip) + logit_new, logit_old, action, value_new, value_old, adv, return_, weight, logit_pretrained = data + policy_data = ppo_policy_data(logit_new, logit_old, action, adv, weight, logit_pretrained) + policy_output, policy_info = ppo_policy_error(policy_data, clip_ratio, dual_clip, kl_type=kl_type) value_data = ppo_value_data(value_new, value_old, return_, weight) value_loss = ppo_value_error(value_data, clip_ratio, use_value_clip) - return ppo_loss(policy_output.policy_loss, value_loss, policy_output.entropy_loss), policy_info + return ppo_loss( + policy_output.policy_loss, value_loss, policy_output.entropy_loss, policy_output.kl_div + ), policy_info -def ppo_policy_error(data: namedtuple, - clip_ratio: float = 0.2, - dual_clip: Optional[float] = None) -> Tuple[namedtuple, namedtuple]: - logit_new, logit_old, action, adv, weight = data +def ppo_policy_error( + data: namedtuple, + clip_ratio: float = 0.2, + dual_clip: Optional[float] = None, + entropy_bonus: bool = True, + kl_type: str = 'k1' +) -> Tuple[namedtuple, namedtuple]: + """ + Overview: + Get PPO policy loss (both for classical RL in control/video games and LLM/VLM RLHF). + Arguments: + - data (:obj:`namedtuple`): Ppo input data with fieids shown in ``ppo_policy_data``. + - clip_ratio (:obj:`float`): Clip value for ratio, defaults to 0.2. + - dual_clip (:obj:`float`): A parameter c mentioned in arXiv:1912.09729 Equ. 5, shoule be in [1, inf), \ + defaults to 5.0, if you don't want to use it, set this parameter to None + - entropy_bonus (:obj:`bool`): Whether to use entropy bonus, defaults to True. LLM RLHF usually does not use it. + - kl_type (:obj:`str`): which kl loss to use, default set to 'k1'. + Returns: + - ppo_policy_loss (:obj:`namedtuple`): the ppo policy loss item, all of them are the differentiable 0-dim tensor + - ppo_info (:obj:`namedtuple`): the ppo optim information for monitoring, all of them are Python scalar + Shapes: + - logit_new (:obj:`torch.FloatTensor`): :math:`(B, N)`, where B is batch size and N is action dim + - logit_old (:obj:`torch.FloatTensor`): :math:`(B, N)` + - action (:obj:`torch.LongTensor`): :math:`(B, )` + - adv (:obj:`torch.FloatTensor`): :math:`(B, )` + - weight (:obj:`torch.FloatTensor` or :obj:`None`): :math:`(B, )` + - policy_loss (:obj:`torch.FloatTensor`): :math:`()`, 0-dim tensor + - entropy_loss (:obj:`torch.FloatTensor`): :math:`()` + Examples: + >>> action_dim = 4 + >>> data = ppo_policy_data( + >>> logit_new=torch.randn(3, action_dim), + >>> logit_old=torch.randn(3, action_dim), + >>> action=torch.randint(0, action_dim, (3,)), + >>> adv=torch.randn(3), + >>> weight=torch.ones(3), + >>> ) + >>> loss, info = ppo_policy_error(data) + + .. note:: + This function can be extended from `B` to more parallel dimensions, like `(B, S)`, where `S` is the + sequence length in LLM/VLM. + + .. note:: + For the action mask often used in LLM/VLM, users can set the `weight` to the action mask. + """ + logit_new, logit_old, action, adv, weight, logit_pretrained = data if weight is None: weight = torch.ones_like(adv) dist_new = torch.distributions.categorical.Categorical(logits=logit_new) dist_old = torch.distributions.categorical.Categorical(logits=logit_old) logp_new = dist_new.log_prob(action) logp_old = dist_old.log_prob(action) - dist_new_entropy = dist_new.entropy() - if dist_new_entropy.shape != weight.shape: - dist_new_entropy = dist_new.entropy().mean(dim=1) - entropy_loss = (dist_new_entropy * weight).mean() + + if entropy_bonus: + dist_new_entropy = dist_new.entropy() + if dist_new_entropy.shape != weight.shape: # for the multi-agent rl case + dist_new_entropy = dist_new.entropy().mean(dim=1) + entropy_loss = (dist_new_entropy * weight).mean() + else: + entropy_loss = torch.tensor(0.0) # policy_loss ratio = torch.exp(logp_new - logp_old) if ratio.shape != adv.shape: @@ -113,7 +218,16 @@ def ppo_policy_error(data: namedtuple, approx_kl = (logp_old - logp_new).mean().item() clipped = ratio.gt(1 + clip_ratio) | ratio.lt(1 - clip_ratio) clipfrac = torch.as_tensor(clipped).float().mean().item() - return ppo_policy_loss(policy_loss, entropy_loss), ppo_info(approx_kl, clipfrac) + + if logit_pretrained is not None: + dist_pretrained = torch.distributions.categorical.Categorical(logits=logit_pretrained) + logp_pretrained = dist_pretrained.log_prob(action) + log_ratio = logp_new - logp_pretrained + kl_div = calculate_kl_div(log_ratio, kl_type) + else: + kl_div = torch.tensor(0., dtype=policy_loss.dtype, device=policy_loss.device) + + return ppo_policy_loss(policy_loss, entropy_loss, kl_div), ppo_info(approx_kl, clipfrac) def ppo_value_error( @@ -121,6 +235,32 @@ def ppo_value_error( clip_ratio: float = 0.2, use_value_clip: bool = True, ) -> torch.Tensor: + ''' + Overview: + Get PPO value loss + Arguments: + - data (:obj:`namedtuple`): ppo input data with fieids shown in ``ppo_value_data`` + - clip_ratio (:obj:`float`): clip value for ratio + - use_value_clip (:obj:`bool`): whether use value clip + Returns: + - value_loss (:obj:`torch.FloatTensor`): the ppo value loss item, \ + all of them are the differentiable 0-dim tensor + Shapes: + - value_new (:obj:`torch.FloatTensor`): :math:`(B, )`, where B is batch size + - value_old (:obj:`torch.FloatTensor`): :math:`(B, )` + - return (:obj:`torch.FloatTensor`): :math:`(B, )` + - weight (:obj:`torch.FloatTensor` or :obj:`None`): :math:`(B, )` + - value_loss (:obj:`torch.FloatTensor`): :math:`()`, 0-dim tensor + Examples: + >>> action_dim = 4 + >>> data = ppo_value_data( + >>> value_new=torch.randn(3), + >>> value_old=torch.randn(3), + >>> return_=torch.randn(3), + >>> weight=torch.ones(3), + >>> ) + >>> loss, info = ppo_value_error(data) + ''' value_new, value_old, return_, weight = data if weight is None: weight = torch.ones_like(value_old) @@ -139,7 +279,8 @@ def ppo_error_continuous( data: namedtuple, clip_ratio: float = 0.2, use_value_clip: bool = True, - dual_clip: Optional[float] = None + dual_clip: Optional[float] = None, + kl_type: str = 'k1' ) -> Tuple[namedtuple, namedtuple]: """ Overview: @@ -150,6 +291,7 @@ def ppo_error_continuous( - use_value_clip (:obj:`bool`): whether to use clip in value loss with the same ratio as policy - dual_clip (:obj:`float`): a parameter c mentioned in arXiv:1912.09729 Equ. 5, shoule be in [1, inf),\ defaults to 5.0, if you don't want to use it, set this parameter to None + - kl_type (:obj:`str`): which kl loss to use, default set to 'k1'. Returns: - ppo_loss (:obj:`namedtuple`): the ppo loss item, all of them are the differentiable 0-dim tensor - ppo_info (:obj:`namedtuple`): the ppo optim information for monitoring, all of them are Python scalar @@ -164,6 +306,20 @@ def ppo_error_continuous( - weight (:obj:`torch.FloatTensor` or :obj:`None`): :math:`(B, )` - policy_loss (:obj:`torch.FloatTensor`): :math:`()`, 0-dim tensor - value_loss (:obj:`torch.FloatTensor`): :math:`()` + - entropy_loss (:obj:`torch.FloatTensor`): :math:`()` + Examples: + >>> action_dim = 4 + >>> data = ppo_data_continuous( + >>> mu_sigma_new= dict(mu=torch.randn(3, action_dim), sigma=torch.randn(3, action_dim)**2), + >>> mu_sigma_old= dict(mu=torch.randn(3, action_dim), sigma=torch.randn(3, action_dim)**2), + >>> action=torch.randn(3, action_dim), + >>> value_new=torch.randn(3), + >>> value_old=torch.randn(3), + >>> adv=torch.randn(3), + >>> return_=torch.randn(3), + >>> weight=torch.ones(3), + >>> ) + >>> loss, info = ppo_error(data) .. note:: @@ -174,7 +330,7 @@ def ppo_error_continuous( assert dual_clip is None or dual_clip > 1.0, "dual_clip value must be greater than 1.0, but get value: {}".format( dual_clip ) - mu_sigma_new, mu_sigma_old, action, value_new, value_old, adv, return_, weight = data + mu_sigma_new, mu_sigma_old, action, value_new, value_old, adv, return_, weight, logit_pretrained = data if weight is None: weight = torch.ones_like(adv) @@ -207,12 +363,23 @@ def ppo_error_continuous( else: value_loss = 0.5 * ((return_ - value_new).pow(2) * weight).mean() - return ppo_loss(policy_loss, value_loss, entropy_loss), ppo_info(approx_kl, clipfrac) + if logit_pretrained is not None: + dist_pretrained = Independent(Normal(logit_pretrained['mu'], logit_pretrained['sigma']), 1) + logp_pretrained = dist_pretrained.log_prob(action) + log_ratio = logp_new - logp_pretrained + kl_div = calculate_kl_div(log_ratio, kl_type) + else: + kl_div = torch.tensor(0., dtype=policy_loss.dtype, device=policy_loss.device) + + return ppo_loss(policy_loss, value_loss, entropy_loss, kl_div), ppo_info(approx_kl, clipfrac) -def ppo_policy_error_continuous(data: namedtuple, - clip_ratio: float = 0.2, - dual_clip: Optional[float] = None) -> Tuple[namedtuple, namedtuple]: +def ppo_policy_error_continuous( + data: namedtuple, + clip_ratio: float = 0.2, + dual_clip: Optional[float] = None, + kl_type: str = 'k1' +) -> Tuple[namedtuple, namedtuple]: """ Overview: Implementation of Proximal Policy Optimization (arXiv:1707.06347) with dual_clip @@ -221,6 +388,7 @@ def ppo_policy_error_continuous(data: namedtuple, - clip_ratio (:obj:`float`): the ppo clip ratio for the constraint of policy update, defaults to 0.2 - dual_clip (:obj:`float`): a parameter c mentioned in arXiv:1912.09729 Equ. 5, shoule be in [1, inf),\ defaults to 5.0, if you don't want to use it, set this parameter to None + - kl_type (:obj:`str`): which kl loss to use, default set to 'k1'. Returns: - ppo_loss (:obj:`namedtuple`): the ppo loss item, all of them are the differentiable 0-dim tensor - ppo_info (:obj:`namedtuple`): the ppo optim information for monitoring, all of them are Python scalar @@ -231,11 +399,22 @@ def ppo_policy_error_continuous(data: namedtuple, - adv (:obj:`torch.FloatTensor`): :math:`(B, )` - weight (:obj:`torch.FloatTensor` or :obj:`None`): :math:`(B, )` - policy_loss (:obj:`torch.FloatTensor`): :math:`()`, 0-dim tensor + - entropy_loss (:obj:`torch.FloatTensor`): :math:`()` + Examples: + >>> action_dim = 4 + >>> data = ppo_policy_data_continuous( + >>> mu_sigma_new=dict(mu=torch.randn(3, action_dim), sigma=torch.randn(3, action_dim)**2), + >>> mu_sigma_old=dict(mu=torch.randn(3, action_dim), sigma=torch.randn(3, action_dim)**2), + >>> action=torch.randn(3, action_dim), + >>> adv=torch.randn(3), + >>> weight=torch.ones(3), + >>> ) + >>> loss, info = ppo_policy_error_continuous(data) """ assert dual_clip is None or dual_clip > 1.0, "dual_clip value must be greater than 1.0, but get value: {}".format( dual_clip ) - mu_sigma_new, mu_sigma_old, action, adv, weight = data + mu_sigma_new, mu_sigma_old, action, adv, weight, logit_pretrained = data if weight is None: weight = torch.ones_like(adv) @@ -259,4 +438,13 @@ def ppo_policy_error_continuous(data: namedtuple, approx_kl = (logp_old - logp_new).mean().item() clipped = ratio.gt(1 + clip_ratio) | ratio.lt(1 - clip_ratio) clipfrac = torch.as_tensor(clipped).float().mean().item() - return ppo_policy_loss(policy_loss, entropy_loss), ppo_info(approx_kl, clipfrac) + + if logit_pretrained is not None: + dist_pretrained = Independent(Normal(logit_pretrained['mu'], logit_pretrained['sigma']), 1) + logp_pretrained = dist_pretrained.log_prob(action) + log_ratio = logp_new - logp_pretrained + kl_div = calculate_kl_div(log_ratio, kl_type) + else: + kl_div = 0 + + return ppo_policy_loss(policy_loss, entropy_loss, kl_div), ppo_info(approx_kl, clipfrac) diff --git a/ding/rl_utils/retrace.py b/ding/rl_utils/retrace.py index 3223be8ffa..01b1f3f0f3 100644 --- a/ding/rl_utils/retrace.py +++ b/ding/rl_utils/retrace.py @@ -1,7 +1,7 @@ import torch import torch.nn.functional as F from collections import namedtuple -from .isw import compute_importance_weights +from ding.rl_utils.isw import compute_importance_weights def compute_q_retraces( @@ -23,6 +23,17 @@ def compute_q_retraces( - weights (:obj:`torch.Tensor`): :math:`(T, B)` - ratio (:obj:`torch.Tensor`): :math:`(T, B, N)` - q_retraces (:obj:`torch.Tensor`): :math:`(T + 1, B, 1)` + Examples: + >>> T=2 + >>> B=3 + >>> N=4 + >>> q_values=torch.randn(T+1, B, N) + >>> v_pred=torch.randn(T+1, B, 1) + >>> rewards=torch.randn(T, B) + >>> actions=torch.randint(0, N, (T, B)) + >>> weights=torch.ones(T, B) + >>> ratio=torch.randn(T, B, N) + >>> q_retraces = compute_q_retraces(q_values, v_pred, rewards, actions, weights, ratio) .. note:: q_retrace operation doesn't need to compute gradient, just executes forward computation. diff --git a/ding/rl_utils/rloo.py b/ding/rl_utils/rloo.py new file mode 100644 index 0000000000..108019419e --- /dev/null +++ b/ding/rl_utils/rloo.py @@ -0,0 +1,69 @@ +from typing import Tuple +from collections import namedtuple +import torch +from .log_prob_utils import efficient_method, naive_method, less_efficient_method, LogProbFunction + +rloo_policy_data = namedtuple('rloo_policy_data', ['logit_new', 'logit_old', 'action', 'reward', 'weight']) +rloo_info = namedtuple('rloo_info', ['approx_kl', 'clipfrac']) + + +def rloo_policy_error( + data: namedtuple, + log_prob_fn: LogProbFunction = efficient_method, # Method to calculate the log probabilities + clip_ratio: float = 0.2, +) -> Tuple[namedtuple, namedtuple]: + """ + Overview: + REINFORCE Leave-One-Out (RLOO) algorithm, see https://arxiv.org/abs/2402.14740. + Arguments: + - data (:obj:`namedtuple`): the rloo input data with fields shown in ``rloo_policy_data``. + - clip_ratio (:obj:`float`): the ppo clip ratio for the constraint of policy update, defaults to 0.2. + - log_prob_fn (:obj:`LogProbFunction`): The method to calculate the log probabilities, \ + defaults to `efficient_method`. + Returns: + - loss (:obj:`torch.FloatTensor`): the rloo policy loss, a differentiable 0-dim tensor. + - rloo_info (:obj:`namedtuple`): the rloo optim information for monitoring, all of them are Python scalar. + Shapes: + - logit_new (:obj:`torch.FloatTensor`): :math:`(B, S, V)`, where B is batch size, S is sequence length,\ + and V is vocabulary size. + - logit_old (:obj:`torch.FloatTensor`): :math:`(B, S, V)`. + - action (:obj:`torch.LongTensor`): :math:`(B, S)`. + - reward (:obj:`torch.FloatTensor`): :math:`(K, B)`, where K is the number of samples per prompt. + - weight (:obj:`torch.FloatTensor` or :obj:`None`): :math:`(B, S)`. + - policy_loss (:obj:`torch.FloatTensor`): :math:`()`, 0-dim tensor. + - mean_ratio (:obj:`float`): mean probability ratio. + - mean_clipped (:obj:`float`): proportion of clipped probability ratios. + - mean_advantage (:obj:`float`): mean advantage value. + """ + + # Calculate advantage of each action + rloo_k = data.reward.size(0) + baseline = (data.reward.sum(0) - data.reward) / (rloo_k - 1) + adv = data.reward - baseline + adv = adv.flatten() + + # Get log probabilities for selected actions + per_token_logps = log_prob_fn(data.logit_new, data.action) + per_token_old_logps = log_prob_fn(data.logit_old, data.action) + + # Calculate policy ratio + ratio = torch.exp(per_token_logps - per_token_old_logps) + ratio_clipped = torch.clamp(ratio, 1 - clip_ratio, 1 + clip_ratio) + + # Calculate loss for each token + advantages = adv.unsqueeze(1) # [B, 1] + per_token_loss_unclipped = ratio * advantages + per_token_loss_clipped = ratio_clipped * advantages + per_token_loss = -torch.min(per_token_loss_unclipped, per_token_loss_clipped) + + # Calculate average loss using weight mask + weight = data.weight if data.weight is not None else (torch.ones_like(per_token_loss)) + loss = ((per_token_loss * weight).sum(dim=1) / weight.sum(dim=1)).mean() + + # Calculate additional metrics + with torch.no_grad(): + approx_kl = (per_token_old_logps - per_token_logps).mean().item() + clipped = ratio.gt(1 + clip_ratio) | ratio.lt(1 - clip_ratio) + clipfrac = torch.as_tensor(clipped).float().mean().item() + + return loss, rloo_info(approx_kl=approx_kl, clipfrac=clipfrac) diff --git a/ding/rl_utils/sampler.py b/ding/rl_utils/sampler.py new file mode 100644 index 0000000000..8afbd605bb --- /dev/null +++ b/ding/rl_utils/sampler.py @@ -0,0 +1,127 @@ +import torch +import treetensor.torch as ttorch +from torch.distributions import Normal, Independent + + +class ArgmaxSampler: + ''' + Overview: + Argmax sampler, return the index of the maximum value + ''' + + def __call__(self, logit: torch.Tensor) -> torch.Tensor: + ''' + Overview: + Return the index of the maximum value + Arguments: + - logit (:obj:`torch.Tensor`): The input tensor + Returns: + - action (:obj:`torch.Tensor`): The index of the maximum value + ''' + return logit.argmax(dim=-1) + + +class MultinomialSampler: + ''' + Overview: + Multinomial sampler, return the index of the sampled value + ''' + + def __call__(self, logit: torch.Tensor) -> torch.Tensor: + ''' + Overview: + Return the index of the sampled value + Arguments: + - logit (:obj:`torch.Tensor`): The input tensor + Returns: + - action (:obj:`torch.Tensor`): The index of the sampled value + ''' + dist = torch.distributions.Categorical(logits=logit) + return dist.sample() + + +class MuSampler: + ''' + Overview: + Mu sampler, return the mu of the input tensor + ''' + + def __call__(self, logit: ttorch.Tensor) -> torch.Tensor: + ''' + Overview: + Return the mu of the input tensor + Arguments: + - logit (:obj:`ttorch.Tensor`): The input tensor + Returns: + - action (:obj:`torch.Tensor`): The mu of the input tensor + ''' + return logit.mu + + +class ReparameterizationSampler: + ''' + Overview: + Reparameterization sampler, return the reparameterized value of the input tensor + ''' + + def __call__(self, logit: ttorch.Tensor) -> torch.Tensor: + ''' + Overview: + Return the reparameterized value of the input tensor + Arguments: + - logit (:obj:`ttorch.Tensor`): The input tensor + Returns: + - action (:obj:`torch.Tensor`): The reparameterized value of the input tensor + ''' + dist = Normal(logit.mu, logit.sigma) + dist = Independent(dist, 1) + return dist.rsample() + + +class HybridStochasticSampler: + ''' + Overview: + Hybrid stochastic sampler, return the sampled action type and the reparameterized action args + ''' + + def __call__(self, logit: ttorch.Tensor) -> ttorch.Tensor: + ''' + Overview: + Return the sampled action type and the reparameterized action args + Arguments: + - logit (:obj:`ttorch.Tensor`): The input tensor + Returns: + - action (:obj:`ttorch.Tensor`): The sampled action type and the reparameterized action args + ''' + dist = torch.distributions.Categorical(logits=logit.action_type) + action_type = dist.sample() + dist = Normal(logit.action_args.mu, logit.action_args.sigma) + dist = Independent(dist, 1) + action_args = dist.rsample() + return ttorch.as_tensor({ + 'action_type': action_type, + 'action_args': action_args, + }) + + +class HybridDeterminsticSampler: + ''' + Overview: + Hybrid deterministic sampler, return the argmax action type and the mu action args + ''' + + def __call__(self, logit: ttorch.Tensor) -> ttorch.Tensor: + ''' + Overview: + Return the argmax action type and the mu action args + Arguments: + - logit (:obj:`ttorch.Tensor`): The input tensor + Returns: + - action (:obj:`ttorch.Tensor`): The argmax action type and the mu action args + ''' + action_type = logit.action_type.argmax(dim=-1) + action_args = logit.action_args.mu + return ttorch.as_tensor({ + 'action_type': action_type, + 'action_args': action_args, + }) diff --git a/ding/rl_utils/td.py b/ding/rl_utils/td.py index 3b01115dbe..a3f32589c4 100644 --- a/ding/rl_utils/td.py +++ b/ding/rl_utils/td.py @@ -28,6 +28,38 @@ def q_1step_td_error( gamma: float, criterion: torch.nn.modules = nn.MSELoss(reduction='none') # noqa ) -> torch.Tensor: + """ + Overview: + 1 step td_error, support single agent case and multi agent case. + Arguments: + - data (:obj:`q_1step_td_data`): The input data, q_1step_td_data to calculate loss + - gamma (:obj:`float`): Discount factor + - criterion (:obj:`torch.nn.modules`): Loss function criterion + Returns: + - loss (:obj:`torch.Tensor`): 1step td error + Shapes: + - data (:obj:`q_1step_td_data`): the q_1step_td_data containing\ + ['q', 'next_q', 'act', 'next_act', 'reward', 'done', 'weight'] + - q (:obj:`torch.FloatTensor`): :math:`(B, N)` i.e. [batch_size, action_dim] + - next_q (:obj:`torch.FloatTensor`): :math:`(B, N)` i.e. [batch_size, action_dim] + - act (:obj:`torch.LongTensor`): :math:`(B, )` + - next_act (:obj:`torch.LongTensor`): :math:`(B, )` + - reward (:obj:`torch.FloatTensor`): :math:`( , B)` + - done (:obj:`torch.BoolTensor`) :math:`(B, )`, whether done in last timestep + - weight (:obj:`torch.FloatTensor` or None): :math:`(B, )`, the training sample weight + Examples: + >>> action_dim = 4 + >>> data = q_1step_td_data( + >>> q=torch.randn(3, action_dim), + >>> next_q=torch.randn(3, action_dim), + >>> act=torch.randint(0, action_dim, (3,)), + >>> next_act=torch.randint(0, action_dim, (3,)), + >>> reward=torch.randn(3), + >>> done=torch.randint(0, 2, (3,)).bool(), + >>> weight=torch.ones(3), + >>> ) + >>> loss = q_1step_td_error(data, 0.99) + """ q, next_q, act, next_act, reward, done, weight = data assert len(act.shape) == 1, act.shape assert len(reward.shape) == 1, reward.shape @@ -40,6 +72,92 @@ def q_1step_td_error( return (criterion(q_s_a, target_q_s_a.detach()) * weight).mean() +m_q_1step_td_data = namedtuple('m_q_1step_td_data', ['q', 'target_q', 'next_q', 'act', 'reward', 'done', 'weight']) + + +def m_q_1step_td_error( + data: namedtuple, + gamma: float, + tau: float, + alpha: float, + criterion: torch.nn.modules = nn.MSELoss(reduction='none') # noqa +) -> torch.Tensor: + """ + Overview: + Munchausen td_error for DQN algorithm, support 1 step td error. + Arguments: + - data (:obj:`m_q_1step_td_data`): The input data, m_q_1step_td_data to calculate loss + - gamma (:obj:`float`): Discount factor + - tau (:obj:`float`): Entropy factor for Munchausen DQN + - alpha (:obj:`float`): Discount factor for Munchausen term + - criterion (:obj:`torch.nn.modules`): Loss function criterion + Returns: + - loss (:obj:`torch.Tensor`): 1step td error, 0-dim tensor + Shapes: + - data (:obj:`m_q_1step_td_data`): the m_q_1step_td_data containing\ + ['q', 'target_q', 'next_q', 'act', 'reward', 'done', 'weight'] + - q (:obj:`torch.FloatTensor`): :math:`(B, N)` i.e. [batch_size, action_dim] + - target_q (:obj:`torch.FloatTensor`): :math:`(B, N)` i.e. [batch_size, action_dim] + - next_q (:obj:`torch.FloatTensor`): :math:`(B, N)` i.e. [batch_size, action_dim] + - act (:obj:`torch.LongTensor`): :math:`(B, )` + - reward (:obj:`torch.FloatTensor`): :math:`( , B)` + - done (:obj:`torch.BoolTensor`) :math:`(B, )`, whether done in last timestep + - weight (:obj:`torch.FloatTensor` or None): :math:`(B, )`, the training sample weight + Examples: + >>> action_dim = 4 + >>> data = m_q_1step_td_data( + >>> q=torch.randn(3, action_dim), + >>> target_q=torch.randn(3, action_dim), + >>> next_q=torch.randn(3, action_dim), + >>> act=torch.randint(0, action_dim, (3,)), + >>> reward=torch.randn(3), + >>> done=torch.randint(0, 2, (3,)), + >>> weight=torch.ones(3), + >>> ) + >>> loss = m_q_1step_td_error(data, 0.99, 0.01, 0.01) + """ + q, target_q, next_q, act, reward, done, weight = data + lower_bound = -1 + assert len(act.shape) == 1, act.shape + assert len(reward.shape) == 1, reward.shape + batch_range = torch.arange(act.shape[0]) + if weight is None: + weight = torch.ones_like(reward) + q_s_a = q[batch_range, act] + # calculate muchausen addon + # replay_log_policy + target_v_s = target_q[batch_range].max(1)[0].unsqueeze(-1) + + logsum = torch.logsumexp((target_q - target_v_s) / tau, 1).unsqueeze(-1) + log_pi = target_q - target_v_s - tau * logsum + act_get = act.unsqueeze(-1) + # same to the last second tau_log_pi_a + munchausen_addon = log_pi.gather(1, act_get) + + muchausen_term = alpha * torch.clamp(munchausen_addon, min=lower_bound, max=1) + + # replay_next_log_policy + target_v_s_next = next_q[batch_range].max(1)[0].unsqueeze(-1) + logsum_next = torch.logsumexp((next_q - target_v_s_next) / tau, 1).unsqueeze(-1) + tau_log_pi_next = next_q - target_v_s_next - tau * logsum_next + # do stable softmax == replay_next_policy + pi_target = F.softmax((next_q - target_v_s_next) / tau) + target_q_s_a = (gamma * (pi_target * (next_q - tau_log_pi_next) * (1 - done.unsqueeze(-1))).sum(1)).unsqueeze(-1) + + target_q_s_a = reward.unsqueeze(-1) + muchausen_term + target_q_s_a + td_error_per_sample = criterion(q_s_a.unsqueeze(-1), target_q_s_a.detach()).squeeze(-1) + + # calculate action_gap and clipfrac + with torch.no_grad(): + top2_q_s = target_q[batch_range].topk(2, dim=1, largest=True, sorted=True)[0] + action_gap = (top2_q_s[:, 0] - top2_q_s[:, 1]).mean() + + clipped = munchausen_addon.gt(1) | munchausen_addon.lt(lower_bound) + clipfrac = torch.as_tensor(clipped).float() + + return (td_error_per_sample * weight).mean(), td_error_per_sample, action_gap, clipfrac + + q_v_1step_td_data = namedtuple('q_v_1step_td_data', ['q', 'v', 'act', 'reward', 'done', 'weight']) @@ -47,6 +165,36 @@ def q_v_1step_td_error( data: namedtuple, gamma: float, criterion: torch.nn.modules = nn.MSELoss(reduction='none') ) -> torch.Tensor: # we will use this function in discrete sac algorithm to calculate td error between q and v value. + """ + Overview: + td_error between q and v value for SAC algorithm, support 1 step td error. + Arguments: + - data (:obj:`q_v_1step_td_data`): The input data, q_v_1step_td_data to calculate loss + - gamma (:obj:`float`): Discount factor + - criterion (:obj:`torch.nn.modules`): Loss function criterion + Returns: + - loss (:obj:`torch.Tensor`): 1step td error, 0-dim tensor + Shapes: + - data (:obj:`q_v_1step_td_data`): the q_v_1step_td_data containing\ + ['q', 'v', 'act', 'reward', 'done', 'weight'] + - q (:obj:`torch.FloatTensor`): :math:`(B, N)` i.e. [batch_size, action_dim] + - v (:obj:`torch.FloatTensor`): :math:`(B, )` + - act (:obj:`torch.LongTensor`): :math:`(B, )` + - reward (:obj:`torch.FloatTensor`): :math:`( , B)` + - done (:obj:`torch.BoolTensor`) :math:`(B, )`, whether done in last timestep + - weight (:obj:`torch.FloatTensor` or None): :math:`(B, )`, the training sample weight + Examples: + >>> action_dim = 4 + >>> data = q_v_1step_td_data( + >>> q=torch.randn(3, action_dim), + >>> v=torch.randn(3), + >>> act=torch.randint(0, action_dim, (3,)), + >>> reward=torch.randn(3), + >>> done=torch.randint(0, 2, (3,)), + >>> weight=torch.ones(3), + >>> ) + >>> loss = q_v_1step_td_error(data, 0.99) + """ q, v, act, reward, done, weight = data if len(act.shape) == 1: assert len(reward.shape) == 1, reward.shape @@ -80,6 +228,31 @@ def view_similar(x: torch.Tensor, target: torch.Tensor) -> torch.Tensor: def nstep_return(data: namedtuple, gamma: Union[float, list], nstep: int, value_gamma: Optional[torch.Tensor] = None): + ''' + Overview: + Calculate nstep return for DQN algorithm, support single agent case and multi agent case. + Arguments: + - data (:obj:`nstep_return_data`): The input data, nstep_return_data to calculate loss + - gamma (:obj:`float`): Discount factor + - nstep (:obj:`int`): nstep num + - value_gamma (:obj:`torch.Tensor`): Discount factor for value + Returns: + - return (:obj:`torch.Tensor`): nstep return + Shapes: + - data (:obj:`nstep_return_data`): the nstep_return_data containing\ + ['reward', 'next_value', 'done'] + - reward (:obj:`torch.FloatTensor`): :math:`(T, B)`, where T is timestep(nstep) + - next_value (:obj:`torch.FloatTensor`): :math:`(, B)` + - done (:obj:`torch.BoolTensor`) :math:`(B, )`, whether done in last timestep + Examples: + >>> data = nstep_return_data( + >>> reward=torch.randn(3, 3), + >>> next_value=torch.randn(3), + >>> done=torch.randint(0, 2, (3,)), + >>> ) + >>> loss = nstep_return(data, 0.99, 3) + ''' + reward, next_value, done = data assert reward.shape[0] == nstep device = reward.device @@ -93,6 +266,10 @@ def nstep_return(data: namedtuple, gamma: Union[float, list], nstep: int, value_ if value_gamma is None: return_ = return_tmp + (gamma ** nstep) * next_value * (1 - done) else: + if np.isscalar(value_gamma): + value_gamma = torch.full_like(next_value, value_gamma) + value_gamma = view_similar(value_gamma, next_value) + done = view_similar(done, next_value) return_ = return_tmp + value_gamma * next_value * (1 - done) elif isinstance(gamma, list): @@ -121,6 +298,37 @@ def dist_1step_td_error( v_max: float, n_atom: int, ) -> torch.Tensor: + """ + Overview: + 1 step td_error for distributed q-learning based algorithm + Arguments: + - data (:obj:`dist_1step_td_data`): The input data, dist_nstep_td_data to calculate loss + - gamma (:obj:`float`): Discount factor + - v_min (:obj:`float`): The min value of support + - v_max (:obj:`float`): The max value of support + - n_atom (:obj:`int`): The num of atom + Returns: + - loss (:obj:`torch.Tensor`): nstep td error, 0-dim tensor + Shapes: + - data (:obj:`dist_1step_td_data`): the dist_1step_td_data containing\ + ['dist', 'next_n_dist', 'act', 'reward', 'done', 'weight'] + - dist (:obj:`torch.FloatTensor`): :math:`(B, N, n_atom)` i.e. [batch_size, action_dim, n_atom] + - next_dist (:obj:`torch.FloatTensor`): :math:`(B, N, n_atom)` + - act (:obj:`torch.LongTensor`): :math:`(B, )` + - next_act (:obj:`torch.LongTensor`): :math:`(B, )` + - reward (:obj:`torch.FloatTensor`): :math:`(, B)` + - done (:obj:`torch.BoolTensor`) :math:`(B, )`, whether done in last timestep + - weight (:obj:`torch.FloatTensor` or None): :math:`(B, )`, the training sample weight + Examples: + >>> dist = torch.randn(4, 3, 51).abs().requires_grad_(True) + >>> next_dist = torch.randn(4, 3, 51).abs() + >>> act = torch.randint(0, 3, (4,)) + >>> next_act = torch.randint(0, 3, (4,)) + >>> reward = torch.randn(4) + >>> done = torch.randint(0, 2, (4,)) + >>> data = dist_1step_td_data(dist, next_dist, act, next_act, reward, done, None) + >>> loss = dist_1step_td_error(data, 0.99, -10.0, 10.0, 51) + """ dist, next_dist, act, next_act, reward, done, weight = data device = reward.device assert len(reward.shape) == 1, reward.shape @@ -211,13 +419,13 @@ def dist_nstep_td_error( nstep: int = 1, value_gamma: Optional[torch.Tensor] = None, ) -> torch.Tensor: - r""" + """ Overview: Multistep (1 step or n step) td_error for distributed q-learning based algorithm, support single\ agent case and multi agent case. Arguments: - - data (:obj:`dist_nstep_td_data`): the input data, dist_nstep_td_data to calculate loss - - gamma (:obj:`float`): discount factor + - data (:obj:`dist_nstep_td_data`): The input data, dist_nstep_td_data to calculate loss + - gamma (:obj:`float`): Discount factor - nstep (:obj:`int`): nstep num, default set to 1 Returns: - loss (:obj:`torch.Tensor`): nstep td error, 0-dim tensor @@ -230,6 +438,15 @@ def dist_nstep_td_error( - next_n_act (:obj:`torch.LongTensor`): :math:`(B, )` - reward (:obj:`torch.FloatTensor`): :math:`(T, B)`, where T is timestep(nstep) - done (:obj:`torch.BoolTensor`) :math:`(B, )`, whether done in last timestep + Examples: + >>> dist = torch.randn(4, 3, 51).abs().requires_grad_(True) + >>> next_n_dist = torch.randn(4, 3, 51).abs() + >>> done = torch.randn(4) + >>> action = torch.randint(0, 3, size=(4, )) + >>> next_action = torch.randint(0, 3, size=(4, )) + >>> reward = torch.randn(5, 4) + >>> data = dist_nstep_td_data(dist, next_n_dist, action, next_action, reward, done, None) + >>> loss, _ = dist_nstep_td_error(data, 0.95, -10.0, 10.0, 51, 5) """ dist, next_n_dist, act, next_n_act, reward, done, weight = data device = reward.device @@ -314,6 +531,31 @@ def v_1step_td_error( gamma: float, criterion: torch.nn.modules = nn.MSELoss(reduction='none') # noqa ) -> torch.Tensor: + ''' + Overview: + 1 step td_error for distributed value based algorithm + Arguments: + - data (:obj:`v_1step_td_data`): The input data, v_1step_td_data to calculate loss + - gamma (:obj:`float`): Discount factor + - criterion (:obj:`torch.nn.modules`): Loss function criterion + Returns: + - loss (:obj:`torch.Tensor`): 1step td error, 0-dim tensor + Shapes: + - data (:obj:`v_1step_td_data`): the v_1step_td_data containing\ + ['v', 'next_v', 'reward', 'done', 'weight'] + - v (:obj:`torch.FloatTensor`): :math:`(B, )` i.e. [batch_size, ] + - next_v (:obj:`torch.FloatTensor`): :math:`(B, )` + - reward (:obj:`torch.FloatTensor`): :math:`(, B)` + - done (:obj:`torch.BoolTensor`) :math:`(B, )`, whether done in last timestep + - weight (:obj:`torch.FloatTensor` or None): :math:`(B, )`, the training sample weight + Examples: + >>> v = torch.randn(5).requires_grad_(True) + >>> next_v = torch.randn(5) + >>> reward = torch.rand(5) + >>> done = torch.zeros(5) + >>> data = v_1step_td_data(v, next_v, reward, done, None) + >>> loss, td_error_per_sample = v_1step_td_error(data, 0.99) + ''' v, next_v, reward, done, weight = data if weight is None: weight = torch.ones_like(v) @@ -340,25 +582,32 @@ def v_nstep_td_error( nstep: int = 1, criterion: torch.nn.modules = nn.MSELoss(reduction='none') # noqa ) -> torch.Tensor: - r""" + """ Overview: Multistep (n step) td_error for distributed value based algorithm Arguments: - - data (:obj:`dist_nstep_td_data`): the input data, v_nstep_td_data to calculate loss - - gamma (:obj:`float`): discount factor + - data (:obj:`dist_nstep_td_data`): The input data, v_nstep_td_data to calculate loss + - gamma (:obj:`float`): Discount factor - nstep (:obj:`int`): nstep num, default set to 1 Returns: - loss (:obj:`torch.Tensor`): nstep td error, 0-dim tensor Shapes: - - data (:obj:`dist_nstep_td_data`): the v_nstep_td_data containing\ + - data (:obj:`dist_nstep_td_data`): The v_nstep_td_data containing \ ['v', 'next_n_v', 'reward', 'done', 'weight', 'value_gamma'] - v (:obj:`torch.FloatTensor`): :math:`(B, )` i.e. [batch_size, ] - next_v (:obj:`torch.FloatTensor`): :math:`(B, )` - reward (:obj:`torch.FloatTensor`): :math:`(T, B)`, where T is timestep(nstep) - done (:obj:`torch.BoolTensor`) :math:`(B, )`, whether done in last timestep - weight (:obj:`torch.FloatTensor` or None): :math:`(B, )`, the training sample weight - - value_gamma (:obj:`torch.Tensor`): If the remaining data in the buffer is less than n_step\ + - value_gamma (:obj:`torch.Tensor`): If the remaining data in the buffer is less than n_step \ we use value_gamma as the gamma discount value for next_v rather than gamma**n_step + Examples: + >>> v = torch.randn(5).requires_grad_(True) + >>> next_v = torch.randn(5) + >>> reward = torch.rand(5, 5) + >>> done = torch.zeros(5) + >>> data = v_nstep_td_data(v, next_v, reward, done, 0.9, 0.99) + >>> loss, td_error_per_sample = v_nstep_td_error(data, 0.99, 5) """ v, next_n_v, reward, done, weight, value_gamma = data if weight is None: @@ -409,17 +658,17 @@ def q_nstep_td_error( Overview: Multistep (1 step or n step) td_error for q-learning based algorithm Arguments: - - data (:obj:`q_nstep_td_data`): the input data, q_nstep_td_data to calculate loss - - gamma (:obj:`float`): discount factor - - cum_reward (:obj:`bool`): whether to use cumulative nstep reward, which is figured out when collecting data - - value_gamma (:obj:`torch.Tensor`): gamma discount value for target q_value - - criterion (:obj:`torch.nn.modules`): loss function criterion + - data (:obj:`q_nstep_td_data`): The input data, q_nstep_td_data to calculate loss + - gamma (:obj:`float`): Discount factor + - cum_reward (:obj:`bool`): Whether to use cumulative nstep reward, which is figured out when collecting data + - value_gamma (:obj:`torch.Tensor`): Gamma discount value for target q_value + - criterion (:obj:`torch.nn.modules`): Loss function criterion - nstep (:obj:`int`): nstep num, default set to 1 Returns: - loss (:obj:`torch.Tensor`): nstep td error, 0-dim tensor - td_error_per_sample (:obj:`torch.Tensor`): nstep td error, 1-dim tensor Shapes: - - data (:obj:`q_nstep_td_data`): the q_nstep_td_data containing\ + - data (:obj:`q_nstep_td_data`): The q_nstep_td_data containing\ ['q', 'next_n_q', 'action', 'reward', 'done'] - q (:obj:`torch.FloatTensor`): :math:`(B, N)` i.e. [batch_size, action_dim] - next_n_q (:obj:`torch.FloatTensor`): :math:`(B, N)` @@ -428,17 +677,102 @@ def q_nstep_td_error( - reward (:obj:`torch.FloatTensor`): :math:`(T, B)`, where T is timestep(nstep) - done (:obj:`torch.BoolTensor`) :math:`(B, )`, whether done in last timestep - td_error_per_sample (:obj:`torch.FloatTensor`): :math:`(B, )` + Examples: + >>> next_q = torch.randn(4, 3) + >>> done = torch.randn(4) + >>> action = torch.randint(0, 3, size=(4, )) + >>> next_action = torch.randint(0, 3, size=(4, )) + >>> nstep =3 + >>> q = torch.randn(4, 3).requires_grad_(True) + >>> reward = torch.rand(nstep, 4) + >>> data = q_nstep_td_data(q, next_q, action, next_action, reward, done, None) + >>> loss, td_error_per_sample = q_nstep_td_error(data, 0.95, nstep=nstep) """ q, next_n_q, action, next_n_action, reward, done, weight = data if weight is None: weight = torch.ones_like(reward) - if len(action.shape) > 1: # MARL case + + if len(action.shape) == 1 or len(action.shape) < len(q.shape): + # we need to unsqueeze action and q to make them have the same shape + # e.g. single agent case: action is [B, ] and q is [B, ] + # e.g. multi agent case: action is [B, agent_num] and q is [B, agent_num, action_shape] + action = action.unsqueeze(-1) + elif len(action.shape) > 1: # MARL case reward = reward.unsqueeze(-1) weight = weight.unsqueeze(-1) done = done.unsqueeze(-1) if value_gamma is not None: value_gamma = value_gamma.unsqueeze(-1) + q_s_a = q.gather(-1, action).squeeze(-1) + + target_q_s_a = next_n_q.gather(-1, next_n_action.unsqueeze(-1)).squeeze(-1) + + if cum_reward: + if value_gamma is None: + target_q_s_a = reward + (gamma ** nstep) * target_q_s_a * (1 - done) + else: + target_q_s_a = reward + value_gamma * target_q_s_a * (1 - done) + else: + target_q_s_a = nstep_return(nstep_return_data(reward, target_q_s_a, done), gamma, nstep, value_gamma) + td_error_per_sample = criterion(q_s_a, target_q_s_a.detach()) + return (td_error_per_sample * weight).mean(), td_error_per_sample + + +def bdq_nstep_td_error( + data: namedtuple, + gamma: Union[float, list], + nstep: int = 1, + cum_reward: bool = False, + value_gamma: Optional[torch.Tensor] = None, + criterion: torch.nn.modules = nn.MSELoss(reduction='none'), +) -> torch.Tensor: + """ + Overview: + Multistep (1 step or n step) td_error for BDQ algorithm, referenced paper "Action Branching Architectures for \ + Deep Reinforcement Learning", link: https://arxiv.org/pdf/1711.08946. + In fact, the original paper only provides the 1-step TD-error calculation method, and here we extend the \ + calculation method of n-step, i.e., TD-error: + Arguments: + - data (:obj:`q_nstep_td_data`): The input data, q_nstep_td_data to calculate loss + - gamma (:obj:`float`): Discount factor + - cum_reward (:obj:`bool`): Whether to use cumulative nstep reward, which is figured out when collecting data + - value_gamma (:obj:`torch.Tensor`): Gamma discount value for target q_value + - criterion (:obj:`torch.nn.modules`): Loss function criterion + - nstep (:obj:`int`): nstep num, default set to 1 + Returns: + - loss (:obj:`torch.Tensor`): nstep td error, 0-dim tensor + - td_error_per_sample (:obj:`torch.Tensor`): nstep td error, 1-dim tensor + Shapes: + - data (:obj:`q_nstep_td_data`): The q_nstep_td_data containing \ + ['q', 'next_n_q', 'action', 'reward', 'done'] + - q (:obj:`torch.FloatTensor`): :math:`(B, D, N)` i.e. [batch_size, branch_num, action_bins_per_branch] + - next_n_q (:obj:`torch.FloatTensor`): :math:`(B, D, N)` + - action (:obj:`torch.LongTensor`): :math:`(B, D)` + - next_n_action (:obj:`torch.LongTensor`): :math:`(B, D)` + - reward (:obj:`torch.FloatTensor`): :math:`(T, B)`, where T is timestep(nstep) + - done (:obj:`torch.BoolTensor`) :math:`(B, )`, whether done in last timestep + - td_error_per_sample (:obj:`torch.FloatTensor`): :math:`(B, )` + Examples: + >>> action_per_branch = 3 + >>> next_q = torch.randn(8, 6, action_per_branch) + >>> done = torch.randn(8) + >>> action = torch.randint(0, action_per_branch, size=(8, 6)) + >>> next_action = torch.randint(0, action_per_branch, size=(8, 6)) + >>> nstep =3 + >>> q = torch.randn(8, 6, action_per_branch).requires_grad_(True) + >>> reward = torch.rand(nstep, 8) + >>> data = q_nstep_td_data(q, next_q, action, next_action, reward, done, None) + >>> loss, td_error_per_sample = bdq_nstep_td_error(data, 0.95, nstep=nstep) + """ + q, next_n_q, action, next_n_action, reward, done, weight = data + if weight is None: + weight = torch.ones_like(reward) + reward = reward.unsqueeze(-1) + done = done.unsqueeze(-1) + if value_gamma is not None: + value_gamma = value_gamma.unsqueeze(-1) + q_s_a = q.gather(-1, action.unsqueeze(-1)).squeeze(-1) target_q_s_a = next_n_q.gather(-1, next_n_action.unsqueeze(-1)).squeeze(-1) @@ -450,6 +784,7 @@ def q_nstep_td_error( else: target_q_s_a = nstep_return(nstep_return_data(reward, target_q_s_a, done), gamma, nstep, value_gamma) td_error_per_sample = criterion(q_s_a, target_q_s_a.detach()) + td_error_per_sample = td_error_per_sample.mean(-1) return (td_error_per_sample * weight).mean(), td_error_per_sample @@ -485,18 +820,18 @@ def q_nstep_td_error_with_rescale( Overview: Multistep (1 step or n step) td_error with value rescaling Arguments: - - data (:obj:`q_nstep_td_data`): the input data, q_nstep_td_data to calculate loss - - gamma (:obj:`float`): discount factor + - data (:obj:`q_nstep_td_data`): The input data, q_nstep_td_data to calculate loss + - gamma (:obj:`float`): Discount factor - nstep (:obj:`int`): nstep num, default set to 1 - - criterion (:obj:`torch.nn.modules`): loss function criterion - - trans_fn (:obj:`Callable`): value transfrom function, default to value_transform\ + - criterion (:obj:`torch.nn.modules`): Loss function criterion + - trans_fn (:obj:`Callable`): Value transfrom function, default to value_transform\ (refer to rl_utils/value_rescale.py) - - inv_trans_fn (:obj:`Callable`): value inverse transfrom function, default to value_inv_transform\ + - inv_trans_fn (:obj:`Callable`): Value inverse transfrom function, default to value_inv_transform\ (refer to rl_utils/value_rescale.py) Returns: - loss (:obj:`torch.Tensor`): nstep td error, 0-dim tensor Shapes: - - data (:obj:`q_nstep_td_data`): the q_nstep_td_data containing\ + - data (:obj:`q_nstep_td_data`): The q_nstep_td_data containing\ ['q', 'next_n_q', 'action', 'reward', 'done'] - q (:obj:`torch.FloatTensor`): :math:`(B, N)` i.e. [batch_size, action_dim] - next_n_q (:obj:`torch.FloatTensor`): :math:`(B, N)` @@ -504,6 +839,16 @@ def q_nstep_td_error_with_rescale( - next_n_action (:obj:`torch.LongTensor`): :math:`(B, )` - reward (:obj:`torch.FloatTensor`): :math:`(T, B)`, where T is timestep(nstep) - done (:obj:`torch.BoolTensor`) :math:`(B, )`, whether done in last timestep + Examples: + >>> next_q = torch.randn(4, 3) + >>> done = torch.randn(4) + >>> action = torch.randint(0, 3, size=(4, )) + >>> next_action = torch.randint(0, 3, size=(4, )) + >>> nstep =3 + >>> q = torch.randn(4, 3).requires_grad_(True) + >>> reward = torch.rand(nstep, 4) + >>> data = q_nstep_td_data(q, next_q, action, next_action, reward, done, None) + >>> loss, _ = q_nstep_td_error_with_rescale(data, 0.95, nstep=nstep) """ q, next_n_q, action, next_n_action, reward, done, weight = data assert len(action.shape) == 1, action.shape @@ -538,11 +883,11 @@ def dqfd_nstep_td_error( Overview: Multistep n step td_error + 1 step td_error + supervised margin loss or dqfd Arguments: - - data (:obj:`dqfd_nstep_td_data`): the input data, dqfd_nstep_td_data to calculate loss + - data (:obj:`dqfd_nstep_td_data`): The input data, dqfd_nstep_td_data to calculate loss - gamma (:obj:`float`): discount factor - - cum_reward (:obj:`bool`): whether to use cumulative nstep reward, which is figured out when collecting data - - value_gamma (:obj:`torch.Tensor`): gamma discount value for target q_value - - criterion (:obj:`torch.nn.modules`): loss function criterion + - cum_reward (:obj:`bool`): Whether to use cumulative nstep reward, which is figured out when collecting data + - value_gamma (:obj:`torch.Tensor`): Gamma discount value for target q_value + - criterion (:obj:`torch.nn.modules`): Loss function criterion - nstep (:obj:`int`): nstep num, default set to 10 Returns: - loss (:obj:`torch.Tensor`): Multistep n step td_error + 1 step td_error + supervised margin loss, 0-dim tensor @@ -562,6 +907,26 @@ def dqfd_nstep_td_error( - new_n_q_one_step (:obj:`torch.FloatTensor`): :math:`(B, N)` - next_n_action_one_step (:obj:`torch.LongTensor`): :math:`(B, )` - is_expert (:obj:`int`) : 0 or 1 + Examples: + >>> next_q = torch.randn(4, 3) + >>> done = torch.randn(4) + >>> done_1 = torch.randn(4) + >>> next_q_one_step = torch.randn(4, 3) + >>> action = torch.randint(0, 3, size=(4, )) + >>> next_action = torch.randint(0, 3, size=(4, )) + >>> next_action_one_step = torch.randint(0, 3, size=(4, )) + >>> is_expert = torch.ones((4)) + >>> nstep = 3 + >>> q = torch.randn(4, 3).requires_grad_(True) + >>> reward = torch.rand(nstep, 4) + >>> data = dqfd_nstep_td_data( + >>> q, next_q, action, next_action, reward, done, done_1, None, + >>> next_q_one_step, next_action_one_step, is_expert + >>> ) + >>> loss, td_error_per_sample, loss_statistics = dqfd_nstep_td_error( + >>> data, 0.95, lambda_n_step_td=1, lambda_supervised_loss=1, + >>> margin_function=0.8, nstep=nstep + >>> ) """ q, next_n_q, action, next_n_action, reward, done, done_one_step, weight, new_n_q_one_step, next_n_action_one_step, \ is_expert = data # set is_expert flag(expert 1, agent 0) @@ -636,18 +1001,18 @@ def dqfd_nstep_td_error_with_rescale( Overview: Multistep n step td_error + 1 step td_error + supervised margin loss or dqfd Arguments: - - data (:obj:`dqfd_nstep_td_data`): the input data, dqfd_nstep_td_data to calculate loss - - gamma (:obj:`float`): discount factor - - cum_reward (:obj:`bool`): whether to use cumulative nstep reward, which is figured out when collecting data - - value_gamma (:obj:`torch.Tensor`): gamma discount value for target q_value - - criterion (:obj:`torch.nn.modules`): loss function criterion + - data (:obj:`dqfd_nstep_td_data`): The input data, dqfd_nstep_td_data to calculate loss + - gamma (:obj:`float`): Discount factor + - cum_reward (:obj:`bool`): Whether to use cumulative nstep reward, which is figured out when collecting data + - value_gamma (:obj:`torch.Tensor`): Gamma discount value for target q_value + - criterion (:obj:`torch.nn.modules`): Loss function criterion - nstep (:obj:`int`): nstep num, default set to 10 Returns: - loss (:obj:`torch.Tensor`): Multistep n step td_error + 1 step td_error + supervised margin loss, 0-dim tensor - td_error_per_sample (:obj:`torch.Tensor`): Multistep n step td_error + 1 step td_error\ + supervised margin loss, 1-dim tensor Shapes: - - data (:obj:`q_nstep_td_data`): the q_nstep_td_data containing\ + - data (:obj:`q_nstep_td_data`): The q_nstep_td_data containing\ ['q', 'next_n_q', 'action', 'next_n_action', 'reward', 'done', 'weight'\ , 'new_n_q_one_step', 'next_n_action_one_step', 'is_expert'] - q (:obj:`torch.FloatTensor`): :math:`(B, N)` i.e. [batch_size, action_dim] @@ -740,13 +1105,13 @@ def qrdqn_nstep_td_error( Overview: Multistep (1 step or n step) td_error with in QRDQN Arguments: - - data (:obj:`iqn_nstep_td_data`): the input data, iqn_nstep_td_data to calculate loss - - gamma (:obj:`float`): discount factor + - data (:obj:`qrdqn_nstep_td_data`): The input data, qrdqn_nstep_td_data to calculate loss + - gamma (:obj:`float`): Discount factor - nstep (:obj:`int`): nstep num, default set to 1 Returns: - loss (:obj:`torch.Tensor`): nstep td error, 0-dim tensor Shapes: - - data (:obj:`q_nstep_td_data`): the q_nstep_td_data containing\ + - data (:obj:`q_nstep_td_data`): The q_nstep_td_data containing\ ['q', 'next_n_q', 'action', 'reward', 'done'] - q (:obj:`torch.FloatTensor`): :math:`(tau, B, N)` i.e. [tau x batch_size, action_dim] - next_n_q (:obj:`torch.FloatTensor`): :math:`(tau', B, N)` @@ -754,6 +1119,16 @@ def qrdqn_nstep_td_error( - next_n_action (:obj:`torch.LongTensor`): :math:`(B, )` - reward (:obj:`torch.FloatTensor`): :math:`(T, B)`, where T is timestep(nstep) - done (:obj:`torch.BoolTensor`) :math:`(B, )`, whether done in last timestep + Examples: + >>> next_q = torch.randn(4, 3, 3) + >>> done = torch.randn(4) + >>> action = torch.randint(0, 3, size=(4, )) + >>> next_action = torch.randint(0, 3, size=(4, )) + >>> nstep = 3 + >>> q = torch.randn(4, 3, 3).requires_grad_(True) + >>> reward = torch.rand(nstep, 4) + >>> data = qrdqn_nstep_td_data(q, next_q, action, next_action, reward, done, 3, None) + >>> loss, td_error_per_sample = qrdqn_nstep_td_error(data, 0.95, nstep=nstep) """ q, next_n_q, action, next_n_action, reward, done, tau, weight = data @@ -810,18 +1185,18 @@ def q_nstep_sql_td_error( Overview: Multistep (1 step or n step) td_error for q-learning based algorithm Arguments: - - data (:obj:`q_nstep_td_data`): the input data, q_nstep_sql_td_data to calculate loss - - gamma (:obj:`float`): discount factor + - data (:obj:`q_nstep_td_data`): The input data, q_nstep_sql_td_data to calculate loss + - gamma (:obj:`float`): Discount factor - Alpha (:obj:｀float`): A parameter to weight entropy term in a policy equation - - cum_reward (:obj:`bool`): whether to use cumulative nstep reward, which is figured out when collecting data - - value_gamma (:obj:`torch.Tensor`): gamma discount value for target soft_q_value - - criterion (:obj:`torch.nn.modules`): loss function criterion + - cum_reward (:obj:`bool`): Whether to use cumulative nstep reward, which is figured out when collecting data + - value_gamma (:obj:`torch.Tensor`): Gamma discount value for target soft_q_value + - criterion (:obj:`torch.nn.modules`): Loss function criterion - nstep (:obj:`int`): nstep num, default set to 1 Returns: - loss (:obj:`torch.Tensor`): nstep td error, 0-dim tensor - td_error_per_sample (:obj:`torch.Tensor`): nstep td error, 1-dim tensor Shapes: - - data (:obj:`q_nstep_td_data`): the q_nstep_td_data containing\ + - data (:obj:`q_nstep_td_data`): The q_nstep_td_data containing\ ['q', 'next_n_q', 'action', 'reward', 'done'] - q (:obj:`torch.FloatTensor`): :math:`(B, N)` i.e. [batch_size, action_dim] - next_n_q (:obj:`torch.FloatTensor`): :math:`(B, N)` @@ -830,6 +1205,16 @@ def q_nstep_sql_td_error( - reward (:obj:`torch.FloatTensor`): :math:`(T, B)`, where T is timestep(nstep) - done (:obj:`torch.BoolTensor`) :math:`(B, )`, whether done in last timestep - td_error_per_sample (:obj:`torch.FloatTensor`): :math:`(B, )` + Examples: + >>> next_q = torch.randn(4, 3) + >>> done = torch.randn(4) + >>> action = torch.randint(0, 3, size=(4, )) + >>> next_action = torch.randint(0, 3, size=(4, )) + >>> nstep = 3 + >>> q = torch.randn(4, 3).requires_grad_(True) + >>> reward = torch.rand(nstep, 4) + >>> data = q_nstep_td_data(q, next_q, action, next_action, reward, done, None) + >>> loss, td_error_per_sample, record_target_v = q_nstep_sql_td_error(data, 0.95, 1.0, nstep=nstep) """ q, next_n_q, action, next_n_action, reward, done, weight = data assert len(action.shape) == 1, action.shape @@ -878,15 +1263,15 @@ def iqn_nstep_td_error( referenced paper Implicit Quantile Networks for Distributional Reinforcement Learning \ Arguments: - - data (:obj:`iqn_nstep_td_data`): the input data, iqn_nstep_td_data to calculate loss - - gamma (:obj:`float`): discount factor + - data (:obj:`iqn_nstep_td_data`): The input data, iqn_nstep_td_data to calculate loss + - gamma (:obj:`float`): Discount factor - nstep (:obj:`int`): nstep num, default set to 1 - - criterion (:obj:`torch.nn.modules`): loss function criterion - - beta_function (:obj:`Callable`): the risk function + - criterion (:obj:`torch.nn.modules`): Loss function criterion + - beta_function (:obj:`Callable`): The risk function Returns: - loss (:obj:`torch.Tensor`): nstep td error, 0-dim tensor Shapes: - - data (:obj:`q_nstep_td_data`): the q_nstep_td_data containing\ + - data (:obj:`q_nstep_td_data`): The q_nstep_td_data containing\ ['q', 'next_n_q', 'action', 'reward', 'done'] - q (:obj:`torch.FloatTensor`): :math:`(tau, B, N)` i.e. [tau x batch_size, action_dim] - next_n_q (:obj:`torch.FloatTensor`): :math:`(tau', B, N)` @@ -894,6 +1279,17 @@ def iqn_nstep_td_error( - next_n_action (:obj:`torch.LongTensor`): :math:`(B, )` - reward (:obj:`torch.FloatTensor`): :math:`(T, B)`, where T is timestep(nstep) - done (:obj:`torch.BoolTensor`) :math:`(B, )`, whether done in last timestep + Examples: + >>> next_q = torch.randn(3, 4, 3) + >>> done = torch.randn(4) + >>> action = torch.randint(0, 3, size=(4, )) + >>> next_action = torch.randint(0, 3, size=(4, )) + >>> nstep = 3 + >>> q = torch.randn(3, 4, 3).requires_grad_(True) + >>> replay_quantile = torch.randn([3, 4, 1]) + >>> reward = torch.rand(nstep, 4) + >>> data = iqn_nstep_td_data(q, next_q, action, next_action, reward, done, replay_quantile, None) + >>> loss, td_error_per_sample = iqn_nstep_td_error(data, 0.95, nstep=nstep) """ q, next_n_q, action, next_n_action, reward, done, replay_quantiles, weight = data @@ -973,15 +1369,15 @@ def fqf_nstep_td_error( referenced paper Fully Parameterized Quantile Function for Distributional Reinforcement Learning \ Arguments: - - data (:obj:`fqf_nstep_td_data`): the input data, fqf_nstep_td_data to calculate loss - - gamma (:obj:`float`): discount factor + - data (:obj:`fqf_nstep_td_data`): The input data, fqf_nstep_td_data to calculate loss + - gamma (:obj:`float`): Discount factor - nstep (:obj:`int`): nstep num, default set to 1 - - criterion (:obj:`torch.nn.modules`): loss function criterion - - beta_function (:obj:`Callable`): the risk function + - criterion (:obj:`torch.nn.modules`): Loss function criterion + - beta_function (:obj:`Callable`): The risk function Returns: - loss (:obj:`torch.Tensor`): nstep td error, 0-dim tensor Shapes: - - data (:obj:`q_nstep_td_data`): the q_nstep_td_data containing\ + - data (:obj:`q_nstep_td_data`): The q_nstep_td_data containing\ ['q', 'next_n_q', 'action', 'reward', 'done'] - q (:obj:`torch.FloatTensor`): :math:`(B, tau, N)` i.e. [batch_size, tau, action_dim] - next_n_q (:obj:`torch.FloatTensor`): :math:`(B, tau', N)` @@ -990,6 +1386,17 @@ def fqf_nstep_td_error( - reward (:obj:`torch.FloatTensor`): :math:`(T, B)`, where T is timestep(nstep) - done (:obj:`torch.BoolTensor`) :math:`(B, )`, whether done in last timestep - quantiles_hats (:obj:`torch.FloatTensor`): :math:`(B, tau)` + Examples: + >>> next_q = torch.randn(4, 3, 3) + >>> done = torch.randn(4) + >>> action = torch.randint(0, 3, size=(4, )) + >>> next_action = torch.randint(0, 3, size=(4, )) + >>> nstep = 3 + >>> q = torch.randn(4, 3, 3).requires_grad_(True) + >>> quantiles_hats = torch.randn([4, 3]) + >>> reward = torch.rand(nstep, 4) + >>> data = fqf_nstep_td_data(q, next_q, action, next_action, reward, done, quantiles_hats, None) + >>> loss, td_error_per_sample = fqf_nstep_td_error(data, 0.95, nstep=nstep) """ q, next_n_q, action, next_n_action, reward, done, quantiles_hats, weight = data @@ -1058,11 +1465,17 @@ def evaluate_quantile_at_action(q_s, actions): def fqf_calculate_fraction_loss(q_tau_i, q_value, quantiles, actions): """ - Shapes: - - q_tau_i (:obj:`torch.FloatTensor`) :math:`(batch_size, num_quantiles-1, action_dim)` - - q_value (:obj:`torch.FloatTensor`) :math:`(batch_size, num_quantiles, action_dim)` - - quantiles (:obj:`torch.FloatTensor`) :math:`(batch_size, num_quantiles+1)` - - actions (:obj:`torch.LongTensor`) :math:`(batch_size, )` + Overview: + Calculate the fraction loss in FQF, \ + referenced paper Fully Parameterized Quantile Function for Distributional Reinforcement Learning \ + + Arguments: + - q_tau_i (:obj:`torch.FloatTensor`): :math:`(batch_size, num_quantiles-1, action_dim)` + - q_value (:obj:`torch.FloatTensor`): :math:`(batch_size, num_quantiles, action_dim)` + - quantiles (:obj:`torch.FloatTensor`): :math:`(batch_size, num_quantiles+1)` + - actions (:obj:`torch.LongTensor`): :math:`(batch_size, )` + Returns: + - fraction_loss (:obj:`torch.Tensor`): fraction loss, 0-dim tensor """ assert q_value.requires_grad @@ -1132,8 +1545,8 @@ def td_lambda_error(data: namedtuple, gamma: float = 0.9, lambda_: float = 0.8) (*including the terminal state*, values[terminal] should also be 0) Arguments: - data (:obj:`namedtuple`): td_lambda input data with fields ['value', 'reward', 'weight'] - - gamma (:obj:`float`): constant discount factor gamma, should be in [0, 1], defaults to 0.9 - - lambda (:obj:`float`): constant lambda, should be in [0, 1], defaults to 0.8 + - gamma (:obj:`float`): Constant discount factor gamma, should be in [0, 1], defaults to 0.9 + - lambda (:obj:`float`): Constant lambda, should be in [0, 1], defaults to 0.8 Returns: - loss (:obj:`torch.Tensor`): Computed MSE loss, averaged over the batch Shapes: @@ -1142,6 +1555,11 @@ def td_lambda_error(data: namedtuple, gamma: float = 0.9, lambda_: float = 0.8) - reward (:obj:`torch.FloatTensor`): :math:`(T, B)`, the returns from time step 0 to T-1 - weight (:obj:`torch.FloatTensor` or None): :math:`(B, )`, the training sample weight - loss (:obj:`torch.FloatTensor`): :math:`()`, 0-dim tensor + Examples: + >>> T, B = 8, 4 + >>> value = torch.randn(T + 1, B).requires_grad_(True) + >>> reward = torch.rand(T, B) + >>> loss = td_lambda_error(td_lambda_data(value, reward, None)) """ value, reward, weight = data if weight is None: @@ -1168,13 +1586,13 @@ def generalized_lambda_returns( Arguments: - bootstrap_values (:obj:`torch.Tensor` or :obj:`float`): estimation of the value at step 0 to *T*, of size [T_traj+1, batchsize] - - rewards (:obj:`torch.Tensor`): the returns from 0 to T-1, of size [T_traj, batchsize] + - rewards (:obj:`torch.Tensor`): The returns from 0 to T-1, of size [T_traj, batchsize] - gammas (:obj:`torch.Tensor` or :obj:`float`): - discount factor for each step (from 0 to T-1), of size [T_traj, batchsize] - - lambda (:obj:`torch.Tensor` or :obj:`float`): determining the mix of bootstrapping + Discount factor for each step (from 0 to T-1), of size [T_traj, batchsize] + - lambda (:obj:`torch.Tensor` or :obj:`float`): Determining the mix of bootstrapping vs further accumulation of multistep returns at each timestep, of size [T_traj, batchsize] - done (:obj:`torch.Tensor` or :obj:`float`): - whether the episode done at current step (from 0 to T-1), of size [T_traj, batchsize] + Whether the episode done at current step (from 0 to T-1), of size [T_traj, batchsize] Returns: - return (:obj:`torch.Tensor`): Computed lambda return value for each state from 0 to T-1, of size [T_traj, batchsize] @@ -1194,27 +1612,25 @@ def multistep_forward_view( lambda_: float, done: Optional[torch.Tensor] = None ) -> torch.Tensor: - r""" + """ Overview: - Same as trfl.sequence_ops.multistep_forward_view - Implementing (12.18) in Sutton & Barto + Same as trfl.sequence_ops.multistep_forward_view, which implements (12.18) in Sutton & Barto. + Assuming the first dim of input tensors correspond to the index in batch. - ``` + .. note:: result[T-1] = rewards[T-1] + gammas[T-1] * bootstrap_values[T] for t in 0...T-2 : result[t] = rewards[t] + gammas[t]*(lambdas[t]*result[t+1] + (1-lambdas[t])*bootstrap_values[t+1]) - ``` - Assuming the first dim of input tensors correspond to the index in batch Arguments: - - bootstrap_values (:obj:`torch.Tensor`): estimation of the value at *step 1 to T*, of size [T_traj, batchsize] - - rewards (:obj:`torch.Tensor`): the returns from 0 to T-1, of size [T_traj, batchsize] - - gammas (:obj:`torch.Tensor`): discount factor for each step (from 0 to T-1), of size [T_traj, batchsize] - - lambda (:obj:`torch.Tensor`): determining the mix of bootstrapping vs further accumulation of \ + - bootstrap_values (:obj:`torch.Tensor`): Estimation of the value at *step 1 to T*, of size [T_traj, batchsize] + - rewards (:obj:`torch.Tensor`): The returns from 0 to T-1, of size [T_traj, batchsize] + - gammas (:obj:`torch.Tensor`): Discount factor for each step (from 0 to T-1), of size [T_traj, batchsize] + - lambda (:obj:`torch.Tensor`): Determining the mix of bootstrapping vs further accumulation of \ multistep returns at each timestep of size [T_traj, batchsize], the element for T-1 is ignored \ and effectively set to 0, as there is no information about future rewards. - done (:obj:`torch.Tensor` or :obj:`float`): - whether the episode done at current step (from 0 to T-1), of size [T_traj, batchsize] + Whether the episode done at current step (from 0 to T-1), of size [T_traj, batchsize] Returns: - ret (:obj:`torch.Tensor`): Computed lambda return value \ for each state from 0 to T-1, of size [T_traj, batchsize] diff --git a/ding/rl_utils/tests/test_a2c.py b/ding/rl_utils/tests/test_a2c.py index e2c635a7f2..e321db28c4 100644 --- a/ding/rl_utils/tests/test_a2c.py +++ b/ding/rl_utils/tests/test_a2c.py @@ -2,7 +2,7 @@ from itertools import product import numpy as np import torch -from ding.rl_utils import a2c_data, a2c_error +from ding.rl_utils import a2c_data, a2c_error, a2c_error_continuous random_weight = torch.rand(4) + 1 weight_args = [None, random_weight] @@ -26,3 +26,28 @@ def test_a2c(weight): total_loss.backward() assert isinstance(logit.grad, torch.Tensor) assert isinstance(value.grad, torch.Tensor) + + +@pytest.mark.unittest +@pytest.mark.parametrize('weight, ', weight_args) +def test_a2c_continuous(weight): + B, N = 4, 32 + logit = { + "mu": torch.randn(B, N).requires_grad_(True), + "sigma": torch.exp(torch.randn(B, N)).requires_grad_(True), + } + action = torch.randn(B, N).requires_grad_(True) + value = torch.randn(B).requires_grad_(True) + adv = torch.rand(B) + return_ = torch.randn(B) * 2 + data = a2c_data(logit, action, value, adv, return_, weight) + loss = a2c_error_continuous(data) + assert all([l.shape == tuple() for l in loss]) + assert logit["mu"].grad is None + assert logit["sigma"].grad is None + assert value.grad is None + total_loss = sum(loss) + total_loss.backward() + assert isinstance(logit["mu"].grad, torch.Tensor) + assert isinstance(logit['sigma'].grad, torch.Tensor) + assert isinstance(value.grad, torch.Tensor) diff --git a/ding/rl_utils/tests/test_grpo_rlhf.py b/ding/rl_utils/tests/test_grpo_rlhf.py new file mode 100644 index 0000000000..b2ac533d95 --- /dev/null +++ b/ding/rl_utils/tests/test_grpo_rlhf.py @@ -0,0 +1,109 @@ +import pytest +import numpy as np +import torch +# Import GRPO related functions +from ding.rl_utils.grpo import grpo_policy_data, grpo_policy_error + + +@pytest.fixture +def batch_size(): + return 4 + + +@pytest.fixture +def seq_length(): + return 8 + + +@pytest.fixture +def dictionary_num(): + return 1000 + + +@pytest.mark.unittest +def test_grpo_policy_loss_with_mask(batch_size: int = 4, seq_length: int = 8, vocab_size: int = 1000): + """Test GRPO policy loss calculation with mask""" + # 1. Create test data + logit_new = (torch.randn(batch_size, seq_length, vocab_size).requires_grad_(True)) + logit_old = logit_new + torch.randn_like(logit_new) * 0.1 + logit_ref = logit_new + torch.randn_like(logit_new) * 0.2 + action = torch.randint(0, vocab_size, (batch_size, seq_length)) + adv = torch.randn(batch_size) + weight = torch.ones(batch_size, seq_length) + weight[:, -2:] = 0 + + # 2. Create grpo_policy_data instance + data = grpo_policy_data( + logit_new=logit_new, # Current policy output + logit_old=logit_old, # Old policy output + logit_ref=logit_ref, # Reference policy output + action=action, # Sampled tokens + adv=adv, # Advantage values + weight=weight # Attention mask + ) + + # 3. Calculate GRPO loss + loss, info = grpo_policy_error( + data=data, + clip_ratio=0.2, # PPO clipping ratio + beta=0.1 # KL divergence weight + ) + + # 4. Verify outputs + assert isinstance(loss, torch.Tensor) + assert loss.shape == torch.Size([]) # Ensure scalar output + assert not torch.isnan(loss) + assert not torch.isinf(loss) + + # 5. Test gradients + assert logit_new.grad is None + loss.backward() + assert isinstance(logit_new.grad, torch.Tensor) + + assert 'approx_kl' in info._asdict() + assert 'clipfrac' in info._asdict() + assert all([np.isscalar(v) for v in info._asdict().values()]) + + +@pytest.mark.unittest +def test_grpo_policy_loss_without_mask(batch_size: int = 4, seq_length: int = 8, vocab_size: int = 1000): + """Test GRPO policy loss calculation without mask""" + # 1. Create test data + logit_new = torch.randn(batch_size, seq_length, vocab_size).requires_grad_(True) + logit_old = logit_new + torch.randn_like(logit_new) * 0.1 + logit_ref = logit_new + torch.randn_like(logit_new) * 0.2 + action = torch.randint(0, vocab_size, (batch_size, seq_length)) + adv = torch.randn(batch_size) + + # 2. Create grpo_policy_data instance + data = grpo_policy_data( + logit_new=logit_new, # Current policy output + logit_old=logit_old, # Old policy output + logit_ref=logit_ref, # Reference policy output + action=action, # Sampled tokens + adv=adv, # Advantage values + weight=None # No mask + ) + + # 3. Calculate GRPO loss + loss, info = grpo_policy_error( + data=data, + clip_ratio=0.2, # PPO clipping ratio + beta=0.1 # KL divergence weight + ) + + # 4. Verify outputs + assert isinstance(loss, torch.Tensor) + assert loss.shape == torch.Size([]) # Ensure scalar output + assert not torch.isnan(loss) + assert not torch.isinf(loss) + + # 5. Test gradients + assert logit_new.grad is None + loss.backward() + assert isinstance(logit_new.grad, torch.Tensor) + + # 6. Verify metrics + assert 'approx_kl' in info._asdict() + assert 'clipfrac' in info._asdict() + assert all([np.isscalar(v) for v in info._asdict().values()]) diff --git a/ding/rl_utils/tests/test_happo.py b/ding/rl_utils/tests/test_happo.py new file mode 100644 index 0000000000..d82e5a37bc --- /dev/null +++ b/ding/rl_utils/tests/test_happo.py @@ -0,0 +1,71 @@ +import pytest +from itertools import product +import numpy as np +import torch + +from ding.rl_utils import happo_data, happo_error, happo_error_continuous +from ding.rl_utils.ppo import shape_fn_ppo + +use_value_clip_args = [True, False] +dual_clip_args = [None, 5.0] +random_weight = torch.rand(4) + 1 +weight_args = [None, random_weight] +factor_args = [torch.rand(4, 1)] +args = [item for item in product(*[use_value_clip_args, dual_clip_args, weight_args, factor_args])] + + +@pytest.mark.unittest +def test_shape_fn_ppo(): + data = happo_data(torch.randn(3, 5, 8), None, None, None, None, None, None, None, None) + shape1 = shape_fn_ppo([data], {}) + shape2 = shape_fn_ppo([], {'data': data}) + assert shape1 == shape2 == (3, 5, 8) + + +@pytest.mark.unittest +@pytest.mark.parametrize('use_value_clip, dual_clip, weight, factor', args) +def test_happo(use_value_clip, dual_clip, weight, factor): + B, N = 4, 32 + logit_new = torch.randn(B, N).requires_grad_(True) + logit_old = logit_new + torch.rand_like(logit_new) * 0.1 + action = torch.randint(0, N, size=(B, )) + value_new = torch.randn(B).requires_grad_(True) + value_old = value_new + torch.rand_like(value_new) * 0.1 + adv = torch.rand(B) + return_ = torch.randn(B) * 2 + data = happo_data(logit_new, logit_old, action, value_new, value_old, adv, return_, weight, factor) + loss, info = happo_error(data, use_value_clip=use_value_clip, dual_clip=dual_clip) + assert all([l.shape == tuple() for l in loss]) + assert all([np.isscalar(i) for i in info]) + assert logit_new.grad is None + assert value_new.grad is None + total_loss = sum(loss) + total_loss.backward() + assert isinstance(logit_new.grad, torch.Tensor) + assert isinstance(value_new.grad, torch.Tensor) + + +@pytest.mark.unittest +@pytest.mark.parametrize('use_value_clip, dual_clip, weight, factor', args) +def test_happo_error_continous(use_value_clip, dual_clip, weight, factor): + B, N = 4, 6 + mu_sigma_new = {'mu': torch.rand(B, N).requires_grad_(True), 'sigma': torch.rand(B, N).requires_grad_(True)} + mu_sigma_old = { + 'mu': mu_sigma_new['mu'] + torch.rand_like(mu_sigma_new['mu']) * 0.1, + 'sigma': mu_sigma_new['sigma'] + torch.rand_like(mu_sigma_new['sigma']) * 0.1 + } + action = torch.rand(B, N) + value_new = torch.randn(B).requires_grad_(True) + value_old = value_new + torch.rand_like(value_new) * 0.1 + adv = torch.rand(B) + return_ = torch.randn(B) * 2 + data = happo_data(mu_sigma_new, mu_sigma_old, action, value_new, value_old, adv, return_, weight, factor) + loss, info = happo_error_continuous(data, use_value_clip=use_value_clip, dual_clip=dual_clip) + assert all([l.shape == tuple() for l in loss]) + assert all([np.isscalar(i) for i in info]) + assert mu_sigma_new['mu'].grad is None + assert value_new.grad is None + total_loss = sum(loss) + total_loss.backward() + assert isinstance(mu_sigma_new['mu'].grad, torch.Tensor) + assert isinstance(value_new.grad, torch.Tensor) diff --git a/ding/rl_utils/tests/test_log_prob_fn.py b/ding/rl_utils/tests/test_log_prob_fn.py new file mode 100644 index 0000000000..81b3d35c5c --- /dev/null +++ b/ding/rl_utils/tests/test_log_prob_fn.py @@ -0,0 +1,150 @@ +import pytest +import numpy as np +import torch +from torch import Tensor +from typing import Dict, List, Tuple +from ding.rl_utils.log_prob_utils import (efficient_method, naive_method, less_efficient_method, LogProbFunction) + + +def get_gpu_memory() -> float: + """获取当前GPU内存使用情况""" + if torch.cuda.is_available(): + return torch.cuda.max_memory_allocated() / 1024 / 1024 # 转换为MB + return 0 + + +@pytest.fixture +def batch_size() -> int: + return 16 + + +@pytest.fixture +def seq_length() -> int: + return 1024 + + +@pytest.fixture +def dictionary_num() -> int: + return 32768 + + +@pytest.mark.gputest +def test_log_prob_methods_float32(batch_size: int, seq_length: int, dictionary_num: int) -> None: + """Benchmark different methods for calculating log probabilities with float32""" + print("\n" + "=" * 50) + print("Testing with float32 precision") + print("=" * 50) + + # 设置参数 + device = "cuda" if torch.cuda.is_available() else "cpu" + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() # 重置内存统计 + + # 生成测试数据 + logits: Tensor = torch.randn(batch_size, seq_length, dictionary_num, device=device, dtype=torch.float32) + input_ids: Tensor = torch.randint(0, dictionary_num, (batch_size, seq_length), device=device) + + # 预热 GPU + for _ in range(3): + _ = naive_method(logits[:2], input_ids[:2]) + torch.cuda.synchronize() + + # 测试每个方法 + results: Dict[str, Tensor] = {} + peak_memory: Dict[str, float] = {} + for method, name in [(naive_method, "Naive"), (efficient_method, "Efficient"), + (less_efficient_method, "Less_Efficient")]: + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() # 重置每个方法的内存统计 + + # 运行多次并计时 + times: List[float] = [] + for _ in range(10): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + torch.cuda.synchronize() + start.record() + result = method(logits, input_ids) + end.record() + torch.cuda.synchronize() + times.append(start.elapsed_time(end)) + if len(times) == 1: + results[name] = result + + # 记录内存使用 + peak_memory[name] = get_gpu_memory() + + # 计算统计信息 + mean_time = np.mean(times) + std_time = np.std(times) + print(f"\n{name}:") + print(f"Time: {mean_time:.2f} ± {std_time:.2f} ms") + print(f"Peak GPU Memory: {peak_memory[name]:.2f} MB") + + # 验证结果正确性 + for name, result in results.items(): + if name != "Naive": + diff = (results["Naive"] - result).abs().max().item() + assert diff < 1e-5, f"Results mismatch between Naive and {name}: {diff}" + + +@pytest.mark.gputest +def test_log_prob_methods_bfloat16(batch_size: int, seq_length: int, dictionary_num: int) -> None: + """Benchmark different methods for calculating log probabilities with bfloat16""" + print("\n" + "=" * 50) + print("Testing with bfloat16 precision") + print("=" * 50) + + # 设置参数 + device = "cuda" if torch.cuda.is_available() else "cpu" + tolerance = 0.1 # bfloat16的容差值要更大一些 + + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() # 重置内存统计 + + # 生成测试数据 + logits: Tensor = torch.randn(batch_size, seq_length, dictionary_num, device=device, dtype=torch.bfloat16) + input_ids: Tensor = torch.randint(0, dictionary_num, (batch_size, seq_length), device=device) + + # 预热 GPU + for _ in range(3): + _ = naive_method(logits[:2], input_ids[:2]) + torch.cuda.synchronize() + + # 测试每个方法 + results: Dict[str, Tensor] = {} + peak_memory: Dict[str, float] = {} + for method, name in [(naive_method, "Naive"), (efficient_method, "Efficient"), + (less_efficient_method, "Less_Efficient")]: + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() # 重置每个方法的内存统计 + + # 运行多次并计时 + times: List[float] = [] + for _ in range(10): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + torch.cuda.synchronize() + start.record() + result = method(logits, input_ids) + end.record() + torch.cuda.synchronize() + times.append(start.elapsed_time(end)) + if len(times) == 1: + results[name] = result + + # 记录内存使用 + peak_memory[name] = get_gpu_memory() + + # 计算统计信息 + mean_time = np.mean(times) + std_time = np.std(times) + print(f"\n{name}:") + print(f"Time: {mean_time:.2f} ± {std_time:.2f} ms") + print(f"Peak GPU Memory: {peak_memory[name]:.2f} MB") + + # 验证结果正确性 + for name, result in results.items(): + if name != "Naive": + diff = (results["Naive"] - result).abs().max().item() + assert diff < tolerance, f"Results mismatch between Naive and {name}: {diff}" diff --git a/ding/rl_utils/tests/test_log_prob_utils.py b/ding/rl_utils/tests/test_log_prob_utils.py new file mode 100644 index 0000000000..4b7d46a465 --- /dev/null +++ b/ding/rl_utils/tests/test_log_prob_utils.py @@ -0,0 +1,150 @@ +import pytest +import numpy as np +import torch +from torch import Tensor +from typing import Dict, List, Tuple +from ding.rl_utils.log_prob_utils import (efficient_method, naive_method, less_efficient_method, LogProbFunction) + + +def get_gpu_memory() -> float: + """Get current GPU memory usage""" + if torch.cuda.is_available(): + return torch.cuda.max_memory_allocated() / 1024 / 1024 # Convert to MB + return 0 + + +@pytest.fixture +def batch_size() -> int: + return 16 + + +@pytest.fixture +def seq_length() -> int: + return 1024 + + +@pytest.fixture +def dictionary_num() -> int: + return 32768 + + +@pytest.mark.gputest +def test_log_prob_methods_float32(batch_size: int, seq_length: int, dictionary_num: int) -> None: + """Benchmark different methods for calculating log probabilities with float32""" + print("\n" + "=" * 50) + print("Testing with float32 precision") + print("=" * 50) + + # Set parameters + device = "cuda" if torch.cuda.is_available() else "cpu" + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() # Reset memory statistics + + # Generate test data + logits: Tensor = torch.randn(batch_size, seq_length, dictionary_num, device=device, dtype=torch.float32) + input_ids: Tensor = torch.randint(0, dictionary_num, (batch_size, seq_length), device=device) + + # Warm up GPU + for _ in range(3): + _ = naive_method(logits[:2], input_ids[:2]) + torch.cuda.synchronize() + + # Test each method + results: Dict[str, Tensor] = {} + peak_memory: Dict[str, float] = {} + for method, name in [(naive_method, "Naive"), (efficient_method, "Efficient"), + (less_efficient_method, "Less_Efficient")]: + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() # Reset memory statistics for each method + + # Run multiple times and measure time + times: List[float] = [] + for _ in range(10): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + torch.cuda.synchronize() + start.record() + result = method(logits, input_ids) + end.record() + torch.cuda.synchronize() + times.append(start.elapsed_time(end)) + if len(times) == 1: + results[name] = result + + # Record memory usage + peak_memory[name] = get_gpu_memory() + + # Calculate statistics + mean_time = np.mean(times) + std_time = np.std(times) + print(f"\n{name}:") + print(f"Time: {mean_time:.2f} ± {std_time:.2f} ms") + print(f"Peak GPU Memory: {peak_memory[name]:.2f} MB") + + # Verify results correctness + for name, result in results.items(): + if name != "Naive": + diff = (results["Naive"] - result).abs().max().item() + assert diff < 1e-5, f"Results mismatch between Naive and {name}: {diff}" + + +@pytest.mark.gputest +def test_log_prob_methods_bfloat16(batch_size: int, seq_length: int, dictionary_num: int) -> None: + """Benchmark different methods for calculating log probabilities with bfloat16""" + print("\n" + "=" * 50) + print("Testing with bfloat16 precision") + print("=" * 50) + + # Set parameters + device = "cuda" if torch.cuda.is_available() else "cpu" + tolerance = 0.1 # Higher tolerance for bfloat16 + + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() # Reset memory statistics + + # Generate test data + logits: Tensor = torch.randn(batch_size, seq_length, dictionary_num, device=device, dtype=torch.bfloat16) + input_ids: Tensor = torch.randint(0, dictionary_num, (batch_size, seq_length), device=device) + + # Warm up GPU + for _ in range(3): + _ = naive_method(logits[:2], input_ids[:2]) + torch.cuda.synchronize() + + # Test each method + results: Dict[str, Tensor] = {} + peak_memory: Dict[str, float] = {} + for method, name in [(naive_method, "Naive"), (efficient_method, "Efficient"), + (less_efficient_method, "Less_Efficient")]: + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() # Reset memory statistics for each method + + # Run multiple times and measure time + times: List[float] = [] + for _ in range(10): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + torch.cuda.synchronize() + start.record() + result = method(logits, input_ids) + end.record() + torch.cuda.synchronize() + times.append(start.elapsed_time(end)) + if len(times) == 1: + results[name] = result + + # Record memory usage + peak_memory[name] = get_gpu_memory() + + # Calculate statistics + mean_time = np.mean(times) + std_time = np.std(times) + print(f"\n{name}:") + print(f"Time: {mean_time:.2f} ± {std_time:.2f} ms") + print(f"Peak GPU Memory: {peak_memory[name]:.2f} MB") + + # Verify results correctness + for name, result in results.items(): + if name != "Naive": + diff = (results["Naive"] - result).abs().max().item() + assert diff < tolerance, f"Results mismatch between Naive and {name}: {diff}" diff --git a/ding/rl_utils/tests/test_muzero_utils.py b/ding/rl_utils/tests/test_muzero_utils.py deleted file mode 100644 index e4cea523c1..0000000000 --- a/ding/rl_utils/tests/test_muzero_utils.py +++ /dev/null @@ -1,24 +0,0 @@ -import random -import numpy as np -import pytest -from ding.rl_utils.efficientzero.utils import get_augmented_data - - -@pytest.mark.unittest -class TestMuZeroUtils(): - - def test_get_augmented_data(self): - num_of_data = 100 - board_size = 15 - state = np.random.randint(0, 3, (board_size, board_size, 3), dtype=np.uint8) - mcts_prob = np.random.randn(board_size, board_size) - winner = np.random.randint(0, 2, 1, dtype=np.uint8) - play_data = [{'state': state, 'mcts_prob': mcts_prob, 'winner': winner} for _ in range(num_of_data)] - - extented_data = get_augmented_data(board_size, play_data) - assert len(extented_data) == num_of_data * 8 - # TODO(pu): extented data shape is not same as original data? - # assert extented_data[0]['state'].shape == state.shape - assert extented_data[0]['state'].flatten().shape == state.flatten().shape - assert extented_data[0]['mcts_prob'].shape == mcts_prob.flatten().shape - assert extented_data[0]['winner'].shape == winner.shape diff --git a/ding/rl_utils/tests/test_ppo.py b/ding/rl_utils/tests/test_ppo.py index a72d0e3b16..bbff8f64ee 100644 --- a/ding/rl_utils/tests/test_ppo.py +++ b/ding/rl_utils/tests/test_ppo.py @@ -15,7 +15,7 @@ @pytest.mark.unittest def test_shape_fn_ppo(): - data = ppo_data(torch.randn(3, 5, 8), None, None, None, None, None, None, None) + data = ppo_data(torch.randn(3, 5, 8), None, None, None, None, None, None, None, None) shape1 = shape_fn_ppo([data], {}) shape2 = shape_fn_ppo([], {'data': data}) assert shape1 == shape2 == (3, 5, 8) @@ -32,7 +32,7 @@ def test_ppo(use_value_clip, dual_clip, weight): value_old = value_new + torch.rand_like(value_new) * 0.1 adv = torch.rand(B) return_ = torch.randn(B) * 2 - data = ppo_data(logit_new, logit_old, action, value_new, value_old, adv, return_, weight) + data = ppo_data(logit_new, logit_old, action, value_new, value_old, adv, return_, weight, None) loss, info = ppo_error(data, use_value_clip=use_value_clip, dual_clip=dual_clip) assert all([l.shape == tuple() for l in loss]) assert all([np.isscalar(i) for i in info]) @@ -54,7 +54,7 @@ def test_mappo(): value_old = value_new + torch.rand_like(value_new) * 0.1 adv = torch.rand(B, A) return_ = torch.randn(B, A) * 2 - data = ppo_data(logit_new, logit_old, action, value_new, value_old, adv, return_, None) + data = ppo_data(logit_new, logit_old, action, value_new, value_old, adv, return_, None, None) loss, info = ppo_error(data) assert all([l.shape == tuple() for l in loss]) assert all([np.isscalar(i) for i in info]) @@ -80,7 +80,7 @@ def test_ppo_error_continous(use_value_clip, dual_clip, weight): value_old = value_new + torch.rand_like(value_new) * 0.1 adv = torch.rand(B) return_ = torch.randn(B) * 2 - data = ppo_data(mu_sigma_new, mu_sigma_old, action, value_new, value_old, adv, return_, weight) + data = ppo_data(mu_sigma_new, mu_sigma_old, action, value_new, value_old, adv, return_, weight, None) loss, info = ppo_error_continuous(data, use_value_clip=use_value_clip, dual_clip=dual_clip) assert all([l.shape == tuple() for l in loss]) assert all([np.isscalar(i) for i in info]) diff --git a/ding/rl_utils/tests/test_ppo_rlhf.py b/ding/rl_utils/tests/test_ppo_rlhf.py new file mode 100644 index 0000000000..4227bc8eb1 --- /dev/null +++ b/ding/rl_utils/tests/test_ppo_rlhf.py @@ -0,0 +1,92 @@ +import pytest +import numpy as np +import torch +from ding.rl_utils import ppo_policy_data, ppo_value_data, ppo_policy_error, ppo_value_error + + +@pytest.fixture +def batch_size(): + return 4 + + +@pytest.fixture +def seq_length(): + return 8 + + +@pytest.fixture +def dictionary_num(): + return 1000 + + +@pytest.mark.unittest +def test_policy_loss_without_mask(batch_size: int, seq_length: int, dictionary_num: int): + # Create test data + logit_new = torch.randn(batch_size, seq_length, dictionary_num).requires_grad_(True) + logit_old = logit_new + torch.randn_like(logit_new) * 0.1 + logit_pretrained = logit_new + torch.randn_like(logit_new) * 0.1 + action = torch.randint(0, 10, (batch_size, seq_length)) + advantages = torch.randn(batch_size, seq_length) + + # Compute loss + data = ppo_policy_data(logit_new, logit_old, action, advantages, weight=None, logit_pretrained=logit_pretrained) + loss, info = ppo_policy_error(data, clip_ratio=0.2, entropy_bonus=False) + + # Verify output + assert isinstance(loss.policy_loss, torch.Tensor) + assert loss.policy_loss.shape == torch.Size([]) # scalar + assert not torch.isnan(loss.policy_loss) + assert not torch.isinf(loss.policy_loss) + assert logit_new.grad is None + loss.policy_loss.backward() + assert isinstance(logit_new.grad, torch.Tensor) + assert all([np.isscalar(i) for i in info]) + + +@pytest.mark.unittest +def test_policy_loss_with_mask(batch_size: int, seq_length: int, dictionary_num: int): + # Create test data + logit_new = torch.randn(batch_size, seq_length, dictionary_num).requires_grad_(True) + logit_old = logit_new + torch.randn_like(logit_new) * 0.1 + logit_pretrained = logit_new + torch.randn_like(logit_new) * 0.1 + action = torch.randint(0, 10, (batch_size, seq_length)) + advantages = torch.randn(batch_size, seq_length) + action_mask = torch.ones(batch_size, seq_length) + action_mask[:, -2:] = 0 # Set last two timesteps as padding + + # Compute loss + data = ppo_policy_data( + logit_new, logit_old, action, advantages, weight=action_mask, logit_pretrained=logit_pretrained + ) + loss, info = ppo_policy_error(data, clip_ratio=0.2, entropy_bonus=False) + + # Verify output + assert isinstance(loss.policy_loss, torch.Tensor) + assert loss.policy_loss.shape == torch.Size([]) # scalar + assert not torch.isnan(loss.policy_loss) + assert not torch.isinf(loss.policy_loss) + assert logit_new.grad is None + loss.policy_loss.backward() + assert isinstance(logit_new.grad, torch.Tensor) + assert all([np.isscalar(i) for i in info]) + + +@pytest.mark.unittest +def test_value_loss(batch_size: int, seq_length: int): + # Create test data + values = torch.randn(batch_size, seq_length).requires_grad_(True) + old_values = values + torch.randn_like(values) * 0.1 + returns = torch.randn(batch_size, seq_length) + + # Compute loss + data = ppo_value_data(values, old_values, returns, weight=None) + value_loss = ppo_value_error(data, clip_ratio=0.2, use_value_clip=True) + + # Verify output + assert isinstance(value_loss, torch.Tensor) + assert value_loss.shape == torch.Size([]) + assert not torch.isnan(value_loss) + assert not torch.isinf(value_loss) + assert values.grad is None + value_loss.backward() + assert isinstance(values.grad, torch.Tensor) diff --git a/ding/rl_utils/tests/test_rloo_rlhf.py b/ding/rl_utils/tests/test_rloo_rlhf.py new file mode 100644 index 0000000000..d7899e7185 --- /dev/null +++ b/ding/rl_utils/tests/test_rloo_rlhf.py @@ -0,0 +1,78 @@ +import pytest +import numpy as np +import torch +from ding.rl_utils.rloo import (rloo_policy_data, rloo_policy_error) + + +@pytest.fixture +def batch_size(): + return 4 + + +@pytest.fixture +def seq_length(): + return 8 + + +@pytest.fixture +def dictionary_num(): + return 1000 + + +@pytest.mark.unittest +def test_rloo_policy_loss_without_mask(batch_size, seq_length, dictionary_num): + """Test RLOO policy loss calculation without mask""" + # Create test data + logit_new = torch.randn(batch_size, seq_length, dictionary_num).requires_grad_(True) + logit_old = logit_new + torch.randn_like(logit_new) * 0.1 + action = torch.randint(0, dictionary_num, (batch_size, seq_length)) + reward = torch.randn(batch_size) + + # Calculate loss + data = rloo_policy_data(logit_new=logit_new, logit_old=logit_old, action=action, reward=reward, weight=None) + loss, info = rloo_policy_error(data, clip_ratio=0.2) + + # Verify outputs + assert isinstance(loss, torch.Tensor) + assert loss.shape == torch.Size([]) # Ensure scalar output + assert not torch.isnan(loss) + assert not torch.isinf(loss) + assert logit_new.grad is None + loss.backward() + assert isinstance(logit_new.grad, torch.Tensor) + assert all([np.isscalar(v) for v in info._asdict().values()]) + + # Verify metrics + assert 'approx_kl' in info._asdict() + assert 'clipfrac' in info._asdict() + assert all([np.isscalar(v) for v in info._asdict().values()]) + + +@pytest.mark.unittest +def test_rloo_policy_loss_with_mask(batch_size, seq_length, dictionary_num): + """Test RLOO policy loss calculation with mask""" + # Create test data + logit_new = torch.randn(batch_size, seq_length, dictionary_num).requires_grad_(True) + logit_old = logit_new + torch.randn_like(logit_new) * 0.1 + action = torch.randint(0, dictionary_num, (batch_size, seq_length)) + reward = torch.randn(batch_size) + action_mask = torch.ones(batch_size, seq_length) + action_mask[:, -2:] = 0 + + # Calculate loss + data = rloo_policy_data(logit_new=logit_new, logit_old=logit_old, action=action, reward=reward, weight=action_mask) + loss, info = rloo_policy_error(data, clip_ratio=0.2) + + # Verify outputs + assert isinstance(loss, torch.Tensor) + assert loss.shape == torch.Size([]) # Ensure scalar output + assert not torch.isnan(loss) + assert not torch.isinf(loss) + assert logit_new.grad is None + loss.backward() + assert isinstance(logit_new.grad, torch.Tensor) + + # Verify metrics + assert 'approx_kl' in info._asdict() + assert 'clipfrac' in info._asdict() + assert all([np.isscalar(v) for v in info._asdict().values()]) diff --git a/ding/rl_utils/tests/test_td.py b/ding/rl_utils/tests/test_td.py index 3b792ddc35..c710695cb2 100644 --- a/ding/rl_utils/tests/test_td.py +++ b/ding/rl_utils/tests/test_td.py @@ -4,7 +4,8 @@ td_lambda_error, q_nstep_td_error_with_rescale, dist_1step_td_data, dist_1step_td_error, dist_nstep_td_data,\ dqfd_nstep_td_data, dqfd_nstep_td_error, dist_nstep_td_error, v_1step_td_data, v_1step_td_error, v_nstep_td_data,\ v_nstep_td_error, q_nstep_sql_td_error, iqn_nstep_td_data, iqn_nstep_td_error,\ - fqf_nstep_td_data, fqf_nstep_td_error, qrdqn_nstep_td_data, qrdqn_nstep_td_error + fqf_nstep_td_data, fqf_nstep_td_error, qrdqn_nstep_td_data, qrdqn_nstep_td_error, bdq_nstep_td_error,\ + m_q_1step_td_data, m_q_1step_td_error from ding.rl_utils.td import shape_fn_dntd, shape_fn_qntd, shape_fn_td_lambda, shape_fn_qntd_rescale @@ -35,6 +36,36 @@ def test_q_nstep_td(): assert isinstance(q.grad, torch.Tensor) +@pytest.mark.unittest +def test_bdq_nstep_td(): + batch_size = 8 + branch_num = 6 + action_per_branch = 3 + next_q = torch.randn(batch_size, branch_num, action_per_branch) + done = torch.randn(batch_size) + action = torch.randint(0, action_per_branch, size=(batch_size, branch_num)) + next_action = torch.randint(0, action_per_branch, size=(batch_size, branch_num)) + for nstep in range(1, 10): + q = torch.randn(batch_size, branch_num, action_per_branch).requires_grad_(True) + reward = torch.rand(nstep, batch_size) + data = q_nstep_td_data(q, next_q, action, next_action, reward, done, None) + loss, td_error_per_sample = bdq_nstep_td_error(data, 0.95, nstep=nstep) + assert td_error_per_sample.shape == (batch_size, ) + assert loss.shape == () + assert q.grad is None + loss.backward() + assert isinstance(q.grad, torch.Tensor) + data = q_nstep_td_data(q, next_q, action, next_action, reward, done, None) + loss, td_error_per_sample = bdq_nstep_td_error(data, 0.95, nstep=nstep, cum_reward=True) + value_gamma = torch.tensor(0.9) + data = q_nstep_td_data(q, next_q, action, next_action, reward, done, None) + loss, td_error_per_sample = bdq_nstep_td_error( + data, 0.95, nstep=nstep, cum_reward=True, value_gamma=value_gamma + ) + loss.backward() + assert isinstance(q.grad, torch.Tensor) + + @pytest.mark.unittest def test_q_nstep_td_ngu(): batch_size = 4 @@ -555,3 +586,25 @@ def test_fn_td_lambda(): assert tmp == reward.shape[0] tmp = shape_fn_td_lambda([data], {}) assert tmp == reward.shape + + +@pytest.mark.unittest +def test_fn_m_q_1step_td_error(): + batch_size = 128 + action_dim = 9 + q = torch.randn(batch_size, action_dim).requires_grad_(True) + target_q_current = torch.randn(batch_size, action_dim).requires_grad_(False) + target_q_next = torch.randn(batch_size, action_dim).requires_grad_(False) + done = torch.randn(batch_size) + action = torch.randint(0, action_dim, size=(batch_size, )) + reward = torch.randn(batch_size) + data = m_q_1step_td_data(q, target_q_current, target_q_next, action, reward, done, None) + loss, td_error_per_sample, action_gap, clip_frac = m_q_1step_td_error(data, 0.99, 0.03, 0.6) + + assert loss.shape == () + assert q.grad is None + loss.backward() + assert isinstance(q.grad, torch.Tensor) + assert clip_frac.mean().item() <= 1 + assert action_gap.item() > 0 + assert td_error_per_sample.shape == (batch_size, ) diff --git a/ding/rl_utils/tests/test_value_rescale.py b/ding/rl_utils/tests/test_value_rescale.py index 86cac06b06..ce5cb0c9b9 100644 --- a/ding/rl_utils/tests/test_value_rescale.py +++ b/ding/rl_utils/tests/test_value_rescale.py @@ -1,6 +1,6 @@ import pytest import torch -from ding.rl_utils.value_rescale import value_inv_transform, value_transform +from ding.rl_utils.value_rescale import value_inv_transform, value_transform, symlog, inv_symlog @pytest.mark.unittest @@ -24,3 +24,26 @@ def test_trans_inverse(self): diff = value_inv_transform(value_transform(t)) - t assert pytest.approx(diff.abs().max().item(), abs=2e-5) == 0 assert pytest.approx(diff.abs().max().item(), abs=2e-5) == 0 + + +@pytest.mark.unittest +class TestSymlog: + + def test_symlog(self): + for _ in range(10): + t = torch.rand((3, 4)) + assert isinstance(symlog(t), torch.Tensor) + assert symlog(t).shape == t.shape + + def test_inv_symlog(self): + for _ in range(10): + t = torch.rand((3, 4)) + assert isinstance(inv_symlog(t), torch.Tensor) + assert inv_symlog(t).shape == t.shape + + def test_trans_inverse(self): + for _ in range(10): + t = torch.rand((4, 16)) + diff = inv_symlog(symlog(t)) - t + assert pytest.approx(diff.abs().max().item(), abs=2e-5) == 0 + assert pytest.approx(diff.abs().max().item(), abs=2e-5) == 0 diff --git a/ding/rl_utils/tests/test_vtrace.py b/ding/rl_utils/tests/test_vtrace.py index d69a5c3c58..a5fa94fa0b 100644 --- a/ding/rl_utils/tests/test_vtrace.py +++ b/ding/rl_utils/tests/test_vtrace.py @@ -1,10 +1,10 @@ import pytest import torch -from ding.rl_utils import vtrace_data, vtrace_error +from ding.rl_utils import vtrace_data, vtrace_error_discrete_action, vtrace_error_continuous_action @pytest.mark.unittest -def test_vtrace(): +def test_vtrace_discrete_action(): T, B, N = 4, 8, 16 value = torch.randn(T + 1, B).requires_grad_(True) reward = torch.rand(T, B) @@ -12,7 +12,7 @@ def test_vtrace(): behaviour_output = torch.randn(T, B, N) action = torch.randint(0, N, size=(T, B)) data = vtrace_data(target_output, behaviour_output, action, value, reward, None) - loss = vtrace_error(data, rho_clip_ratio=1.1) + loss = vtrace_error_discrete_action(data, rho_clip_ratio=1.1) assert all([l.shape == tuple() for l in loss]) assert target_output.grad is None assert value.grad is None @@ -20,3 +20,28 @@ def test_vtrace(): loss.backward() assert isinstance(target_output, torch.Tensor) assert isinstance(value, torch.Tensor) + + +@pytest.mark.unittest +def test_vtrace_continuous_action(): + T, B, N = 4, 8, 16 + value = torch.randn(T + 1, B).requires_grad_(True) + reward = torch.rand(T, B) + target_output = {} + target_output['mu'] = torch.randn(T, B, N).requires_grad_(True) + target_output['sigma'] = torch.exp(torch.randn(T, B, N).requires_grad_(True)) + behaviour_output = {} + behaviour_output['mu'] = torch.randn(T, B, N) + behaviour_output['sigma'] = torch.exp(torch.randn(T, B, N)) + action = torch.randn((T, B, N)) + data = vtrace_data(target_output, behaviour_output, action, value, reward, None) + loss = vtrace_error_continuous_action(data, rho_clip_ratio=1.1) + assert all([l.shape == tuple() for l in loss]) + assert target_output['mu'].grad is None + assert target_output['sigma'].grad is None + assert value.grad is None + loss = sum(loss) + loss.backward() + assert isinstance(target_output['mu'], torch.Tensor) + assert isinstance(target_output['sigma'], torch.Tensor) + assert isinstance(value, torch.Tensor) diff --git a/ding/rl_utils/upgo.py b/ding/rl_utils/upgo.py index 366d67ada4..e0f54159ea 100644 --- a/ding/rl_utils/upgo.py +++ b/ding/rl_utils/upgo.py @@ -5,6 +5,21 @@ def tb_cross_entropy(logit, label, mask=None): + """ + Overview: + Compute the cross entropy loss for label and logit, with mask support + Arguments: + - logit (:obj:`torch.Tensor`): the logit tensor, of size [T, B, N] or [T, B, N, N2] + - label (:obj:`torch.Tensor`): the label tensor, of size [T, B] or [T, B, N2] + - mask (:obj:`torch.Tensor` or :obj:`None`): the mask tensor, of size [T, B] or [T, B, N2] + Returns: + - ce (:obj:`torch.Tensor`): the computed cross entropy, of size [T, B] + Examples: + >>> T, B, N, N2 = 4, 8, 5, 7 + >>> logit = torch.randn(T, B, N, N2).softmax(-1).requires_grad_(True) + >>> action = logit.argmax(-1).detach() + >>> ce = tb_cross_entropy(logit, action) + """ assert (len(label.shape) >= 2) T, B = label.shape[:2] # Special 2D case @@ -29,7 +44,7 @@ def tb_cross_entropy(logit, label, mask=None): def upgo_returns(rewards: torch.Tensor, bootstrap_values: torch.Tensor) -> torch.Tensor: - r""" + """ Overview: Computing UPGO return targets. Also notice there is no special handling for the terminal state. Arguments: @@ -40,6 +55,11 @@ def upgo_returns(rewards: torch.Tensor, bootstrap_values: torch.Tensor) -> torch Returns: - ret (:obj:`torch.Tensor`): Computed lambda return value for each state from 0 to T-1, \ of size [T_traj, batchsize] + Examples: + >>> T, B, N, N2 = 4, 8, 5, 7 + >>> rewards = torch.randn(T, B) + >>> bootstrap_values = torch.randn(T + 1, B).requires_grad_(True) + >>> returns = upgo_returns(rewards, bootstrap_values) """ # UPGO can be viewed as a lambda return! The trace continues for V_t (i.e. lambda = 1.0) if r_tp1 + V_tp2 > V_tp1. # as the lambdas[-1, :] is ignored in generalized_lambda_returns, we don't care about bootstrap_values_tp2[-1] @@ -62,7 +82,7 @@ def upgo_loss( bootstrap_values: torch.Tensor, mask=None ) -> torch.Tensor: - r""" + """ Overview: Computing UPGO loss given constant gamma and lambda. There is no special handling for terminal state value, if the last state in trajectory is the terminal, just pass a 0 as bootstrap_terminal_value. @@ -76,6 +96,10 @@ def upgo_loss( of size [T_traj+1, batchsize] Returns: - loss (:obj:`torch.Tensor`): Computed importance sampled UPGO loss, averaged over the samples, of size [] + Examples: + >>> T, B, N, N2 = 4, 8, 5, 7 + >>> rhos = torch.randn(T, B) + >>> loss = upgo_loss(logit, rhos, action, rewards, bootstrap_values) """ # discard the value at T as it should be considered in the next slice with torch.no_grad(): diff --git a/ding/rl_utils/value_rescale.py b/ding/rl_utils/value_rescale.py index 1120fae883..e51aed4d83 100644 --- a/ding/rl_utils/value_rescale.py +++ b/ding/rl_utils/value_rescale.py @@ -1,20 +1,63 @@ -""" -Referenced papar -""" import torch def value_transform(x: torch.Tensor, eps: float = 1e-2) -> torch.Tensor: - r""" + """ Overview: - :math: `h(x) = sign(x)(\sqrt{(abs(x)+1)} - 1) + \eps * x` . + A function to reduce the scale of the action-value function. + :math: `h(x) = sign(x)(\sqrt{(abs(x)+1)} - 1) + \epsilon * x` . + Arguments: + - x: (:obj:`torch.Tensor`) The input tensor to be normalized. + - eps: (:obj:`float`) The coefficient of the additive regularization term \ + to ensure inverse function is Lipschitz continuous + Returns: + - (:obj:`torch.Tensor`) Normalized tensor. + + .. note:: + Observe and Look Further: Achieving Consistent Performance on Atari (https://arxiv.org/abs/1805.11593). """ return torch.sign(x) * (torch.sqrt(torch.abs(x) + 1) - 1) + eps * x def value_inv_transform(x: torch.Tensor, eps: float = 1e-2) -> torch.Tensor: - r""" + """ Overview: - :math: `h^{-1}(x) = sign(x)({(\frac{\sqrt{1+4\eps(|x|+1+\eps)}-1}{2\eps})}^2-1)` . + The inverse form of value rescale. + :math: `h^{-1}(x) = sign(x)({(\frac{\sqrt{1+4\epsilon(|x|+1+\epsilon)}-1}{2\epsilon})}^2-1)` . + Arguments: + - x: (:obj:`torch.Tensor`) The input tensor to be unnormalized. + - eps: (:obj:`float`) The coefficient of the additive regularization term \ + to ensure inverse function is Lipschitz continuous + Returns: + - (:obj:`torch.Tensor`) Unnormalized tensor. """ return torch.sign(x) * (((torch.sqrt(1 + 4 * eps * (torch.abs(x) + 1 + eps)) - 1) / (2 * eps)) ** 2 - 1) + + +def symlog(x: torch.Tensor) -> torch.Tensor: + """ + Overview: + A function to normalize the targets. + :math: `symlog(x) = sign(x)(\ln{|x|+1})` . + Arguments: + - x: (:obj:`torch.Tensor`) The input tensor to be normalized. + Returns: + - (:obj:`torch.Tensor`) Normalized tensor. + + .. note:: + Mastering Diverse Domains through World Models (https://arxiv.org/abs/2301.04104) + """ + return torch.sign(x) * (torch.log(torch.abs(x) + 1)) + + +def inv_symlog(x: torch.Tensor) -> torch.Tensor: + """ + Overview: + The inverse form of symlog. + :math: `symexp(x) = sign(x)(\exp{|x|}-1)` . + Arguments: + - x: (:obj:`torch.Tensor`) The input tensor to be unnormalized. + Returns: + - (:obj:`torch.Tensor`) Unnormalized tensor. + """ + return torch.sign(x) * (torch.exp(torch.abs(x)) - 1) diff --git a/ding/rl_utils/vtrace.py b/ding/rl_utils/vtrace.py index b8c262e1a0..44728fd6cd 100644 --- a/ding/rl_utils/vtrace.py +++ b/ding/rl_utils/vtrace.py @@ -1,5 +1,6 @@ import torch import torch.nn.functional as F +from torch.distributions import Categorical, Independent, Normal from collections import namedtuple from .isw import compute_importance_weights from ding.hpc_rl import hpc_wrapper @@ -14,7 +15,7 @@ def vtrace_nstep_return(clipped_rhos, clipped_cs, reward, bootstrap_values, gamm Shapes: - clipped_rhos (:obj:`torch.FloatTensor`): :math:`(T, B)`, where T is timestep, B is batch size - clipped_cs (:obj:`torch.FloatTensor`): :math:`(T, B)` - - reward: (:obj:`torch.FloatTensor`): :math:`(T, B)` + - reward (:obj:`torch.FloatTensor`): :math:`(T, B)` - bootstrap_values (:obj:`torch.FloatTensor`): :math:`(T+1, B)` - vtrace_return (:obj:`torch.FloatTensor`): :math:`(T, B)` """ @@ -36,10 +37,10 @@ def vtrace_advantage(clipped_pg_rhos, reward, return_, bootstrap_values, gamma): - vtrace_advantage (:obj:`namedtuple`): the vtrace loss item, all of them are the differentiable 0-dim tensor Shapes: - clipped_pg_rhos (:obj:`torch.FloatTensor`): :math:`(T, B)`, where T is timestep, B is batch size - - reward: (:obj:`torch.FloatTensor`): :math:`(T, B)` - - return_ (:obj:`torch.FloatTensor`): :math:`(T, B)` + - reward (:obj:`torch.FloatTensor`): :math:`(T, B)` + - return (:obj:`torch.FloatTensor`): :math:`(T, B)` - bootstrap_values (:obj:`torch.FloatTensor`): :math:`(T, B)` - - vtrace_advantage (:obj:`torch.FloatTensor`): :math:`(T, B)` + - vtrace_advantage (:obj:`torch.FloatTensor`): :math:`(T, B)` """ return clipped_pg_rhos * (reward + gamma * return_ - bootstrap_values) @@ -48,7 +49,7 @@ def vtrace_advantage(clipped_pg_rhos, reward, return_, bootstrap_values, gamma): vtrace_loss = namedtuple('vtrace_loss', ['policy_loss', 'value_loss', 'entropy_loss']) -def shape_fn_vtrace(args, kwargs): +def shape_fn_vtrace_discrete_action(args, kwargs): r""" Overview: Return shape of vtrace for hpc @@ -63,12 +64,12 @@ def shape_fn_vtrace(args, kwargs): @hpc_wrapper( - shape_fn=shape_fn_vtrace, + shape_fn=shape_fn_vtrace_discrete_action, namedtuple_data=True, include_args=[0, 1, 2, 3, 4, 5], include_kwargs=['data', 'gamma', 'lambda_', 'rho_clip_ratio', 'c_clip_ratio', 'rho_pg_clip_ratio'] ) -def vtrace_error( +def vtrace_error_discrete_action( data: namedtuple, gamma: float = 0.99, lambda_: float = 0.95, @@ -106,10 +107,19 @@ def vtrace_error( - value (:obj:`torch.FloatTensor`): :math:`(T+1, B)` - reward (:obj:`torch.LongTensor`): :math:`(T, B)` - weight (:obj:`torch.LongTensor`): :math:`(T, B)` + Examples: + >>> T, B, N = 4, 8, 16 + >>> value = torch.randn(T + 1, B).requires_grad_(True) + >>> reward = torch.rand(T, B) + >>> target_output = torch.randn(T, B, N).requires_grad_(True) + >>> behaviour_output = torch.randn(T, B, N) + >>> action = torch.randint(0, N, size=(T, B)) + >>> data = vtrace_data(target_output, behaviour_output, action, value, reward, None) + >>> loss = vtrace_error_discrete_action(data, rho_clip_ratio=1.1) """ target_output, behaviour_output, action, value, reward, weight = data with torch.no_grad(): - IS = compute_importance_weights(target_output, behaviour_output, action) + IS = compute_importance_weights(target_output, behaviour_output, action, 'discrete') rhos = torch.clamp(IS, max=rho_clip_ratio) cs = torch.clamp(IS, max=c_clip_ratio) return_ = vtrace_nstep_return(rhos, cs, reward, value, gamma, lambda_) @@ -119,7 +129,83 @@ def vtrace_error( if weight is None: weight = torch.ones_like(reward) - dist_target = torch.distributions.Categorical(logits=target_output) + dist_target = Categorical(logits=target_output) + pg_loss = -(dist_target.log_prob(action) * adv * weight).mean() + value_loss = (F.mse_loss(value[:-1], return_, reduction='none') * weight).mean() + entropy_loss = (dist_target.entropy() * weight).mean() + return vtrace_loss(pg_loss, value_loss, entropy_loss) + + +def vtrace_error_continuous_action( + data: namedtuple, + gamma: float = 0.99, + lambda_: float = 0.95, + rho_clip_ratio: float = 1.0, + c_clip_ratio: float = 1.0, + rho_pg_clip_ratio: float = 1.0 +): + """ + Overview: + Implementation of vtrace(IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner\ + Architectures), (arXiv:1802.01561) + Arguments: + - data (:obj:`namedtuple`): input data with fields shown in ``vtrace_data`` + - target_output (:obj:`dict{key:torch.Tensor}`): the output taking the action \ + by the current policy network, usually this output is network output, \ + which represents the distribution by reparameterization trick. + - behaviour_output (:obj:`dict{key:torch.Tensor}`): the output taking the action \ + by the behaviour policy network, usually this output is network output logit, \ + which represents the distribution by reparameterization trick. + - action (:obj:`torch.Tensor`): the chosen action(index for the discrete action space) in trajectory, \ + i.e.: behaviour_action + - gamma: (:obj:`float`): the future discount factor, defaults to 0.95 + - lambda: (:obj:`float`): mix factor between 1-step (lambda_=0) and n-step, defaults to 1.0 + - rho_clip_ratio (:obj:`float`): the clipping threshold for importance weights (rho) when calculating\ + the baseline targets (vs) + - c_clip_ratio (:obj:`float`): the clipping threshold for importance weights (c) when calculating\ + the baseline targets (vs) + - rho_pg_clip_ratio (:obj:`float`): the clipping threshold for importance weights (rho) when calculating\ + the policy gradient advantage + Returns: + - trace_loss (:obj:`namedtuple`): the vtrace loss item, all of them are the differentiable 0-dim tensor + Shapes: + - target_output (:obj:`dict{key:torch.FloatTensor}`): :math:`(T, B, N)`, \ + where T is timestep, B is batch size and \ + N is action dim. The keys are usually parameters of reparameterization trick. + - behaviour_output (:obj:`dict{key:torch.FloatTensor}`): :math:`(T, B, N)` + - action (:obj:`torch.LongTensor`): :math:`(T, B)` + - value (:obj:`torch.FloatTensor`): :math:`(T+1, B)` + - reward (:obj:`torch.LongTensor`): :math:`(T, B)` + - weight (:obj:`torch.LongTensor`): :math:`(T, B)` + Examples: + >>> T, B, N = 4, 8, 16 + >>> value = torch.randn(T + 1, B).requires_grad_(True) + >>> reward = torch.rand(T, B) + >>> target_output = dict( + >>> 'mu': torch.randn(T, B, N).requires_grad_(True), + >>> 'sigma': torch.exp(torch.randn(T, B, N).requires_grad_(True)), + >>> ) + >>> behaviour_output = dict( + >>> 'mu': torch.randn(T, B, N), + >>> 'sigma': torch.exp(torch.randn(T, B, N)), + >>> ) + >>> action = torch.randn((T, B, N)) + >>> data = vtrace_data(target_output, behaviour_output, action, value, reward, None) + >>> loss = vtrace_error_continuous_action(data, rho_clip_ratio=1.1) + """ + target_output, behaviour_output, action, value, reward, weight = data + with torch.no_grad(): + IS = compute_importance_weights(target_output, behaviour_output, action, 'continuous') + rhos = torch.clamp(IS, max=rho_clip_ratio) + cs = torch.clamp(IS, max=c_clip_ratio) + return_ = vtrace_nstep_return(rhos, cs, reward, value, gamma, lambda_) + pg_rhos = torch.clamp(IS, max=rho_pg_clip_ratio) + return_t_plus_1 = torch.cat([return_[1:], value[-1:]], 0) + adv = vtrace_advantage(pg_rhos, reward, return_t_plus_1, value[:-1], gamma) + + if weight is None: + weight = torch.ones_like(reward) + dist_target = Independent(Normal(loc=target_output['mu'], scale=target_output['sigma']), 1) pg_loss = -(dist_target.log_prob(action) * adv * weight).mean() value_loss = (F.mse_loss(value[:-1], return_, reduction='none') * weight).mean() entropy_loss = (dist_target.entropy() * weight).mean() diff --git a/ding/torch_utils/__init__.py b/ding/torch_utils/__init__.py old mode 100644 new mode 100755 index 9bd2e63a3f..151b4da7e1 --- a/ding/torch_utils/__init__.py +++ b/ding/torch_utils/__init__.py @@ -1,6 +1,7 @@ from .checkpoint_helper import build_checkpoint_helper, CountVar, auto_checkpoint from .data_helper import to_device, to_tensor, to_ndarray, to_list, to_dtype, same_shape, tensor_to_list, \ - build_log_buffer, CudaFetcher, get_tensor_data, unsqueeze, get_null_data + build_log_buffer, CudaFetcher, get_tensor_data, unsqueeze, squeeze, get_null_data, get_shape0, to_item, \ + zeros_like from .distribution import CategoricalPd, CategoricalPdPytorch from .metric import levenshtein_distance, hamming_distance from .network import * @@ -10,3 +11,4 @@ from .math_helper import cov from .dataparallel import DataParallel from .reshape_helper import fold_batch, unfold_batch, unsqueeze_repeat +from .parameter import NonegativeParameter, TanhParameter diff --git a/ding/torch_utils/backend_helper.py b/ding/torch_utils/backend_helper.py new file mode 100644 index 0000000000..b7346b2d79 --- /dev/null +++ b/ding/torch_utils/backend_helper.py @@ -0,0 +1,12 @@ +import torch + + +def enable_tf32() -> None: + """ + Overview: + Enable tf32 on matmul and cudnn for faster computation. This only works on Ampere GPU devices. \ + For detailed information, please refer to: \ + https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices. + """ + torch.backends.cuda.matmul.allow_tf32 = True # allow tf32 on matmul + torch.backends.cudnn.allow_tf32 = True # allow tf32 on cudnn diff --git a/ding/torch_utils/checkpoint_helper.py b/ding/torch_utils/checkpoint_helper.py index d51d15d07e..5d600b556e 100644 --- a/ding/torch_utils/checkpoint_helper.py +++ b/ding/torch_utils/checkpoint_helper.py @@ -11,7 +11,7 @@ def build_checkpoint_helper(cfg): - r""" + """ Overview: Use config to build checkpoint helper. Arguments: @@ -23,18 +23,18 @@ def build_checkpoint_helper(cfg): class CheckpointHelper: - r""" + """ Overview: Help to save or load checkpoint by give args. - Interface: - save, load + Interfaces: + ``__init__``, ``save``, ``load``, ``_remove_prefix``, ``_add_prefix``, ``_load_matched_model_state_dict`` """ def __init__(self): pass def _remove_prefix(self, state_dict: dict, prefix: str = 'module.') -> dict: - r""" + """ Overview: Remove prefix in state_dict Arguments: @@ -53,7 +53,7 @@ def _remove_prefix(self, state_dict: dict, prefix: str = 'module.') -> dict: return new_state_dict def _add_prefix(self, state_dict: dict, prefix: str = 'module.') -> dict: - r""" + """ Overview: Add prefix in state_dict Arguments: @@ -77,7 +77,7 @@ def save( prefix_op: str = None, prefix: str = None, ) -> None: - r""" + """ Overview: Save checkpoint by given args Arguments: @@ -119,7 +119,7 @@ def save( logger.info('save checkpoint in {}'.format(path)) def _load_matched_model_state_dict(self, model: torch.nn.Module, ckpt_state_dict: dict) -> None: - r""" + """ Overview: Load matched model state_dict, and show mismatch keys between model's state_dict and checkpoint's state_dict Arguments: @@ -169,7 +169,7 @@ def load( logger_prefix: str = '', state_dict_mask: list = [], ): - r""" + """ Overview: Load checkpoint by given path Arguments: @@ -254,22 +254,36 @@ def load( class CountVar(object): - r""" + """ Overview: Number counter - Interface: - val, update, add + Interfaces: + ``__init__``, ``update``, ``add`` + Properties: + - val (:obj:`int`): the value of the counter """ def __init__(self, init_val: int) -> None: + """ + Overview: + Init the var counter + Arguments: + - init_val (:obj:`int`): the init value of the counter + """ + self._val = init_val @property def val(self) -> int: + """ + Overview: + Get the var counter + """ + return self._val def update(self, val: int) -> None: - r""" + """ Overview: Update the var counter Arguments: @@ -278,7 +292,7 @@ def update(self, val: int) -> None: self._val = val def add(self, add_num: int): - r""" + """ Overview: Add the number to counter Arguments: @@ -288,7 +302,7 @@ def add(self, add_num: int): def auto_checkpoint(func: Callable) -> Callable: - r""" + """ Overview: Create a wrapper to wrap function, and the wrapper will call the save_checkpoint method whenever an exception happens. diff --git a/ding/torch_utils/data_helper.py b/ding/torch_utils/data_helper.py index 9e0b8e7861..b906279b1e 100644 --- a/ding/torch_utils/data_helper.py +++ b/ding/torch_utils/data_helper.py @@ -8,18 +8,32 @@ import numpy as np import torch +import treetensor.torch as ttorch + +from ding.utils.default_helper import get_shape0 def to_device(item: Any, device: str, ignore_keys: list = []) -> Any: """ Overview: - Transfer data to certain device + Transfer data to certain device. Arguments: - - item (:obj:`Any`): the item to be transferred - - device (:obj:`str`): the device wanted - - ignore_keys (:obj:`list`): the keys to be ignored in transfer, defalut set to empty + - item (:obj:`Any`): The item to be transferred. + - device (:obj:`str`): The device wanted. + - ignore_keys (:obj:`list`): The keys to be ignored in transfer, default set to empty. Returns: - - item (:obj:`Any`): the transferred item + - item (:obj:`Any`): The transferred item. + Examples: + >>> setup_data_dict['module'] = nn.Linear(3, 5) + >>> device = 'cuda' + >>> cuda_d = to_device(setup_data_dict, device, ignore_keys=['module']) + >>> assert cuda_d['module'].weight.device == torch.device('cpu') + + Examples: + >>> setup_data_dict['module'] = nn.Linear(3, 5) + >>> device = 'cuda' + >>> cuda_d = to_device(setup_data_dict, device) + >>> assert cuda_d['module'].weight.device == torch.device('cuda:0') .. note: @@ -29,6 +43,15 @@ def to_device(item: Any, device: str, ignore_keys: list = []) -> Any: """ if isinstance(item, torch.nn.Module): return item.to(device) + elif isinstance(item, ttorch.Tensor): + if 'prev_state' in item: + prev_state = to_device(item.prev_state, device) + del item.prev_state + item = item.to(device) + item.prev_state = prev_state + return item + else: + return item.to(device) elif isinstance(item, torch.Tensor): return item.to(device) elif isinstance(item, Sequence): @@ -50,23 +73,39 @@ def to_device(item: Any, device: str, ignore_keys: list = []) -> Any: return item elif item is None or isinstance(item, str): return item + elif isinstance(item, torch.distributions.Distribution): # for compatibility + return item else: raise TypeError("not support item type: {}".format(type(item))) def to_dtype(item: Any, dtype: type) -> Any: - r""" + """ Overview: - Change data to certain dtype + Change data to certain dtype. Arguments: - - item (:obj:`Any`): the item to be dtype changed - - dtype (:obj:`type`): the type wanted + - item (:obj:`Any`): The item for changing the dtype. + - dtype (:obj:`type`): The type wanted. Returns: - - item (:obj:`object`): the dtype changed item + - item (:obj:`object`): The item with changed dtype. + Examples (tensor): + >>> t = torch.randint(0, 10, (3, 5)) + >>> tfloat = to_dtype(t, torch.float) + >>> assert tfloat.dtype == torch.float + + Examples (list): + >>> tlist = [torch.randint(0, 10, (3, 5))] + >>> tlfloat = to_dtype(tlist, torch.float) + >>> assert tlfloat[0].dtype == torch.float + + Examples (dict): + >>> tdict = {'t': torch.randint(0, 10, (3, 5))} + >>> tdictf = to_dtype(tdict, torch.float) + >>> assert tdictf['t'].dtype == torch.float .. note: - Now supports item type: :obj:`torch.Tensor`, :obj:`Sequence`, :obj:`dict` + Now supports item type: :obj:`torch.Tensor`, :obj:`Sequence`, :obj:`dict`. """ if isinstance(item, torch.Tensor): return item.to(dtype=dtype) @@ -79,23 +118,44 @@ def to_dtype(item: Any, dtype: type) -> Any: def to_tensor( - item: Any, - dtype: Optional[torch.dtype] = None, - ignore_keys: list = [], - transform_scalar: bool = True -) -> torch.Tensor: - r""" + item: Any, dtype: Optional[torch.dtype] = None, ignore_keys: list = [], transform_scalar: bool = True +) -> Any: + """ Overview: - Change `numpy.ndarray`, sequence of scalars to torch.Tensor, and keep other data types unchanged. + Convert ``numpy.ndarray`` object to ``torch.Tensor``. Arguments: - - item (:obj:`Any`): the item to be changed - - dtype (:obj:`type`): the type of wanted tensor + - item (:obj:`Any`): The ``numpy.ndarray`` objects to be converted. It can be exactly a ``numpy.ndarray`` \ + object or a container (list, tuple or dict) that contains several ``numpy.ndarray`` objects. + - dtype (:obj:`torch.dtype`): The type of wanted tensor. If set to ``None``, its dtype will be unchanged. + - ignore_keys (:obj:`list`): If the ``item`` is a dict, values whose keys are in ``ignore_keys`` will not \ + be converted. + - transform_scalar (:obj:`bool`): If set to ``True``, a scalar will be also converted to a tensor object. Returns: - - item (:obj:`torch.Tensor`): the change tensor + - item (:obj:`Any`): The converted tensors. + + Examples (scalar): + >>> i = 10 + >>> t = to_tensor(i) + >>> assert t.item() == i + + Examples (dict): + >>> d = {'i': i} + >>> dt = to_tensor(d, torch.int) + >>> assert dt['i'].item() == i + + Examples (named tuple): + >>> data_type = namedtuple('data_type', ['x', 'y']) + >>> inputs = data_type(np.random.random(3), 4) + >>> outputs = to_tensor(inputs, torch.float32) + >>> assert type(outputs) == data_type + >>> assert isinstance(outputs.x, torch.Tensor) + >>> assert isinstance(outputs.y, torch.Tensor) + >>> assert outputs.x.dtype == torch.float32 + >>> assert outputs.y.dtype == torch.float32 .. note: - Now supports item type: :obj:`dict`, :obj:`list`, :obj:`tuple` and :obj:`None` + Now supports item type: :obj:`dict`, :obj:`list`, :obj:`tuple` and :obj:`None`. """ def transform(d): @@ -153,19 +213,33 @@ def transform(d): raise TypeError("not support item type: {}".format(type(item))) -def to_ndarray(item: Any, dtype: np.dtype = None) -> np.ndarray: - r""" +def to_ndarray(item: Any, dtype: np.dtype = None) -> Any: + """ Overview: - Change `torch.Tensor`, sequence of scalars to ndarray, and keep other data types unchanged. + Convert ``torch.Tensor`` to ``numpy.ndarray``. Arguments: - - item (:obj:`object`): the item to be changed - - dtype (:obj:`type`): the type of wanted ndarray + - item (:obj:`Any`): The ``torch.Tensor`` objects to be converted. It can be exactly a ``torch.Tensor`` \ + object or a container (list, tuple or dict) that contains several ``torch.Tensor`` objects. + - dtype (:obj:`np.dtype`): The type of wanted array. If set to ``None``, its dtype will be unchanged. Returns: - - item (:obj:`object`): the changed ndarray + - item (:obj:`object`): The changed arrays. + + Examples (ndarray): + >>> t = torch.randn(3, 5) + >>> tarray1 = to_ndarray(t) + >>> assert tarray1.shape == (3, 5) + >>> assert isinstance(tarray1, np.ndarray) + + Examples (list): + >>> t = [torch.randn(5, ) for i in range(3)] + >>> tarray1 = to_ndarray(t, np.float32) + >>> assert isinstance(tarray1, list) + >>> assert tarray1[0].shape == (5, ) + >>> assert isinstance(tarray1[0], np.ndarray) .. note: - Now supports item type: :obj:`torch.Tensor`, :obj:`dict`, :obj:`list`, :obj:`tuple` and :obj:`None` + Now supports item type: :obj:`torch.Tensor`, :obj:`dict`, :obj:`list`, :obj:`tuple` and :obj:`None`. """ def transform(d): @@ -204,26 +278,43 @@ def transform(d): elif isinstance(item, bool) or isinstance(item, str): return item elif np.isscalar(item): - return np.array(item) + if dtype is None: + return np.array(item) + else: + return np.array(item, dtype=dtype) elif item is None: return None else: raise TypeError("not support item type: {}".format(type(item))) -def to_list(item: Any) -> list: - r""" +def to_list(item: Any) -> Any: + """ Overview: - Transform `torch.Tensor`, `numpy.ndarray` to `list`, keep other data types unchanged + Convert ``torch.Tensor``, ``numpy.ndarray`` objects to ``list`` objects, and keep their dtypes unchanged. Arguments: - - item (:obj:`Any`): the item to be transformed + - item (:obj:`Any`): The item to be converted. Returns: - - item (:obj:`list`): the list after transformation + - item (:obj:`Any`): The list after conversion. + + Examples: + >>> data = { \ + 'tensor': torch.randn(4), \ + 'list': [True, False, False], \ + 'tuple': (4, 5, 6), \ + 'bool': True, \ + 'int': 10, \ + 'float': 10., \ + 'array': np.random.randn(4), \ + 'str': "asdf", \ + 'none': None, \ + } \ + >>> transformed_data = to_list(data) .. note:: Now supports item type: :obj:`torch.Tensor`, :obj:`numpy.ndarray`, :obj:`dict`, :obj:`list`, \ - :obj:`tuple` and :obj:`None` + :obj:`tuple` and :obj:`None`. """ if item is None: return item @@ -241,18 +332,41 @@ def to_list(item: Any) -> list: raise TypeError("not support item type: {}".format(type(item))) -def tensor_to_list(item): - r""" +def tensor_to_list(item: Any) -> Any: + """ Overview: - Transform `torch.Tensor` to `list`, keep other data types unchanged + Convert ``torch.Tensor`` objects to ``list``, and keep their dtypes unchanged. Arguments: - - item (:obj:`Any`): the item to be transformed + - item (:obj:`Any`): The item to be converted. Returns: - - item (:obj:`list`): the list after transformation + - item (:obj:`Any`): The lists after conversion. + + Examples (2d-tensor): + >>> t = torch.randn(3, 5) + >>> tlist1 = tensor_to_list(t) + >>> assert len(tlist1) == 3 + >>> assert len(tlist1[0]) == 5 + + Examples (1d-tensor): + >>> t = torch.randn(3, ) + >>> tlist1 = tensor_to_list(t) + >>> assert len(tlist1) == 3 + + Examples (list) + >>> t = [torch.randn(5, ) for i in range(3)] + >>> tlist1 = tensor_to_list(t) + >>> assert len(tlist1) == 3 + >>> assert len(tlist1[0]) == 5 + + Examples (dict): + >>> td = {'t': torch.randn(3, 5)} + >>> tdlist1 = tensor_to_list(td) + >>> assert len(tdlist1['t']) == 3 + >>> assert len(tdlist1['t'][0]) == 5 .. note:: - Now supports item type: :obj:`torch.Tensor`, :obj:`dict`, :obj:`list`, :obj:`tuple` and :obj:`None` + Now supports item type: :obj:`torch.Tensor`, :obj:`dict`, :obj:`list`, :obj:`tuple` and :obj:`None`. """ if item is None: return item @@ -268,14 +382,78 @@ def tensor_to_list(item): raise TypeError("not support item type: {}".format(type(item))) +def to_item(data: Any, ignore_error: bool = True) -> Any: + """ + Overview: + Convert data to python native scalar (i.e. data item), and keep their dtypes unchanged. + Arguments: + - data (:obj:`Any`): The data that needs to be converted. + - ignore_error (:obj:`bool`): Whether to ignore the error when the data type is not supported. That is to \ + say, only the data can be transformed into a python native scalar will be returned. + Returns: + - data (:obj:`Any`): Converted data. + + Examples: + >>>> data = { \ + 'tensor': torch.randn(1), \ + 'list': [True, False, torch.randn(1)], \ + 'tuple': (4, 5, 6), \ + 'bool': True, \ + 'int': 10, \ + 'float': 10., \ + 'array': np.random.randn(1), \ + 'str': "asdf", \ + 'none': None, \ + } + >>>> new_data = to_item(data) + >>>> assert np.isscalar(new_data['tensor']) + >>>> assert np.isscalar(new_data['array']) + >>>> assert np.isscalar(new_data['list'][-1]) + + .. note:: + + Now supports item type: :obj:`torch.Tensor`, :obj:`torch.Tensor`, :obj:`ttorch.Tensor`, \ + :obj:`bool`, :obj:`str`, :obj:`dict`, :obj:`list`, :obj:`tuple` and :obj:`None`. + """ + if data is None: + return data + elif isinstance(data, bool) or isinstance(data, str): + return data + elif np.isscalar(data): + return data + elif isinstance(data, np.ndarray) or isinstance(data, torch.Tensor) or isinstance(data, ttorch.Tensor): + return data.item() + elif isinstance(data, list) or isinstance(data, tuple): + return [to_item(d) for d in data] + elif isinstance(data, dict): + new_data = {} + for k, v in data.items(): + if ignore_error: + try: + new_data[k] = to_item(v) + except (ValueError, RuntimeError): + pass + else: + new_data[k] = to_item(v) + return new_data + else: + raise TypeError("not support data type: {}".format(data)) + + def same_shape(data: list) -> bool: - r""" + """ Overview: - Judge whether all data elements in a list have the same shape. + Judge whether all data elements in a list have the same shapes. Arguments: - - data (:obj:`list`): the list of data + - data (:obj:`list`): The list of data. Returns: - - same (:obj:`bool`): whether the list of data all have the same shape + - same (:obj:`bool`): Whether the list of data all have the same shape. + + Examples: + >>> tlist = [torch.randn(3, 5) for i in range(5)] + >>> assert same_shape(tlist) + >>> tlist = [torch.randn(3, 5), torch.randn(4, 5)] + >>> assert not same_shape(tlist) """ assert (isinstance(data, list)) shapes = [t.shape for t in data] @@ -283,33 +461,61 @@ def same_shape(data: list) -> bool: class LogDict(dict): - ''' + """ Overview: - Derived from ``dict``; Would transform ``torch.Tensor`` to ``list`` for convenient logging. - ''' + Derived from ``dict``. Would convert ``torch.Tensor`` to ``list`` for convenient logging. + Interfaces: + ``_transform``, ``__setitem__``, ``update``. + """ - def _transform(self, data): + def _transform(self, data: Any) -> None: + """ + Overview: + Convert tensor objects to lists for better logging. + Arguments: + - data (:obj:`Any`): The input data to be converted. + """ if isinstance(data, torch.Tensor): new_data = data.tolist() else: new_data = data return new_data - def __setitem__(self, key, value): + def __setitem__(self, key: Any, value: Any) -> None: + """ + Overview: + Override the ``__setitem__`` function of built-in dict. + Arguments: + - key (:obj:`Any`): The key of the data item. + - value (:obj:`Any`): The value of the data item. + """ new_value = self._transform(value) super().__setitem__(key, new_value) - def update(self, data): + def update(self, data: dict) -> None: + """ + Overview: + Override the ``update`` function of built-in dict. + Arguments: + - data (:obj:`dict`): The dict for updating current object. + """ for k, v in data.items(): self.__setitem__(k, v) -def build_log_buffer(): - r""" +def build_log_buffer() -> LogDict: + """ Overview: - Builg log buffer, a subclass of dict, which can transform the input data into log format. + Build log buffer, a subclass of dict, which can convert the input data into log format. Returns: - - log_buffer (:obj:`LogDict`): Log buffer dict + - log_buffer (:obj:`LogDict`): Log buffer dict. + Examples: + >>> log_buffer = build_log_buffer() + >>> log_buffer['not_tensor'] = torch.randn(3) + >>> assert isinstance(log_buffer['not_tensor'], list) + >>> assert len(log_buffer['not_tensor']) == 3 + >>> log_buffer.update({'not_tensor': 4, 'a': 5}) + >>> assert log_buffer['not_tensor'] == 4 """ return LogDict() @@ -317,12 +523,21 @@ def build_log_buffer(): class CudaFetcher(object): """ Overview: - Fetch data from source, and transfer it to specified device. + Fetch data from source, and transfer it to a specified device. Interfaces: - run, close + ``__init__``, ``__next__``, ``run``, ``close``. """ def __init__(self, data_source: Iterable, device: str, queue_size: int = 4, sleep: float = 0.1) -> None: + """ + Overview: + Initialize the CudaFetcher object using the given arguments. + Arguments: + - data_source (:obj:`Iterable`): The iterable data source. + - device (:obj:`str`): The device to put data to, such as "cuda:0". + - queue_size (:obj:`int`): The internal size of queue, such as 4. + - sleep (:obj:`float`): Sleeping time when the internal queue is full. + """ self._source = data_source self._queue = Queue(maxsize=queue_size) self._stream = torch.cuda.Stream() @@ -331,13 +546,25 @@ def __init__(self, data_source: Iterable, device: str, queue_size: int = 4, slee self._device = device def __next__(self) -> Any: + """ + Overview: + Response to the request for data. Return one data item from the internal queue. + Returns: + - item (:obj:`Any`): The data item on the required device. + """ return self._queue.get() def run(self) -> None: """ Overview: - Start ``producer`` thread: Keep fetching data from source, - change the device, and put into ``queue`` for request. + Start ``producer`` thread: Keep fetching data from source, change the device, and put into \ + ``queue`` for request. + Examples: + >>> timer = EasyTimer() + >>> dataloader = iter([torch.randn(3, 3) for _ in range(10)]) + >>> dataloader = CudaFetcher(dataloader, device='cuda', sleep=0.1) + >>> dataloader.run() + >>> data = next(dataloader) """ self._end_flag = False self._producer_thread.start() @@ -350,6 +577,11 @@ def close(self) -> None: self._end_flag = True def _producer(self) -> None: + """ + Overview: + Keep fetching data from source, change the device, and put into ``queue`` for request. + """ + with torch.cuda.stream(self._stream): while not self._end_flag: if self._queue.full(): @@ -363,7 +595,21 @@ def _producer(self) -> None: def get_tensor_data(data: Any) -> Any: """ Overview: - Get pure tensor data from the given data(without disturbing grad computation graph) + Get pure tensor data from the given data (without disturbing grad computation graph). + Arguments: + - data (:obj:`Any`): The original data. It can be exactly a tensor or a container (Sequence or dict). + Returns: + - output (:obj:`Any`): The output data. + Examples: + >>> a = { \ + 'tensor': torch.tensor([1, 2, 3.], requires_grad=True), \ + 'list': [torch.tensor([1, 2, 3.], requires_grad=True) for _ in range(2)], \ + 'none': None \ + } + >>> tensor_a = get_tensor_data(a) + >>> assert not tensor_a['tensor'].requires_grad + >>> for t in tensor_a['list']: + >>> assert not t.requires_grad """ if isinstance(data, torch.Tensor): return data.data.clone() @@ -378,6 +624,30 @@ def get_tensor_data(data: Any) -> Any: def unsqueeze(data: Any, dim: int = 0) -> Any: + """ + Overview: + Unsqueeze the tensor data. + Arguments: + - data (:obj:`Any`): The original data. It can be exactly a tensor or a container (Sequence or dict). + - dim (:obj:`int`): The dimension to be unsqueezed. + Returns: + - output (:obj:`Any`): The output data. + + Examples (tensor): + >>> t = torch.randn(3, 3) + >>> tt = unsqueeze(t, dim=0) + >>> assert tt.shape == torch.Shape([1, 3, 3]) + + Examples (list): + >>> t = [torch.randn(3, 3)] + >>> tt = unsqueeze(t, dim=0) + >>> assert tt[0].shape == torch.Shape([1, 3, 3]) + + Examples (dict): + >>> t = {"t": torch.randn(3, 3)} + >>> tt = unsqueeze(t, dim=0) + >>> assert tt["t"].shape == torch.Shape([1, 3, 3]) + """ if isinstance(data, torch.Tensor): return data.unsqueeze(dim) elif isinstance(data, Sequence): @@ -388,7 +658,57 @@ def unsqueeze(data: Any, dim: int = 0) -> Any: raise TypeError("not support type in unsqueeze: {}".format(type(data))) +def squeeze(data: Any, dim: int = 0) -> Any: + """ + Overview: + Squeeze the tensor data. + Arguments: + - data (:obj:`Any`): The original data. It can be exactly a tensor or a container (Sequence or dict). + - dim (:obj:`int`): The dimension to be Squeezed. + Returns: + - output (:obj:`Any`): The output data. + + Examples (tensor): + >>> t = torch.randn(1, 3, 3) + >>> tt = squeeze(t, dim=0) + >>> assert tt.shape == torch.Shape([3, 3]) + + Examples (list): + >>> t = [torch.randn(1, 3, 3)] + >>> tt = squeeze(t, dim=0) + >>> assert tt[0].shape == torch.Shape([3, 3]) + + Examples (dict): + >>> t = {"t": torch.randn(1, 3, 3)} + >>> tt = squeeze(t, dim=0) + >>> assert tt["t"].shape == torch.Shape([3, 3]) + """ + if isinstance(data, torch.Tensor): + return data.squeeze(dim) + elif isinstance(data, Sequence): + return [squeeze(d) for d in data] + elif isinstance(data, dict): + return {k: squeeze(v, 0) for k, v in data.items()} + else: + raise TypeError("not support type in squeeze: {}".format(type(data))) + + def get_null_data(template: Any, num: int) -> List[Any]: + """ + Overview: + Get null data given an input template. + Arguments: + - template (:obj:`Any`): The template data. + - num (:obj:`int`): The number of null data items to generate. + Returns: + - output (:obj:`List[Any]`): The generated null data. + + Examples: + >>> temp = {'obs': [1, 2, 3], 'action': 1, 'done': False, 'reward': torch.tensor(1.)} + >>> null_data = get_null_data(temp, 2) + >>> assert len(null_data) ==2 + >>> assert null_data[0]['null'] and null_data[0]['done'] + """ ret = [] for _ in range(num): data = copy.deepcopy(template) @@ -397,3 +717,40 @@ def get_null_data(template: Any, num: int) -> List[Any]: data['reward'].zero_() ret.append(data) return ret + + +def zeros_like(h: Any) -> Any: + """ + Overview: + Generate zero-tensors like the input data. + Arguments: + - h (:obj:`Any`): The original data. It can be exactly a tensor or a container (Sequence or dict). + Returns: + - output (:obj:`Any`): The output zero-tensors. + + Examples (tensor): + >>> t = torch.randn(3, 3) + >>> tt = zeros_like(t) + >>> assert tt.shape == torch.Shape([3, 3]) + >>> assert torch.sum(torch.abs(tt)) < 1e-8 + + Examples (list): + >>> t = [torch.randn(3, 3)] + >>> tt = zeros_like(t) + >>> assert tt[0].shape == torch.Shape([3, 3]) + >>> assert torch.sum(torch.abs(tt[0])) < 1e-8 + + Examples (dict): + >>> t = {"t": torch.randn(3, 3)} + >>> tt = zeros_like(t) + >>> assert tt["t"].shape == torch.Shape([3, 3]) + >>> assert torch.sum(torch.abs(tt["t"])) < 1e-8 + """ + if isinstance(h, torch.Tensor): + return torch.zeros_like(h) + elif isinstance(h, (list, tuple)): + return [zeros_like(t) for t in h] + elif isinstance(h, dict): + return {k: zeros_like(v) for k, v in h.items()} + else: + raise TypeError("not support type: {}".format(h)) diff --git a/ding/torch_utils/dataparallel.py b/ding/torch_utils/dataparallel.py index 654ceaa7d1..f4ea14f767 100644 --- a/ding/torch_utils/dataparallel.py +++ b/ding/torch_utils/dataparallel.py @@ -3,10 +3,33 @@ class DataParallel(nn.DataParallel): + """ + Overview: + A wrapper class for nn.DataParallel. + Interfaces: + ``__init__``, ``parameters`` + """ def __init__(self, module, device_ids=None, output_device=None, dim=0): + """ + Overview: + Initialize the DataParallel object. + Arguments: + - module (:obj:`nn.Module`): The module to be parallelized. + - device_ids (:obj:`list`): The list of GPU ids. + - output_device (:obj:`int`): The output GPU id. + - dim (:obj:`int`): The dimension to be parallelized. + """ super().__init__(module, device_ids=None, output_device=None, dim=0) self.module = module def parameters(self, recurse: bool = True): + """ + Overview: + Return the parameters of the module. + Arguments: + - recurse (:obj:`bool`): Whether to return the parameters of the submodules. + Returns: + - params (:obj:`generator`): The generator of the parameters. + """ return self.module.parameters(recurse=True) diff --git a/dizoo/board_games/go/envs/__init__.py b/ding/torch_utils/diffusion_SDE/__init__.py similarity index 100% rename from dizoo/board_games/go/envs/__init__.py rename to ding/torch_utils/diffusion_SDE/__init__.py diff --git a/ding/torch_utils/diffusion_SDE/dpm_solver_pytorch.py b/ding/torch_utils/diffusion_SDE/dpm_solver_pytorch.py new file mode 100644 index 0000000000..4f43c10510 --- /dev/null +++ b/ding/torch_utils/diffusion_SDE/dpm_solver_pytorch.py @@ -0,0 +1,1288 @@ +############################################################# +# This DPM-Solver snippet is from https://github.com/ChenDRAG/CEP-energy-guided-diffusion +# wich is based on https://github.com/LuChengTHU/dpm-solver +############################################################# + +import torch +import math + + +class NoiseScheduleVP: + + def __init__( + self, + schedule='discrete', + betas=None, + alphas_cumprod=None, + continuous_beta_0=0.1, + continuous_beta_1=20., + ): + """Create a wrapper class for the forward SDE (VP type). + + *** + Update: We support discrete-time diffusion models by implementing a picewise linear interpolation \ + for log_alpha_t. We recommend to use schedule='discrete' for the discrete-time diffusion models, \ + especially for high-resolution images. + *** + + The forward SDE ensures that the condition distribution \ + q_{t|0}(x_t | x_0) = N ( alpha_t * x_0, sigma_t^2 * I ). + We further define lambda_t = log(alpha_t) - log(sigma_t), \ + which is the half-logSNR (described in the DPM-Solver paper). + Therefore, we implement the functions for computing alpha_t, sigma_t and lambda_t. + For t in [0, T], we have: + log_alpha_t = self.marginal_log_mean_coeff(t) + sigma_t = self.marginal_std(t) + lambda_t = self.marginal_lambda(t) + + Moreover, as lambda(t) is an invertible function, we also support its inverse function: + + t = self.inverse_lambda(lambda_t) + + =============================================================== + + We support both discrete-time DPMs (trained on n = 0, 1, ..., N-1) and \ + continuous-time DPMs (trained on t in [t_0, T]). + + 1. For discrete-time DPMs: + + For discrete-time DPMs trained on n = 0, 1, ..., N-1, \ + we convert the discrete steps to continuous time steps by: + t_i = (i + 1) / N + e.g. for N = 1000, we have t_0 = 1e-3 and T = t_{N-1} = 1. + We solve the corresponding diffusion ODE from time T = 1 to time t_0 = 1e-3. + + Args: + betas: A `torch.Tensor`. The beta array for the discrete-time DPM. \ + (See the original DDPM paper for details) + alphas_cumprod: A `torch.Tensor`. The cumprod alphas for the discrete-time DPM. \ + (See the original DDPM paper for details) + + Note that we always have alphas_cumprod = cumprod(betas). \ + Therefore, we only need to set one of `betas` and `alphas_cumprod`. + + **Important**: Please pay special attention for the args for `alphas_cumprod`: + The `alphas_cumprod` is the \hat{alpha_n} arrays in the notations of DDPM. \ + Specifically, DDPMs assume that + q_{t_n | 0}(x_{t_n} | x_0) = N ( \sqrt{\hat{alpha_n}} * x_0, (1 - \hat{alpha_n}) * I ). + Therefore, the notation \hat{alpha_n} is different from the notation alpha_t in DPM-Solver. + In fact, we have + alpha_{t_n} = \sqrt{\hat{alpha_n}}, + and + log(alpha_{t_n}) = 0.5 * log(\hat{alpha_n}). + + + 2. For continuous-time DPMs: + + We support two types of VPSDEs: linear (DDPM) and cosine (improved-DDPM). + The hyperparameters for the noise schedule are the default settings in DDPM and improved-DDPM: + + Args: + beta_min: A `float` number. The smallest beta for the linear schedule. + beta_max: A `float` number. The largest beta for the linear schedule. + cosine_s: A `float` number. The hyperparameter in the cosine schedule. + cosine_beta_max: A `float` number. The hyperparameter in the cosine schedule. + T: A `float` number. The ending time of the forward process. + + =============================================================== + + Args: + schedule: A `str`. The noise schedule of the forward SDE. 'discrete' for discrete-time DPMs, + 'linear' or 'cosine' for continuous-time DPMs. + Returns: + A wrapper object of the forward SDE (VP type). + + =============================================================== + + Example: + + # For discrete-time DPMs, given betas (the beta array for n = 0, 1, ..., N - 1): + >>> ns = NoiseScheduleVP('discrete', betas=betas) + + # For discrete-time DPMs, given alphas_cumprod (the \hat{alpha_n} array \ + for n = 0, 1, ..., N - 1): + >>> ns = NoiseScheduleVP('discrete', alphas_cumprod=alphas_cumprod) + + # For continuous-time DPMs (VPSDE), linear schedule: + >>> ns = NoiseScheduleVP('linear', continuous_beta_0=0.1, continuous_beta_1=20.) + + """ + + if schedule not in ['discrete', 'linear', 'cosine']: + raise ValueError( + "Unsupported noise schedule {}. \ + The schedule needs to be 'discrete' or 'linear' or 'cosine'".format(schedule) + ) + + self.schedule = schedule + if schedule == 'discrete': + if betas is not None: + log_alphas = 0.5 * torch.log(1 - betas).cumsum(dim=0) + else: + assert alphas_cumprod is not None + log_alphas = 0.5 * torch.log(alphas_cumprod) + self.total_N = len(log_alphas) + self.T = 1. + self.t_array = torch.linspace(0., 1., self.total_N + 1)[1:].reshape((1, -1)) + self.log_alpha_array = log_alphas.reshape(( + 1, + -1, + )) + else: + self.total_N = 1000 + self.beta_0 = continuous_beta_0 + self.beta_1 = continuous_beta_1 + self.cosine_s = 0.008 + self.cosine_beta_max = 999. + self.cosine_t_max = math.atan(self.cosine_beta_max * (1. + self.cosine_s) / math.pi + ) * 2. * (1. + self.cosine_s) / math.pi - self.cosine_s + self.cosine_log_alpha_0 = math.log(math.cos(self.cosine_s / (1. + self.cosine_s) * math.pi / 2.)) + self.schedule = schedule + if schedule == 'cosine': + # For the cosine schedule, T = 1 will have numerical issues. So we manually set the ending time T. + # Note that T = 0.9946 may be not the optimal setting. However, we find it works well. + self.T = 0.9946 + else: + self.T = 1. + + def marginal_log_mean_coeff(self, t): + """ + Compute log(alpha_t) of a given continuous-time label t in [0, T]. + """ + if self.schedule == 'discrete': + return interpolate_fn(t.reshape((-1, 1)), self.t_array.to(t.device), + self.log_alpha_array.to(t.device)).reshape((-1)) + elif self.schedule == 'linear': + return -0.25 * t ** 2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0 + elif self.schedule == 'cosine': + log_alpha_fn = lambda s: torch.log(torch.cos((s + self.cosine_s) / (1. + self.cosine_s) * math.pi / 2.)) + log_alpha_t = log_alpha_fn(t) - self.cosine_log_alpha_0 + return log_alpha_t + + def marginal_alpha(self, t): + """ + Compute alpha_t of a given continuous-time label t in [0, T]. + """ + return torch.exp(self.marginal_log_mean_coeff(t)) + + def marginal_std(self, t): + """ + Compute sigma_t of a given continuous-time label t in [0, T]. + """ + return torch.sqrt(1. - torch.exp(2. * self.marginal_log_mean_coeff(t))) + + def marginal_lambda(self, t): + """ + Compute lambda_t = log(alpha_t) - log(sigma_t) of a given continuous-time label t in [0, T]. + """ + log_mean_coeff = self.marginal_log_mean_coeff(t) + log_std = 0.5 * torch.log(1. - torch.exp(2. * log_mean_coeff)) + return log_mean_coeff - log_std + + def inverse_lambda(self, lamb): + """ + Compute the continuous-time label t in [0, T] of a given half-logSNR lambda_t. + """ + if self.schedule == 'linear': + tmp = 2. * (self.beta_1 - self.beta_0) * torch.logaddexp(-2. * lamb, torch.zeros((1, )).to(lamb)) + Delta = self.beta_0 ** 2 + tmp + return tmp / (torch.sqrt(Delta) + self.beta_0) / (self.beta_1 - self.beta_0) + elif self.schedule == 'discrete': + log_alpha = -0.5 * torch.logaddexp(torch.zeros((1, )).to(lamb.device), -2. * lamb) + t = interpolate_fn( + log_alpha.reshape((-1, 1)), torch.flip(self.log_alpha_array.to(lamb.device), [1]), + torch.flip(self.t_array.to(lamb.device), [1]) + ) + return t.reshape((-1, )) + else: + log_alpha = -0.5 * torch.logaddexp(-2. * lamb, torch.zeros((1, )).to(lamb)) + t_fn = lambda log_alpha_t: torch.arccos(torch.exp(log_alpha_t + self.cosine_log_alpha_0) + ) * 2. * (1. + self.cosine_s) / math.pi - self.cosine_s + t = t_fn(log_alpha) + return t + + +def model_wrapper( + model, + noise_schedule, + model_type="noise", + model_kwargs={}, + guidance_type="uncond", + condition=None, + unconditional_condition=None, + guidance_scale=1., + classifier_fn=None, + classifier_kwargs={}, +): + """Create a wrapper function for the noise prediction model. + + DPM-Solver needs to solve the continuous-time diffusion ODEs. For DPMs trained on discrete-time labels, we need to + firstly wrap the model function to a noise prediction model that accepts the continuous time as the input. + + We support four types of the diffusion model by setting `model_type`: + + 1. "noise": noise prediction model. (Trained by predicting noise). + + 2. "x_start": data prediction model. (Trained by predicting the data x_0 at time 0). + + 3. "v": velocity prediction model. (Trained by predicting the velocity). + The "v" prediction is derivation detailed in Appendix D of [1], and is used in Imagen-Video [2]. + + [1] Salimans, Tim, and Jonathan Ho. "Progressive distillation for fast sampling of diffusion models." + arXiv preprint arXiv:2202.00512 (2022). + [2] Ho, Jonathan, et al. "Imagen Video: High Definition Video Generation with Diffusion Models." + arXiv preprint arXiv:2210.02303 (2022). + + 4. "score": marginal score function. (Trained by denoising score matching). + Note that the score function and the noise prediction model follows a simple relationship: + ``` + noise(x_t, t) = -sigma_t * score(x_t, t) + ``` + + We support three types of guided sampling by DPMs by setting `guidance_type`: + 1. "uncond": unconditional sampling by DPMs. + The input `model` has the following format: + `` + model(x, t_input, **model_kwargs) -> noise | x_start | v | score + `` + + 2. "classifier": classifier guidance sampling [3] by DPMs and another classifier. + The input `model` has the following format: + `` + model(x, t_input, **model_kwargs) -> noise | x_start | v | score + `` + + The input `classifier_fn` has the following format: + `` + classifier_fn(x, t_input, cond, **classifier_kwargs) -> logits(x, t_input, cond) + `` + + [3] P. Dhariwal and A. Q. Nichol, "Diffusion models beat GANs on image synthesis," + in Advances in Neural Information Processing Systems, vol. 34, 2021, pp. 8780-8794. + + 3. "classifier-free": classifier-free guidance sampling by conditional DPMs. + The input `model` has the following format: + `` + model(x, t_input, cond, **model_kwargs) -> noise | x_start | v | score + `` + And if cond == `unconditional_condition`, the model output is the unconditional DPM output. + + [4] Ho, Jonathan, and Tim Salimans. "Classifier-free diffusion guidance." + arXiv preprint arXiv:2207.12598 (2022). + + + The `t_input` is the time label of the model, which may be discrete-time labels (i.e. 0 to 999) + or continuous-time labels (i.e. epsilon to T). + + We wrap the model function to accept only `x` and `t_continuous` as inputs, and outputs the predicted noise: + `` + def model_fn(x, t_continuous) -> noise: + t_input = get_model_input_time(t_continuous) + return noise_pred(model, x, t_input, **model_kwargs) + `` + where `t_continuous` is the continuous time labels (i.e. epsilon to T). And we use `model_fn` for DPM-Solver. + + =============================================================== + + Args: + model: A diffusion model with the corresponding format described above. + noise_schedule: A noise schedule object, such as NoiseScheduleVP. + model_type: A `str`. The parameterization type of the diffusion model. + "noise" or "x_start" or "v" or "score". + model_kwargs: A `dict`. A dict for the other inputs of the model function. + guidance_type: A `str`. The type of the guidance for sampling. + "uncond" or "classifier" or "classifier-free". + condition: A pytorch tensor. The condition for the guided sampling. + Only used for "classifier" or "classifier-free" guidance type. + unconditional_condition: A pytorch tensor. The condition for the unconditional sampling. + Only used for "classifier-free" guidance type. + guidance_scale: A `float`. The scale for the guided sampling. + classifier_fn: A classifier function. Only used for the classifier guidance. + classifier_kwargs: A `dict`. A dict for the other inputs of the classifier function. + Returns: + A noise prediction model that accepts the noised data and the continuous time as the inputs. + """ + + def get_model_input_time(t_continuous): + """ + Convert the continuous-time `t_continuous` (in [epsilon, T]) to the model input time. + For discrete-time DPMs, we convert `t_continuous` in [1 / N, 1] to `t_input` in [0, 1000 * (N - 1) / N]. + For continuous-time DPMs, we just use `t_continuous`. + """ + if noise_schedule.schedule == 'discrete': + return (t_continuous - 1. / noise_schedule.total_N) * 1000. + else: + return t_continuous + + def noise_pred_fn(x, t_continuous, cond=None): + if t_continuous.reshape((-1, )).shape[0] == 1: + t_continuous = t_continuous.expand((x.shape[0])) + t_input = get_model_input_time(t_continuous) + if cond is None: + output = model(x, t_input, **model_kwargs) + else: + output = model(x, t_input, cond, **model_kwargs) + if model_type == "noise": + return output + elif model_type == "x_start": + alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous) + dims = x.dim() + return (x - expand_dims(alpha_t, dims) * output) / expand_dims(sigma_t, dims) + elif model_type == "v": + alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous) + dims = x.dim() + return expand_dims(alpha_t, dims) * output + expand_dims(sigma_t, dims) * x + elif model_type == "score": + sigma_t = noise_schedule.marginal_std(t_continuous) + dims = x.dim() + return -expand_dims(sigma_t, dims) * output + + def cond_grad_fn(x, t_input): + """ + Compute the gradient of the classifier, i.e. nabla_{x} log p_t(cond | x_t). + """ + with torch.enable_grad(): + x_in = x.detach().requires_grad_(True) + log_prob = classifier_fn(x_in, t_input, condition, **classifier_kwargs) + return torch.autograd.grad(log_prob.sum(), x_in)[0] + + def model_fn(x, t_continuous): + """ + The noise predicition model function that is used for DPM-Solver. + """ + if t_continuous.reshape((-1, )).shape[0] == 1: + t_continuous = t_continuous.expand((x.shape[0])) + if guidance_type == "uncond": + return noise_pred_fn(x, t_continuous) + elif guidance_type == "classifier": + assert classifier_fn is not None + t_input = get_model_input_time(t_continuous) + cond_grad = cond_grad_fn(x, t_input) + sigma_t = noise_schedule.marginal_std(t_continuous) + noise = noise_pred_fn(x, t_continuous) + return noise - guidance_scale * expand_dims(sigma_t, dims=cond_grad.dim()) * cond_grad + elif guidance_type == "classifier-free": + if guidance_scale == 1. or unconditional_condition is None: + return noise_pred_fn(x, t_continuous, cond=condition) + else: + x_in = torch.cat([x] * 2) + t_in = torch.cat([t_continuous] * 2) + c_in = torch.cat([unconditional_condition, condition]) + noise_uncond, noise = noise_pred_fn(x_in, t_in, cond=c_in).chunk(2) + return noise_uncond + guidance_scale * (noise - noise_uncond) + + assert model_type in ["noise", "x_start", "v"] + assert guidance_type in ["uncond", "classifier", "classifier-free"] + return model_fn + + +class DPM_Solver: + + def __init__(self, model_fn, noise_schedule, predict_x0=False, thresholding=False, max_val=1.): + """Construct a DPM-Solver. + + We support both the noise prediction model ("predicting epsilon") and \ + the data prediction model ("predicting x0"). + If `predict_x0` is False, we use the solver for the noise prediction model (DPM-Solver). + If `predict_x0` is True, we use the solver for the data prediction model (DPM-Solver++). + In such case, we further support the "dynamic thresholding" in [1] when `thresholding` is True. + The "dynamic thresholding" can greatly improve the sample quality for pixel-space DPMs \ + with large guidance scales. + + Args: + model_fn: A noise prediction model function which accepts the continuous-time input (t in [epsilon, T]): + `` + def model_fn(x, t_continuous): + return noise + `` + noise_schedule: A noise schedule object, such as NoiseScheduleVP. + predict_x0: A `bool`. If true, use the data prediction model; else, use the noise prediction model. + thresholding: A `bool`. Valid when `predict_x0` is True. Whether to use the "dynamic thresholding" in [1]. + max_val: A `float`. Valid when both `predict_x0` and `thresholding` are True. \ + The max value for thresholding. + + [1] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, \ + Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S Sara Mahdavi, Rapha Gontijo Lopes, et al. \ + Photorealistic text-to-image diffusion models with deep language understanding. \ + arXiv preprint arXiv:2205.11487, 2022b. + """ + self.model = model_fn + self.noise_schedule = noise_schedule + self.predict_x0 = predict_x0 + self.thresholding = thresholding + self.max_val = max_val + + def noise_prediction_fn(self, x, t): + """ + Return the noise prediction model. + """ + return self.model(x, t) + + def data_prediction_fn(self, x, t): + """ + Return the data prediction model (with thresholding). + """ + noise = self.noise_prediction_fn(x, t) + dims = x.dim() + alpha_t, sigma_t = self.noise_schedule.marginal_alpha(t), self.noise_schedule.marginal_std(t) + x0 = (x - expand_dims(sigma_t, dims) * noise) / expand_dims(alpha_t, dims) + if self.thresholding: + p = 0.995 # A hyperparameter in the paper of "Imagen" [1]. + s = torch.quantile(torch.abs(x0).reshape((x0.shape[0], -1)), p, dim=1) + s = expand_dims(torch.maximum(s, torch.ones_like(s).to(s.device)), dims) + x0 = torch.clamp(x0, -s, s) / (s / self.max_val) + return x0 + + def model_fn(self, x, t): + """ + Convert the model to the noise prediction model or the data prediction model. + """ + if self.predict_x0: + return self.data_prediction_fn(x, t) + else: + return self.noise_prediction_fn(x, t) + + def get_time_steps(self, skip_type, t_T, t_0, N, device): + """Compute the intermediate time steps for sampling. + + Args: + skip_type: A `str`. The type for the spacing of the time steps. We support three types: + - 'logSNR': uniform logSNR for the time steps. + - 'time_uniform': uniform time for the time steps. \ + (**Recommended for high-resolutional data**.) + - 'time_quadratic': quadratic time for the time steps. \ + (Used in DDIM for low-resolutional data.) + t_T: A `float`. The starting time of the sampling (default is T). + t_0: A `float`. The ending time of the sampling (default is epsilon). + N: A `int`. The total number of the spacing of the time steps. + device: A torch device. + Returns: + A pytorch tensor of the time steps, with the shape (N + 1,). + """ + if skip_type == 'logSNR': + lambda_T = self.noise_schedule.marginal_lambda(torch.tensor(t_T).to(device)) + lambda_0 = self.noise_schedule.marginal_lambda(torch.tensor(t_0).to(device)) + logSNR_steps = torch.linspace(lambda_T.cpu().item(), lambda_0.cpu().item(), N + 1).to(device) + return self.noise_schedule.inverse_lambda(logSNR_steps) + elif skip_type == 'time_uniform': + return torch.linspace(t_T, t_0, N + 1).to(device) + elif skip_type == 'time_quadratic': + t_order = 2 + t = torch.linspace(t_T ** (1. / t_order), t_0 ** (1. / t_order), N + 1).pow(t_order).to(device) + return t + else: + raise ValueError( + "Unsupported skip_type {}, need to be 'logSNR' or 'time_uniform' or 'time_quadratic'".format(skip_type) + ) + + def get_orders_for_singlestep_solver(self, steps, order): + """ + Get the order of each step for sampling by the singlestep DPM-Solver. + + We combine both DPM-Solver-1,2,3 to use all the function evaluations, which is named as "DPM-Solver-fast". + Given a fixed number of function evaluations by `steps`, the sampling procedure by DPM-Solver-fast is: + - If order == 1: + We take `steps` of DPM-Solver-1 (i.e. DDIM). + - If order == 2: + - Denote K = (steps // 2). We take K or (K + 1) intermediate time steps for sampling. + - If steps % 2 == 0, we use K steps of DPM-Solver-2. + - If steps % 2 == 1, we use K steps of DPM-Solver-2 and 1 step of DPM-Solver-1. + - If order == 3: + - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling. + - If steps % 3 == 0, we use (K - 2) steps of DPM-Solver-3, \ + and 1 step of DPM-Solver-2 and 1 step of DPM-Solver-1. + - If steps % 3 == 1, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-1. + - If steps % 3 == 2, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-2. + + ============================================ + Args: + order: A `int`. The max order for the solver (2 or 3). + steps: A `int`. The total number of function evaluations (NFE). + Returns: + orders: A list of the solver order of each step. + """ + if order == 3: + K = steps // 3 + 1 + if steps % 3 == 0: + orders = [ + 3, + ] * (K - 2) + [2, 1] + elif steps % 3 == 1: + orders = [ + 3, + ] * (K - 1) + [1] + else: + orders = [ + 3, + ] * (K - 1) + [2] + return orders + elif order == 2: + K = steps // 2 + if steps % 2 == 0: + # orders = [2,] * K + K = steps // 2 + 1 + orders = [ + 2, + ] * (K - 2) + [ + 1, + ] * 2 + else: + orders = [ + 2, + ] * K + [1] + return orders + elif order == 1: + return [ + 1, + ] * steps + else: + raise ValueError("'order' must be '1' or '2' or '3'.") + + def denoise_fn(self, x, s): + """ + Denoise at the final step, which is equivalent to solve the ODE \ + from lambda_s to infty by first-order discretization. + """ + return self.data_prediction_fn(x, s) + + def dpm_solver_first_update(self, x, s, t, model_s=None, return_intermediate=False): + """ + DPM-Solver-1 (equivalent to DDIM) from time `s` to time `t`. + + Args: + x: A pytorch tensor. The initial value at time `s`. + s: A pytorch tensor. The starting time, with the shape (x.shape[0],). + t: A pytorch tensor. The ending time, with the shape (x.shape[0],). + model_s: A pytorch tensor. The model function evaluated at time `s`. + If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it. + return_intermediate: A `bool`. If true, also return the model value at time `s`. + Returns: + x_t: A pytorch tensor. The approximated solution at time `t`. + """ + ns = self.noise_schedule + dims = x.dim() + lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t) + h = lambda_t - lambda_s + log_alpha_s, log_alpha_t = ns.marginal_log_mean_coeff(s), ns.marginal_log_mean_coeff(t) + sigma_s, sigma_t = ns.marginal_std(s), ns.marginal_std(t) + alpha_t = torch.exp(log_alpha_t) + + if self.predict_x0: + phi_1 = torch.expm1(-h) + if model_s is None: + model_s = self.model_fn(x, s) + x_t = (expand_dims(sigma_t / sigma_s, dims) * x - expand_dims(alpha_t * phi_1, dims) * model_s) + if return_intermediate: + return x_t, {'model_s': model_s} + else: + return x_t + else: + phi_1 = torch.expm1(h) + if model_s is None: + model_s = self.model_fn(x, s) + x_t = ( + expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x - + expand_dims(sigma_t * phi_1, dims) * model_s + ) + if return_intermediate: + return x_t, {'model_s': model_s} + else: + return x_t + + def singlestep_dpm_solver_second_update( + self, x, s, t, r1=0.5, model_s=None, return_intermediate=False, solver_type='dpm_solver' + ): + """ + Singlestep solver DPM-Solver-2 from time `s` to time `t`. + + Args: + x: A pytorch tensor. The initial value at time `s`. + s: A pytorch tensor. The starting time, with the shape (x.shape[0],). + t: A pytorch tensor. The ending time, with the shape (x.shape[0],). + r1: A `float`. The hyperparameter of the second-order solver. + model_s: A pytorch tensor. The model function evaluated at time `s`. + If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it. + return_intermediate: A `bool`. If true, also return the model value at time `s` and `s1` \ + (the intermediate time). + solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers. + The type slightly impacts the performance. We recommend to use 'dpm_solver' type. + Returns: + x_t: A pytorch tensor. The approximated solution at time `t`. + """ + if solver_type not in ['dpm_solver', 'taylor']: + raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type)) + if r1 is None: + r1 = 0.5 + ns = self.noise_schedule + dims = x.dim() + lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t) + h = lambda_t - lambda_s + lambda_s1 = lambda_s + r1 * h + s1 = ns.inverse_lambda(lambda_s1) + log_alpha_s, log_alpha_s1, log_alpha_t = ns.marginal_log_mean_coeff(s), ns.marginal_log_mean_coeff( + s1 + ), ns.marginal_log_mean_coeff(t) + sigma_s, sigma_s1, sigma_t = ns.marginal_std(s), ns.marginal_std(s1), ns.marginal_std(t) + alpha_s1, alpha_t = torch.exp(log_alpha_s1), torch.exp(log_alpha_t) + + if self.predict_x0: + phi_11 = torch.expm1(-r1 * h) + phi_1 = torch.expm1(-h) + + if model_s is None: + model_s = self.model_fn(x, s) + x_s1 = (expand_dims(sigma_s1 / sigma_s, dims) * x - expand_dims(alpha_s1 * phi_11, dims) * model_s) + model_s1 = self.model_fn(x_s1, s1) + if solver_type == 'dpm_solver': + x_t = ( + expand_dims(sigma_t / sigma_s, dims) * x - expand_dims(alpha_t * phi_1, dims) * model_s - + (0.5 / r1) * expand_dims(alpha_t * phi_1, dims) * (model_s1 - model_s) + ) + elif solver_type == 'taylor': + x_t = ( + expand_dims(sigma_t / sigma_s, dims) * x - expand_dims(alpha_t * phi_1, dims) * model_s + + (1. / r1) * expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * (model_s1 - model_s) + ) + else: + phi_11 = torch.expm1(r1 * h) + phi_1 = torch.expm1(h) + + if model_s is None: + model_s = self.model_fn(x, s) + x_s1 = ( + expand_dims(torch.exp(log_alpha_s1 - log_alpha_s), dims) * x - + expand_dims(sigma_s1 * phi_11, dims) * model_s + ) + model_s1 = self.model_fn(x_s1, s1) + if solver_type == 'dpm_solver': + x_t = ( + expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x - + expand_dims(sigma_t * phi_1, dims) * model_s - (0.5 / r1) * expand_dims(sigma_t * phi_1, dims) * + (model_s1 - model_s) + ) + elif solver_type == 'taylor': + x_t = ( + expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x - + expand_dims(sigma_t * phi_1, dims) * model_s - + (1. / r1) * expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * (model_s1 - model_s) + ) + if return_intermediate: + return x_t, {'model_s': model_s, 'model_s1': model_s1} + else: + return x_t + + def singlestep_dpm_solver_third_update( + self, + x, + s, + t, + r1=1. / 3., + r2=2. / 3., + model_s=None, + model_s1=None, + return_intermediate=False, + solver_type='dpm_solver' + ): + """ + Singlestep solver DPM-Solver-3 from time `s` to time `t`. + + Args: + x: A pytorch tensor. The initial value at time `s`. + s: A pytorch tensor. The starting time, with the shape (x.shape[0],). + t: A pytorch tensor. The ending time, with the shape (x.shape[0],). + r1: A `float`. The hyperparameter of the third-order solver. + r2: A `float`. The hyperparameter of the third-order solver. + model_s: A pytorch tensor. The model function evaluated at time `s`. + If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it. + model_s1: A pytorch tensor. The model function evaluated at time `s1` (the intermediate time given by `r1`). + If `model_s1` is None, we evaluate the model at `s1`; otherwise we directly use it. + return_intermediate: A `bool`. If true, also return the model value at time `s`, `s1` and `s2` \ + (the intermediate times). + solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers. + The type slightly impacts the performance. We recommend to use 'dpm_solver' type. + Returns: + x_t: A pytorch tensor. The approximated solution at time `t`. + """ + if solver_type not in ['dpm_solver', 'taylor']: + raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type)) + if r1 is None: + r1 = 1. / 3. + if r2 is None: + r2 = 2. / 3. + ns = self.noise_schedule + dims = x.dim() + lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t) + h = lambda_t - lambda_s + lambda_s1 = lambda_s + r1 * h + lambda_s2 = lambda_s + r2 * h + s1 = ns.inverse_lambda(lambda_s1) + s2 = ns.inverse_lambda(lambda_s2) + log_alpha_s, log_alpha_s1, log_alpha_s2, log_alpha_t = ns.marginal_log_mean_coeff( + s + ), ns.marginal_log_mean_coeff(s1), ns.marginal_log_mean_coeff(s2), ns.marginal_log_mean_coeff(t) + sigma_s, sigma_s1, sigma_s2, sigma_t = ns.marginal_std(s), ns.marginal_std(s1), ns.marginal_std( + s2 + ), ns.marginal_std(t) + alpha_s1, alpha_s2, alpha_t = torch.exp(log_alpha_s1), torch.exp(log_alpha_s2), torch.exp(log_alpha_t) + + if self.predict_x0: + phi_11 = torch.expm1(-r1 * h) + phi_12 = torch.expm1(-r2 * h) + phi_1 = torch.expm1(-h) + phi_22 = torch.expm1(-r2 * h) / (r2 * h) + 1. + phi_2 = phi_1 / h + 1. + phi_3 = phi_2 / h - 0.5 + + if model_s is None: + model_s = self.model_fn(x, s) + if model_s1 is None: + x_s1 = (expand_dims(sigma_s1 / sigma_s, dims) * x - expand_dims(alpha_s1 * phi_11, dims) * model_s) + model_s1 = self.model_fn(x_s1, s1) + x_s2 = ( + expand_dims(sigma_s2 / sigma_s, dims) * x - expand_dims(alpha_s2 * phi_12, dims) * model_s + + r2 / r1 * expand_dims(alpha_s2 * phi_22, dims) * (model_s1 - model_s) + ) + model_s2 = self.model_fn(x_s2, s2) + if solver_type == 'dpm_solver': + x_t = ( + expand_dims(sigma_t / sigma_s, dims) * x - expand_dims(alpha_t * phi_1, dims) * model_s + + (1. / r2) * expand_dims(alpha_t * phi_2, dims) * (model_s2 - model_s) + ) + elif solver_type == 'taylor': + D1_0 = (1. / r1) * (model_s1 - model_s) + D1_1 = (1. / r2) * (model_s2 - model_s) + D1 = (r2 * D1_0 - r1 * D1_1) / (r2 - r1) + D2 = 2. * (D1_1 - D1_0) / (r2 - r1) + x_t = ( + expand_dims(sigma_t / sigma_s, dims) * x - expand_dims(alpha_t * phi_1, dims) * model_s + + expand_dims(alpha_t * phi_2, dims) * D1 - expand_dims(alpha_t * phi_3, dims) * D2 + ) + else: + phi_11 = torch.expm1(r1 * h) + phi_12 = torch.expm1(r2 * h) + phi_1 = torch.expm1(h) + phi_22 = torch.expm1(r2 * h) / (r2 * h) - 1. + phi_2 = phi_1 / h - 1. + phi_3 = phi_2 / h - 0.5 + + if model_s is None: + model_s = self.model_fn(x, s) + if model_s1 is None: + x_s1 = ( + expand_dims(torch.exp(log_alpha_s1 - log_alpha_s), dims) * x - + expand_dims(sigma_s1 * phi_11, dims) * model_s + ) + model_s1 = self.model_fn(x_s1, s1) + x_s2 = ( + expand_dims(torch.exp(log_alpha_s2 - log_alpha_s), dims) * x - + expand_dims(sigma_s2 * phi_12, dims) * model_s - r2 / r1 * expand_dims(sigma_s2 * phi_22, dims) * + (model_s1 - model_s) + ) + model_s2 = self.model_fn(x_s2, s2) + if solver_type == 'dpm_solver': + x_t = ( + expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x - + expand_dims(sigma_t * phi_1, dims) * model_s - (1. / r2) * expand_dims(sigma_t * phi_2, dims) * + (model_s2 - model_s) + ) + elif solver_type == 'taylor': + D1_0 = (1. / r1) * (model_s1 - model_s) + D1_1 = (1. / r2) * (model_s2 - model_s) + D1 = (r2 * D1_0 - r1 * D1_1) / (r2 - r1) + D2 = 2. * (D1_1 - D1_0) / (r2 - r1) + x_t = ( + expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x - + expand_dims(sigma_t * phi_1, dims) * model_s - expand_dims(sigma_t * phi_2, dims) * D1 - + expand_dims(sigma_t * phi_3, dims) * D2 + ) + + if return_intermediate: + return x_t, {'model_s': model_s, 'model_s1': model_s1, 'model_s2': model_s2} + else: + return x_t + + def multistep_dpm_solver_second_update(self, x, model_prev_list, t_prev_list, t, solver_type="dpm_solver"): + """ + Multistep solver DPM-Solver-2 from time `t_prev_list[-1]` to time `t`. + + Args: + x: A pytorch tensor. The initial value at time `s`. + model_prev_list: A list of pytorch tensor. The previous computed model values. + t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],) + t: A pytorch tensor. The ending time, with the shape (x.shape[0],). + solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers. + The type slightly impacts the performance. We recommend to use 'dpm_solver' type. + Returns: + x_t: A pytorch tensor. The approximated solution at time `t`. + """ + if solver_type not in ['dpm_solver', 'taylor']: + raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type)) + ns = self.noise_schedule + dims = x.dim() + model_prev_1, model_prev_0 = model_prev_list + t_prev_1, t_prev_0 = t_prev_list + lambda_prev_1, lambda_prev_0, lambda_t = ns.marginal_lambda(t_prev_1), ns.marginal_lambda( + t_prev_0 + ), ns.marginal_lambda(t) + log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t) + sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t) + alpha_t = torch.exp(log_alpha_t) + + h_0 = lambda_prev_0 - lambda_prev_1 + h = lambda_t - lambda_prev_0 + r0 = h_0 / h + D1_0 = expand_dims(1. / r0, dims) * (model_prev_0 - model_prev_1) + if self.predict_x0: + if solver_type == 'dpm_solver': + x_t = ( + expand_dims(sigma_t / sigma_prev_0, dims) * x - + expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0 - + 0.5 * expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * D1_0 + ) + elif solver_type == 'taylor': + x_t = ( + expand_dims(sigma_t / sigma_prev_0, dims) * x - + expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0 + + expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * D1_0 + ) + else: + if solver_type == 'dpm_solver': + x_t = ( + expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x - + expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0 - + 0.5 * expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * D1_0 + ) + elif solver_type == 'taylor': + x_t = ( + expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x - + expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0 - + expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * D1_0 + ) + return x_t + + def multistep_dpm_solver_third_update(self, x, model_prev_list, t_prev_list, t, solver_type='dpm_solver'): + """ + Multistep solver DPM-Solver-3 from time `t_prev_list[-1]` to time `t`. + + Args: + x: A pytorch tensor. The initial value at time `s`. + model_prev_list: A list of pytorch tensor. The previous computed model values. + t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],) + t: A pytorch tensor. The ending time, with the shape (x.shape[0],). + solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers. + The type slightly impacts the performance. We recommend to use 'dpm_solver' type. + Returns: + x_t: A pytorch tensor. The approximated solution at time `t`. + """ + ns = self.noise_schedule + dims = x.dim() + model_prev_2, model_prev_1, model_prev_0 = model_prev_list + t_prev_2, t_prev_1, t_prev_0 = t_prev_list + lambda_prev_2, lambda_prev_1, lambda_prev_0, lambda_t = ns.marginal_lambda(t_prev_2), ns.marginal_lambda( + t_prev_1 + ), ns.marginal_lambda(t_prev_0), ns.marginal_lambda(t) + log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t) + sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t) + alpha_t = torch.exp(log_alpha_t) + + h_1 = lambda_prev_1 - lambda_prev_2 + h_0 = lambda_prev_0 - lambda_prev_1 + h = lambda_t - lambda_prev_0 + r0, r1 = h_0 / h, h_1 / h + D1_0 = expand_dims(1. / r0, dims) * (model_prev_0 - model_prev_1) + D1_1 = expand_dims(1. / r1, dims) * (model_prev_1 - model_prev_2) + D1 = D1_0 + expand_dims(r0 / (r0 + r1), dims) * (D1_0 - D1_1) + D2 = expand_dims(1. / (r0 + r1), dims) * (D1_0 - D1_1) + if self.predict_x0: + x_t = ( + expand_dims(sigma_t / sigma_prev_0, dims) * x - expand_dims(alpha_t * + (torch.exp(-h) - 1.), dims) * model_prev_0 + + expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * D1 - + expand_dims(alpha_t * ((torch.exp(-h) - 1. + h) / h ** 2 - 0.5), dims) * D2 + ) + else: + x_t = ( + expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x - + expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0 - + expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * D1 - + expand_dims(sigma_t * ((torch.exp(h) - 1. - h) / h ** 2 - 0.5), dims) * D2 + ) + return x_t + + def singlestep_dpm_solver_update( + self, x, s, t, order, return_intermediate=False, solver_type='dpm_solver', r1=None, r2=None + ): + """ + Singlestep DPM-Solver with the order `order` from time `s` to time `t`. + + Args: + x: A pytorch tensor. The initial value at time `s`. + s: A pytorch tensor. The starting time, with the shape (x.shape[0],). + t: A pytorch tensor. The ending time, with the shape (x.shape[0],). + order: A `int`. The order of DPM-Solver. We only support order == 1 or 2 or 3. + return_intermediate: A `bool`. If true, also return the model value at time `s`, \ + `s1` and `s2` (the intermediate times). + solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers. + The type slightly impacts the performance. We recommend to use 'dpm_solver' type. + r1: A `float`. The hyperparameter of the second-order or third-order solver. + r2: A `float`. The hyperparameter of the third-order solver. + Returns: + x_t: A pytorch tensor. The approximated solution at time `t`. + """ + if order == 1: + return self.dpm_solver_first_update(x, s, t, return_intermediate=return_intermediate) + elif order == 2: + return self.singlestep_dpm_solver_second_update( + x, s, t, return_intermediate=return_intermediate, solver_type=solver_type, r1=r1 + ) + elif order == 3: + return self.singlestep_dpm_solver_third_update( + x, s, t, return_intermediate=return_intermediate, solver_type=solver_type, r1=r1, r2=r2 + ) + else: + raise ValueError("Solver order must be 1 or 2 or 3, got {}".format(order)) + + def multistep_dpm_solver_update(self, x, model_prev_list, t_prev_list, t, order, solver_type='dpm_solver'): + """ + Multistep DPM-Solver with the order `order` from time `t_prev_list[-1]` to time `t`. + + Args: + x: A pytorch tensor. The initial value at time `s`. + model_prev_list: A list of pytorch tensor. The previous computed model values. + t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],) + t: A pytorch tensor. The ending time, with the shape (x.shape[0],). + order: A `int`. The order of DPM-Solver. We only support order == 1 or 2 or 3. + solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers. + The type slightly impacts the performance. We recommend to use 'dpm_solver' type. + Returns: + x_t: A pytorch tensor. The approximated solution at time `t`. + """ + if order == 1: + return self.dpm_solver_first_update(x, t_prev_list[-1], t, model_s=model_prev_list[-1]) + elif order == 2: + return self.multistep_dpm_solver_second_update(x, model_prev_list, t_prev_list, t, solver_type=solver_type) + elif order == 3: + return self.multistep_dpm_solver_third_update(x, model_prev_list, t_prev_list, t, solver_type=solver_type) + else: + raise ValueError("Solver order must be 1 or 2 or 3, got {}".format(order)) + + def dpm_solver_adaptive( + self, x, order, t_T, t_0, h_init=0.05, atol=0.0078, rtol=0.05, theta=0.9, t_err=1e-5, solver_type='dpm_solver' + ): + """ + The adaptive step size solver based on singlestep DPM-Solver. + + Args: + x: A pytorch tensor. The initial value at time `t_T`. + order: A `int`. The (higher) order of the solver. We only support order == 2 or 3. + t_T: A `float`. The starting time of the sampling (default is T). + t_0: A `float`. The ending time of the sampling (default is epsilon). + h_init: A `float`. The initial step size (for logSNR). + atol: A `float`. The absolute tolerance of the solver. \ + For image data, the default setting is 0.0078, followed [1]. + rtol: A `float`. The relative tolerance of the solver. The default setting is 0.05. + theta: A `float`. The safety hyperparameter for adapting the step size. \ + The default setting is 0.9, followed [1]. + t_err: A `float`. The tolerance for the time. \ + We solve the diffusion ODE until the absolute error between the \ + current time and `t_0` is less than `t_err`. The default setting is 1e-5. + solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers. + The type slightly impacts the performance. We recommend to use 'dpm_solver' type. + Returns: + x_0: A pytorch tensor. The approximated solution at time `t_0`. + + [1] A. Jolicoeur-Martineau, K. Li, R. Piché-Taillefer, T. Kachman, and I. Mitliagkas, \ + "Gotta go fast when generating data with score-based models," \ + arXiv preprint arXiv:2105.14080, 2021. + """ + ns = self.noise_schedule + s = t_T * torch.ones((x.shape[0], )).to(x) + lambda_s = ns.marginal_lambda(s) + lambda_0 = ns.marginal_lambda(t_0 * torch.ones_like(s).to(x)) + h = h_init * torch.ones_like(s).to(x) + x_prev = x + nfe = 0 + if order == 2: + r1 = 0.5 + lower_update = lambda x, s, t: self.dpm_solver_first_update(x, s, t, return_intermediate=True) + higher_update = lambda x, s, t, **kwargs: self.singlestep_dpm_solver_second_update( + x, s, t, r1=r1, solver_type=solver_type, **kwargs + ) + elif order == 3: + r1, r2 = 1. / 3., 2. / 3. + lower_update = lambda x, s, t: self.singlestep_dpm_solver_second_update( + x, s, t, r1=r1, return_intermediate=True, solver_type=solver_type + ) + higher_update = lambda x, s, t, **kwargs: self.singlestep_dpm_solver_third_update( + x, s, t, r1=r1, r2=r2, solver_type=solver_type, **kwargs + ) + else: + raise ValueError("For adaptive step size solver, order must be 2 or 3, got {}".format(order)) + while torch.abs((s - t_0)).mean() > t_err: + t = ns.inverse_lambda(lambda_s + h) + x_lower, lower_noise_kwargs = lower_update(x, s, t) + x_higher = higher_update(x, s, t, **lower_noise_kwargs) + delta = torch.max(torch.ones_like(x).to(x) * atol, rtol * torch.max(torch.abs(x_lower), torch.abs(x_prev))) + norm_fn = lambda v: torch.sqrt(torch.square(v.reshape((v.shape[0], -1))).mean(dim=-1, keepdim=True)) + E = norm_fn((x_higher - x_lower) / delta).max() + if torch.all(E <= 1.): + x = x_higher + s = t + x_prev = x_lower + lambda_s = ns.marginal_lambda(s) + h = torch.min(theta * h * torch.float_power(E, -1. / order).float(), lambda_0 - lambda_s) + nfe += order + print('adaptive solver nfe', nfe) + return x + + def sample( + self, + x, + steps=20, + t_start=None, + t_end=None, + order=3, + skip_type='time_uniform', + method='singlestep', + denoise=False, + solver_type='dpm_solver', + atol=0.0078, + rtol=0.05, + ): + """ + Compute the sample at time `t_end` by DPM-Solver, given the initial `x` at time `t_start`. + + ===================================================== + + We support the following algorithms for both noise prediction model and data prediction model: + - 'singlestep': + Singlestep DPM-Solver (i.e. "DPM-Solver-fast" in the paper), \ + which combines different orders of singlestep DPM-Solver. + We combine all the singlestep solvers with order <= `order` \ + to use up all the function evaluations (steps). + The total number of function evaluations (NFE) == `steps`. + Given a fixed NFE == `steps`, the sampling procedure is: + - If `order` == 1: + - Denote K = steps. We use K steps of DPM-Solver-1 (i.e. DDIM). + - If `order` == 2: + - Denote K = (steps // 2) + (steps % 2). We take K intermediate time steps for sampling. + - If steps % 2 == 0, we use K steps of singlestep DPM-Solver-2. + - If steps % 2 == 1, we use (K - 1) steps of singlestep DPM-Solver-2 \ + and 1 step of DPM-Solver-1. + - If `order` == 3: + - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling. + - If steps % 3 == 0, we use (K - 2) steps of singlestep DPM-Solver-3, \ + and 1 step of singlestep DPM-Solver-2 \ + and 1 step of DPM-Solver-1. + - If steps % 3 == 1, we use (K - 1) steps of singlestep DPM-Solver-3 \ + and 1 step of DPM-Solver-1. + - If steps % 3 == 2, we use (K - 1) steps of singlestep DPM-Solver-3 \ + and 1 step of singlestep DPM-Solver-2. + - 'multistep': + Multistep DPM-Solver with the order of `order`. + The total number of function evaluations (NFE) == `steps`. + We initialize the first `order` values by lower order multistep solvers. + Given a fixed NFE == `steps`, the sampling procedure is: + Denote K = steps. + - If `order` == 1: + - We use K steps of DPM-Solver-1 (i.e. DDIM). + - If `order` == 2: + - We firstly use 1 step of DPM-Solver-1, then use (K - 1) step of multistep DPM-Solver-2. + - If `order` == 3: + - We firstly use 1 step of DPM-Solver-1, then 1 step of multistep DPM-Solver-2, \ + then (K - 2) step of multistep DPM-Solver-3. + - 'singlestep_fixed': + Fixed order singlestep DPM-Solver \ + (i.e. DPM-Solver-1 or singlestep DPM-Solver-2 or singlestep DPM-Solver-3). + We use singlestep DPM-Solver-`order` for `order`=1 or 2 or 3, \ + with total [`steps` // `order`] * `order` NFE. + - 'adaptive': + Adaptive step size DPM-Solver (i.e. "DPM-Solver-12" and "DPM-Solver-23" in the paper). + We ignore `steps` and use adaptive step size DPM-Solver with a higher order of `order`. + You can adjust the absolute tolerance `atol` and the relative tolerance `rtol` \ + to balance the computatation costs (NFE) and the sample quality. + - If `order` == 2, we use DPM-Solver-12 which combines DPM-Solver-1 and singlestep DPM-Solver-2. + - If `order` == 3, we use DPM-Solver-23 which combines singlestep DPM-Solver-2 \ + and singlestep DPM-Solver-3. + + ===================================================== + + Some advices for choosing the algorithm: + - For **unconditional sampling** or **guided sampling with small guidance scale** by DPMs: + Use singlestep DPM-Solver ("DPM-Solver-fast" in the paper) with `order = 3`. + e.g. + >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, predict_x0=False) + >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=3, + skip_type='time_uniform', method='singlestep') + - For **guided sampling with large guidance scale** by DPMs: + Use multistep DPM-Solver with `predict_x0 = True` and `order = 2`. + e.g. + >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, predict_x0=True) + >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=2, + skip_type='time_uniform', method='multistep') + + We support three types of `skip_type`: + - 'logSNR': uniform logSNR for the time steps. **Recommended for low-resolutional images** + - 'time_uniform': uniform time for the time steps. **Recommended for high-resolutional images**. + - 'time_quadratic': quadratic time for the time steps. + + ===================================================== + Args: + x: A pytorch tensor. The initial value at time `t_start` + e.g. if `t_start` == T, then `x` is a sample from the standard normal distribution. + steps: A `int`. The total number of function evaluations (NFE). + t_start: A `float`. The starting time of the sampling. + If `T` is None, we use self.noise_schedule.T (default is 1.0). + t_end: A `float`. The ending time of the sampling. + If `t_end` is None, we use 1. / self.noise_schedule.total_N. + e.g. if total_N == 1000, we have `t_end` == 1e-3. + For discrete-time DPMs: + - We recommend `t_end` == 1. / self.noise_schedule.total_N. + For continuous-time DPMs: + - We recommend `t_end` == 1e-3 when `steps` <= 15; and `t_end` == 1e-4 when `steps` > 15. + order: A `int`. The order of DPM-Solver. + skip_type: A `str`. The type for the spacing of the time steps. \ + 'time_uniform' or 'logSNR' or 'time_quadratic'. + method: A `str`. The method for sampling. 'singlestep' or 'multistep' or 'singlestep_fixed' or 'adaptive'. + denoise: A `bool`. Whether to denoise at the final step. Default is False. + If `denoise` is True, the total NFE is (`steps` + 1). + solver_type: A `str`. The taylor expansion type for the solver. \ + `dpm_solver` or `taylor`. We recommend `dpm_solver`. + atol: A `float`. The absolute tolerance of the adaptive step size solver. Valid when `method` == 'adaptive'. + rtol: A `float`. The relative tolerance of the adaptive step size solver. Valid when `method` == 'adaptive'. + Returns: + x_end: A pytorch tensor. The approximated solution at time `t_end`. + + """ + t_0 = 1. / self.noise_schedule.total_N if t_end is None else t_end + t_T = self.noise_schedule.T if t_start is None else t_start + device = x.device + if method == 'adaptive': + with torch.no_grad(): + x = self.dpm_solver_adaptive( + x, order=order, t_T=t_T, t_0=t_0, atol=atol, rtol=rtol, solver_type=solver_type + ) + elif method == 'multistep': + assert steps >= order + timesteps = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=steps, device=device) + assert timesteps.shape[0] - 1 == steps + with torch.no_grad(): + vec_t = timesteps[0].expand((x.shape[0])) + model_prev_list = [self.model_fn(x, vec_t)] + t_prev_list = [vec_t] + # Init the first `order` values by lower order multistep DPM-Solver. + for init_order in range(1, order): + vec_t = timesteps[init_order].expand(x.shape[0]) + x = self.multistep_dpm_solver_update( + x, model_prev_list, t_prev_list, vec_t, init_order, solver_type=solver_type + ) + model_prev_list.append(self.model_fn(x, vec_t)) + t_prev_list.append(vec_t) + # Compute the remaining values by `order`-th order multistep DPM-Solver. + for step in range(order, steps + 1): + vec_t = timesteps[step].expand(x.shape[0]) + x = self.multistep_dpm_solver_update( + x, model_prev_list, t_prev_list, vec_t, order, solver_type=solver_type + ) + for i in range(order - 1): + t_prev_list[i] = t_prev_list[i + 1] + model_prev_list[i] = model_prev_list[i + 1] + t_prev_list[-1] = vec_t + # We do not need to evaluate the final model value. + if step < steps: + model_prev_list[-1] = self.model_fn(x, vec_t) + elif method in ['singlestep', 'singlestep_fixed']: + if method == 'singlestep': + orders = self.get_orders_for_singlestep_solver(steps=steps, order=order) + timesteps = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=steps, device=device) + elif method == 'singlestep_fixed': + K = steps // order + orders = [ + order, + ] * K + timesteps = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=(K * order), device=device) + with torch.no_grad(): + i = 0 + for order in orders: + vec_s, vec_t = timesteps[i].expand(x.shape[0]), timesteps[i + order].expand(x.shape[0]) + h = self.noise_schedule.marginal_lambda(timesteps[i + order] + ) - self.noise_schedule.marginal_lambda(timesteps[i]) + r1 = None if order <= 1 else ( + self.noise_schedule.marginal_lambda(timesteps[i + 1]) - + self.noise_schedule.marginal_lambda(timesteps[i]) + ) / h + r2 = None if order <= 2 else ( + self.noise_schedule.marginal_lambda(timesteps[i + 2]) - + self.noise_schedule.marginal_lambda(timesteps[i]) + ) / h + x = self.singlestep_dpm_solver_update(x, vec_s, vec_t, order, solver_type=solver_type, r1=r1, r2=r2) + i += order + if denoise: + x = self.denoise_fn(x, torch.ones((x.shape[0], )).to(device) * t_0) + return x + + +############################################################# +# other utility functions +############################################################# + + +def interpolate_fn(x, xp, yp): + """ + A piecewise linear function y = f(x), using xp and yp as keypoints. + We implement f(x) in a differentiable way (i.e. applicable for autograd). + The function f(x) is well-defined for all x-axis. \ + (For x beyond the bounds of xp, we use the outmost points of xp to define the linear function.) + + Args: + x: PyTorch tensor with shape [N, C], where N is the batch size, C is the number of channels + (we use C = 1 for DPM-Solver). + xp: PyTorch tensor with shape [C, K], where K is the number of keypoints. + yp: PyTorch tensor with shape [C, K]. + Returns: + The function values f(x), with shape [N, C]. + """ + N, K = x.shape[0], xp.shape[1] + all_x = torch.cat([x.unsqueeze(2), xp.unsqueeze(0).repeat((N, 1, 1))], dim=2) + sorted_all_x, x_indices = torch.sort(all_x, dim=2) + x_idx = torch.argmin(x_indices, dim=2) + cand_start_idx = x_idx - 1 + start_idx = torch.where( + torch.eq(x_idx, 0), + torch.tensor(1, device=x.device), + torch.where( + torch.eq(x_idx, K), + torch.tensor(K - 2, device=x.device), + cand_start_idx, + ), + ) + end_idx = torch.where(torch.eq(start_idx, cand_start_idx), start_idx + 2, start_idx + 1) + start_x = torch.gather(sorted_all_x, dim=2, index=start_idx.unsqueeze(2)).squeeze(2) + end_x = torch.gather(sorted_all_x, dim=2, index=end_idx.unsqueeze(2)).squeeze(2) + start_idx2 = torch.where( + torch.eq(x_idx, 0), + torch.tensor(0, device=x.device), + torch.where( + torch.eq(x_idx, K), + torch.tensor(K - 2, device=x.device), + cand_start_idx, + ), + ) + y_positions_expanded = yp.unsqueeze(0).expand(N, -1, -1) + start_y = torch.gather(y_positions_expanded, dim=2, index=start_idx2.unsqueeze(2)).squeeze(2) + end_y = torch.gather(y_positions_expanded, dim=2, index=(start_idx2 + 1).unsqueeze(2)).squeeze(2) + cand = start_y + (x - start_x) * (end_y - start_y) / (end_x - start_x) + return cand + + +def expand_dims(v, dims): + """ + Expand the tensor `v` to the dim `dims`. + + Args: + `v`: a PyTorch tensor with shape [N]. + `dim`: a `int`. + Returns: + a PyTorch tensor with shape [N, 1, 1, ..., 1] and the total dimension is `dims`. + """ + return v[(..., ) + (None, ) * (dims - 1)] diff --git a/ding/torch_utils/distribution.py b/ding/torch_utils/distribution.py index dc34fb492a..f68ef6fba0 100644 --- a/ding/torch_utils/distribution.py +++ b/ding/torch_utils/distribution.py @@ -9,11 +9,11 @@ class Pd(object): - r""" + """ Overview: Abstract class for parameterizable probability distributions and sampling functions. - Interface: - neglogp, entropy, noise_mode, mode, sample + Interfaces: + ``neglogp``, ``entropy``, ``noise_mode``, ``mode``, ``sample`` .. tip:: @@ -21,7 +21,7 @@ class Pd(object): """ def neglogp(self, x: torch.Tensor) -> torch.Tensor: - r""" + """ Overview: Calculate cross_entropy between input x and logits Arguments: @@ -32,7 +32,7 @@ def neglogp(self, x: torch.Tensor) -> torch.Tensor: raise NotImplementedError def entropy(self) -> torch.Tensor: - r""" + """ Overview: Calculate the softmax entropy of logits Arguments: @@ -43,21 +43,21 @@ def entropy(self) -> torch.Tensor: raise NotImplementedError def noise_mode(self): - r""" + """ Overview: Add noise to logits. This method is designed for randomness """ raise NotImplementedError def mode(self): - r""" + """ Overview: Return logits argmax result. This method is designed for deterministic. """ raise NotImplementedError def sample(self): - r""" + """ Overview: Sample from logits's distribution by using softmax. This method is designed for multinomial. """ @@ -65,15 +65,15 @@ def sample(self): class CategoricalPd(Pd): - r""" + """ Overview: Catagorical probility distribution sampler - Interface: - update_logits, neglogp, entropy, noise_mode, mode, sample + Interfaces: + ``__init__``, ``neglogp``, ``entropy``, ``noise_mode``, ``mode``, ``sample`` """ def __init__(self, logits: torch.Tensor = None) -> None: - r""" + """ Overview: Init the Pd with logits Arguments: @@ -82,7 +82,7 @@ def __init__(self, logits: torch.Tensor = None) -> None: self.update_logits(logits) def update_logits(self, logits: torch.Tensor) -> None: - r""" + """ Overview: Updata logits Arguments: @@ -91,7 +91,7 @@ def update_logits(self, logits: torch.Tensor) -> None: self.logits = logits def neglogp(self, x, reduction: str = 'mean') -> torch.Tensor: - r""" + """ Overview: Calculate cross_entropy between input x and logits Arguments: @@ -103,7 +103,7 @@ def neglogp(self, x, reduction: str = 'mean') -> torch.Tensor: return F.cross_entropy(self.logits, x, reduction=reduction) def entropy(self, reduction: str = 'mean') -> torch.Tensor: - r""" + """ Overview: Calculate the softmax entropy of logits Arguments: @@ -191,16 +191,22 @@ class CategoricalPdPytorch(torch.distributions.Categorical): Overview: Wrapped ``torch.distributions.Categorical`` - Interface: - update_logits, update_probs, sample, neglogp, mode, entropy + Interfaces: + ``__init__``, ``update_logits``, ``update_probs``, ``sample``, ``neglogp``, ``mode``, ``entropy`` """ def __init__(self, probs: torch.Tensor = None) -> None: + """ + Overview: + Initialize the CategoricalPdPytorch object. + Arguments: + - probs (:obj:`torch.Tensor`): The tensor of probabilities. + """ if probs is not None: self.update_probs(probs) def update_logits(self, logits: torch.Tensor) -> None: - r""" + """ Overview: Updata logits Arguments: @@ -209,7 +215,7 @@ def update_logits(self, logits: torch.Tensor) -> None: super().__init__(logits=logits) def update_probs(self, probs: torch.Tensor) -> None: - r""" + """ Overview: Updata probs Arguments: @@ -218,7 +224,7 @@ def update_probs(self, probs: torch.Tensor) -> None: super().__init__(probs=probs) def sample(self) -> torch.Tensor: - r""" + """ Overview: Sample from logits's distribution by using softmax Return: @@ -227,7 +233,7 @@ def sample(self) -> torch.Tensor: return super().sample() def neglogp(self, actions: torch.Tensor, reduction: str = 'mean') -> torch.Tensor: - r""" + """ Overview: Calculate cross_entropy between input x and logits Arguments: @@ -244,7 +250,7 @@ def neglogp(self, actions: torch.Tensor, reduction: str = 'mean') -> torch.Tenso return neglogp.mean(dim=0) def mode(self) -> torch.Tensor: - r""" + """ Overview: Return logits argmax result Return: @@ -253,7 +259,7 @@ def mode(self) -> torch.Tensor: return self.probs.argmax(dim=-1) def entropy(self, reduction: str = None) -> torch.Tensor: - r""" + """ Overview: Calculate the softmax entropy of logits Arguments: diff --git a/ding/torch_utils/loss/contrastive_loss.py b/ding/torch_utils/loss/contrastive_loss.py index ed17eebcb1..94ef62f8de 100644 --- a/ding/torch_utils/loss/contrastive_loss.py +++ b/ding/torch_utils/loss/contrastive_loss.py @@ -8,10 +8,11 @@ class ContrastiveLoss(nn.Module): """ - The class for contrastive learning losses. - Only InfoNCE loss supported currently. - Code Reference: https://github.com/rdevon/DIM. - paper: https://arxiv.org/abs/1808.06670. + Overview: + The class for contrastive learning losses. Only InfoNCE loss is supported currently. \ + Code Reference: https://github.com/rdevon/DIM. Paper Reference: https://arxiv.org/abs/1808.06670. + Interfaces: + ``__init__``, ``forward``. """ def __init__( @@ -24,13 +25,18 @@ def __init__( temperature: float = 1.0, ) -> None: """ - Args: - x_size: input shape for x, both the obs shape and the encoding shape are supported. - y_size: input shape for y, both the obs shape and the encoding shape are supported. - heads: a list of 2 int elems, heads[0] for x and head[1] for y. + Overview: + Initialize the ContrastiveLoss object using the given arguments. + Arguments: + - x_size (:obj:`Union[int, SequenceType]`): input shape for x, both the obs shape and the encoding shape \ + are supported. + - y_size (:obj:`Union[int, SequenceType]`): Input shape for y, both the obs shape and the encoding shape \ + are supported. + - heads (:obj:`SequenceType`): A list of 2 int elems, ``heads[0]`` for x and ``head[1]`` for y. \ Used in multi-head, global-local, local-local MI maximization process. - loss_type: only the InfoNCE loss is available now. - temperature: the parameter to adjust the log_softmax. + - encoder_shape (:obj:`Union[int, SequenceType]`): The dimension of encoder hidden state. + - loss_type: Only the InfoNCE loss is available now. + - temperature: The parameter to adjust the ``log_softmax``. """ super(ContrastiveLoss, self).__init__() assert len(heads) == 2, "Expected length of 2, but got: {}".format(len(heads)) @@ -39,36 +45,68 @@ def __init__( self._type = loss_type.lower() self._encode_shape = encode_shape self._heads = heads - self._x_encoder = self._get_encoder(x_size, heads[0]) - self._y_encoder = self._get_encoder(y_size, heads[1]) + self._x_encoder = self._create_encoder(x_size, heads[0]) + self._y_encoder = self._create_encoder(y_size, heads[1]) self._temperature = temperature - def _get_encoder(self, obs: Union[int, SequenceType], heads: int): + def _create_encoder(self, obs_size: Union[int, SequenceType], heads: int) -> nn.Module: + """ + Overview: + Create the encoder for the input obs. + Arguments: + - obs_size (:obj:`Union[int, SequenceType]`): input shape for x, both the obs shape and the encoding shape \ + are supported. If the obs_size is an int, it means the obs is a 1D vector. If the obs_size is a list \ + such as [1, 16, 16], it means the obs is a 3D image with shape [1, 16, 16]. + - heads (:obj:`int`): The number of heads. + Returns: + - encoder (:obj:`nn.Module`): The encoder module. + Examples: + >>> obs_size = 16 + or + >>> obs_size = [1, 16, 16] + >>> heads = 1 + >>> encoder = self._create_encoder(obs_size, heads) + """ from ding.model import ConvEncoder, FCEncoder - if isinstance(obs, int): - obs = [obs] - assert len(obs) in [1, 3] + if isinstance(obs_size, int): + obs_size = [obs_size] + assert len(obs_size) in [1, 3] - if len(obs) == 1: + if len(obs_size) == 1: hidden_size_list = [128, 128, self._encode_shape * heads] - encoder = FCEncoder(obs[0], hidden_size_list) + encoder = FCEncoder(obs_size[0], hidden_size_list) else: hidden_size_list = [32, 64, 64, self._encode_shape * heads] - if obs[-1] >= 36: - encoder = ConvEncoder(obs, hidden_size_list) + if obs_size[-1] >= 36: + encoder = ConvEncoder(obs_size, hidden_size_list) else: - encoder = ConvEncoder(obs, hidden_size_list, kernel_size=[4, 3, 2], stride=[2, 1, 1]) + encoder = ConvEncoder(obs_size, hidden_size_list, kernel_size=[4, 3, 2], stride=[2, 1, 1]) return encoder - def forward(self, x: torch.Tensor, y: torch.Tensor): + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: """ - Computes the noise contrastive estimation-based loss, a.k.a. infoNCE. - Args: - x: the input x, both raw obs and encoding are supported. - y: the input y, both raw obs and encoding are supported. + Overview: + Computes the noise contrastive estimation-based loss, a.k.a. infoNCE. + Arguments: + - x (:obj:`torch.Tensor`): The input x, both raw obs and encoding are supported. + - y (:obj:`torch.Tensor`): The input y, both raw obs and encoding are supported. Returns: - torch.Tensor: loss value. + loss (:obj:`torch.Tensor`): The calculated loss value. + Examples: + >>> x_dim = [3, 16] + >>> encode_shape = 16 + >>> x = np.random.normal(0, 1, size=x_dim) + >>> y = x ** 2 + 0.01 * np.random.normal(0, 1, size=x_dim) + >>> estimator = ContrastiveLoss(dims, dims, encode_shape=encode_shape) + >>> loss = estimator.forward(x, y) + Examples: + >>> x_dim = [3, 1, 16, 16] + >>> encode_shape = 16 + >>> x = np.random.normal(0, 1, size=x_dim) + >>> y = x ** 2 + 0.01 * np.random.normal(0, 1, size=x_dim) + >>> estimator = ContrastiveLoss(dims, dims, encode_shape=encode_shape) + >>> loss = estimator.forward(x, y) """ N = x.size(0) @@ -79,7 +117,7 @@ def forward(self, x: torch.Tensor, y: torch.Tensor): x_n = x.view(-1, self._encode_shape) y_n = y.view(-1, self._encode_shape) - # Use inner product to obtain postive samples. + # Use inner product to obtain positive samples. # [N, x_heads, encode_dim] * [N, encode_dim, y_heads] -> [N, x_heads, y_heads] u_pos = torch.matmul(x, y.permute(0, 2, 1)).unsqueeze(2) # Use outer product to obtain all sample permutations. @@ -92,7 +130,7 @@ def forward(self, x: torch.Tensor, y: torch.Tensor): u_neg = (n_mask * u_all) - (10. * (1 - n_mask)) u_neg = u_neg.view(N, N * x_heads, y_heads).unsqueeze(dim=1).expand(-1, x_heads, -1, -1) - # Concatenate postive and negative samples and apply log softmax. + # Concatenate positive and negative samples and apply log softmax. pred_lgt = torch.cat([u_pos, u_neg], dim=2) pred_log = F.log_softmax(pred_lgt * self._temperature, dim=2) diff --git a/ding/torch_utils/loss/cross_entropy_loss.py b/ding/torch_utils/loss/cross_entropy_loss.py index f8365645c1..3cdcb969d7 100644 --- a/ding/torch_utils/loss/cross_entropy_loss.py +++ b/ding/torch_utils/loss/cross_entropy_loss.py @@ -5,19 +5,26 @@ class LabelSmoothCELoss(nn.Module): - r""" + """ Overview: Label smooth cross entropy loss. Interfaces: - forward + ``__init__``, ``forward``. """ def __init__(self, ratio: float) -> None: + """ + Overview: + Initialize the LabelSmoothCELoss object using the given arguments. + Arguments: + - ratio (:obj:`float`): The ratio of label-smoothing (the value is in 0-1). If the ratio is larger, the \ + extent of label smoothing is larger. + """ super().__init__() self.ratio = ratio def forward(self, logits: torch.Tensor, labels: torch.LongTensor) -> torch.Tensor: - r""" + """ Overview: Calculate label smooth cross entropy loss. Arguments: @@ -35,22 +42,37 @@ def forward(self, logits: torch.Tensor, labels: torch.LongTensor) -> torch.Tenso class SoftFocalLoss(nn.Module): - r""" + """ Overview: Soft focal loss. Interfaces: - forward + ``__init__``, ``forward``. """ def __init__( self, gamma: int = 2, weight: Any = None, size_average: bool = True, reduce: Optional[bool] = None ) -> None: + """ + Overview: + Initialize the SoftFocalLoss object using the given arguments. + Arguments: + - gamma (:obj:`int`): The extent of focus on hard samples. A smaller ``gamma`` will lead to more focus on \ + easy samples, while a larger ``gamma`` will lead to more focus on hard samples. + - weight (:obj:`Any`): The weight for loss of each class. + - size_average (:obj:`bool`): By default, the losses are averaged over each loss element in the batch. \ + Note that for some losses, there are multiple elements per sample. If the field ``size_average`` is \ + set to ``False``, the losses are instead summed for each minibatch. Ignored when ``reduce`` is \ + ``False``. + - reduce (:obj:`Optional[bool]`): By default, the losses are averaged or summed over observations for \ + each minibatch depending on size_average. When ``reduce`` is ``False``, returns a loss for each batch \ + element instead and ignores ``size_average``. + """ super().__init__() self.gamma = gamma self.nll_loss = torch.nn.NLLLoss2d(weight, size_average, reduce=reduce) def forward(self, inputs: torch.Tensor, targets: torch.LongTensor) -> torch.Tensor: - r""" + """ Overview: Calculate soft focal loss. Arguments: @@ -63,11 +85,14 @@ def forward(self, inputs: torch.Tensor, targets: torch.LongTensor) -> torch.Tens def build_ce_criterion(cfg: dict) -> nn.Module: - r""" + """ Overview: - Get a cross enntropy loss instance according to given config. + Get a cross entropy loss instance according to given config. Arguments: - - cfg (:obj:`dict`) + - cfg (:obj:`dict`) : Config dict. It contains: + - type (:obj:`str`): Type of loss function, now supports ['cross_entropy', 'label_smooth_ce', \ + 'soft_focal_loss']. + - kwargs (:obj:`dict`): Arguments for the corresponding loss function. Returns: - loss (:obj:`nn.Module`): loss function instance """ diff --git a/ding/torch_utils/loss/multi_logits_loss.py b/ding/torch_utils/loss/multi_logits_loss.py index 666a29df3a..293d74e114 100644 --- a/ding/torch_utils/loss/multi_logits_loss.py +++ b/ding/torch_utils/loss/multi_logits_loss.py @@ -6,36 +6,22 @@ from ding.torch_utils.network import one_hot -def get_distance_matrix(lx, ly, mat, M: int) -> np.ndarray: - nlx = np.broadcast_to(lx, [M, M]).T - nly = np.broadcast_to(ly, [M, M]) - nret = nlx + nly - mat - - # ret = [] - # for i in range(M): - # ret.append(lx[i] + ly - mat[i]) - # ret = np.stack(ret) - # assert ret.shape == (M, M) - # assert np.all(nret == ret) - return nret - - class MultiLogitsLoss(nn.Module): - ''' + """ Overview: Base class for supervised learning on linklink, including basic processes. - Interface: - forward - ''' + Interfaces: + ``__init__``, ``forward``. + """ def __init__(self, criterion: str = None, smooth_ratio: float = 0.1) -> None: - ''' + """ Overview: - initialization method, use cross_entropy as default criterion + Initialization method, use cross_entropy as default criterion. Arguments: - - criterion (:obj:`str`): criterion type, supports ['cross_entropy', 'label_smooth_ce'] - - smooth_ratio (:obs:`float`): smooth_ratio for label smooth - ''' + - criterion (:obj:`str`): Criterion type, supports ['cross_entropy', 'label_smooth_ce']. + - smooth_ratio (:obj:`float`): Smoothing ratio for label smoothing. + """ super(MultiLogitsLoss, self).__init__() if criterion is None: criterion = 'cross_entropy' @@ -45,6 +31,15 @@ def __init__(self, criterion: str = None, smooth_ratio: float = 0.1) -> None: self.ratio = smooth_ratio def _label_process(self, logits: torch.Tensor, labels: torch.LongTensor) -> torch.LongTensor: + """ + Overview: + Process the label according to the criterion. + Arguments: + - logits (:obj:`torch.Tensor`): Predicted logits. + - labels (:obj:`torch.LongTensor`): Ground truth. + Returns: + - ret (:obj:`torch.LongTensor`): Processed label. + """ N = logits.shape[1] if self.criterion == 'cross_entropy': return one_hot(labels, num=N) @@ -55,10 +50,28 @@ def _label_process(self, logits: torch.Tensor, labels: torch.LongTensor) -> torc return ret def _nll_loss(self, nlls: torch.Tensor, labels: torch.LongTensor) -> torch.Tensor: + """ + Overview: + Calculate the negative log likelihood loss. + Arguments: + - nlls (:obj:`torch.Tensor`): Negative log likelihood loss. + - labels (:obj:`torch.LongTensor`): Ground truth. + Returns: + - ret (:obj:`torch.Tensor`): Calculated loss. + """ ret = (-nlls * (labels.detach())) return ret.sum(dim=1) def _get_metric_matrix(self, logits: torch.Tensor, labels: torch.LongTensor) -> torch.Tensor: + """ + Overview: + Calculate the metric matrix. + Arguments: + - logits (:obj:`torch.Tensor`): Predicted logits. + - labels (:obj:`torch.LongTensor`): Ground truth. + Returns: + - metric (:obj:`torch.Tensor`): Calculated metric matrix. + """ M, N = logits.shape labels = self._label_process(logits, labels) logits = F.log_softmax(logits, dim=1) @@ -70,14 +83,22 @@ def _get_metric_matrix(self, logits: torch.Tensor, labels: torch.LongTensor) -> return torch.stack(metric, dim=0) def _match(self, matrix: torch.Tensor): + """ + Overview: + Match the metric matrix. + Arguments: + - matrix (:obj:`torch.Tensor`): Metric matrix. + Returns: + - index (:obj:`np.ndarray`): Matched index. + """ mat = matrix.clone().detach().to('cpu').numpy() mat = -mat # maximize M = mat.shape[0] index = np.full(M, -1, dtype=np.int32) # -1 note not find link lx = mat.max(axis=1) ly = np.zeros(M, dtype=np.float32) - visx = np.zeros(M, dtype=np.bool) - visy = np.zeros(M, dtype=np.bool) + visx = np.zeros(M, dtype=np.bool_) + visy = np.zeros(M, dtype=np.bool_) def has_augmented_path(t, binary_distance_matrix): # What is changed? visx, visy, distance_matrix, index @@ -95,7 +116,7 @@ def has_augmented_path(t, binary_distance_matrix): while True: visx.fill(False) visy.fill(False) - distance_matrix = get_distance_matrix(lx, ly, mat, M) + distance_matrix = self._get_distance_matrix(lx, ly, mat, M) binary_distance_matrix = np.abs(distance_matrix) < 1e-4 if has_augmented_path(i, binary_distance_matrix): break @@ -108,8 +129,24 @@ def has_augmented_path(t, binary_distance_matrix): ly[visy] += d return index + @staticmethod + def _get_distance_matrix(lx: np.ndarray, ly: np.ndarray, mat: np.ndarray, M: int) -> np.ndarray: + """ + Overview: + Get distance matrix. + Arguments: + - lx (:obj:`np.ndarray`): lx. + - ly (:obj:`np.ndarray`): ly. + - mat (:obj:`np.ndarray`): mat. + - M (:obj:`int`): M. + """ + nlx = np.broadcast_to(lx, [M, M]).T + nly = np.broadcast_to(ly, [M, M]) + nret = nlx + nly - mat + return nret + def forward(self, logits: torch.Tensor, labels: torch.LongTensor) -> torch.Tensor: - r""" + """ Overview: Calculate multiple logits loss. Arguments: diff --git a/ding/torch_utils/loss/tests/test_contrastive_loss.py b/ding/torch_utils/loss/tests/test_contrastive_loss.py index 785f703bc2..41506bc476 100644 --- a/ding/torch_utils/loss/tests/test_contrastive_loss.py +++ b/ding/torch_utils/loss/tests/test_contrastive_loss.py @@ -5,7 +5,7 @@ from ding.torch_utils.loss.contrastive_loss import ContrastiveLoss -@pytest.mark.benchmark +@pytest.mark.unittest @pytest.mark.parametrize('noise', [0.1, 1.0]) @pytest.mark.parametrize('dims', [16, [1, 16, 16], [1, 40, 40]]) def test_infonce_loss(noise, dims): diff --git a/ding/torch_utils/lr_scheduler.py b/ding/torch_utils/lr_scheduler.py new file mode 100644 index 0000000000..09d341430a --- /dev/null +++ b/ding/torch_utils/lr_scheduler.py @@ -0,0 +1,60 @@ +from functools import partial +import math + +import torch.optim +from torch.optim.lr_scheduler import LambdaLR + + +def get_lr_ratio(epoch: int, warmup_epochs: int, learning_rate: float, lr_decay_epochs: int, min_lr: float) -> float: + """ + Overview: + Get learning rate ratio for each epoch. + Arguments: + - epoch (:obj:`int`): Current epoch. + - warmup_epochs (:obj:`int`): Warmup epochs. + - learning_rate (:obj:`float`): Learning rate. + - lr_decay_epochs (:obj:`int`): Learning rate decay epochs. + - min_lr (:obj:`float`): Minimum learning rate. + """ + + # 1) linear warmup for warmup_epochs. + if epoch < warmup_epochs: + return epoch / warmup_epochs + # 2) if epoch> lr_decay_epochs, return min learning rate + if epoch > lr_decay_epochs: + return min_lr / learning_rate + # 3) in between, use cosine decay down to min learning rate + decay_ratio = (epoch - warmup_epochs) / (lr_decay_epochs - warmup_epochs) + assert 0 <= decay_ratio <= 1 + coefficient = 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) + return (min_lr + coefficient * (learning_rate - min_lr)) / learning_rate + + +def cos_lr_scheduler( + optimizer: torch.optim.Optimizer, + learning_rate: float, + warmup_epochs: float = 5, + lr_decay_epochs: float = 100, + min_lr: float = 6e-5 +) -> torch.optim.lr_scheduler.LambdaLR: + """ + Overview: + Cosine learning rate scheduler. + Arguments: + - optimizer (:obj:`torch.optim.Optimizer`): Optimizer. + - learning_rate (:obj:`float`): Learning rate. + - warmup_epochs (:obj:`float`): Warmup epochs. + - lr_decay_epochs (:obj:`float`): Learning rate decay epochs. + - min_lr (:obj:`float`): Minimum learning rate. + """ + + return LambdaLR( + optimizer, + partial( + get_lr_ratio, + warmup_epochs=warmup_epochs, + lr_decay_epochs=lr_decay_epochs, + min_lr=min_lr, + learning_rate=learning_rate + ) + ) diff --git a/ding/torch_utils/math_helper.py b/ding/torch_utils/math_helper.py index bd94b36558..751c807597 100644 --- a/ding/torch_utils/math_helper.py +++ b/ding/torch_utils/math_helper.py @@ -9,17 +9,24 @@ def cov( ddof: Optional[int] = None, aweights: Optional[torch.Tensor] = None ) -> torch.Tensor: - r""" + """ Overview: - Estimates covariance matrix like ``numpy.cov`` + Estimates covariance matrix like ``numpy.cov``. Arguments: - - x (:obj:`torch.Tensor`) - - rowvar (:obj:`bool`) - - bias (:obj:`bool`) - - ddof (:obj:`Optional[int]`) - - aweights (:obj:`Optional[torch.Tensor]`) + - x (:obj:`torch.Tensor`): A 1-D or 2-D tensor containing multiple variables and observations. Each row of \ + ``x`` represents a variable, and each column a single observation of all those variables. + - rowvar (:obj:`bool`): If ``rowvar`` is True by default, and each column is a single observation of all those \ + variables. Otherwise, each column represents a variable, while the rows contain observations. + - bias (:obj:`bool`): Default normalization (False) is by dividing ``N - 1``, where ``N`` is the number of \ + observations given (unbiased estimate). If ``bias`` is ``True``, then normalization is by ``N``. + - ddof (:obj:`Optional[int]`): If ``ddof`` is not ``None``, it implies that the argument ``bias`` is \ + overridden. Note that ``ddof=1`` will return the unbiased estimate (equals to ``bias=False``), and \ + ``ddof=0`` will return the biased estimation (equals to ``bias=True``). + - aweights (:obj:`Optional[torch.Tensor]`): 1-D tensor of observation vector weights. These relative weights \ + are typically large for observations considered “important” and smaller for observations considered less \ + “important”. If ``ddof=0``, the tensor of weights can be used to assign weights to observation vectors. Returns: - - cov_mat (:obj:`torch.Tensor`): Covariance matrix + - cov_mat (:obj:`torch.Tensor`): Covariance matrix calculated. """ if x.dim() == 1 and rowvar: raise NotImplementedError diff --git a/ding/torch_utils/metric.py b/ding/torch_utils/metric.py index 7fc2edf230..75554c211d 100644 --- a/ding/torch_utils/metric.py +++ b/ding/torch_utils/metric.py @@ -9,17 +9,20 @@ def levenshtein_distance( target_extra: Optional[torch.Tensor] = None, extra_fn: Optional[Callable] = None ) -> torch.FloatTensor: - r""" + """ Overview: Levenshtein Distance, i.e. Edit Distance. Arguments: - - pred (:obj:`torch.LongTensor`): shape: (N1, ) (N1 >= 0) - - target (:obj:`torch.LongTensor`): shape: (N2, ) (N2 >= 0) - - pred_extra (:obj:`Optional[torch.Tensor]`) - - target_extra (:obj:`Optional[torch.Tensor]`) - - extra_fn (:obj:`Optional[Callable]`): if specified, the distance metric of the extra input data + - pred (:obj:`torch.LongTensor`): The first tensor to calculate the distance, shape: (N1, ) (N1 >= 0). + - target (:obj:`torch.LongTensor`): The second tensor to calculate the distance, shape: (N2, ) (N2 >= 0). + - pred_extra (:obj:`Optional[torch.Tensor]`): Extra tensor to calculate the distance, only works when \ + ``extra_fn`` is not ``None``. + - target_extra (:obj:`Optional[torch.Tensor]`): Extra tensor to calculate the distance, only works when \ + ``extra_fn`` is not ``None``. + - extra_fn (:obj:`Optional[Callable]`): The distance function for ``pred_extra`` and \ + ``target_extra``. If set to ``None``, this distance will not be considered. Returns: - - distance (:obj:`torch.FloatTensor`): distance(scalar), shape: (1, ) + - distance (:obj:`torch.FloatTensor`): distance(scalar), shape: (1, ). """ assert (isinstance(pred, torch.Tensor) and isinstance(target, torch.Tensor)) assert (pred.dtype == torch.long and target.dtype == torch.long), '{}\t{}'.format(pred.dtype, target.dtype) @@ -57,19 +60,19 @@ def levenshtein_distance( def hamming_distance(pred: torch.LongTensor, target: torch.LongTensor, weight=1.) -> torch.LongTensor: - r''' + """ Overview: - Hamming Distance + Hamming Distance. Arguments: - - pred (:obj:`torch.LongTensor`): pred input, boolean vector(0 or 1) - - target (:obj:`torch.LongTensor`): target input, boolean vector(0 or 1) - - weight (:obj:`torch.LongTensor`): weight to multiply + - pred (:obj:`torch.LongTensor`): Pred input, boolean vector(0 or 1). + - target (:obj:`torch.LongTensor`): Target input, boolean vector(0 or 1). + - weight (:obj:`torch.LongTensor`): Weight to multiply. Returns: - - distance(:obj:`torch.LongTensor`): distance(scalar), shape (1, ) + - distance(:obj:`torch.LongTensor`): Distance (scalar), shape (1, ). Shapes: - pred & target (:obj:`torch.LongTensor`): shape :math:`(B, N)`, \ while B is the batch size, N is the dimension - ''' + """ assert (isinstance(pred, torch.Tensor) and isinstance(target, torch.Tensor)) assert (pred.dtype == torch.long and target.dtype == torch.long) assert (pred.device == target.device) diff --git a/ding/torch_utils/model_helper.py b/ding/torch_utils/model_helper.py new file mode 100644 index 0000000000..05a83ab773 --- /dev/null +++ b/ding/torch_utils/model_helper.py @@ -0,0 +1,17 @@ +import torch + + +def get_num_params(model: torch.nn.Module) -> int: + """ + Overview: + Return the number of parameters in the model. + Arguments: + - model (:obj:`torch.nn.Module`): The model object to calculate the parameter number. + Returns: + - n_params (:obj:`int`): The calculated number of parameters. + Examples: + >>> model = torch.nn.Linear(3, 5) + >>> num = get_num_params(model) + >>> assert num == 15 + """ + return sum(p.numel() for p in model.parameters()) diff --git a/ding/torch_utils/network/__init__.py b/ding/torch_utils/network/__init__.py index 7f520f7d19..dda50c339e 100644 --- a/ding/torch_utils/network/__init__.py +++ b/ding/torch_utils/network/__init__.py @@ -1,7 +1,7 @@ from .activation import build_activation, Swish from .res_block import ResBlock, ResFCBlock from .nn_module import fc_block, conv2d_block, one_hot, deconv2d_block, BilinearUpsample, NearestUpsample, \ - binary_encode, NoiseLinearLayer, noise_block, MLP, Flatten, normed_linear, normed_conv2d + binary_encode, NoiseLinearLayer, noise_block, MLP, Flatten, normed_linear, normed_conv2d, conv1d_block from .normalization import build_normalization from .rnn import get_lstm, sequence_mask from .soft_argmax import SoftArgmax @@ -10,3 +10,6 @@ from .resnet import resnet18, ResNet from .gumbel_softmax import GumbelSoftmax from .gtrxl import GTrXL, GRUGatingUnit +from .popart import PopArt +#from .dreamer import Conv2dSame, DreamerLayerNorm, ActionHead, DenseHead +from .merge import GatingType, SumMerge, VectorMerge diff --git a/ding/torch_utils/network/activation.py b/ding/torch_utils/network/activation.py index 56afa03ce3..e97d6645b2 100644 --- a/ding/torch_utils/network/activation.py +++ b/ding/torch_utils/network/activation.py @@ -1,41 +1,57 @@ +import math +from collections.abc import Callable + import torch import torch.nn as nn -import torch.nn.functional as F -class GLU(nn.Module): - r""" +class Lambda(nn.Module): + """ Overview: - Gating Linear Unit. - This class does a thing like this: - - .. code:: python - - # Inputs: input, context, output_size - # The gate value is a learnt function of the input. - gate = sigmoid(linear(input.size)(context)) - # Gate the input and return an output of desired size. - gated_input = gate * input - output = linear(output_size)(gated_input) - return output + A custom lambda module for constructing custom layers. Interfaces: - forward + ``__init__``, ``forward``. + """ + + def __init__(self, f: Callable): + """ + Overview: + Initialize the lambda module with a given function. + Arguments: + - f (:obj:`Callable`): a python function + """ + super(Lambda, self).__init__() + self.f = f + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Overview: + Compute the function of the input tensor. + Arguments: + - x (:obj:`torch.Tensor`): The input tensor. + """ + return self.f(x) - .. tip:: - This module also supports 2D convolution, in which case, the input and context must have the same shape. +class GLU(nn.Module): + """ + Overview: + Gating Linear Unit (GLU), a specific type of activation function, which is first proposed in + [Language Modeling with Gated Convolutional Networks](https://arxiv.org/pdf/1612.08083.pdf). + Interfaces: + ``__init__``, ``forward``. """ def __init__(self, input_dim: int, output_dim: int, context_dim: int, input_type: str = 'fc') -> None: - r""" + """ Overview: - Init GLU + Initialize the GLU module. Arguments: - - input_dim (:obj:`int`): the input dimension - - output_dim (:obj:`int`): the output dimension - - context_dim (:obj:`int`): the context dimension - - input_type (:obj:`str`): the type of input, now support ['fc', 'conv2d'] - """ + - input_dim (:obj:`int`): The dimension of the input tensor. + - output_dim (:obj:`int`): The dimension of the output tensor. + - context_dim (:obj:`int`): The dimension of the context tensor. + - input_type (:obj:`str`): The type of input, now supports ['fc', 'conv2d'] + """ super(GLU, self).__init__() assert (input_type in ['fc', 'conv2d']) if input_type == 'fc': @@ -46,14 +62,14 @@ def __init__(self, input_dim: int, output_dim: int, context_dim: int, input_type self.layer2 = nn.Conv2d(input_dim, output_dim, 1, 1, 0) def forward(self, x: torch.Tensor, context: torch.Tensor) -> torch.Tensor: - r""" + """ Overview: - Return GLU computed tensor + Compute the GLU transformation of the input tensor. Arguments: - - x (:obj:`torch.Tensor`) : the input tensor - - context (:obj:`torch.Tensor`) : the context tensor + - x (:obj:`torch.Tensor`): The input tensor. + - context (:obj:`torch.Tensor`): The context tensor. Returns: - - x (:obj:`torch.Tensor`): the computed tensor + - x (:obj:`torch.Tensor`): The output tensor after GLU transformation. """ gate = self.layer1(context) gate = torch.sigmoid(gate) @@ -63,31 +79,91 @@ def forward(self, x: torch.Tensor, context: torch.Tensor) -> torch.Tensor: class Swish(nn.Module): + """ + Overview: + Swish activation function, which is a smooth, non-monotonic activation function. For more details, please refer + to [Searching for Activation Functions](https://arxiv.org/pdf/1710.05941.pdf). + Interfaces: + ``__init__``, ``forward``. + """ def __init__(self): + """ + Overview: + Initialize the Swish module. + """ super(Swish, self).__init__() - def forward(self, x): - x = x * torch.sigmoid(x) - return x + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Overview: + Compute the Swish transformation of the input tensor. + Arguments: + - x (:obj:`torch.Tensor`): The input tensor. + Returns: + - x (:obj:`torch.Tensor`): The output tensor after Swish transformation. + """ + return x * torch.sigmoid(x) + + +class GELU(nn.Module): + """ + Overview: + Gaussian Error Linear Units (GELU) activation function, which is widely used in NLP models like GPT, BERT. + For more details, please refer to the original paper: https://arxiv.org/pdf/1606.08415.pdf. + Interfaces: + ``__init__``, ``forward``. + """ + + def __init__(self): + """ + Overview: + Initialize the GELU module. + """ + super(GELU, self).__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Overview: + Compute the GELU transformation of the input tensor. + Arguments: + - x (:obj:`torch.Tensor`): The input tensor. + Returns: + - x (:obj:`torch.Tensor`): The output tensor after GELU transformation. + """ + return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0)))) def build_activation(activation: str, inplace: bool = None) -> nn.Module: - r""" + """ Overview: - Return the activation module according to the given type. + Build and return the activation module according to the given type. Arguments: - - actvation (:obj:`str`): the type of activation module, now supports ['relu', 'glu', 'prelu'] - - inplace (:obj:`bool`): can optionally do the operation in-place in relu. Default ``None`` + - activation (:obj:`str`): The type of activation module, now supports \ + ['relu', 'glu', 'prelu', 'swish', 'gelu', 'tanh', 'sigmoid', 'softplus', 'elu', 'square', 'identity']. + - inplace (Optional[:obj:`bool`): Execute the operation in-place in activation, defaults to None. Returns: - - act_func (:obj:`nn.module`): the corresponding activation module + - act_func (:obj:`nn.module`): The corresponding activation module. """ if inplace is not None: assert activation == 'relu', 'inplace argument is not compatible with {}'.format(activation) else: inplace = False - act_func = {'relu': nn.ReLU(inplace=inplace), 'glu': GLU, 'prelu': nn.PReLU(), 'swish': Swish()} - if activation in act_func.keys(): + act_func = { + 'relu': nn.ReLU(inplace=inplace), + 'glu': GLU, + 'prelu': nn.PReLU(), + 'swish': Swish(), + 'gelu': GELU(), + "tanh": nn.Tanh(), + "sigmoid": nn.Sigmoid(), + "softplus": nn.Softplus(), + "elu": nn.ELU(), + "silu": torch.nn.SiLU(inplace=inplace), + "square": Lambda(lambda x: x ** 2), + "identity": Lambda(lambda x: x), + } + if activation.lower() in act_func.keys(): return act_func[activation] else: raise KeyError("invalid key for activation: {}".format(activation)) diff --git a/ding/torch_utils/network/diffusion.py b/ding/torch_utils/network/diffusion.py new file mode 100755 index 0000000000..deb95c9022 --- /dev/null +++ b/ding/torch_utils/network/diffusion.py @@ -0,0 +1,661 @@ +from typing import Union, List, Dict +from collections import namedtuple +import numpy as np +import math +import torch +import torch.nn as nn +import torch.nn.functional as F +from ding.utils import list_split, MODEL_REGISTRY, squeeze, SequenceType + + +def extract(a, t, x_shape): + """ + Overview: + extract output from a through index t. + Arguments: + - a (:obj:`torch.Tensor`): input tensor + - t (:obj:`torch.Tensor`): index tensor + - x_shape (:obj:`torch.Tensor`): shape of x + """ + b, *_ = t.shape + out = a.gather(-1, t) + return out.reshape(b, *((1, ) * (len(x_shape) - 1))) + + +def cosine_beta_schedule(timesteps: int, s: float = 0.008, dtype=torch.float32): + """ + Overview: + cosine schedule + as proposed in https://openreview.net/forum?id=-NEXDKk8gZ + Arguments: + - timesteps (:obj:`int`): timesteps of diffusion step + - s (:obj:`float`): s + - dtype (:obj:`torch.dtype`): dtype of beta + Return: + Tensor of beta [timesteps,], computing by cosine. + """ + steps = timesteps + 1 + x = np.linspace(0, steps, steps) + alphas_cumprod = np.cos(((x / steps) + s) / (1 + s) * np.pi * 0.5) ** 2 + alphas_cumprod = alphas_cumprod / alphas_cumprod[0] + betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1]) + betas_clipped = np.clip(betas, a_min=0, a_max=0.999) + return torch.tensor(betas_clipped, dtype=dtype) + + +def apply_conditioning(x, conditions, action_dim): + """ + Overview: + add condition into x + Arguments: + - x (:obj:`torch.Tensor`): input tensor + - conditions (:obj:`dict`): condition dict, key is timestep, value is condition + - action_dim (:obj:`int`): action dim + """ + for t, val in conditions.items(): + x[:, t, action_dim:] = val.clone() + return x + + +class DiffusionConv1d(nn.Module): + """ + Overview: + Conv1d with activation and normalization for diffusion models. + Interfaces: + ``__init__``, ``forward`` + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + padding: int, + activation: nn.Module = None, + n_groups: int = 8 + ) -> None: + """ + Overview: + Create a 1-dim convlution layer with activation and normalization. This Conv1d have GropuNorm. + And need add 1-dim when compute norm + Arguments: + - in_channels (:obj:`int`): Number of channels in the input tensor + - out_channels (:obj:`int`): Number of channels in the output tensor + - kernel_size (:obj:`int`): Size of the convolving kernel + - padding (:obj:`int`): Zero-padding added to both sides of the input + - activation (:obj:`nn.Module`): the optional activation function + """ + super().__init__() + self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size=kernel_size, padding=padding) + self.norm = nn.GroupNorm(n_groups, out_channels) + self.act = activation + + def forward(self, inputs) -> torch.Tensor: + """ + Overview: + compute conv1d for inputs. + Arguments: + - inputs (:obj:`torch.Tensor`): input tensor + Return: + - out (:obj:`torch.Tensor`): output tensor + """ + x = self.conv1(inputs) + # [batch, channels, horizon] -> [batch, channels, 1, horizon] + x = x.unsqueeze(-2) + x = self.norm(x) + # [batch, channels, 1, horizon] -> [batch, channels, horizon] + x = x.squeeze(-2) + out = self.act(x) + return out + + +class SinusoidalPosEmb(nn.Module): + """ + Overview: + class for computing sin position embeding + Interfaces: + ``__init__``, ``forward`` + """ + + def __init__(self, dim: int) -> None: + """ + Overview: + Initialization of SinusoidalPosEmb class + Arguments: + - dim (:obj:`int`): dimension of embeding + """ + + super().__init__() + self.dim = dim + + def forward(self, x) -> torch.Tensor: + """ + Overview: + compute sin position embeding + Arguments: + - x (:obj:`torch.Tensor`): input tensor + Return: + - emb (:obj:`torch.Tensor`): output tensor + """ + + device = x.device + half_dim = self.dim // 2 + emb = math.log(10000) / (half_dim - 1) + emb = torch.exp(torch.arange(half_dim, device=device) * -emb) + emb = x[:, None] * emb[None, :] + emb = torch.cat((emb.sin(), emb.cos()), dim=1) + return emb + + +class Residual(nn.Module): + """ + Overview: + Basic Residual block + Interfaces: + ``__init__``, ``forward`` + """ + + def __init__(self, fn): + """ + Overview: + Initialization of Residual class + Arguments: + - fn (:obj:`nn.Module`): function of residual block + """ + + super().__init__() + self.fn = fn + + def forward(self, x, *arg, **kwargs): + """ + Overview: + compute residual block + Arguments: + - x (:obj:`torch.Tensor`): input tensor + """ + + return self.fn(x, *arg, **kwargs) + x + + +class LayerNorm(nn.Module): + """ + Overview: + LayerNorm, compute dim = 1, because Temporal input x [batch, dim, horizon] + Interfaces: + ``__init__``, ``forward`` + """ + + def __init__(self, dim, eps=1e-5) -> None: + """ + Overview: + Initialization of LayerNorm class + Arguments: + - dim (:obj:`int`): dimension of input + - eps (:obj:`float`): eps of LayerNorm + """ + + super().__init__() + self.eps = eps + self.g = nn.Parameter(torch.ones(1, dim, 1)) + self.b = nn.Parameter(torch.zeros(1, dim, 1)) + + def forward(self, x): + """ + Overview: + compute LayerNorm + Arguments: + - x (:obj:`torch.Tensor`): input tensor + """ + + print('x.shape:', x.shape) + var = torch.var(x, dim=1, unbiased=False, keepdim=True) + mean = torch.mean(x, dim=1, keepdim=True) + return (x - mean) / (var + self.eps).sqrt() * self.g + self.b + + +class PreNorm(nn.Module): + """ + Overview: + PreNorm, compute dim = 1, because Temporal input x [batch, dim, horizon] + Interfaces: + ``__init__``, ``forward`` + """ + + def __init__(self, dim, fn) -> None: + """ + Overview: + Initialization of PreNorm class + Arguments: + - dim (:obj:`int`): dimension of input + - fn (:obj:`nn.Module`): function of residual block + """ + + super().__init__() + self.fn = fn + self.norm = LayerNorm(dim) + + def forward(self, x): + """ + Overview: + compute PreNorm + Arguments: + - x (:obj:`torch.Tensor`): input tensor + """ + x = self.norm(x) + return self.fn(x) + + +class LinearAttention(nn.Module): + """ + Overview: + Linear Attention head + Interfaces: + ``__init__``, ``forward`` + """ + + def __init__(self, dim, heads=4, dim_head=32) -> None: + """ + Overview: + Initialization of LinearAttention class + Arguments: + - dim (:obj:`int`): dimension of input + - heads (:obj:`int`): heads of attention + - dim_head (:obj:`int`): dim of head + """ + super().__init__() + self.scale = dim_head ** -0.5 + self.heads = heads + hidden_dim = dim_head * heads + self.to_qkv = nn.Conv1d(dim, hidden_dim * 3, 1, bias=False) + self.to_out = nn.Conv1d(hidden_dim, dim, 1) + + def forward(self, x): + """ + Overview: + compute LinearAttention + Arguments: + - x (:obj:`torch.Tensor`): input tensor + """ + qkv = self.to_qkv(x).chunk(3, dim=1) + q, k, v = map(lambda t: t.reshape(t.shape[0], self.heads, -1, t.shape[-1]), qkv) + q = q * self.scale + k = k.softmax(dim=-1) + context = torch.einsum('b h d n, b h e n -> b h d e', k, v) + + out = torch.einsum('b h d e, b h d n -> b h e n', context, q) + out = out.reshape(out.shape[0], -1, out.shape[-1]) + return self.to_out(out) + + +class ResidualTemporalBlock(nn.Module): + """ + Overview: + Residual block of temporal + Interfaces: + ``__init__``, ``forward`` + """ + + def __init__( + self, in_channels: int, out_channels: int, embed_dim: int, kernel_size: int = 5, mish: bool = True + ) -> None: + """ + Overview: + Initialization of ResidualTemporalBlock class + Arguments: + - in_channels (:obj:'int'): dim of in_channels + - out_channels (:obj:'int'): dim of out_channels + - embed_dim (:obj:'int'): dim of embeding layer + - kernel_size (:obj:'int'): kernel_size of conv1d + - mish (:obj:'bool'): whether use mish as a activate function + """ + super().__init__() + if mish: + act = nn.Mish() + else: + act = nn.SiLU() + self.blocks = nn.ModuleList( + [ + DiffusionConv1d(in_channels, out_channels, kernel_size, kernel_size // 2, act), + DiffusionConv1d(out_channels, out_channels, kernel_size, kernel_size // 2, act), + ] + ) + self.time_mlp = nn.Sequential( + act, + nn.Linear(embed_dim, out_channels), + ) + self.residual_conv = nn.Conv1d(in_channels, out_channels, 1) \ + if in_channels != out_channels else nn.Identity() + + def forward(self, x, t): + """ + Overview: + compute residual block + Arguments: + - x (:obj:'tensor'): input tensor + - t (:obj:'tensor'): time tensor + """ + out = self.blocks[0](x) + self.time_mlp(t).unsqueeze(-1) + out = self.blocks[1](out) + return out + self.residual_conv(x) + + +class DiffusionUNet1d(nn.Module): + """ + Overview: + Diffusion unet for 1d vector data + Interfaces: + ``__init__``, ``forward``, ``get_pred`` + """ + + def __init__( + self, + transition_dim: int, + dim: int = 32, + dim_mults: SequenceType = [1, 2, 4, 8], + returns_condition: bool = False, + condition_dropout: float = 0.1, + calc_energy: bool = False, + kernel_size: int = 5, + attention: bool = False, + ) -> None: + """ + Overview: + Initialization of DiffusionUNet1d class + Arguments: + - transition_dim (:obj:'int'): dim of transition, it is obs_dim + action_dim + - dim (:obj:'int'): dim of layer + - dim_mults (:obj:'SequenceType'): mults of dim + - returns_condition (:obj:'bool'): whether use return as a condition + - condition_dropout (:obj:'float'): dropout of returns condition + - calc_energy (:obj:'bool'): whether use calc_energy + - kernel_size (:obj:'int'): kernel_size of conv1d + - attention (:obj:'bool'): whether use attention + """ + super().__init__() + dims = [transition_dim, *map(lambda m: dim * m, dim_mults)] + in_out = list(zip(dims[:-1], dims[1:])) + + if calc_energy: + mish = False + act = nn.SiLU() + else: + mish = True + act = nn.Mish() + + self.time_dim = dim + self.returns_dim = dim + + self.time_mlp = nn.Sequential( + SinusoidalPosEmb(dim), + nn.Linear(dim, dim * 4), + act, + nn.Linear(dim * 4, dim), + ) + + self.returns_condition = returns_condition + self.condition_dropout = condition_dropout + self.cale_energy = calc_energy + + if self.returns_condition: + self.returns_mlp = nn.Sequential( + nn.Linear(1, dim), + act, + nn.Linear(dim, dim * 4), + act, + nn.Linear(dim * 4, dim), + ) + self.mask_dist = torch.distributions.Bernoulli(probs=1 - self.condition_dropout) + embed_dim = 2 * dim + else: + embed_dim = dim + + self.downs = nn.ModuleList([]) + self.ups = nn.ModuleList([]) + num_resolution = len(in_out) + + for ind, (dim_in, dim_out) in enumerate(in_out): + is_last = ind >= (num_resolution - 1) + self.downs.append( + nn.ModuleList( + [ + ResidualTemporalBlock(dim_in, dim_out, embed_dim, kernel_size, mish=mish), + ResidualTemporalBlock(dim_out, dim_out, embed_dim, kernel_size, mish=mish), + Residual(PreNorm(dim_out, LinearAttention(dim_out))) if attention else nn.Identity(), + nn.Conv1d(dim_out, dim_out, 3, 2, 1) if not is_last else nn.Identity() + ] + ) + ) + + mid_dim = dims[-1] + self.mid_block1 = ResidualTemporalBlock(mid_dim, mid_dim, embed_dim, kernel_size, mish) + self.mid_atten = Residual(PreNorm(mid_dim, LinearAttention(mid_dim))) if attention else nn.Identity() + self.mid_block2 = ResidualTemporalBlock(mid_dim, mid_dim, embed_dim, kernel_size, mish) + + for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])): + is_last = ind >= (num_resolution - 1) + self.ups.append( + nn.ModuleList( + [ + ResidualTemporalBlock(dim_out * 2, dim_in, embed_dim, kernel_size, mish=mish), + ResidualTemporalBlock(dim_in, dim_in, embed_dim, kernel_size, mish=mish), + Residual(PreNorm(dim_in, LinearAttention(dim_in))) if attention else nn.Identity(), + nn.ConvTranspose1d(dim_in, dim_in, 4, 2, 1) if not is_last else nn.Identity() + ] + ) + ) + + self.final_conv = nn.Sequential( + DiffusionConv1d(dim, dim, kernel_size=kernel_size, padding=kernel_size // 2, activation=act), + nn.Conv1d(dim, transition_dim, 1), + ) + + def forward(self, x, cond, time, returns=None, use_dropout: bool = True, force_dropout: bool = False): + """ + Overview: + compute diffusion unet forward + Arguments: + - x (:obj:'tensor'): noise trajectory + - cond (:obj:'tuple'): [ (time, state), ... ] state is init state of env, time = 0 + - time (:obj:'int'): timestep of diffusion step + - returns (:obj:'tensor'): condition returns of trajectory, returns is normal return + - use_dropout (:obj:'bool'): Whether use returns condition mask + - force_dropout (:obj:'bool'): Whether use returns condition + """ + if self.cale_energy: + x_inp = x + + # [batch, horizon, transition ] -> [batch, transition , horizon] + x = x.transpose(1, 2) + t = self.time_mlp(time) + + if self.returns_condition: + assert returns is not None + returns_embed = self.returns_mlp(returns) + if use_dropout: + mask = self.mask_dist.sample(sample_shape=(returns_embed.size(0), 1)).to(returns_embed.device) + returns_embed = mask * returns_embed + if force_dropout: + returns_embed = 0 * returns_embed + t = torch.cat([t, returns_embed], dim=-1) + + h = [] + + for resnet, resnet2, atten, downsample in self.downs: + x = resnet(x, t) + x = resnet2(x, t) + x = atten(x) + h.append(x) + x = downsample(x) + + x = self.mid_block1(x, t) + x = self.mid_atten(x) + x = self.mid_block2(x, t) + + for resnet, resnet2, atten, upsample in self.ups: + x = torch.cat((x, h.pop()), dim=1) + x = resnet(x, t) + x = resnet2(x, t) + x = atten(x) + x = upsample(x) + + x = self.final_conv(x) + # [batch, transition , horizon] -> [batch, horizon, transition ] + x = x.transpose(1, 2) + + if self.cale_energy: + # Energy function + energy = ((x - x_inp) ** 2).mean() + grad = torch.autograd.grad(outputs=energy, inputs=x_inp, create_graph=True) + return grad[0] + else: + return x + + def get_pred(self, x, cond, time, returns: bool = None, use_dropout: bool = True, force_dropout: bool = False): + """ + Overview: + compute diffusion unet forward + Arguments: + - x (:obj:'tensor'): noise trajectory + - cond (:obj:'tuple'): [ (time, state), ... ] state is init state of env, time = 0 + - time (:obj:'int'): timestep of diffusion step + - returns (:obj:'tensor'): condition returns of trajectory, returns is normal return + - use_dropout (:obj:'bool'): Whether use returns condition mask + - force_dropout (:obj:'bool'): Whether use returns condition + """ + # [batch, horizon, transition ] -> [batch, transition , horizon] + x = x.transpose(1, 2) + t = self.time_mlp(time) + + if self.returns_condition: + assert returns is not None + returns_embed = self.returns_mlp(returns) + if use_dropout: + mask = self.mask_dist.sample(sample_shape=(returns_embed.size(0), 1)).to(returns_embed.device) + returns_embed = mask * returns_embed + if force_dropout: + returns_embed = 0 * returns_embed + t = torch.cat([t, returns_embed], dim=-1) + + h = [] + + for resnet, resnet2, downsample in self.downs: + x = resnet(x, t) + x = resnet2(x, t) + h.append(x) + x = downsample(x) + + x = self.mid_block1(x, t) + x = self.mid_block2(x, t) + + for resnet, resnet2, upsample in self.ups: + x = torch.cat((x, h.pop()), dim=1) + x = resnet(x, t) + x = resnet2(x, t) + x = upsample(x) + + x = self.final_conv(x) + # [batch, transition , horizon] -> [batch, horizon, transition ] + x = x.transpose(1, 2) + return x + + +class TemporalValue(nn.Module): + """ + Overview: + temporal net for value function + Interfaces: + ``__init__``, ``forward`` + """ + + def __init__( + self, + horizon: int, + transition_dim: int, + dim: int = 32, + time_dim: int = None, + out_dim: int = 1, + kernel_size: int = 5, + dim_mults: SequenceType = [1, 2, 4, 8], + ) -> None: + """ + Overview: + Initialization of TemporalValue class + Arguments: + - horizon (:obj:'int'): horizon of trajectory + - transition_dim (:obj:'int'): dim of transition, it is obs_dim + action_dim + - dim (:obj:'int'): dim of layer + - time_dim (:obj:'int'): dim of time + - out_dim (:obj:'int'): dim of output + - kernel_size (:obj:'int'): kernel_size of conv1d + - dim_mults (:obj:'SequenceType'): mults of dim + """ + super().__init__() + dims = [transition_dim, *map(lambda m: dim * m, dim_mults)] + in_out = list(zip(dims[:-1], dims[1:])) + + time_dim = time_dim or dim + self.time_mlp = nn.Sequential( + SinusoidalPosEmb(dim), + nn.Linear(dim, dim * 4), + nn.Mish(), + nn.Linear(dim * 4, dim), + ) + self.blocks = nn.ModuleList([]) + + for ind, (dim_in, dim_out) in enumerate(in_out): + self.blocks.append( + nn.ModuleList( + [ + ResidualTemporalBlock(dim_in, dim_out, kernel_size=kernel_size, embed_dim=time_dim), + ResidualTemporalBlock(dim_out, dim_out, kernel_size=kernel_size, embed_dim=time_dim), + nn.Conv1d(dim_out, dim_out, 3, 2, 1) + ] + ) + ) + + horizon = horizon // 2 + + mid_dim = dims[-1] + mid_dim_2 = mid_dim // 2 + mid_dim_3 = mid_dim // 4 + + self.mid_block1 = ResidualTemporalBlock(mid_dim, mid_dim_2, kernel_size=kernel_size, embed_dim=time_dim) + self.mid_down1 = nn.Conv1d(mid_dim_2, mid_dim_2, 3, 2, 1) + + horizon = horizon // 2 + self.mid_block2 = ResidualTemporalBlock(mid_dim_2, mid_dim_3, kernel_size=kernel_size, embed_dim=time_dim) + self.mid_down2 = nn.Conv1d(mid_dim_3, mid_dim_3, 3, 2, 1) + horizon = horizon // 2 + + fc_dim = mid_dim_3 * max(horizon, 1) + self.final_block = nn.Sequential( + nn.Linear(fc_dim + time_dim, fc_dim // 2), + nn.Mish(), + nn.Linear(fc_dim // 2, out_dim), + ) + + def forward(self, x, cond, time, *args): + """ + Overview: + compute temporal value forward + Arguments: + - x (:obj:'tensor'): noise trajectory + - cond (:obj:'tuple'): [ (time, state), ... ] state is init state of env, time = 0 + - time (:obj:'int'): timestep of diffusion step + """ + # [batch, horizon, transition ] -> [batch, transition , horizon] + x = x.transpose(1, 2) + t = self.time_mlp(time) + for resnet, resnet2, downsample in self.blocks: + x = resnet(x, t) + x = resnet2(x, t) + x = downsample(x) + + x = self.mid_block1(x, t) + x = self.mid_down1(x) + + x = self.mid_block2(x, t) + x = self.mid_down2(x) + x = x.view(len(x), -1) + out = self.final_block(torch.cat([x, t], dim=-1)) + return out diff --git a/ding/torch_utils/network/dreamer.py b/ding/torch_utils/network/dreamer.py new file mode 100644 index 0000000000..b7ae67b57c --- /dev/null +++ b/ding/torch_utils/network/dreamer.py @@ -0,0 +1,937 @@ +import math +import numpy as np + +import torch +from torch import nn +import torch.nn.functional as F +from torch import distributions as torchd +from ding.torch_utils import MLP +from ding.rl_utils import symlog, inv_symlog + + +class Conv2dSame(torch.nn.Conv2d): + """ + Overview: + Conv2dSame Network for dreamerv3. + Interfaces: + ``__init__``, ``forward`` + """ + + def calc_same_pad(self, i, k, s, d): + """ + Overview: + Calculate the same padding size. + Arguments: + - i (:obj:`int`): Input size. + - k (:obj:`int`): Kernel size. + - s (:obj:`int`): Stride size. + - d (:obj:`int`): Dilation size. + """ + return max((math.ceil(i / s) - 1) * s + (k - 1) * d + 1 - i, 0) + + def forward(self, x): + """ + Overview: + compute the forward of Conv2dSame. + Arguments: + - x (:obj:`torch.Tensor`): Input tensor. + """ + ih, iw = x.size()[-2:] + pad_h = self.calc_same_pad(i=ih, k=self.kernel_size[0], s=self.stride[0], d=self.dilation[0]) + pad_w = self.calc_same_pad(i=iw, k=self.kernel_size[1], s=self.stride[1], d=self.dilation[1]) + + if pad_h > 0 or pad_w > 0: + x = F.pad(x, [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2]) + + ret = F.conv2d( + x, + self.weight, + self.bias, + self.stride, + self.padding, + self.dilation, + self.groups, + ) + return ret + + +class DreamerLayerNorm(nn.Module): + """ + Overview: + DreamerLayerNorm Network for dreamerv3. + Interfaces: + ``__init__``, ``forward`` + """ + + def __init__(self, ch, eps=1e-03): + """ + Overview: + Init the DreamerLayerNorm class. + Arguments: + - ch (:obj:`int`): Input channel. + - eps (:obj:`float`): Epsilon. + """ + + super(DreamerLayerNorm, self).__init__() + self.norm = torch.nn.LayerNorm(ch, eps=eps) + + def forward(self, x): + """ + Overview: + compute the forward of DreamerLayerNorm. + Arguments: + - x (:obj:`torch.Tensor`): Input tensor. + """ + + x = x.permute(0, 2, 3, 1) + x = self.norm(x) + x = x.permute(0, 3, 1, 2) + return x + + +class DenseHead(nn.Module): + """ + Overview: + DenseHead Network for value head, reward head, and discount head of dreamerv3. + Interfaces: + ``__init__``, ``forward`` + """ + + def __init__( + self, + inp_dim, + shape, # (255,) + layer_num, + units, # 512 + act='SiLU', + norm='LN', + dist='normal', + std=1.0, + outscale=1.0, + device='cpu', + ): + """ + Overview: + Init the DenseHead class. + Arguments: + - inp_dim (:obj:`int`): Input dimension. + - shape (:obj:`tuple`): Output shape. + - layer_num (:obj:`int`): Number of layers. + - units (:obj:`int`): Number of units. + - act (:obj:`str`): Activation function. + - norm (:obj:`str`): Normalization function. + - dist (:obj:`str`): Distribution function. + - std (:obj:`float`): Standard deviation. + - outscale (:obj:`float`): Output scale. + - device (:obj:`str`): Device. + """ + + super(DenseHead, self).__init__() + self._shape = (shape, ) if isinstance(shape, int) else shape + if len(self._shape) == 0: + self._shape = (1, ) + self._layer_num = layer_num + self._units = units + self._act = getattr(torch.nn, act)() + self._norm = norm + self._dist = dist + self._std = std + self._device = device + + self.mlp = MLP( + inp_dim, + self._units, + self._units, + self._layer_num, + layer_fn=nn.Linear, + activation=self._act, + norm_type=self._norm + ) + self.mlp.apply(weight_init) + + self.mean_layer = nn.Linear(self._units, np.prod(self._shape)) + self.mean_layer.apply(uniform_weight_init(outscale)) + + if self._std == "learned": + self.std_layer = nn.Linear(self._units, np.prod(self._shape)) + self.std_layer.apply(uniform_weight_init(outscale)) + + def forward(self, features): + """ + Overview: + compute the forward of DenseHead. + Arguments: + - features (:obj:`torch.Tensor`): Input tensor. + """ + + x = features + out = self.mlp(x) # (batch, time, _units=512) + mean = self.mean_layer(out) # (batch, time, 255) + if self._std == "learned": + std = self.std_layer(out) + else: + std = self._std + if self._dist == "normal": + return ContDist(torchd.independent.Independent(torchd.normal.Normal(mean, std), len(self._shape))) + elif self._dist == "huber": + return ContDist(torchd.independent.Independent(UnnormalizedHuber(mean, std, 1.0), len(self._shape))) + elif self._dist == "binary": + return Bernoulli(torchd.independent.Independent(torchd.bernoulli.Bernoulli(logits=mean), len(self._shape))) + elif self._dist == "twohot_symlog": + return TwoHotDistSymlog(logits=mean, low=-1., high=1., device=self._device) + raise NotImplementedError(self._dist) + + +class ActionHead(nn.Module): + """ + Overview: + ActionHead Network for action head of dreamerv3. + Interfaces: + ``__init__``, ``forward`` + """ + + def __init__( + self, + inp_dim, + size, + layers, + units, + act=nn.ELU, + norm=nn.LayerNorm, + dist="trunc_normal", + init_std=0.0, + min_std=0.1, + max_std=1.0, + temp=0.1, + outscale=1.0, + unimix_ratio=0.01, + ): + """ + Overview: + Initialize the ActionHead class. + Arguments: + - inp_dim (:obj:`int`): Input dimension. + - size (:obj:`int`): Output size. + - layers (:obj:`int`): Number of layers. + - units (:obj:`int`): Number of units. + - act (:obj:`str`): Activation function. + - norm (:obj:`str`): Normalization function. + - dist (:obj:`str`): Distribution function. + - init_std (:obj:`float`): Initial standard deviation. + - min_std (:obj:`float`): Minimum standard deviation. + - max_std (:obj:`float`): Maximum standard deviation. + - temp (:obj:`float`): Temperature. + - outscale (:obj:`float`): Output scale. + - unimix_ratio (:obj:`float`): Unimix ratio. + """ + super(ActionHead, self).__init__() + self._size = size + self._layers = layers + self._units = units + self._dist = dist + self._act = getattr(torch.nn, act) + self._norm = getattr(torch.nn, norm) + self._min_std = min_std + self._max_std = max_std + self._init_std = init_std + self._unimix_ratio = unimix_ratio + self._temp = temp() if callable(temp) else temp + + pre_layers = [] + for index in range(self._layers): + pre_layers.append(nn.Linear(inp_dim, self._units, bias=False)) + pre_layers.append(self._norm(self._units, eps=1e-03)) + pre_layers.append(self._act()) + if index == 0: + inp_dim = self._units + self._pre_layers = nn.Sequential(*pre_layers) + self._pre_layers.apply(weight_init) + + if self._dist in ["tanh_normal", "tanh_normal_5", "normal", "trunc_normal"]: + self._dist_layer = nn.Linear(self._units, 2 * self._size) + self._dist_layer.apply(uniform_weight_init(outscale)) + + elif self._dist in ["normal_1", "onehot", "onehot_gumbel"]: + self._dist_layer = nn.Linear(self._units, self._size) + self._dist_layer.apply(uniform_weight_init(outscale)) + + def forward(self, features): + """ + Overview: + compute the forward of ActionHead. + Arguments: + - features (:obj:`torch.Tensor`): Input tensor. + """ + + x = features + x = self._pre_layers(x) + if self._dist == "tanh_normal": + x = self._dist_layer(x) + mean, std = torch.split(x, 2, -1) + mean = torch.tanh(mean) + std = F.softplus(std + self._init_std) + self._min_std + dist = torchd.normal.Normal(mean, std) + dist = torchd.transformed_distribution.TransformedDistribution(dist, TanhBijector()) + dist = torchd.independent.Independent(dist, 1) + dist = SampleDist(dist) + elif self._dist == "tanh_normal_5": + x = self._dist_layer(x) + mean, std = torch.split(x, 2, -1) + mean = 5 * torch.tanh(mean / 5) + std = F.softplus(std + 5) + 5 + dist = torchd.normal.Normal(mean, std) + dist = torchd.transformed_distribution.TransformedDistribution(dist, TanhBijector()) + dist = torchd.independent.Independent(dist, 1) + dist = SampleDist(dist) + elif self._dist == "normal": + x = self._dist_layer(x) + mean, std = torch.split(x, [self._size] * 2, -1) + std = (self._max_std - self._min_std) * torch.sigmoid(std + 2.0) + self._min_std + dist = torchd.normal.Normal(torch.tanh(mean), std) + dist = ContDist(torchd.independent.Independent(dist, 1)) + elif self._dist == "normal_1": + x = self._dist_layer(x) + dist = torchd.normal.Normal(mean, 1) + dist = ContDist(torchd.independent.Independent(dist, 1)) + elif self._dist == "trunc_normal": + x = self._dist_layer(x) + mean, std = torch.split(x, [self._size] * 2, -1) + mean = torch.tanh(mean) + std = 2 * torch.sigmoid(std / 2) + self._min_std + dist = SafeTruncatedNormal(mean, std, -1, 1) + dist = ContDist(torchd.independent.Independent(dist, 1)) + elif self._dist == "onehot": + x = self._dist_layer(x) + dist = OneHotDist(x, unimix_ratio=self._unimix_ratio) + elif self._dist == "onehot_gumble": + x = self._dist_layer(x) + temp = self._temp + dist = ContDist(torchd.gumbel.Gumbel(x, 1 / temp)) + else: + raise NotImplementedError(self._dist) + return dist + + +class SampleDist: + """ + Overview: + A kind of sample Dist for ActionHead of dreamerv3. + Interfaces: + ``__init__``, ``mean``, ``mode``, ``entropy`` + """ + + def __init__(self, dist, samples=100): + """ + Overview: + Initialize the SampleDist class. + Arguments: + - dist (:obj:`torch.Tensor`): Distribution. + - samples (:obj:`int`): Number of samples. + """ + + self._dist = dist + self._samples = samples + + def mean(self): + """ + Overview: + Calculate the mean of the distribution. + """ + + samples = self._dist.sample(self._samples) + return torch.mean(samples, 0) + + def mode(self): + """ + Overview: + Calculate the mode of the distribution. + """ + + sample = self._dist.sample(self._samples) + logprob = self._dist.log_prob(sample) + return sample[torch.argmax(logprob)][0] + + def entropy(self): + """ + Overview: + Calculate the entropy of the distribution. + """ + + sample = self._dist.sample(self._samples) + logprob = self.log_prob(sample) + return -torch.mean(logprob, 0) + + +class OneHotDist(torchd.one_hot_categorical.OneHotCategorical): + """ + Overview: + A kind of onehot Dist for dreamerv3. + Interfaces: + ``__init__``, ``mode``, ``sample`` + """ + + def __init__(self, logits=None, probs=None, unimix_ratio=0.0): + """ + Overview: + Initialize the OneHotDist class. + Arguments: + - logits (:obj:`torch.Tensor`): Logits. + - probs (:obj:`torch.Tensor`): Probabilities. + - unimix_ratio (:obj:`float`): Unimix ratio. + """ + + if logits is not None and unimix_ratio > 0.0: + probs = F.softmax(logits, dim=-1) + probs = probs * (1.0 - unimix_ratio) + unimix_ratio / probs.shape[-1] + logits = torch.log(probs) + super().__init__(logits=logits, probs=None) + else: + super().__init__(logits=logits, probs=probs) + + def mode(self): + """ + Overview: + Calculate the mode of the distribution. + """ + + _mode = F.one_hot(torch.argmax(super().logits, axis=-1), super().logits.shape[-1]) + return _mode.detach() + super().logits - super().logits.detach() + + def sample(self, sample_shape=(), seed=None): + """ + Overview: + Sample from the distribution. + Arguments: + - sample_shape (:obj:`tuple`): Sample shape. + - seed (:obj:`int`): Seed. + """ + + if seed is not None: + raise ValueError('need to check') + sample = super().sample(sample_shape) + probs = super().probs + while len(probs.shape) < len(sample.shape): + probs = probs[None] + sample += probs - probs.detach() + return sample + + +class TwoHotDistSymlog: + """ + Overview: + A kind of twohotsymlog Dist for dreamerv3. + Interfaces: + ``__init__``, ``mode``, ``mean``, ``log_prob``, ``log_prob_target`` + """ + + def __init__(self, logits=None, low=-20.0, high=20.0, device='cpu'): + """ + Overview: + Initialize the TwoHotDistSymlog class. + Arguments: + - logits (:obj:`torch.Tensor`): Logits. + - low (:obj:`float`): Low. + - high (:obj:`float`): High. + - device (:obj:`str`): Device. + """ + + self.logits = logits + self.probs = torch.softmax(logits, -1) + self.buckets = torch.linspace(low, high, steps=255).to(device) + self.width = (self.buckets[-1] - self.buckets[0]) / 255 + + def mean(self): + """ + Overview: + Calculate the mean of the distribution. + """ + + _mean = self.probs * self.buckets + return inv_symlog(torch.sum(_mean, dim=-1, keepdim=True)) + + def mode(self): + """ + Overview: + Calculate the mode of the distribution. + """ + + _mode = self.probs * self.buckets + return inv_symlog(torch.sum(_mode, dim=-1, keepdim=True)) + + # Inside OneHotCategorical, log_prob is calculated using only max element in targets + def log_prob(self, x): + """ + Overview: + Calculate the log probability of the distribution. + Arguments: + - x (:obj:`torch.Tensor`): Input tensor. + """ + + x = symlog(x) + # x(time, batch, 1) + below = torch.sum((self.buckets <= x[..., None]).to(torch.int32), dim=-1) - 1 + above = len(self.buckets) - torch.sum((self.buckets > x[..., None]).to(torch.int32), dim=-1) + below = torch.clip(below, 0, len(self.buckets) - 1) + above = torch.clip(above, 0, len(self.buckets) - 1) + equal = (below == above) + + dist_to_below = torch.where(equal, torch.tensor(1).to(x), torch.abs(self.buckets[below] - x)) + dist_to_above = torch.where(equal, torch.tensor(1).to(x), torch.abs(self.buckets[above] - x)) + total = dist_to_below + dist_to_above + weight_below = dist_to_above / total + weight_above = dist_to_below / total + target = ( + F.one_hot(below, num_classes=len(self.buckets)) * weight_below[..., None] + + F.one_hot(above, num_classes=len(self.buckets)) * weight_above[..., None] + ) + log_pred = self.logits - torch.logsumexp(self.logits, -1, keepdim=True) + target = target.squeeze(-2) + + return (target * log_pred).sum(-1) + + def log_prob_target(self, target): + """ + Overview: + Calculate the log probability of the target. + Arguments: + - target (:obj:`torch.Tensor`): Target tensor. + """ + + log_pred = super().logits - torch.logsumexp(super().logits, -1, keepdim=True) + return (target * log_pred).sum(-1) + + +class SymlogDist: + """ + Overview: + A kind of Symlog Dist for dreamerv3. + Interfaces: + ``__init__``, ``entropy``, ``mode``, ``mean``, ``log_prob`` + """ + + def __init__(self, mode, dist='mse', aggregation='sum', tol=1e-8, dim_to_reduce=[-1, -2, -3]): + """ + Overview: + Initialize the SymlogDist class. + Arguments: + - mode (:obj:`torch.Tensor`): Mode. + - dist (:obj:`str`): Distribution function. + - aggregation (:obj:`str`): Aggregation function. + - tol (:obj:`float`): Tolerance. + - dim_to_reduce (:obj:`list`): Dimension to reduce. + """ + self._mode = mode + self._dist = dist + self._aggregation = aggregation + self._tol = tol + self._dim_to_reduce = dim_to_reduce + + def mode(self): + """ + Overview: + Calculate the mode of the distribution. + """ + + return inv_symlog(self._mode) + + def mean(self): + """ + Overview: + Calculate the mean of the distribution. + """ + + return inv_symlog(self._mode) + + def log_prob(self, value): + """ + Overview: + Calculate the log probability of the distribution. + Arguments: + - value (:obj:`torch.Tensor`): Input tensor. + """ + + assert self._mode.shape == value.shape + if self._dist == 'mse': + distance = (self._mode - symlog(value)) ** 2.0 + distance = torch.where(distance < self._tol, 0, distance) + elif self._dist == 'abs': + distance = torch.abs(self._mode - symlog(value)) + distance = torch.where(distance < self._tol, 0, distance) + else: + raise NotImplementedError(self._dist) + if self._aggregation == 'mean': + loss = distance.mean(self._dim_to_reduce) + elif self._aggregation == 'sum': + loss = distance.sum(self._dim_to_reduce) + else: + raise NotImplementedError(self._aggregation) + return -loss + + +class ContDist: + """ + Overview: + A kind of ordinary Dist for dreamerv3. + Interfaces: + ``__init__``, ``entropy``, ``mode``, ``sample``, ``log_prob`` + """ + + def __init__(self, dist=None): + """ + Overview: + Initialize the ContDist class. + Arguments: + - dist (:obj:`torch.Tensor`): Distribution. + """ + + super().__init__() + self._dist = dist + self.mean = dist.mean + + def __getattr__(self, name): + """ + Overview: + Get attribute. + Arguments: + - name (:obj:`str`): Attribute name. + """ + + return getattr(self._dist, name) + + def entropy(self): + """ + Overview: + Calculate the entropy of the distribution. + """ + + return self._dist.entropy() + + def mode(self): + """ + Overview: + Calculate the mode of the distribution. + """ + + return self._dist.mean + + def sample(self, sample_shape=()): + """ + Overview: + Sample from the distribution. + Arguments: + - sample_shape (:obj:`tuple`): Sample shape. + """ + + return self._dist.rsample(sample_shape) + + def log_prob(self, x): + return self._dist.log_prob(x) + + +class Bernoulli: + """ + Overview: + A kind of Bernoulli Dist for dreamerv3. + Interfaces: + ``__init__``, ``entropy``, ``mode``, ``sample``, ``log_prob`` + """ + + def __init__(self, dist=None): + """ + Overview: + Initialize the Bernoulli distribution. + Arguments: + - dist (:obj:`torch.Tensor`): Distribution. + """ + + super().__init__() + self._dist = dist + self.mean = dist.mean + + def __getattr__(self, name): + """ + Overview: + Get attribute. + Arguments: + - name (:obj:`str`): Attribute name. + """ + + return getattr(self._dist, name) + + def entropy(self): + """ + Overview: + Calculate the entropy of the distribution. + """ + return self._dist.entropy() + + def mode(self): + """ + Overview: + Calculate the mode of the distribution. + """ + + _mode = torch.round(self._dist.mean) + return _mode.detach() + self._dist.mean - self._dist.mean.detach() + + def sample(self, sample_shape=()): + """ + Overview: + Sample from the distribution. + Arguments: + - sample_shape (:obj:`tuple`): Sample shape. + """ + + return self._dist.rsample(sample_shape) + + def log_prob(self, x): + """ + Overview: + Calculate the log probability of the distribution. + Arguments: + - x (:obj:`torch.Tensor`): Input tensor. + """ + + _logits = self._dist.base_dist.logits + log_probs0 = -F.softplus(_logits) + log_probs1 = -F.softplus(-_logits) + + return log_probs0 * (1 - x) + log_probs1 * x + + +class UnnormalizedHuber(torchd.normal.Normal): + """ + Overview: + A kind of UnnormalizedHuber Dist for dreamerv3. + Interfaces: + ``__init__``, ``mode``, ``log_prob`` + """ + + def __init__(self, loc, scale, threshold=1, **kwargs): + """ + Overview: + Initialize the UnnormalizedHuber class. + Arguments: + - loc (:obj:`torch.Tensor`): Location. + - scale (:obj:`torch.Tensor`): Scale. + - threshold (:obj:`float`): Threshold. + """ + super().__init__(loc, scale, **kwargs) + self._threshold = threshold + + def log_prob(self, event): + """ + Overview: + Calculate the log probability of the distribution. + Arguments: + - event (:obj:`torch.Tensor`): Event. + """ + + return -(torch.sqrt((event - self.mean) ** 2 + self._threshold ** 2) - self._threshold) + + def mode(self): + """ + Overview: + Calculate the mode of the distribution. + """ + + return self.mean + + +class SafeTruncatedNormal(torchd.normal.Normal): + """ + Overview: + A kind of SafeTruncatedNormal Dist for dreamerv3. + Interfaces: + ``__init__``, ``sample`` + """ + + def __init__(self, loc, scale, low, high, clip=1e-6, mult=1): + """ + Overview: + Initialize the SafeTruncatedNormal class. + Arguments: + - loc (:obj:`torch.Tensor`): Location. + - scale (:obj:`torch.Tensor`): Scale. + - low (:obj:`float`): Low. + - high (:obj:`float`): High. + - clip (:obj:`float`): Clip. + - mult (:obj:`float`): Mult. + """ + + super().__init__(loc, scale) + self._low = low + self._high = high + self._clip = clip + self._mult = mult + + def sample(self, sample_shape): + """ + Overview: + Sample from the distribution. + Arguments: + - sample_shape (:obj:`tuple`): Sample shape. + """ + + event = super().sample(sample_shape) + if self._clip: + clipped = torch.clip(event, self._low + self._clip, self._high - self._clip) + event = event - event.detach() + clipped.detach() + if self._mult: + event *= self._mult + return event + + +class TanhBijector(torchd.Transform): + """ + Overview: + A kind of TanhBijector Dist for dreamerv3. + Interfaces: + ``__init__``, ``_forward``, ``_inverse``, ``_forward_log_det_jacobian`` + """ + + def __init__(self, validate_args=False, name='tanh'): + """ + Overview: + Initialize the TanhBijector class. + Arguments: + - validate_args (:obj:`bool`): Validate arguments. + - name (:obj:`str`): Name. + """ + + super().__init__() + + def _forward(self, x): + """ + Overview: + Calculate the forward of the distribution. + Arguments: + - x (:obj:`torch.Tensor`): Input tensor. + """ + + return torch.tanh(x) + + def _inverse(self, y): + """ + Overview: + Calculate the inverse of the distribution. + Arguments: + - y (:obj:`torch.Tensor`): Input tensor. + """ + + y = torch.where((torch.abs(y) <= 1.), torch.clamp(y, -0.99999997, 0.99999997), y) + y = torch.atanh(y) + return y + + def _forward_log_det_jacobian(self, x): + """ + Overview: + Calculate the forward log det jacobian of the distribution. + Arguments: + - x (:obj:`torch.Tensor`): Input tensor. + """ + + log2 = torch.math.log(2.0) + return 2.0 * (log2 - x - torch.softplus(-2.0 * x)) + + +def static_scan(fn, inputs, start): + """ + Overview: + Static scan function. + Arguments: + - fn (:obj:`function`): Function. + - inputs (:obj:`tuple`): Inputs. + - start (:obj:`torch.Tensor`): Start tensor. + """ + + last = start # {logit, stoch, deter:[batch_size, self._deter]} + indices = range(inputs[0].shape[0]) + flag = True + for index in indices: + inp = lambda x: (_input[x] for _input in inputs) # inputs:(action:(time, batch, 6), embed:(time, batch, 4096)) + last = fn(last, *inp(index)) # post, prior + if flag: + if isinstance(last, dict): + outputs = {key: value.clone().unsqueeze(0) for key, value in last.items()} + else: + outputs = [] + for _last in last: + if isinstance(_last, dict): + outputs.append({key: value.clone().unsqueeze(0) for key, value in _last.items()}) + else: + outputs.append(_last.clone().unsqueeze(0)) + flag = False + else: + if isinstance(last, dict): + for key in last.keys(): + outputs[key] = torch.cat([outputs[key], last[key].unsqueeze(0)], dim=0) + else: + for j in range(len(outputs)): + if isinstance(last[j], dict): + for key in last[j].keys(): + outputs[j][key] = torch.cat([outputs[j][key], last[j][key].unsqueeze(0)], dim=0) + else: + outputs[j] = torch.cat([outputs[j], last[j].unsqueeze(0)], dim=0) + if isinstance(last, dict): + outputs = [outputs] + return outputs + + +def weight_init(m): + """ + Overview: + weight_init for Linear, Conv2d, ConvTranspose2d, and LayerNorm. + Arguments: + - m (:obj:`torch.nn`): Module. + """ + + if isinstance(m, nn.Linear): + in_num = m.in_features + out_num = m.out_features + denoms = (in_num + out_num) / 2.0 + scale = 1.0 / denoms + std = np.sqrt(scale) / 0.87962566103423978 + nn.init.trunc_normal_(m.weight.data, mean=0.0, std=std, a=-2.0, b=2.0) + if hasattr(m.bias, 'data'): + m.bias.data.fill_(0.0) + elif isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d): + space = m.kernel_size[0] * m.kernel_size[1] + in_num = space * m.in_channels + out_num = space * m.out_channels + denoms = (in_num + out_num) / 2.0 + scale = 1.0 / denoms + std = np.sqrt(scale) / 0.87962566103423978 + nn.init.trunc_normal_(m.weight.data, mean=0.0, std=std, a=-2.0, b=2.0) + if hasattr(m.bias, 'data'): + m.bias.data.fill_(0.0) + elif isinstance(m, nn.LayerNorm): + m.weight.data.fill_(1.0) + if hasattr(m.bias, 'data'): + m.bias.data.fill_(0.0) + + +def uniform_weight_init(given_scale): + """ + Overview: + weight_init for Linear and LayerNorm. + Arguments: + - given_scale (:obj:`float`): Given scale. + """ + + def f(m): + if isinstance(m, nn.Linear): + in_num = m.in_features + out_num = m.out_features + denoms = (in_num + out_num) / 2.0 + scale = given_scale / denoms + limit = np.sqrt(3 * scale) + nn.init.uniform_(m.weight.data, a=-limit, b=limit) + if hasattr(m.bias, 'data'): + m.bias.data.fill_(0.0) + elif isinstance(m, nn.LayerNorm): + m.weight.data.fill_(1.0) + if hasattr(m.bias, 'data'): + m.bias.data.fill_(0.0) + + return f diff --git a/ding/torch_utils/network/gtrxl.py b/ding/torch_utils/network/gtrxl.py index 6c9f4d2a36..3672fc977d 100644 --- a/ding/torch_utils/network/gtrxl.py +++ b/ding/torch_utils/network/gtrxl.py @@ -1,3 +1,8 @@ +""" +Overview: + This file implements the core modules of GTrXL Transformer as described in + "Stabilizing Transformer for Reinforcement Learning" (https://arxiv.org/abs/1910.06764). +""" from typing import Optional, Dict, List import warnings import numpy as np @@ -9,31 +14,40 @@ class PositionalEmbedding(nn.Module): """ Overview: - Positional Embedding used in vanilla Transformer + The PositionalEmbedding module implements the positional embedding used in the vanilla Transformer model. + Interfaces: + ``__init__``, ``forward`` + .. note:: - Adapted from https://github.com/kimiyoung/transformer-xl/blob/master/pytorch/mem_transformer.py + This implementation is adapted from https://github.com/kimiyoung/transformer-xl/blob/ \ + master/pytorch/mem_transformer.py """ def __init__(self, embedding_dim: int): """ + Overview: + Initialize the PositionalEmbedding module. Arguments: - - embedding_dim: (:obj:`int`): dimension of embedding + - embedding_dim: (:obj:`int`): The dimensionality of the embeddings. """ + super(PositionalEmbedding, self).__init__() self.embedding_dim = embedding_dim inv_freq = 1 / (10000 ** (torch.arange(0.0, embedding_dim, 2.0) / embedding_dim)) # (embedding_dim / 2) self.register_buffer('inv_freq', inv_freq) - def forward(self, pos_seq: torch.Tensor): + def forward(self, pos_seq: torch.Tensor) -> torch.Tensor: """ Overview: - Compute positional embedding + Compute positional embedding given a sequence of positions. Arguments: - - pos_seq: (:obj:`torch.Tensor`): positional sequence, - usually a 1D integer sequence as [seq_len-1, seq_len-2, ..., 1, 0], + - pos_seq (:obj:`torch.Tensor`): The positional sequence, \ + typically a 1D tensor of integers in the form of [seq_len-1, seq_len-2, ..., 1, 0], Returns: - - pos_embedding: (:obj:`torch.Tensor`): positional embedding. Shape (seq_len, 1, embedding_dim) + - pos_embedding (:obj:`torch.Tensor`): The computed positional embeddings. \ + The shape of the tensor is (seq_len, 1, embedding_dim). """ + sinusoid_inp = torch.outer(pos_seq, self.inv_freq) # For position embedding, the order of sin/cos is negligible. # This is because tokens are consumed by the matrix multiplication which is permutation-invariant. @@ -44,17 +58,22 @@ def forward(self, pos_seq: torch.Tensor): class GRUGatingUnit(torch.nn.Module): """ Overview: - GRU Gating Unit used in GTrXL. + The GRUGatingUnit module implements the GRU gating mechanism used in the GTrXL model. + Interfaces: + ``__init__``, ``forward`` """ def __init__(self, input_dim: int, bg: float = 2.): """ + Overview: + Initialize the GRUGatingUnit module. Arguments: - - input_dim: (:obj:`int`): dimension of input. - - bg (:obj:`bg`): gate bias. By setting bg > 0 we can explicitly initialize the gating mechanism to - be close to the identity map. This can greatly improve the learning speed and stability since it - initializes the agent close to a Markovian policy (ignore attention at the beginning). + - input_dim (:obj:`int`): The dimensionality of the input. + - bg (:obj:`bg`): The gate bias. By setting bg > 0 we can explicitly initialize the gating mechanism to \ + be close to the identity map. This can greatly improve the learning speed and stability since it \ + initializes the agent close to a Markovian policy (ignore attention at the beginning). """ + super(GRUGatingUnit, self).__init__() self.Wr = torch.nn.Linear(input_dim, input_dim, bias=False) self.Ur = torch.nn.Linear(input_dim, input_dim, bias=False) @@ -69,14 +88,16 @@ def __init__(self, input_dim: int, bg: float = 2.): def forward(self, x: torch.Tensor, y: torch.Tensor): """ Overview: - Compute output value with gating mechanism + Compute the output value using the GRU gating mechanism. Arguments: - - x: (:obj:`torch.Tensor`): first input. - - y: (:obj:`torch.Tensor`): second input. - x and y have same shape and last shape is input_dim. + - x: (:obj:`torch.Tensor`): The first input tensor. + - y: (:obj:`torch.Tensor`): The second input tensor. \ + x and y should have the same shape and their last dimension should match the input_dim. Returns: - - g: (:obj:`torch.Tensor`): output of GRU. Same shape of x and y. + - g: (:obj:`torch.Tensor`): The output of the GRU gating mechanism. \ + The shape of g matches the shapes of x and y. """ + r = self.sigmoid(self.Wr(y) + self.Ur(x)) z = self.sigmoid(self.Wz(y) + self.Uz(x) - self.bg) h = self.tanh(self.Wg(y) + self.Ug(torch.mul(r, x))) # element wise multiplication @@ -87,9 +108,12 @@ def forward(self, x: torch.Tensor, y: torch.Tensor): class Memory: """ Overview: - Stores the context used to add memory to Transformer. + A class that stores the context used to add memory to Transformer. + Interfaces: + ``__init__``, ``init``, ``update``, ``get``, ``to`` + .. note:: - For details refer to Transformer-XL: https://arxiv.org/abs/1901.02860 + For details, refer to Transformer-XL: https://arxiv.org/abs/1901.02860 """ def __init__( @@ -101,12 +125,17 @@ def __init__( memory: Optional[torch.Tensor] = None ) -> None: """ + Overview: + Initialize the Memory module. Arguments: - - memory_len (:obj:`int`): dimension of memory (how many past observations to use as memory) - - batch_size (:obj:`int`): dimension of each batch - - embedding_dim (:obj:`int`): dimension of embedding (dimension of a single observation after embedding) - - layer_num (:obj:`int`): number of transformer layers + - memory_len (:obj:`int`): The dimension of memory, i.e., how many past observations to use as memory. + - batch_size (:obj:`int`): The dimension of each batch. + - embedding_dim (:obj:`int`): The dimension of embedding, which is the dimension of a single observation \ + after embedding. + - layer_num (:obj:`int`): The number of transformer layers. + - memory (:obj:`Optional[torch.Tensor]`): The initial memory. Default is None. """ + super(Memory, self).__init__() self.embedding_dim = embedding_dim self.bs = batch_size @@ -118,12 +147,13 @@ def __init__( def init(self, memory: Optional[torch.Tensor] = None): """ Overview: - Init memory with an input list of tensors or create it automatically given its dimensions. + Initialize memory with an input list of tensors or create it automatically given its dimensions. Arguments: - - memory: (:obj:`Optional[torch.Tensor]`): memory input. - Shape is (layer_num, memory_len, bs, embedding_dim). - memory_len is length of memory, bs is batch size and embedding_dim is the dimension of embedding. + - memory (:obj:`Optional[torch.Tensor]`): Input memory tensor with shape \ + (layer_num, memory_len, bs, embedding_dim). Its shape is (layer_num, memory_len, bs, embedding_dim), \ + where memory_len is length of memory, bs is batch size and embedding_dim is the dimension of embedding. """ + if memory is not None: self.memory = memory layer_num_plus1, self.memory_len, self.bs, self.embedding_dim = memory.shape @@ -137,20 +167,20 @@ def update(self, hidden_state: List[torch.Tensor]): """ Overview: Update the memory given a sequence of hidden states. - Example for single layer: + Example for single layer: (memory_len=3, hidden_size_len=2, bs=3) - memory_len=3, hidden_size_len=2, bs=3 - - m00 m01 m02 h00 h01 h02 m20 m21 m22 - m = m10 m11 m12 h = h10 h11 h12 => new_m = h00 h01 h02 - m20 m21 m22 h10 h11 h12 + m00 m01 m02 h00 h01 h02 m20 m21 m22 + m = m10 m11 m12 h = h10 h11 h12 => new_m = h00 h01 h02 + m20 m21 m22 h10 h11 h12 Arguments: - - hidden_state: (:obj:`List[torch.Tensor]`): hidden states to update the memory. - Shape is (cur_seq, bs, embedding_dim) for each layer. cur_seq is length of sequence. + - hidden_state: (:obj:`List[torch.Tensor]`): The hidden states to update the memory. \ + Each tensor in the list has shape (cur_seq, bs, embedding_dim), where cur_seq \ + is the length of the sequence. Returns: - - memory: (:obj:`Optional[torch.Tensor]`): output memory. - Shape is (layer_num, memory_len, bs, embedding_dim). + - memory: (:obj:`Optional[torch.Tensor]`): The updated memory, with shape \ + (layer_num, memory_len, bs, embedding_dim). """ + if self.memory is None or hidden_state is None: raise ValueError('Failed to update memory! Memory would be None') # TODO add support of no memory sequence_len = hidden_state[0].shape[0] @@ -170,32 +200,44 @@ def update(self, hidden_state: List[torch.Tensor]): def get(self): """ Overview: - Memory getter method. + Get the current memory. Returns: - - memory: (:obj:`Optional[torch.Tensor]`): output memory. - Shape is (layer_num, memory_len, bs, embedding_dim). + - memory: (:obj:`Optional[torch.Tensor]`): The current memory, \ + with shape (layer_num, memory_len, bs, embedding_dim). """ + return self.memory def to(self, device: str = 'cpu'): + """ + Overview: + Move the current memory to the specified device. + Arguments: + device (:obj:`str`): The device to move the memory to. Default is 'cpu'. + """ + self.memory = self.memory.to(device) class AttentionXL(torch.nn.Module): """ Overview: - Attention of TransformerXL. + An implementation of the Attention mechanism used in the TransformerXL model. + Interfaces: + ``__init__``, ``forward`` """ def __init__(self, input_dim: int, head_dim: int, head_num: int, dropout: nn.Module) -> None: - """Overview: - Init AttentionXL. + """ + Overview: + Initialize the AttentionXL module. Arguments: - - input_dim (:obj:`int`): dimension of input - - head_dim (:obj:`int`): dimension of each head - - head_num (:obj:`int`): number of heads for multihead attention - - dropout (:obj:`nn.Module`): dropout function + - input_dim (:obj:`int`): The dimensionality of the input features. + - head_dim (:obj:`int`): The dimensionality of each attention head. + - head_num (:obj:`int`): The number of attention heads. + - dropout (:obj:`nn.Module`): The dropout layer to use """ + super(AttentionXL, self).__init__() self.head_num = head_num self.head_dim = head_dim @@ -206,29 +248,31 @@ def __init__(self, input_dim: int, head_dim: int, head_num: int, dropout: nn.Mod self.project_pos = fc_block(input_dim, head_dim * head_num) # project the positional embedding self.scale = 1 / (head_dim ** 0.5) # for scaled dot product attention - def _rel_shift(self, x: torch.Tensor, zero_upper: bool = False): + def _rel_shift(self, x: torch.Tensor, zero_upper: bool = False) -> torch.Tensor: """ Overview: - Relatively shift the attention score matrix. - Example: - a00 a01 a02 0 a00 a01 a02 0 a00 a01 a02 0 a10 a02 0 0 - a10 a11 a12 => 0 a10 a11 a12 => a02 0 a10 => a11 a12 0 => a11 a12 0 - a20 a21 a22 0 a20 a21 a22 a11 a12 0 a20 a21 a22 a20 a21 a22 - a20 a21 a22 - 1) Append one "column" of zeros to the left - 2) Reshape the matrix from [3 x 4] into [4 x 3] - 3) Remove the first "row" - 4) Mask out the upper triangle (optional) + Perform a relative shift operation on the attention score matrix. + Example: + a00 a01 a02 0 a00 a01 a02 0 a00 a01 a02 0 a10 a02 0 0 + a10 a11 a12 => 0 a10 a11 a12 => a02 0 a10 => a11 a12 0 => a11 a12 0 + a20 a21 a22 0 a20 a21 a22 a11 a12 0 a20 a21 a22 a20 a21 a22 + a20 a21 a22 + 1) Append one "column" of zeros to the left + 2) Reshape the matrix from [3 x 4] into [4 x 3] + 3) Remove the first "row" + 4) Mask out the upper triangle (optional) + .. note:: - See the following material for better understanding: - https://github.com/kimiyoung/transformer-xl/issues/8 - https://arxiv.org/pdf/1901.02860.pdf (Appendix B) + See the following material for better understanding: https://github.com/kimiyoung/transformer-xl/issues/8 \ + https://arxiv.org/pdf/1901.02860.pdf (Appendix B) Arguments: - - x (:obj:`torch.Tensor`): input tensor of shape (cur_seq, full_seq, bs, head_num). - - zero_upper (:obj:`bool`): if True set the upper-right triangle to zero. + - x (:obj:`torch.Tensor`): The input tensor with shape (cur_seq, full_seq, bs, head_num). + - zero_upper (:obj:`bool`): If True, the upper-right triangle of the matrix is set to zero. Returns: - - x (:obj:`torch.Tensor`): input after relative shift. Shape (cur_seq, full_seq, bs, head_num). + - x (:obj:`torch.Tensor`): The input tensor after the relative shift operation, \ + with shape (cur_seq, full_seq, bs, head_num). """ + x_padded = F.pad(x, [1, 0]) # step 1 x_padded = x_padded.view(x.size(0), x.size(1), x.size(3) + 1, x.size(2)) # step 2 x = x_padded[:, :, 1:].view_as(x) # step 3 @@ -246,19 +290,22 @@ def forward( v: torch.nn.Parameter, mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: - """Overview: - Compute AttentionXL. + """ + Overview: + Compute the forward pass for the AttentionXL module. Arguments: - - inputs (:obj:`torch.Tensor`): attention input of shape (cur_seq, bs, input_dim) - - pos_embedding (:obj:`torch.Tensor`): positional embedding of shape (full_seq, 1, full_seq) - - full_input (:obj:`torch.Tensor`): memory + input concatenation of shape (full_seq, bs, input_dim) - - u (:obj:`torch.nn.Parameter`): content parameter of shape (head_num, head_dim) - - v (:obj:`torch.nn.Parameter`): position parameter of shape (head_num, head_dim) - - mask (:obj:`Optional[torch.Tensor]`): attention mask of shape (cur_seq, full_seq, 1) - full_seq = prev_seq + cur_seq + - inputs (:obj:`torch.Tensor`): The attention input with shape (cur_seq, bs, input_dim). + - pos_embedding (:obj:`torch.Tensor`): The positional embedding with shape (full_seq, 1, full_seq). + - full_input (:obj:`torch.Tensor`): The concatenated memory and input tensor with shape \ + (full_seq, bs, input_dim). + - u (:obj:`torch.nn.Parameter`): The content parameter with shape (head_num, head_dim). + - v (:obj:`torch.nn.Parameter`): The position parameter with shape (head_num, head_dim). + - mask (:obj:`Optional[torch.Tensor]`): The attention mask with shape (cur_seq, full_seq, 1). \ + If None, no masking is applied. Returns: - - output (:obj:`torch.Tensor`): attention output of shape (cur_seq, bs, input_dim) + - output (:obj:`torch.Tensor`): The output of the attention mechanism with shape (cur_seq, bs, input_dim). """ + bs, cur_seq, full_seq = inputs.shape[1], inputs.shape[0], full_input.shape[0] prev_seq = full_seq - cur_seq @@ -306,7 +353,9 @@ def forward( class GatedTransformerXLLayer(torch.nn.Module): """ Overview: - Attention layer of GTrXL + This class implements the attention layer of GTrXL (Gated Transformer-XL). + Interfaces: + ``__init__``, ``forward`` """ def __init__( @@ -322,17 +371,21 @@ def __init__( gru_bias: float = 2. ) -> None: """ + Overview: + Initialize GatedTransformerXLLayer. Arguments: - - input_dim (:obj:`int`): dimension of input - - head_dim (:obj:`int`): dimension of each head - - hidden_dim (:obj:`int`): dimension of hidden layer in mlp - - head_num (:obj:`int`): number of heads for multihead attention - - mlp_num (:obj:`int`): number of mlp layers in attention layer - - dropout (:obj:`nn.Module`): dropout - - activation (:obj:`nn.Module`): activation function - - gru_gating (:obj:`bool`): if False replace GRU gates with residual connections - - gru_bias (:obj:`float`): GRU gate bias + - input_dim (:obj:`int`): The dimension of the input tensor. + - head_dim (:obj:`int`): The dimension of each head in the multi-head attention. + - hidden_dim (:obj:`int`): The dimension of the hidden layer in the MLP. + - head_num (:obj:`int`): The number of heads for the multi-head attention. + - mlp_num (:obj:`int`): The number of MLP layers in the attention layer. + - dropout (:obj:`nn.Module`): The dropout module used in the MLP and attention layers. + - activation (:obj:`nn.Module`): The activation function to be used in the MLP layers. + - gru_gating (:obj:`bool`, optional): Whether to use GRU gates. If False, replace GRU gates with \ + residual connections. Default is True. + - gru_bias (:obj:`float`, optional): The bias of the GRU gate. Default is 2. """ + super(GatedTransformerXLLayer, self).__init__() self.dropout = dropout self.gating = gru_gating @@ -366,19 +419,21 @@ def forward( memory: torch.Tensor, mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: - """Overview: + """ + Overview: Compute forward pass of GTrXL layer. Arguments: - - inputs (:obj:`torch.Tensor`): attention input of shape (cur_seq, bs, input_dim) - - pos_embedding (:obj:`torch.Tensor`): positional embedding of shape (full_seq, 1, full_seq) - - u (:obj:`torch.nn.Parameter`): content parameter of shape (head_num, head_dim) - - v (:obj:`torch.nn.Parameter`): position parameter of shape (head_num, head_dim) - - memory (:obj:`Optional[torch.Tensor]`): memory of shape (prev_seq, bs, input_dim) - - mask (:obj:`Optional[torch.Tensor]`): attention mask of shape (cur_seq, full_seq, 1) - full_seq = prev_seq + cur_seq + - inputs (:obj:`torch.Tensor`): The attention input tensor of shape (cur_seq, bs, input_dim). + - pos_embedding (:obj:`torch.Tensor`): The positional embedding tensor of shape (full_seq, 1, full_seq). + - u (:obj:`torch.nn.Parameter`): The content parameter tensor of shape (head_num, head_dim). + - v (:obj:`torch.nn.Parameter`): The position parameter tensor of shape (head_num, head_dim). + - memory (:obj:`torch.Tensor`): The memory tensor of shape (prev_seq, bs, input_dim). + - mask (:obj:`Optional[torch.Tensor]`): The attention mask tensor of shape (cur_seq, full_seq, 1). + Default is None. Returns: - output (:obj:`torch.Tensor`): layer output of shape (cur_seq, bs, input_dim) """ + # concat memory with input across sequence dimension full_input = torch.cat([memory, inputs], dim=0) # full_seq x bs x input_dim x1 = self.layernorm1(full_input) @@ -394,10 +449,10 @@ def forward( class GTrXL(nn.Module): """ Overview: - GTrXL Transformer implementation. - - .. note:: - For details refer to Stabilizing Transformer for Reinforcement Learning: https://arxiv.org/abs/1910.06764 + GTrXL Transformer implementation as described in "Stabilizing Transformer for Reinforcement Learning" + (https://arxiv.org/abs/1910.06764). + Interfaces: + ``__init__``, ``forward``, ``reset_memory``, ``get_memory`` """ def __init__( @@ -416,21 +471,25 @@ def __init__( use_embedding_layer: bool = True, ) -> None: """Overview: - Init GTrXL Model + Init GTrXL Model. Arguments: - - input_dim (:obj:`int`): dimension of input (dimension of a single observation) - - head_dim (:obj:`int`): dimension of each head - - hidden_dim (:obj:`int`): dimension of hidden layer in mlp - - embedding_dim (:obj:`int`): dimension of embedding (dimension of a single observation after embedding) - - head_num (:obj:`int`): number of heads for multihead attention - - mlp_num (:obj:`int`): number of mlp layers in attention layer - - layer_num (:obj:`int`): number of transformer layers - - dropout_ratio (:obj:`float`): dropout ratio - - activation (:obj:`nn.Module`): activation function - - gru_gating (:obj:`bool`): if False replace GRU gates with residual connections - - gru_bias (:obj:`float`): GRU gate bias - - use_embedding_layer (:obj:`bool`): default True. If False, don't use input embedding layer. + - input_dim (:obj:`int`): The dimension of the input observation. + - head_dim (:obj:`int`, optional): The dimension of each head. Default is 128. + - embedding_dim (:obj:`int`, optional): The dimension of the embedding. Default is 256. + - head_num (:obj:`int`, optional): The number of heads for multi-head attention. Default is 2. + - mlp_num (:obj:`int`, optional): The number of MLP layers in the attention layer. Default is 2. + - layer_num (:obj:`int`, optional): The number of transformer layers. Default is 3. + - memory_len (:obj:`int`, optional): The length of memory. Default is 64. + - dropout_ratio (:obj:`float`, optional): The dropout ratio. Default is 0. + - activation (:obj:`nn.Module`, optional): The activation function. Default is nn.ReLU(). + - gru_gating (:obj:`bool`, optional): If False, replace GRU gates with residual connections. \ + Default is True. + - gru_bias (:obj:`float`, optional): The GRU gate bias. Default is 2.0. + - use_embedding_layer (:obj:`bool`, optional): If False, don't use input embedding layer. Default is True. + Raises: + - AssertionError: If `embedding_dim` is not an even number. """ + super(GTrXL, self).__init__() assert embedding_dim % 2 == 0, 'embedding_dim={} should be even'.format(input_dim) self.head_num = head_num @@ -473,9 +532,11 @@ def reset_memory(self, batch_size: Optional[int] = None, state: Optional[torch.T Overview: Clear or set the memory of GTrXL. Arguments: - - batch_size (:obj:`Optional[int]`): batch size - - state (:obj:`Optional[torch.Tensor]`): input memory. Shape is (layer_num, memory_len, bs, embedding_dim). + - batch_size (:obj:`Optional[int]`): The batch size. Default is None. + - state (:obj:`Optional[torch.Tensor]`): The input memory with shape \ + (layer_num, memory_len, bs, embedding_dim). Default is None. """ + self.memory = Memory(memory_len=self.memory_len, layer_num=self.layer_num, embedding_dim=self.embedding_dim) if batch_size is not None: self.memory = Memory(self.memory_len, batch_size, self.embedding_dim, self.layer_num) @@ -485,11 +546,12 @@ def reset_memory(self, batch_size: Optional[int] = None, state: Optional[torch.T def get_memory(self): """ Overview: - Returns memory of GTrXL. + Returns the memory of GTrXL. Returns: - - memory: (:obj:`Optional[torch.Tensor]`): output memory or None if memory has not been initialized. \ - Shape is (layer_num, memory_len, bs, embedding_dim). + - memory (:obj:`Optional[torch.Tensor]`): The output memory or None if memory has not been initialized. \ + The shape is (layer_num, memory_len, bs, embedding_dim). """ + if self.memory is None: return None else: @@ -498,17 +560,18 @@ def get_memory(self): def forward(self, x: torch.Tensor, batch_first: bool = False, return_mem: bool = True) -> Dict[str, torch.Tensor]: """ Overview: - GTrXL forward pass. + Performs a forward pass on the GTrXL. Arguments: - - x (:obj:`torch.Tensor`): input tensor. Shape (seq_len, bs, input_size). - - batch_first (:obj:`bool`): if the input data has shape (bs, seq_len, input_size), set this param to \ - ``True`` in order to transpose along the first and second dimension and obtain shape \ - (seq_len, bs, input_size). This param doesn't affects the output memory. - - return_mem (:obj:`bool`): if this param is False, return only the output tensor without dict. + - x (:obj:`torch.Tensor`): The input tensor with shape (seq_len, bs, input_size). + - batch_first (:obj:`bool`, optional): If the input data has shape (bs, seq_len, input_size), \ + set this parameter to True to transpose along the first and second dimension and obtain shape \ + (seq_len, bs, input_size). This does not affect the output memory. Default is False. \ + - return_mem (:obj:`bool`, optional): If False, return only the output tensor without dict. Default is True. Returns: - - x (:obj:`Dict[str, torch.Tensor]`): dict containing transformer output of shape \ - (seq_len, bs, embedding_size) and memory of shape (layer_num, seq_len, bs, embedding_size) + - x (:obj:`Dict[str, torch.Tensor]`): A dictionary containing the transformer output of shape \ + (seq_len, bs, embedding_size) and memory of shape (layer_num, seq_len, bs, embedding_size). """ + if batch_first: x = torch.transpose(x, 1, 0) # bs x cur_seq x input_dim -> cur_seq x bs x input_dim cur_seq, bs = x.shape[:2] diff --git a/ding/torch_utils/network/gumbel_softmax.py b/ding/torch_utils/network/gumbel_softmax.py index 7e6177c179..fea7612103 100644 --- a/ding/torch_utils/network/gumbel_softmax.py +++ b/ding/torch_utils/network/gumbel_softmax.py @@ -4,44 +4,53 @@ class GumbelSoftmax(nn.Module): - r""" + """ Overview: - An nn.Module that computes GumbelSoftmax - Interface: - __init__, forward - - .. note: - - For more gumbelsoftmax info, you can refer to - the paper + An `nn.Module` that computes GumbelSoftmax. + Interfaces: + ``__init__``, ``forward``, ``gumbel_softmax_sample`` + .. note:: + For more information on GumbelSoftmax, refer to the paper [Categorical Reparameterization \ + with Gumbel-Softmax](https://arxiv.org/abs/1611.01144). """ def __init__(self) -> None: - r""" - Overview: - Initialize the GumbelSoftmax module """ + Overview: + Initialize the `GumbelSoftmax` module. + """ super(GumbelSoftmax, self).__init__() - def gumbel_softmax_sample(self, x: torch.Tensor, temperature, eps=1e-8): - """ Draw a sample from GumbelSoftmax distribution""" + def gumbel_softmax_sample(self, x: torch.Tensor, temperature: float, eps: float = 1e-8) -> torch.Tensor: + """ + Overview: + Draw a sample from the Gumbel-Softmax distribution. + Arguments: + - x (:obj:`torch.Tensor`): Input tensor. + - temperature (:obj:`float`): Non-negative scalar controlling the sharpness of the distribution. + - eps (:obj:`float`): Small number to prevent division by zero, default is `1e-8`. + Returns: + - output (:obj:`torch.Tensor`): Sample from Gumbel-Softmax distribution. + """ U = torch.rand(x.shape) U = U.to(x.device) y = x - torch.log(-torch.log(U + eps) + eps) return F.softmax(y / temperature, dim=1) def forward(self, x: torch.Tensor, temperature: float = 1.0, hard: bool = False) -> torch.Tensor: - r""" + """ + Overview: + Forward pass for the `GumbelSoftmax` module. Arguments: - - x (:obj:`torch.Tensor`): unnormalized log-probs - - temperature(:obj:`float`): non-negative scalar - - hard(:obj:`bool`): if true return one-hot label + - x (:obj:`torch.Tensor`): Unnormalized log-probabilities. + - temperature (:obj:`float`): Non-negative scalar controlling the sharpness of the distribution. + - hard (:obj:`bool`): If `True`, returns one-hot encoded labels. Default is `False`. Returns: - - output (:obj:`torch.Tensor`): sample from GumbelSoftmax distribution + - output (:obj:`torch.Tensor`): Sample from Gumbel-Softmax distribution. Shapes: - - x: :math:`(B, N)`, while B is the batch size, N is number of classes - - output: :math:`(B, N)`, while B is the batch size, N is number of classes + - x: its shape is :math:`(B, N)`, where `B` is the batch size and `N` is the number of classes. + - y: its shape is :math:`(B, N)`, where `B` is the batch size and `N` is the number of classes. """ y = self.gumbel_softmax_sample(x, temperature) if hard: diff --git a/ding/torch_utils/network/merge.py b/ding/torch_utils/network/merge.py new file mode 100644 index 0000000000..25d89885dd --- /dev/null +++ b/ding/torch_utils/network/merge.py @@ -0,0 +1,400 @@ +""" +This file provides an implementation of several different neural network modules that are used for merging and +transforming input data in various ways. The following components can be used when we are dealing with +data from multiple modes, or when we need to merge multiple intermediate embedded representations in +the forward process of a model. + +The main classes defined in this code are: + + - BilinearGeneral: This class implements a bilinear transformation layer that applies a bilinear transformation to + incoming data, as described in the "Multiplicative Interactions and Where to Find Them", published at ICLR 2020, + https://openreview.net/forum?id=rylnK6VtDH. The transformation involves two input features and an output + feature, and also includes an optional bias term. + + - TorchBilinearCustomized: This class implements a bilinear layer similar to the one provided by PyTorch + (torch.nn.Bilinear), but with additional customizations. This class can be used as an alternative to the + BilinearGeneral class. + + - TorchBilinear: This class is a simple wrapper around the PyTorch's built-in nn.Bilinear module. It provides the + same functionality as PyTorch's nn.Bilinear but within the structure of the current module. + + - FiLM: This class implements a Feature-wise Linear Modulation (FiLM) layer. FiLM layers apply an affine + transformation to the input data, conditioned on some additional context information. + + - GatingType: This is an enumeration class that defines different types of gating mechanisms that can be used in + the modules. + + - SumMerge: This class provides a simple summing mechanism to merge input streams. + + - VectorMerge: This class implements a more complex merging mechanism for vector streams. + The streams are first transformed using layer normalization, a ReLU activation, and a linear layer. + Then they are merged either by simple summing or by using a gating mechanism. + +The implementation of these classes involves PyTorch and Numpy libraries, and the classes use PyTorch's nn.Module as +the base class, making them compatible with PyTorch's neural network modules and functionalities. +These modules can be useful building blocks in more complex deep learning architectures. +""" + +import enum +import math +from collections import OrderedDict +from typing import List, Dict, Tuple + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor + + +class BilinearGeneral(nn.Module): + """ + Overview: + Bilinear implementation as in: Multiplicative Interactions and Where to Find Them, + ICLR 2020, https://openreview.net/forum?id=rylnK6VtDH. + Interfaces: + ``__init__``, ``forward`` + """ + + def __init__(self, in1_features: int, in2_features: int, out_features: int): + """ + Overview: + Initialize the Bilinear layer. + Arguments: + - in1_features (:obj:`int`): The size of each first input sample. + - in2_features (:obj:`int`): The size of each second input sample. + - out_features (:obj:`int`): The size of each output sample. + """ + + super(BilinearGeneral, self).__init__() + # Initialize the weight matrices W and U, and the bias vectors V and b + self.W = nn.Parameter(torch.Tensor(out_features, in1_features, in2_features)) + self.U = nn.Parameter(torch.Tensor(out_features, in2_features)) + self.V = nn.Parameter(torch.Tensor(out_features, in1_features)) + self.b = nn.Parameter(torch.Tensor(out_features)) + self.in1_features = in1_features + self.in2_features = in2_features + self.out_features = out_features + self.reset_parameters() + + def reset_parameters(self): + """ + Overview: + Initialize the parameters of the Bilinear layer. + """ + + stdv = 1. / np.sqrt(self.in1_features) + self.W.data.uniform_(-stdv, stdv) + self.U.data.uniform_(-stdv, stdv) + self.V.data.uniform_(-stdv, stdv) + self.b.data.uniform_(-stdv, stdv) + + def forward(self, x: torch.Tensor, z: torch.Tensor): + """ + Overview: + compute the bilinear function. + Arguments: + - x (:obj:`torch.Tensor`): The first input tensor. + - z (:obj:`torch.Tensor`): The second input tensor. + """ + + # Compute the bilinear function + # x^TWz + out_W = torch.einsum('bi,kij,bj->bk', x, self.W, z) + # x^TU + out_U = z.matmul(self.U.t()) + # Vz + out_V = x.matmul(self.V.t()) + # x^TWz + x^TU + Vz + b + out = out_W + out_U + out_V + self.b + return out + + +class TorchBilinearCustomized(nn.Module): + """ + Overview: + Customized Torch Bilinear implementation. + Interfaces: + ``__init__``, ``forward`` + """ + + def __init__(self, in1_features: int, in2_features: int, out_features: int): + """ + Overview: + Initialize the Bilinear layer. + Arguments: + - in1_features (:obj:`int`): The size of each first input sample. + - in2_features (:obj:`int`): The size of each second input sample. + - out_features (:obj:`int`): The size of each output sample. + """ + + super(TorchBilinearCustomized, self).__init__() + self.in1_features = in1_features + self.in2_features = in2_features + self.out_features = out_features + self.weight = nn.Parameter(torch.Tensor(out_features, in1_features, in2_features)) + self.bias = nn.Parameter(torch.Tensor(out_features)) + self.reset_parameters() + + def reset_parameters(self): + """ + Overview: + Initialize the parameters of the Bilinear layer. + """ + + bound = 1 / math.sqrt(self.in1_features) + nn.init.uniform_(self.weight, -bound, bound) + nn.init.uniform_(self.bias, -bound, bound) + + def forward(self, x, z): + """ + Overview: + Compute the bilinear function. + Arguments: + - x (:obj:`torch.Tensor`): The first input tensor. + - z (:obj:`torch.Tensor`): The second input tensor. + """ + + # Using torch.einsum for the bilinear operation + out = torch.einsum('bi,oij,bj->bo', x, self.weight, z) + self.bias + return out.squeeze(-1) + + +""" +Overview: + Implementation of the Bilinear layer as in PyTorch: + https://pytorch.org/docs/stable/generated/torch.nn.Bilinear.html#torch.nn.Bilinear +Arguments: + - in1_features (:obj:`int`): The size of each first input sample. + - in2_features (:obj:`int`): The size of each second input sample. + - out_features (:obj:`int`): The size of each output sample. + - bias (:obj:`bool`): If set to False, the layer will not learn an additive bias. Default: ``True``. +""" +TorchBilinear = nn.Bilinear + + +class FiLM(nn.Module): + """ + Overview: + Feature-wise Linear Modulation (FiLM) Layer. + This layer applies feature-wise affine transformation based on context. + Interfaces: + ``__init__``, ``forward`` + """ + + def __init__(self, feature_dim: int, context_dim: int): + """ + Overview: + Initialize the FiLM layer. + Arguments: + - feature_dim (:obj:`int`). The dimension of the input feature vector. + - context_dim (:obj:`int`). The dimension of the input context vector. + """ + + super(FiLM, self).__init__() + # Define the fully connected layer for context + # The output dimension is twice the feature dimension for gamma and beta + self.context_layer = nn.Linear(context_dim, 2 * feature_dim) + + def forward(self, feature: torch.Tensor, context: torch.Tensor): + """ + Overview: + Forward propagation. + Arguments: + - feature (:obj:`torch.Tensor`). The input feature, shape (batch_size, feature_dim). + - context (:obj:`torch.Tensor`). The input context, shape (batch_size, context_dim). + Returns: + - conditioned_feature : torch.Tensor. The output feature after FiLM, shape (batch_size, feature_dim). + """ + + # Pass context through the fully connected layer + out = self.context_layer(context) + # Split the output into two parts: gamma and beta + # The dimension for splitting is 1 (feature dimension) + gamma, beta = torch.split(out, out.shape[1] // 2, dim=1) + # Apply feature-wise affine transformation + conditioned_feature = gamma * feature + beta + return conditioned_feature + + +class GatingType(enum.Enum): + """ + Overview: + Enum class defining different types of tensor gating and aggregation in modules. + """ + NONE = 'none' + GLOBAL = 'global' + POINTWISE = 'pointwise' + + +class SumMerge(nn.Module): + """ + Overview: + A PyTorch module that merges a list of tensors by computing their sum. All input tensors must have the same + size. This module can work with any type of tensor (vector, units or visual). + Interfaces: + ``__init__``, ``forward`` + """ + + def forward(self, tensors: List[Tensor]) -> Tensor: + """ + Overview: + Forward pass of the SumMerge module, which sums the input tensors. + Arguments: + - tensors (:obj:`List[Tensor]`): List of input tensors to be summed. All tensors must have the same size. + Returns: + - summed (:obj:`Tensor`): Tensor resulting from the sum of all input tensors. + """ + # stack the tensors along the first dimension + stacked = torch.stack(tensors, dim=0) + + # compute the sum along the first dimension + summed = torch.sum(stacked, dim=0) + # summed = sum(tensors) + return summed + + +class VectorMerge(nn.Module): + """ + Overview: + Merges multiple vector streams. Streams are first transformed through layer normalization, relu, and linear + layers, then summed. They don't need to have the same size. Gating can also be used before the sum. + Interfaces: + ``__init__``, ``encode``, ``_compute_gate``, ``forward`` + + .. note:: + For more details about the gating types, please refer to the GatingType enum class. + """ + + def __init__( + self, + input_sizes: Dict[str, int], + output_size: int, + gating_type: GatingType = GatingType.NONE, + use_layer_norm: bool = True, + ): + """ + Overview: + Initialize the `VectorMerge` module. + Arguments: + - input_sizes (:obj:`Dict[str, int]`): A dictionary mapping input names to their sizes. \ + The size is a single integer for 1D inputs, or `None` for 0D inputs. \ + If an input size is `None`, we assume it's `()`. + - output_size (:obj:`int`): The size of the output vector. + - gating_type (:obj:`GatingType`): The type of gating mechanism to use. Default is `GatingType.NONE`. + - use_layer_norm (:obj:`bool`): Whether to use layer normalization. Default is `True`. + """ + super().__init__() + self._input_sizes = OrderedDict(input_sizes) + self._output_size = output_size + self._gating_type = gating_type + self._use_layer_norm = use_layer_norm + + if self._use_layer_norm: + self._layer_norms = nn.ModuleDict() + else: + self._layer_norms = None + + self._linears = nn.ModuleDict() + for name, size in self._input_sizes.items(): + linear_input_size = size if size > 0 else 1 + if self._use_layer_norm: + self._layer_norms[name] = nn.LayerNorm(linear_input_size) + self._linears[name] = nn.Linear(linear_input_size, self._output_size) + + self._gating_linears = nn.ModuleDict() + if self._gating_type is GatingType.GLOBAL: + self.gate_size = 1 + elif self._gating_type is GatingType.POINTWISE: + self.gate_size = self._output_size + elif self._gating_type is GatingType.NONE: + self._gating_linears = None + else: + raise ValueError(f'Gating type {self._gating_type} is not supported') + + if self._gating_linears is not None: + if len(self._input_sizes) == 2: + # more efficient than the general version below + for name, size in self._input_sizes.items(): + gate_input_size = size if size > 0 else 1 + gating_layer = nn.Linear(gate_input_size, self.gate_size) + torch.nn.init.normal_(gating_layer.weight, std=0.005) + torch.nn.init.constant_(gating_layer.bias, 0.0) + self._gating_linears[name] = gating_layer + else: + for name, size in self._input_sizes.items(): + gate_input_size = size if size > 0 else 1 + gating_layer = nn.Linear(gate_input_size, len(self._input_sizes) * self.gate_size) + torch.nn.init.normal_(gating_layer.weight, std=0.005) + torch.nn.init.constant_(gating_layer.bias, 0.0) + self._gating_linears[name] = gating_layer + + def encode(self, inputs: Dict[str, Tensor]) -> Tuple[List[Tensor], List[Tensor]]: + """ + Overview: + Encode the input tensors using layer normalization, relu, and linear transformations. + Arguments: + - inputs (:obj:`Dict[str, Tensor]`): The input tensors. + Returns: + - gates (:obj:`List[Tensor]`): The gate tensors after transformations. + - outputs (:obj:`List[Tensor]`): The output tensors after transformations. + """ + gates, outputs = [], [] + for name, size in self._input_sizes.items(): + feature = inputs[name] + if size <= 0 and feature.dim() == 1: + feature = feature.unsqueeze(-1) + feature = feature.to(torch.float32) + if self._use_layer_norm and name in self._layer_norms: + feature = self._layer_norms[name](feature) + feature = F.relu(feature) + gates.append(feature) + outputs.append(self._linears[name](feature)) + return gates, outputs + + def _compute_gate( + self, + init_gate: List[Tensor], + ) -> List[Tensor]: + """ + Overview: + Compute the gate values based on the initial gate values. + Arguments: + - init_gate (:obj:`List[Tensor]`): The initial gate values. + Returns: + - gate (:obj:`List[Tensor]`): The computed gate values. + """ + if len(self._input_sizes) == 2: + gate = [self._gating_linears[name](y) for name, y in zip(self._input_sizes.keys(), init_gate)] + gate = sum(gate) + sigmoid = torch.sigmoid(gate) + gate = [sigmoid, 1.0 - sigmoid] + else: + gate = [self._gating_linears[name](y) for name, y in zip(self._input_sizes.keys(), init_gate)] + gate = sum(gate) + gate = gate.reshape([-1, len(self._input_sizes), self.gate_size]) + gate = F.softmax(gate, dim=1) + assert gate.shape[1] == len(self._input_sizes) + gate = [gate[:, i] for i in range(len(self._input_sizes))] + return gate + + def forward(self, inputs: Dict[str, Tensor]) -> Tensor: + """ + Overview: + Forward pass through the VectorMerge module. + Arguments: + - inputs (:obj:`Dict[str, Tensor]`): The input tensors. + Returns: + - output (:obj:`Tensor`): The output tensor after passing through the module. + """ + gates, outputs = self.encode(inputs) + if len(outputs) == 1: + # Special case of 1-D inputs that do not need any gating. + output = outputs[0] + elif self._gating_type is GatingType.NONE: + output = sum(outputs) + else: + gate = self._compute_gate(gates) + data = [g * d for g, d in zip(gate, outputs)] + output = sum(data) + return output diff --git a/ding/torch_utils/network/nn_module.py b/ding/torch_utils/network/nn_module.py index 784fbe3d90..c8bdf44301 100644 --- a/ding/torch_utils/network/nn_module.py +++ b/ding/torch_utils/network/nn_module.py @@ -10,13 +10,14 @@ def weight_init_(weight: torch.Tensor, init_type: str = "xavier", activation: str = None) -> None: - r""" + """ Overview: - Init weight according to the specified type. + Initialize weight according to the specified type. Arguments: - - weight (:obj:`torch.Tensor`): the weight that needed to init - - init_type (:obj:`str`): the type of init to implement, supports ["xavier", "kaiming", "orthogonal"] - - activation (:obj:`str`): the activation function name, recommend that use only with \ + - weight (:obj:`torch.Tensor`): The weight that needs to be initialized. + - init_type (:obj:`str`, optional): The type of initialization to implement, \ + supports ["xavier", "kaiming", "orthogonal"]. + - activation (:obj:`str`, optional): The activation function name. Recommended to use only with \ ['relu', 'leaky_relu']. """ @@ -42,16 +43,16 @@ def orthogonal_init(weight, *args): raise KeyError("Invalid Value in init type: {}".format(init_type)) -def sequential_pack(layers: list) -> nn.Sequential: - r""" +def sequential_pack(layers: List[nn.Module]) -> nn.Sequential: + """ Overview: Pack the layers in the input list to a `nn.Sequential` module. If there is a convolutional layer in module, an extra attribute `out_channels` will be added to the module and set to the out_channel of the conv layer. Arguments: - - layers (:obj:`list`): the input list + - layers (:obj:`List[nn.Module]`): The input list of layers. Returns: - - seq (:obj:`nn.Sequential`): packed sequential container + - seq (:obj:`nn.Sequential`): Packed sequential container. """ assert isinstance(layers, list) seq = nn.Sequential(*layers) @@ -76,24 +77,25 @@ def conv1d_block( activation: nn.Module = None, norm_type: str = None ) -> nn.Sequential: - r""" + """ Overview: - Create a 1-dim convlution layer with activation and normalization. + Create a 1-dimensional convolution layer with activation and normalization. Arguments: - - in_channels (:obj:`int`): Number of channels in the input tensor - - out_channels (:obj:`int`): Number of channels in the output tensor - - kernel_size (:obj:`int`): Size of the convolving kernel - - stride (:obj:`int`): Stride of the convolution - - padding (:obj:`int`): Zero-padding added to both sides of the input - - dilation (:obj:`int`): Spacing between kernel elements - - groups (:obj:`int`): Number of blocked connections from input channels to output channels - - activation (:obj:`nn.Module`): the optional activation function - - norm_type (:obj:`str`): type of the normalization + - in_channels (:obj:`int`): Number of channels in the input tensor. + - out_channels (:obj:`int`): Number of channels in the output tensor. + - kernel_size (:obj:`int`): Size of the convolving kernel. + - stride (:obj:`int`, optional): Stride of the convolution. Default is 1. + - padding (:obj:`int`, optional): Zero-padding added to both sides of the input. Default is 0. + - dilation (:obj:`int`, optional): Spacing between kernel elements. Default is 1. + - groups (:obj:`int`, optional): Number of blocked connections from input channels to output channels. \ + Default is 1. + - activation (:obj:`nn.Module`, optional): The optional activation function. + - norm_type (:obj:`str`, optional): Type of the normalization. Returns: - - block (:obj:`nn.Sequential`): a sequential list containing the torch layers of the 1 dim convlution layer + - block (:obj:`nn.Sequential`): A sequential list containing the torch layers of the 1-dimensional \ + convolution layer. .. note:: - Conv1d (https://pytorch.org/docs/stable/generated/torch.nn.Conv1d.html#torch.nn.Conv1d) """ block = [] @@ -115,27 +117,34 @@ def conv2d_block( groups: int = 1, pad_type: str = 'zero', activation: nn.Module = None, - norm_type: str = None + norm_type: str = None, + num_groups_for_gn: int = 1, + bias: bool = True ) -> nn.Sequential: - r""" + """ Overview: - Create a 2-dim convlution layer with activation and normalization. + Create a 2-dimensional convolution layer with activation and normalization. Arguments: - - in_channels (:obj:`int`): Number of channels in the input tensor - - out_channels (:obj:`int`): Number of channels in the output tensor - - kernel_size (:obj:`int`): Size of the convolving kernel - - stride (:obj:`int`): Stride of the convolution - - padding (:obj:`int`): Zero-padding added to both sides of the input - - dilation (:obj:`int`): Spacing between kernel elements - - groups (:obj:`int`): Number of blocked connections from input channels to output channels - - pad_type (:obj:`str`): the way to add padding, include ['zero', 'reflect', 'replicate'], default: None - - activation (:obj:`nn.Module`): the optional activation function - - norm_type (:obj:`str`): type of the normalization, default set to None, now support ['BN', 'IN', 'SyncBN'] + - in_channels (:obj:`int`): Number of channels in the input tensor. + - out_channels (:obj:`int`): Number of channels in the output tensor. + - kernel_size (:obj:`int`): Size of the convolving kernel. + - stride (:obj:`int`, optional): Stride of the convolution. Default is 1. + - padding (:obj:`int`, optional): Zero-padding added to both sides of the input. Default is 0. + - dilation (:obj:`int`): Spacing between kernel elements. + - groups (:obj:`int`, optional): Number of blocked connections from input channels to output channels. \ + Default is 1. + - pad_type (:obj:`str`, optional): The way to add padding, include ['zero', 'reflect', 'replicate']. \ + Default is 'zero'. + - activation (:obj:`nn.Module`): the optional activation function. + - norm_type (:obj:`str`): The type of the normalization, now support ['BN', 'LN', 'IN', 'GN', 'SyncBN'], \ + default set to None, which means no normalization. + - num_groups_for_gn (:obj:`int`): Number of groups for GroupNorm. + - bias (:obj:`bool`): whether to add a learnable bias to the nn.Conv2d. Default is True. Returns: - - block (:obj:`nn.Sequential`): a sequential list containing the torch layers of the 2 dim convlution layer + - block (:obj:`nn.Sequential`): A sequential list containing the torch layers of the 2-dimensional \ + convolution layer. .. note:: - Conv2d (https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html#torch.nn.Conv2d) """ block = [] @@ -149,10 +158,31 @@ def conv2d_block( block.append(nn.ReplicationPad2d(padding)) padding = 0 block.append( - nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding=padding, dilation=dilation, groups=groups) + nn.Conv2d( + in_channels, + out_channels, + kernel_size, + stride, + padding=padding, + dilation=dilation, + groups=groups, + bias=bias + ) ) if norm_type is not None: - block.append(build_normalization(norm_type, dim=2)(out_channels)) + if norm_type == 'LN': + # LN is implemented as GroupNorm with 1 group. + block.append(nn.GroupNorm(1, out_channels)) + elif norm_type == 'GN': + block.append(nn.GroupNorm(num_groups_for_gn, out_channels)) + elif norm_type in ['BN', 'IN', 'SyncBN']: + block.append(build_normalization(norm_type, dim=2)(out_channels)) + else: + raise KeyError( + "Invalid value in norm_type: {}. The valid norm_type are " + "BN, LN, IN, GN and SyncBN.".format(norm_type) + ) + if activation is not None: block.append(activation) return sequential_pack(block) @@ -169,28 +199,29 @@ def deconv2d_block( activation: int = None, norm_type: int = None ) -> nn.Sequential: - r""" + """ Overview: - Create a 2-dim transopse convlution layer with activation and normalization + Create a 2-dimensional transpose convolution layer with activation and normalization. Arguments: - - in_channels (:obj:`int`): Number of channels in the input tensor - - out_channels (:obj:`int`): Number of channels in the output tensor - - kernel_size (:obj:`int`): Size of the convolving kernel - - stride (:obj:`int`): Stride of the convolution - - padding (:obj:`int`): Zero-padding added to both sides of the input - - pad_type (:obj:`str`): the way to add padding, include ['zero', 'reflect', 'replicate'] - - activation (:obj:`nn.Module`): the optional activation function - - norm_type (:obj:`str`): type of the normalization + - in_channels (:obj:`int`): Number of channels in the input tensor. + - out_channels (:obj:`int`): Number of channels in the output tensor. + - kernel_size (:obj:`int`): Size of the convolving kernel. + - stride (:obj:`int`, optional): Stride of the convolution. Default is 1. + - padding (:obj:`int`, optional): Zero-padding added to both sides of the input. Default is 0. + - output_padding (:obj:`int`, optional): Additional size added to one side of the output shape. Default is 0. + - groups (:obj:`int`, optional): Number of blocked connections from input channels to output channels. \ + Default is 1. + - activation (:obj:`int`, optional): The optional activation function. + - norm_type (:obj:`int`, optional): Type of the normalization. Returns: - - block (:obj:`nn.Sequential`): a sequential list containing the torch layers of the 2-dim \ - transpose convlution layer + - block (:obj:`nn.Sequential`): A sequential list containing the torch layers of the 2-dimensional \ + transpose convolution layer. .. note:: ConvTranspose2d (https://pytorch.org/docs/master/generated/torch.nn.ConvTranspose2d.html) """ - block = [] - block.append( + block = [ nn.ConvTranspose2d( in_channels=in_channels, out_channels=out_channels, @@ -200,7 +231,7 @@ def deconv2d_block( output_padding=output_padding, groups=groups ) - ) + ] if norm_type is not None: block.append(build_normalization(norm_type, dim=2)(out_channels)) if activation is not None: @@ -216,24 +247,25 @@ def fc_block( use_dropout: bool = False, dropout_probability: float = 0.5 ) -> nn.Sequential: - r""" + """ Overview: - Create a fully-connected block with activation, normalization and dropout. - Optional normalization can be done to the dim 1 (across the channels) + Create a fully-connected block with activation, normalization, and dropout. + Optional normalization can be done to the dim 1 (across the channels). x -> fc -> norm -> act -> dropout -> out Arguments: - - in_channels (:obj:`int`): Number of channels in the input tensor - - out_channels (:obj:`int`): Number of channels in the output tensor - - activation (:obj:`nn.Module`): the optional activation function - - norm_type (:obj:`str`): type of the normalization - - use_dropout (:obj:`bool`) : whether to use dropout in the fully-connected block - - dropout_probability (:obj:`float`) : probability of an element to be zeroed in the dropout. Default: 0.5 + - in_channels (:obj:`int`): Number of channels in the input tensor. + - out_channels (:obj:`int`): Number of channels in the output tensor. + - activation (:obj:`nn.Module`, optional): The optional activation function. + - norm_type (:obj:`str`, optional): Type of the normalization. + - use_dropout (:obj:`bool`, optional): Whether to use dropout in the fully-connected block. Default is False. + - dropout_probability (:obj:`float`, optional): Probability of an element to be zeroed in the dropout. \ + Default is 0.5. Returns: - - block (:obj:`nn.Sequential`): a sequential list containing the torch layers of the fully-connected block + - block (:obj:`nn.Sequential`): A sequential list containing the torch layers of the fully-connected block. .. note:: - you can refer to nn.linear (https://pytorch.org/docs/master/generated/torch.nn.Linear.html) + You can refer to nn.linear (https://pytorch.org/docs/master/generated/torch.nn.Linear.html). """ block = [] block.append(nn.Linear(in_channels, out_channels)) @@ -246,9 +278,26 @@ def fc_block( return sequential_pack(block) -def normed_linear(in_features, out_features, bias: bool = True, device=None, dtype=None, scale=1.0): +def normed_linear( + in_features: int, + out_features: int, + bias: bool = True, + device=None, + dtype=None, + scale: float = 1.0 +) -> nn.Linear: """ - nn.Linear but with normalized fan-in init + Overview: + Create a nn.Linear module but with normalized fan-in init. + Arguments: + - in_features (:obj:`int`): Number of features in the input tensor. + - out_features (:obj:`int`): Number of features in the output tensor. + - bias (:obj:`bool`, optional): Whether to add a learnable bias to the nn.Linear. Default is True. + - device (:obj:`torch.device`, optional): The device to put the created module on. Default is None. + - dtype (:obj:`torch.dtype`, optional): The desired data type of created module. Default is None. + - scale (:obj:`float`, optional): The scale factor for initialization. Default is 1.0. + Returns: + - out (:obj:`nn.Linear`): A nn.Linear module with normalized fan-in init. """ out = nn.Linear(in_features, out_features, bias) @@ -260,22 +309,41 @@ def normed_linear(in_features, out_features, bias: bool = True, device=None, dty def normed_conv2d( - in_channels, - out_channels, - kernel_size, - stride=1, - padding=0, - dilation=1, - groups=1, - bias: bool = True, - padding_mode='zeros', - device=None, - dtype=None, - scale=1 -): + in_channels: int, + out_channels: int, + kernel_size: Union[int, Tuple[int, int]], + stride: Union[int, Tuple[int, int]] = 1, + padding: Union[int, Tuple[int, int]] = 0, + dilation: Union[int, Tuple[int, int]] = 1, + groups: int = 1, + bias: bool = True, + padding_mode: str = 'zeros', + device=None, + dtype=None, + scale: float = 1 +) -> nn.Conv2d: """ - nn.Conv2d but with normalized fan-in init + Overview: + Create a nn.Conv2d module but with normalized fan-in init. + Arguments: + - in_channels (:obj:`int`): Number of channels in the input tensor. + - out_channels (:obj:`int`): Number of channels in the output tensor. + - kernel_size (:obj:`Union[int, Tuple[int, int]]`): Size of the convolving kernel. + - stride (:obj:`Union[int, Tuple[int, int]]`, optional): Stride of the convolution. Default is 1. + - padding (:obj:`Union[int, Tuple[int, int]]`, optional): Zero-padding added to both sides of the input. \ + Default is 0. + - dilation (:`Union[int, Tuple[int, int]]`, optional): Spacing between kernel elements. Default is 1. + - groups (:obj:`int`, optional): Number of blocked connections from input channels to output channels. \ + Default is 1. + - bias (:obj:`bool`, optional): Whether to add a learnable bias to the nn.Conv2d. Default is True. + - padding_mode (:obj:`str`, optional): The type of padding algorithm to use. Default is 'zeros'. + - device (:obj:`torch.device`, optional): The device to put the created module on. Default is None. + - dtype (:obj:`torch.dtype`, optional): The desired data type of created module. Default is None. + - scale (:obj:`float`, optional): The scale factor for initialization. Default is 1. + Returns: + - out (:obj:`nn.Conv2d`): A nn.Conv2d module with normalized fan-in init. """ + out = nn.Conv2d( in_channels, out_channels, @@ -302,29 +370,39 @@ def MLP( activation: nn.Module = None, norm_type: str = None, use_dropout: bool = False, - dropout_probability: float = 0.5 + dropout_probability: float = 0.5, + output_activation: bool = True, + output_norm: bool = True, + last_linear_layer_init_zero: bool = False ): - r""" + """ Overview: - create a multi-layer perceptron using fully-connected blocks with activation, normalization and dropout, - optional normalization can be done to the dim 1 (across the channels) + Create a multi-layer perceptron using fully-connected blocks with activation, normalization, and dropout, + optional normalization can be done to the dim 1 (across the channels). x -> fc -> norm -> act -> dropout -> out Arguments: - - in_channels (:obj:`int`): Number of channels in the input tensor - - hidden_channels (:obj:`int`): Number of channels in the hidden tensor - - out_channels (:obj:`int`): Number of channels in the output tensor - - layer_num (:obj:`int`): Number of layers - - layer_fn (:obj:`Callable`): layer function - - activation (:obj:`nn.Module`): the optional activation function - - norm_type (:obj:`str`): type of the normalization - - use_dropout (:obj:`bool`): whether to use dropout in the fully-connected block - - dropout_probability (:obj:`float`): probability of an element to be zeroed in the dropout. Default: 0.5 + - in_channels (:obj:`int`): Number of channels in the input tensor. + - hidden_channels (:obj:`int`): Number of channels in the hidden tensor. + - out_channels (:obj:`int`): Number of channels in the output tensor. + - layer_num (:obj:`int`): Number of layers. + - layer_fn (:obj:`Callable`, optional): Layer function. + - activation (:obj:`nn.Module`, optional): The optional activation function. + - norm_type (:obj:`str`, optional): The type of the normalization. + - use_dropout (:obj:`bool`, optional): Whether to use dropout in the fully-connected block. Default is False. + - dropout_probability (:obj:`float`, optional): Probability of an element to be zeroed in the dropout. \ + Default is 0.5. + - output_activation (:obj:`bool`, optional): Whether to use activation in the output layer. If True, \ + we use the same activation as front layers. Default is True. + - output_norm (:obj:`bool`, optional): Whether to use normalization in the output layer. If True, \ + we use the same normalization as front layers. Default is True. + - last_linear_layer_init_zero (:obj:`bool`, optional): Whether to use zero initializations for the last \ + linear layer (including w and b), which can provide stable zero outputs in the beginning, \ + usually used in the policy network in RL settings. Returns: - - block (:obj:`nn.Sequential`): a sequential list containing the torch layers of the fully-connected block + - block (:obj:`nn.Sequential`): A sequential list containing the torch layers of the multi-layer perceptron. .. note:: - - you can refer to nn.linear (https://pytorch.org/docs/master/generated/torch.nn.Linear.html) + you can refer to nn.linear (https://pytorch.org/docs/master/generated/torch.nn.Linear.html). """ assert layer_num >= 0, layer_num if layer_num == 0: @@ -334,7 +412,7 @@ def MLP( if layer_fn is None: layer_fn = nn.Linear block = [] - for i, (in_channels, out_channels) in enumerate(zip(channels[:-1], channels[1:])): + for i, (in_channels, out_channels) in enumerate(zip(channels[:-2], channels[1:-1])): block.append(layer_fn(in_channels, out_channels)) if norm_type is not None: block.append(build_normalization(norm_type, dim=1)(out_channels)) @@ -342,39 +420,65 @@ def MLP( block.append(activation) if use_dropout: block.append(nn.Dropout(dropout_probability)) + + # The last layer + in_channels = channels[-2] + out_channels = channels[-1] + block.append(layer_fn(in_channels, out_channels)) + """ + In the final layer of a neural network, whether to use normalization and activation are typically determined + based on user specifications. These specifications depend on the problem at hand and the desired properties of + the model's output. + """ + if output_norm is True: + # The last layer uses the same norm as front layers. + if norm_type is not None: + block.append(build_normalization(norm_type, dim=1)(out_channels)) + if output_activation is True: + # The last layer uses the same activation as front layers. + if activation is not None: + block.append(activation) + if use_dropout: + block.append(nn.Dropout(dropout_probability)) + + if last_linear_layer_init_zero: + # Locate the last linear layer and initialize its weights and biases to 0. + for _, layer in enumerate(reversed(block)): + if isinstance(layer, nn.Linear): + nn.init.zeros_(layer.weight) + nn.init.zeros_(layer.bias) + break + return sequential_pack(block) class ChannelShuffle(nn.Module): - r""" + """ Overview: - Apply channelShuffle to the input tensor - Interface: - forward - - .. note:: - - You can see the original paper shuffle net in https://arxiv.org/abs/1707.01083 + Apply channel shuffle to the input tensor. For more details about the channel shuffle, + please refer to the 'ShuffleNet' paper: https://arxiv.org/abs/1707.01083 + Interfaces: + ``__init__``, ``forward`` """ def __init__(self, group_num: int) -> None: - r""" + """ Overview: - Init class ChannelShuffle + Initialize the ChannelShuffle class. Arguments: - - group_num (:obj:`int`): the number of groups to exchange + - group_num (:obj:`int`): The number of groups to exchange. """ super().__init__() self.group_num = group_num def forward(self, x: torch.Tensor) -> torch.Tensor: - r""" + """ Overview: - Return the upsampled input + Forward pass through the ChannelShuffle module. Arguments: - - x (:obj:`torch.Tensor`): the input tensor + - x (:obj:`torch.Tensor`): The input tensor. Returns: - - x (:obj:`torch.Tensor`): the shuffled input tensor + - x (:obj:`torch.Tensor`): The shuffled input tensor. """ b, c, h, w = x.shape g = self.group_num @@ -384,17 +488,17 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: def one_hot(val: torch.LongTensor, num: int, num_first: bool = False) -> torch.FloatTensor: - r""" + """ Overview: - Convert a ``torch.LongTensor`` to one hot encoding. - This implementation can be slightly faster than ``torch.nn.functional.one_hot`` + Convert a torch.LongTensor to one-hot encoding. This implementation can be slightly faster than + ``torch.nn.functional.one_hot``. Arguments: - - val (:obj:`torch.LongTensor`): each element contains the state to be encoded, the range should be [0, num-1] - - num (:obj:`int`): number of states of the one hot encoding - - num_first (:obj:`bool`): If ``num_first`` is False, the one hot encoding is added as the last; \ - Otherwise as the first dimension. + - val (:obj:`torch.LongTensor`): Each element contains the state to be encoded, the range should be [0, num-1] + - num (:obj:`int`): Number of states of the one-hot encoding + - num_first (:obj:`bool`, optional): If False, the one-hot encoding is added as the last dimension; otherwise, \ + it is added as the first dimension. Default is False. Returns: - - one_hot (:obj:`torch.FloatTensor`) + - one_hot (:obj:`torch.FloatTensor`): The one-hot encoded tensor. Example: >>> one_hot(2*torch.ones([2,2]).long(),3) tensor([[[0., 0., 1.], @@ -436,75 +540,74 @@ def one_hot(val: torch.LongTensor, num: int, num_first: bool = False) -> torch.F class NearestUpsample(nn.Module): - r""" + """ Overview: - Upsamples the input to the given member varible scale_factor using mode nearest - Interface: - forward + This module upsamples the input to the given scale_factor using the nearest mode. + Interfaces: + ``__init__``, ``forward`` """ def __init__(self, scale_factor: Union[float, List[float]]) -> None: - r""" + """ Overview: - Init class NearestUpsample + Initialize the NearestUpsample class. Arguments: - - scale_factor (:obj:`Union[float, List[float]]`): multiplier for spatial size + - scale_factor (:obj:`Union[float, List[float]]`): The multiplier for the spatial size. """ super(NearestUpsample, self).__init__() self.scale_factor = scale_factor def forward(self, x: torch.Tensor) -> torch.Tensor: - r""" - Overview: - Return the upsampled input - Arguments: - - x (:obj:`torch.Tensor`): the input tensor - Returns: - - upsample(:obj:`torch.Tensor`): the upsampled input tensor """ + Overview: + Return the upsampled input tensor. + Arguments: + - x (:obj:`torch.Tensor`): The input tensor. + Returns: + - upsample(:obj:`torch.Tensor`): The upsampled input tensor. + """ return F.interpolate(x, scale_factor=self.scale_factor, mode='nearest') class BilinearUpsample(nn.Module): - r""" + """ Overview: - Upsamples the input to the given member varible scale_factor using mode biliner - Interface: - forward + This module upsamples the input to the given scale_factor using the bilinear mode. + Interfaces: + ``__init__``, ``forward`` """ def __init__(self, scale_factor: Union[float, List[float]]) -> None: - r""" + """ Overview: - Init class BilinearUpsample - + Initialize the BilinearUpsample class. Arguments: - - scale_factor (:obj:`Union[float, List[float]]`): multiplier for spatial size + - scale_factor (:obj:`Union[float, List[float]]`): The multiplier for the spatial size. """ super(BilinearUpsample, self).__init__() self.scale_factor = scale_factor def forward(self, x: torch.Tensor) -> torch.Tensor: - r""" + """ Overview: - Return the upsampled input + Return the upsampled input tensor. Arguments: - - x (:obj:`torch.Tensor`): the input tensor + - x (:obj:`torch.Tensor`): The input tensor. Returns: - - upsample(:obj:`torch.Tensor`): the upsampled input tensor + - upsample(:obj:`torch.Tensor`): The upsampled input tensor. """ return F.interpolate(x, scale_factor=self.scale_factor, mode='bilinear', align_corners=False) def binary_encode(y: torch.Tensor, max_val: torch.Tensor) -> torch.Tensor: - r""" + """ Overview: - Convert elements in a tensor to its binary representation + Convert elements in a tensor to its binary representation. Arguments: - - y (:obj:`torch.Tensor`): the tensor to be transferred into its binary representation - - max_val (:obj:`torch.Tensor`): the max value of the elements in tensor + - y (:obj:`torch.Tensor`): The tensor to be converted into its binary representation. + - max_val (:obj:`torch.Tensor`): The maximum value of the elements in the tensor. Returns: - - binary (:obj:`torch.Tensor`): the input tensor in its binary representation + - binary (:obj:`torch.Tensor`): The input tensor in its binary representation. Example: >>> binary_encode(torch.tensor([3,2]),torch.tensor(8)) tensor([[0, 0, 1, 1],[0, 0, 1, 0]]) @@ -524,14 +627,26 @@ def binary_encode(y: torch.Tensor, max_val: torch.Tensor) -> torch.Tensor: class NoiseLinearLayer(nn.Module): - r""" + """ Overview: - Linear layer with random noise. - Interface: - reset_noise, reset_parameters, forward + This is a linear layer with random noise. + Interfaces: + ``__init__``, ``reset_noise``, ``reset_parameters``, ``forward`` """ def __init__(self, in_channels: int, out_channels: int, sigma0: int = 0.4) -> None: + """ + Overview: + Initialize the NoiseLinearLayer class. The 'enable_noise' attribute enables external control over whether \ + noise is applied. + - If enable_noise is True, the layer adds noise even if the module is in evaluation mode. + - If enable_noise is False, no noise is added regardless of self.training. + Arguments: + - in_channels (:obj:`int`): Number of channels in the input tensor. + - out_channels (:obj:`int`): Number of channels in the output tensor. + - sigma0 (:obj:`int`, optional): Default noise volume when initializing NoiseLinearLayer. \ + Default is 0.4. + """ super(NoiseLinearLayer, self).__init__() self.in_channels = in_channels self.out_channels = out_channels @@ -542,19 +657,27 @@ def __init__(self, in_channels: int, out_channels: int, sigma0: int = 0.4) -> No self.register_buffer("weight_eps", torch.empty(out_channels, in_channels)) self.register_buffer("bias_eps", torch.empty(out_channels)) self.sigma0 = sigma0 + self.enable_noise = False self.reset_parameters() self.reset_noise() def _scale_noise(self, size: Union[int, Tuple]): + """ + Overview: + Scale the noise. + Arguments: + - size (:obj:`Union[int, Tuple]`): The size of the noise. + """ + x = torch.randn(size) x = x.sign().mul(x.abs().sqrt()) return x def reset_noise(self): - r""" - Overview: - Reset noise settinngs in the layer. """ + Overview: + Reset the noise settings in the layer. + """ is_cuda = self.weight_mu.is_cuda in_noise = self._scale_noise(self.in_channels).to(torch.device("cuda" if is_cuda else "cpu")) out_noise = self._scale_noise(self.out_channels).to(torch.device("cuda" if is_cuda else "cpu")) @@ -562,9 +685,9 @@ def reset_noise(self): self.bias_eps = out_noise def reset_parameters(self): - r""" + """ Overview: - Reset parameters in the layer. + Reset the parameters in the layer. """ stdv = 1. / math.sqrt(self.in_channels) self.weight_mu.data.uniform_(-stdv, stdv) @@ -576,15 +699,16 @@ def reset_parameters(self): self.bias_sigma.data.fill_(std_bias) def forward(self, x: torch.Tensor): - r""" + """ Overview: - Layer forward with noise. + Perform the forward pass with noise. Arguments: - - x (:obj:`torch.Tensor`): the input tensor + - x (:obj:`torch.Tensor`): The input tensor. Returns: - - output (:obj:`torch.Tensor`): the output with noise + - output (:obj:`torch.Tensor`): The output tensor with noise. """ - if self.training: + # Determine whether to add noise: + if self.enable_noise: return F.linear( x, self.weight_mu + self.weight_sigma * self.weight_eps, @@ -603,28 +727,24 @@ def noise_block( dropout_probability: float = 0.5, sigma0: float = 0.4 ): - r""" + """ Overview: - Create a fully-connected block with activation, normalization and dropout - Optional normalization can be done to the dim 1 (across the channels) - x -> fc -> norm -> act -> dropout -> out + Create a fully-connected noise layer with activation, normalization, and dropout. + Optional normalization can be done to the dim 1 (across the channels). Arguments: - - in_channels (:obj:`int`): Number of channels in the input tensor - - out_channels (:obj:`int`): Number of channels in the output tensor - - activation (:obj:`str`): the optional activation function - - norm_type (:obj:`str`): type of the normalization - - use_dropout (:obj:`bool`) : whether to use dropout in the fully-connected block - - dropout_probability (:obj:`float`) : probability of an element to be zeroed in the dropout. Default: 0.5 - - simga0 (:obj:`float`): the sigma0 is the defalut noise volumn when init NoiseLinearLayer + - in_channels (:obj:`int`): Number of channels in the input tensor. + - out_channels (:obj:`int`): Number of channels in the output tensor. + - activation (:obj:`str`, optional): The optional activation function. Default is None. + - norm_type (:obj:`str`, optional): Type of normalization. Default is None. + - use_dropout (:obj:`bool`, optional): Whether to use dropout in the fully-connected block. + - dropout_probability (:obj:`float`, optional): Probability of an element to be zeroed in the dropout. \ + Default is 0.5. + - sigma0 (:obj:`float`, optional): The sigma0 is the default noise volume when initializing NoiseLinearLayer. \ + Default is 0.4. Returns: - - block (:obj:`nn.Sequential`): a sequential list containing the torch layers of the fully-connected block - - .. note:: - - you can refer to nn.linear (https://pytorch.org/docs/master/generated/torch.nn.Linear.html) + - block (:obj:`nn.Sequential`): A sequential list containing the torch layers of the fully-connected block. """ - block = [] - block.append(NoiseLinearLayer(in_channels, out_channels, sigma0=sigma0)) + block = [NoiseLinearLayer(in_channels, out_channels, sigma0=sigma0)] if norm_type is not None: block.append(build_normalization(norm_type, dim=1)(out_channels)) if activation is not None: @@ -635,13 +755,34 @@ def noise_block( class NaiveFlatten(nn.Module): + """ + Overview: + This module is a naive implementation of the flatten operation. + Interfaces: + ``__init__``, ``forward`` + """ def __init__(self, start_dim: int = 1, end_dim: int = -1) -> None: + """ + Overview: + Initialize the NaiveFlatten class. + Arguments: + - start_dim (:obj:`int`, optional): The first dimension to flatten. Default is 1. + - end_dim (:obj:`int`, optional): The last dimension to flatten. Default is -1. + """ super(NaiveFlatten, self).__init__() self.start_dim = start_dim self.end_dim = end_dim def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Overview: + Perform the flatten operation on the input tensor. + Arguments: + - x (:obj:`torch.Tensor`): The input tensor. + Returns: + - output (:obj:`torch.Tensor`): The flattened output tensor. + """ if self.end_dim != -1: return x.view(*x.shape[:self.start_dim], -1, *x.shape[self.end_dim + 1:]) else: diff --git a/ding/torch_utils/network/normalization.py b/ding/torch_utils/network/normalization.py index 054a2a8197..a6d4c8df32 100644 --- a/ding/torch_utils/network/normalization.py +++ b/ding/torch_utils/network/normalization.py @@ -3,24 +3,22 @@ def build_normalization(norm_type: str, dim: Optional[int] = None) -> nn.Module: - r""" + """ Overview: - Build the corresponding normalization module + Construct the corresponding normalization module. For beginners, + refer to [this article](https://zhuanlan.zhihu.com/p/34879333) to learn more about batch normalization. Arguments: - - norm_type (:obj:`str`): type of the normaliztion, now support ['BN', 'IN', 'SyncBN', 'AdaptiveIN'] - - dim (:obj:`int`): dimension of the normalization, when norm_type is in [BN, IN] + - norm_type (:obj:`str`): Type of the normalization. Currently supports ['BN', 'LN', 'IN', 'SyncBN']. + - dim (:obj:`Optional[int]`): Dimension of the normalization, applicable when norm_type is in ['BN', 'IN']. Returns: - - norm_func (:obj:`nn.Module`): the corresponding batch normalization function - - .. note:: - For beginers, you can refer to to learn more about batch normalization. + - norm_func (:obj:`nn.Module`): The corresponding batch normalization function. """ if dim is None: key = norm_type else: - if norm_type in ['BN', 'IN', 'SyncBN']: + if norm_type in ['BN', 'IN']: key = norm_type + str(dim) - elif norm_type in ['LN']: + elif norm_type in ['LN', 'SyncBN']: key = norm_type else: raise NotImplementedError("not support indicated dim when creates {}".format(norm_type)) @@ -28,7 +26,9 @@ def build_normalization(norm_type: str, dim: Optional[int] = None) -> nn.Module: 'BN1': nn.BatchNorm1d, 'BN2': nn.BatchNorm2d, 'LN': nn.LayerNorm, + 'IN1': nn.InstanceNorm1d, 'IN2': nn.InstanceNorm2d, + 'SyncBN': nn.SyncBatchNorm, } if key in norm_func.keys(): return norm_func[key] diff --git a/ding/torch_utils/network/popart.py b/ding/torch_utils/network/popart.py new file mode 100644 index 0000000000..5a0351fb20 --- /dev/null +++ b/ding/torch_utils/network/popart.py @@ -0,0 +1,125 @@ +""" +Implementation of ``POPART`` algorithm for reward rescale. + + +POPART is an adaptive normalization algorithm to normalize the targets used in the learning updates. +The two main components in POPART are: +**ART**: to update scale and shift such that the return is appropriately normalized, +**POP**: to preserve the outputs of the unnormalized function when we change the scale and shift. + +""" +from typing import Optional, Union, Dict +import math +import torch +import torch.nn as nn + + +class PopArt(nn.Module): + """ + Overview: + A linear layer with popart normalization. This class implements a linear transformation followed by + PopArt normalization, which is a method to automatically adapt the contribution of each task to the agent's + updates in multi-task learning, as described in the paper . + + Interfaces: + ``__init__``, ``reset_parameters``, ``forward``, ``update_parameters`` + """ + + def __init__( + self, + input_features: Union[int, None] = None, + output_features: Union[int, None] = None, + beta: float = 0.5 + ) -> None: + """ + Overview: + Initialize the class with input features, output features, and the beta parameter. + Arguments: + - input_features (:obj:`Union[int, None]`): The size of each input sample. + - output_features (:obj:`Union[int, None]`): The size of each output sample. + - beta (:obj:`float`): The parameter for moving average. + """ + super(PopArt, self).__init__() + + self.beta = beta + self.input_features = input_features + self.output_features = output_features + # Initialize the linear layer parameters, weight and bias. + self.weight = nn.Parameter(torch.Tensor(output_features, input_features)) + self.bias = nn.Parameter(torch.Tensor(output_features)) + # Register a buffer for normalization parameters which can not be considered as model parameters. + # The normalization parameters will be used later to save the target value's scale and shift. + self.register_buffer('mu', torch.zeros(output_features, requires_grad=False)) + self.register_buffer('sigma', torch.ones(output_features, requires_grad=False)) + self.register_buffer('v', torch.ones(output_features, requires_grad=False)) + + self.reset_parameters() + + def reset_parameters(self): + """ + Overview: + Reset the parameters including weights and bias using ``kaiming_uniform_`` and ``uniform_`` initialization. + """ + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + if self.bias is not None: + fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) + bound = 1 / math.sqrt(fan_in) + nn.init.uniform_(self.bias, -bound, bound) + + def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]: + """ + Overview: + Implement the forward computation of the linear layer and return both the output and the + normalized output of the layer. + Arguments: + - x (:obj:`torch.Tensor`): Input tensor which is to be normalized. + Returns: + - output (:obj:`Dict[str, torch.Tensor]`): A dictionary contains 'pred' and 'unnormalized_pred'. + """ + normalized_output = x.mm(self.weight.t()) + normalized_output += self.bias.unsqueeze(0).expand_as(normalized_output) + # The unnormalization of output + with torch.no_grad(): + output = normalized_output * self.sigma + self.mu + + return {'pred': normalized_output.squeeze(1), 'unnormalized_pred': output.squeeze(1)} + + def update_parameters(self, value: torch.Tensor) -> Dict[str, torch.Tensor]: + """ + Overview: + Update the normalization parameters based on the given value and return the new mean and + standard deviation after the update. + Arguments: + - value (:obj:`torch.Tensor`): The tensor to be used for updating parameters. + Returns: + - update_results (:obj:`Dict[str, torch.Tensor]`): A dictionary contains 'new_mean' and 'new_std'. + """ + # Tensor device conversion of the normalization parameters. + self.mu = self.mu.to(value.device) + self.sigma = self.sigma.to(value.device) + self.v = self.v.to(value.device) + + old_mu = self.mu + old_std = self.sigma + + # Calculate the first and second moments (mean and variance) of the target value: + batch_mean = torch.mean(value, 0) + batch_v = torch.mean(torch.pow(value, 2), 0) + batch_mean[torch.isnan(batch_mean)] = self.mu[torch.isnan(batch_mean)] + batch_v[torch.isnan(batch_v)] = self.v[torch.isnan(batch_v)] + batch_mean = (1 - self.beta) * self.mu + self.beta * batch_mean + batch_v = (1 - self.beta) * self.v + self.beta * batch_v + batch_std = torch.sqrt(batch_v - (batch_mean ** 2)) + # Clip the standard deviation to reject the outlier data. + batch_std = torch.clamp(batch_std, min=1e-4, max=1e+6) + # Replace the nan value with old value. + batch_std[torch.isnan(batch_std)] = self.sigma[torch.isnan(batch_std)] + + self.mu = batch_mean + self.v = batch_v + self.sigma = batch_std + # Update weight and bias with mean and standard deviation to preserve unnormalised outputs + self.weight.data = (self.weight.data.t() * old_std / self.sigma).t() + self.bias.data = (old_std * self.bias.data + old_mu - self.mu) / self.sigma + + return {'new_mean': batch_mean, 'new_std': batch_std} diff --git a/ding/torch_utils/network/res_block.py b/ding/torch_utils/network/res_block.py index af89bddd7e..82908ce509 100644 --- a/ding/torch_utils/network/res_block.py +++ b/ding/torch_utils/network/res_block.py @@ -1,13 +1,15 @@ -import torch.nn as nn +from typing import Union + import torch +import torch.nn as nn from .nn_module import conv2d_block, fc_block class ResBlock(nn.Module): - r''' + """ Overview: - Residual Block with 2D convolution layers, including 2 types: + Residual Block with 2D convolution layers, including 3 types: basic block: input channel: C x -> 3*3*C -> norm -> act -> 3*3*C -> norm -> act -> out @@ -15,94 +17,183 @@ class ResBlock(nn.Module): bottleneck block: x -> 1*1*(1/4*C) -> norm -> act -> 3*3*(1/4*C) -> norm -> act -> 1*1*C -> norm -> act -> out \_____________________________________________________________________________/+ + downsample block: used in EfficientZero + input channel: C + x -> 3*3*C -> norm -> act -> 3*3*C -> norm -> act -> out + \__________________ 3*3*C ____________________/+ + + .. note:: + You can refer to `Deep Residual Learning for Image Recognition `_ for more \ + details. + Interfaces: - forward - ''' + ``__init__``, ``forward`` + """ def __init__( - self, - in_channels: int, - activation: nn.Module = nn.ReLU(), - norm_type: str = 'BN', - res_type: str = 'basic' + self, + in_channels: int, + activation: nn.Module = nn.ReLU(), + norm_type: str = 'BN', + res_type: str = 'basic', + bias: bool = True, + out_channels: Union[int, None] = None, ) -> None: - r""" + """ Overview: - Init the Residual Block + Init the 2D convolution residual block. Arguments: - - in_channels (:obj:`int`): Number of channels in the input tensor - - activation (:obj:`nn.Module`): the optional activation function - - norm_type (:obj:`str`): type of the normalization, defalut set to 'BN'(Batch Normalization), \ - supports ['BN', 'IN', 'SyncBN', None]. - - res_type (:obj:`str`): type of residual block, supports ['basic', 'bottleneck'] + - in_channels (:obj:`int`): Number of channels in the input tensor. + - activation (:obj:`nn.Module`): The optional activation function. + - norm_type (:obj:`str`): Type of the normalization, default set to 'BN'(Batch Normalization), \ + supports ['BN', 'LN', 'IN', 'GN', 'SyncBN', None]. + - res_type (:obj:`str`): Type of residual block, supports ['basic', 'bottleneck', 'downsample'] + - bias (:obj:`bool`): Whether to add a learnable bias to the conv2d_block. default set to True. + - out_channels (:obj:`int`): Number of channels in the output tensor, default set to None, \ + which means out_channels = in_channels. """ super(ResBlock, self).__init__() self.act = activation - assert res_type in ['basic', - 'bottleneck'], 'residual type only support basic and bottleneck, not:{}'.format(res_type) + assert res_type in ['basic', 'bottleneck', + 'downsample'], 'residual type only support basic and bottleneck, not:{}'.format(res_type) self.res_type = res_type + if out_channels is None: + out_channels = in_channels if self.res_type == 'basic': - self.conv1 = conv2d_block(in_channels, in_channels, 3, 1, 1, activation=self.act, norm_type=norm_type) - self.conv2 = conv2d_block(in_channels, in_channels, 3, 1, 1, activation=None, norm_type=norm_type) + self.conv1 = conv2d_block( + in_channels, out_channels, 3, 1, 1, activation=self.act, norm_type=norm_type, bias=bias + ) + self.conv2 = conv2d_block( + out_channels, out_channels, 3, 1, 1, activation=None, norm_type=norm_type, bias=bias + ) elif self.res_type == 'bottleneck': - self.conv1 = conv2d_block(in_channels, in_channels, 1, 1, 0, activation=self.act, norm_type=norm_type) - self.conv2 = conv2d_block(in_channels, in_channels, 3, 1, 1, activation=self.act, norm_type=norm_type) - self.conv3 = conv2d_block(in_channels, in_channels, 1, 1, 0, activation=None, norm_type=norm_type) + self.conv1 = conv2d_block( + in_channels, out_channels, 1, 1, 0, activation=self.act, norm_type=norm_type, bias=bias + ) + self.conv2 = conv2d_block( + out_channels, out_channels, 3, 1, 1, activation=self.act, norm_type=norm_type, bias=bias + ) + self.conv3 = conv2d_block( + out_channels, out_channels, 1, 1, 0, activation=None, norm_type=norm_type, bias=bias + ) + elif self.res_type == 'downsample': + self.conv1 = conv2d_block( + in_channels, out_channels, 3, 2, 1, activation=self.act, norm_type=norm_type, bias=bias + ) + self.conv2 = conv2d_block( + out_channels, out_channels, 3, 1, 1, activation=None, norm_type=norm_type, bias=bias + ) + self.conv3 = conv2d_block(in_channels, out_channels, 3, 2, 1, activation=None, norm_type=None, bias=bias) def forward(self, x: torch.Tensor) -> torch.Tensor: - r""" + """ Overview: - Return the residual block output + Return the redisual block output. Arguments: - - x (:obj:`torch.Tensor`): the input tensor + - x (:obj:`torch.Tensor`): The input tensor. Returns: - - x(:obj:`torch.Tensor`): the resblock output tensor + - x (:obj:`torch.Tensor`): The resblock output tensor. """ - residual = x + identity = x x = self.conv1(x) x = self.conv2(x) if self.res_type == 'bottleneck': x = self.conv3(x) - x = self.act(x + residual) + elif self.res_type == 'downsample': + identity = self.conv3(identity) + x = self.act(x + identity) return x class ResFCBlock(nn.Module): - r''' + """ Overview: - Residual Block with 2 fully connected block + Residual Block with 2 fully connected layers. x -> fc1 -> norm -> act -> fc2 -> norm -> act -> out \_____________________________________/+ Interfaces: - forward - ''' + ``__init__``, ``forward`` + """ - def __init__(self, in_channels: int, activation: nn.Module = nn.ReLU(), norm_type: str = 'BN'): - r""" + def __init__( + self, in_channels: int, activation: nn.Module = nn.ReLU(), norm_type: str = 'BN', dropout: float = None + ): + """ Overview: - Init the Residual Block + Init the fully connected layer residual block. Arguments: - - in_channels (:obj:`int`): Number of channels in the input tensor - - activation (:obj:`nn.Module`): the optional activation function - - norm_type (:obj:`str`): type of the normalization, defalut set to 'BN' + - in_channels (:obj:`int`): The number of channels in the input tensor. + - activation (:obj:`nn.Module`): The optional activation function. + - norm_type (:obj:`str`): The type of the normalization, default set to 'BN'. + - dropout (:obj:`float`): The dropout rate, default set to None. """ super(ResFCBlock, self).__init__() self.act = activation + if dropout is not None: + self.dropout = nn.Dropout(dropout) + else: + self.dropout = None self.fc1 = fc_block(in_channels, in_channels, activation=self.act, norm_type=norm_type) self.fc2 = fc_block(in_channels, in_channels, activation=None, norm_type=norm_type) def forward(self, x: torch.Tensor) -> torch.Tensor: - r""" + """ Overview: - Return the redisual block output + Return the output of the redisual block. Arguments: - - x (:obj:`torch.Tensor`): the input tensor + - x (:obj:`torch.Tensor`): The input tensor. Returns: - - x(:obj:`torch.Tensor`): the resblock output tensor + - x (:obj:`torch.Tensor`): The resblock output tensor. """ - residual = x + identity = x x = self.fc1(x) x = self.fc2(x) - x = self.act(x + residual) + x = self.act(x + identity) + if self.dropout is not None: + x = self.dropout(x) return x + + +class TemporalSpatialResBlock(nn.Module): + """ + Overview: + Residual Block using MLP layers for both temporal and spatial input. + t → time_mlp → h1 → dense2 → h2 → out + ↗+ ↗+ + x → dense1 → ↗ + ↘ ↗ + → modify_x → → → → + """ + + def __init__(self, input_dim, output_dim, t_dim=128, activation=torch.nn.SiLU()): + """ + Overview: + Init the temporal spatial residual block. + Arguments: + - input_dim (:obj:`int`): The number of channels in the input tensor. + - output_dim (:obj:`int`): The number of channels in the output tensor. + - t_dim (:obj:`int`): The dimension of the temporal input. + - activation (:obj:`nn.Module`): The optional activation function. + """ + super().__init__() + # temporal input is the embedding of time, which is a Gaussian Fourier Feature tensor + self.time_mlp = nn.Sequential( + activation, + nn.Linear(t_dim, output_dim), + ) + self.dense1 = nn.Sequential(nn.Linear(input_dim, output_dim), activation) + self.dense2 = nn.Sequential(nn.Linear(output_dim, output_dim), activation) + self.modify_x = nn.Linear(input_dim, output_dim) if input_dim != output_dim else nn.Identity() + + def forward(self, x, t) -> torch.Tensor: + """ + Overview: + Return the redisual block output. + Arguments: + - x (:obj:`torch.Tensor`): The input tensor. + - t (:obj:`torch.Tensor`): The temporal input tensor. + """ + h1 = self.dense1(x) + self.time_mlp(t) + h2 = self.dense2(h1) + return h2 + self.modify_x(x) diff --git a/ding/torch_utils/network/resnet.py b/ding/torch_utils/network/resnet.py index 42320e53c3..9bfcbeb360 100644 --- a/ding/torch_utils/network/resnet.py +++ b/ding/torch_utils/network/resnet.py @@ -1,15 +1,25 @@ """ This implementation of ResNet is a bit modification version of `https://github.com/rwightman/pytorch-image-models.git` """ -from typing import List +from typing import List, Callable, Optional, Tuple, Type, Dict, Union import math import numpy as np +import torch import torch.nn as nn import torch.nn.functional as F + from .nn_module import Flatten -def to_2tuple(item): +def to_2tuple(item: int) -> tuple: + """ + Overview: + Convert a scalar to a 2-tuple or return the item if it's not a scalar. + Arguments: + - item (:obj:`int`): An item to be converted to a 2-tuple. + Returns: + - (:obj:`tuple`): A 2-tuple of the item. + """ if np.isscalar(item): return (item, item) else: @@ -17,12 +27,35 @@ def to_2tuple(item): # Calculate asymmetric TensorFlow-like 'SAME' padding for a convolution -def get_same_padding(x: int, k: int, s: int, d: int): +def get_same_padding(x: int, k: int, s: int, d: int) -> int: + """ + Overview: + Calculate asymmetric TensorFlow-like 'SAME' padding for a convolution. + Arguments: + - x (:obj:`int`): The size of the input. + - k (:obj:`int`): The size of the kernel. + - s (:obj:`int`): The stride of the convolution. + - d (:obj:`int`): The dilation of the convolution. + Returns: + - (:obj:`int`): The size of the padding. + """ return max((math.ceil(x / s) - 1) * s + (k - 1) * d + 1 - x, 0) # Dynamically pad input x with 'SAME' padding for conv with specified args def pad_same(x, k: List[int], s: List[int], d: List[int] = (1, 1), value: float = 0): + """ + Overview: + Dynamically pad input x with 'SAME' padding for conv with specified args. + Arguments: + - x (:obj:`Tensor`): The input tensor. + - k (:obj:`List[int]`): The size of the kernel. + - s (:obj:`List[int]`): The stride of the convolution. + - d (:obj:`List[int]`): The dilation of the convolution. + - value (:obj:`float`): Value to fill the padding. + Returns: + - (:obj:`Tensor`): The padded tensor. + """ ih, iw = x.size()[-2:] pad_h, pad_w = get_same_padding(ih, k[0], s[0], d[0]), get_same_padding(iw, k[1], s[1], d[1]) if pad_h > 0 or pad_w > 0: @@ -38,29 +71,85 @@ def avg_pool2d_same( ceil_mode: bool = False, count_include_pad: bool = True ): + """ + Overview: + Apply average pooling with 'SAME' padding on the input tensor. + Arguments: + - x (:obj:`Tensor`): The input tensor. + - kernel_size (:obj:`List[int]`): The size of the kernel. + - stride (:obj:`List[int]`): The stride of the convolution. + - padding (:obj:`List[int]`): The size of the padding. + - ceil_mode (:obj:`bool`): When True, will use ceil instead of floor to compute the output shape. + - count_include_pad (:obj:`bool`): When True, will include the zero-padding in the averaging calculation. + Returns: + - (:obj:`Tensor`): The tensor after average pooling. + """ # FIXME how to deal with count_include_pad vs not for external padding? x = pad_same(x, kernel_size, stride) return F.avg_pool2d(x, kernel_size, stride, (0, 0), ceil_mode, count_include_pad) class AvgPool2dSame(nn.AvgPool2d): - """ Tensorflow like 'SAME' wrapper for 2D average pooling + """ + Overview: + Tensorflow-like 'SAME' wrapper for 2D average pooling. + Interfaces: + ``__init__``, ``forward`` """ - def __init__(self, kernel_size: int, stride=None, padding=0, ceil_mode=False, count_include_pad=True): + def __init__( + self, + kernel_size: int, + stride: Optional[Tuple[int, int]] = None, + padding: int = 0, + ceil_mode: bool = False, + count_include_pad: bool = True + ) -> None: + """ + Overview: + Initialize the AvgPool2dSame with given arguments. + Arguments: + - kernel_size (:obj:`int`): The size of the window to take an average over. + - stride (:obj:`Optional[Tuple[int, int]]`): The stride of the window. If None, default to kernel_size. + - padding (:obj:`int`): Implicit zero padding to be added on both sides. + - ceil_mode (:obj:`bool`): When True, will use `ceil` instead of `floor` to compute the output shape. + - count_include_pad (:obj:`bool`): When True, will include the zero-padding in the averaging calculation. + """ kernel_size = to_2tuple(kernel_size) stride = to_2tuple(stride) super(AvgPool2dSame, self).__init__(kernel_size, stride, (0, 0), ceil_mode, count_include_pad) - def forward(self, x): + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Overview: + Forward pass of the AvgPool2dSame. + Argument: + - x (:obj:`torch.Tensor`): Input tensor. + Returns: + - (:obj:`torch.Tensor`): Output tensor after average pooling. + """ x = pad_same(x, self.kernel_size, self.stride) return F.avg_pool2d(x, self.kernel_size, self.stride, self.padding, self.ceil_mode, self.count_include_pad) -def _create_pool(num_features, num_classes, pool_type='avg', use_conv=False): +def _create_pool(num_features: int, + num_classes: int, + pool_type: str = 'avg', + use_conv: bool = False) -> Tuple[nn.Module, int]: + """ + Overview: + Create a global pooling layer based on the given arguments. + Arguments: + - num_features (:obj:`int`): Number of input features. + - num_classes (:obj:`int`): Number of output classes. + - pool_type (:obj:`str`): Type of the pooling operation. Defaults to 'avg'. + - use_conv (:obj:`bool`): Whether to use convolutional layer after pooling. Defaults to False. + Returns: + - (:obj:`Tuple[nn.Module, int]`): The created global pooling layer and the number of pooled features. + """ flatten_in_pool = not use_conv # flatten when we use a Linear layer after pooling if not pool_type: - assert num_classes == 0 or use_conv,\ + assert num_classes == 0 or use_conv, \ 'Pooling can only be disabled if classifier is also removed or conv classifier is used' flatten_in_pool = False # disable flattening if pooling is pass-through (no pooling) assert flatten_in_pool @@ -69,7 +158,17 @@ def _create_pool(num_features, num_classes, pool_type='avg', use_conv=False): return global_pool, num_pooled_features -def _create_fc(num_features, num_classes, use_conv=False): +def _create_fc(num_features: int, num_classes: int, use_conv: bool = False) -> nn.Module: + """ + Overview: + Create a fully connected layer based on the given arguments. + Arguments: + - num_features (:obj:`int`): Number of input features. + - num_classes (:obj:`int`): Number of output classes. + - use_conv (:obj:`bool`): Whether to use convolutional layer. Defaults to False. + Returns: + - (:obj:`nn.Module`): The created fully connected layer. + """ if num_classes <= 0: fc = nn.Identity() # pass-through (no classifier) elif use_conv: @@ -80,7 +179,22 @@ def _create_fc(num_features, num_classes, use_conv=False): return fc -def create_classifier(num_features, num_classes, pool_type='avg', use_conv=False): +def create_classifier(num_features: int, + num_classes: int, + pool_type: str = 'avg', + use_conv: bool = False) -> Tuple[nn.Module, nn.Module]: + """ + Overview: + Create a classifier with global pooling layer and fully connected layer. + Arguments: + - num_features (:obj:`int`): The number of features. + - num_classes (:obj:`int`): The number of classes for the final classification. + - pool_type (:obj:`str`): The type of pooling to use; 'avg' for Average Pooling. + - use_conv (:obj:`bool`): Whether to use convolution or not. + Returns: + - global_pool (:obj:`nn.Module`): The created global pooling layer. + - fc (:obj:`nn.Module`): The created fully connected layer. + """ assert pool_type == 'avg' global_pool, num_pooled_features = _create_pool(num_features, num_classes, pool_type, use_conv=use_conv) fc = _create_fc(num_pooled_features, num_classes, use_conv=use_conv) @@ -88,16 +202,46 @@ def create_classifier(num_features, num_classes, pool_type='avg', use_conv=False class ClassifierHead(nn.Module): - """Classifier head w/ configurable global pooling and dropout.""" + """ + Overview: + Classifier head with configurable global pooling and dropout. + Interfaces: + ``__init__``, ``forward`` + """ - def __init__(self, in_chs, num_classes, pool_type='avg', drop_rate=0., use_conv=False): + def __init__( + self, + in_chs: int, + num_classes: int, + pool_type: str = 'avg', + drop_rate: float = 0., + use_conv: bool = False + ) -> None: + """ + Overview: + Initialize the ClassifierHead with given arguments. + Arguments: + - in_chs (:obj:`int`): Number of input channels. + - num_classes (:obj:`int`): Number of classes for the final classification. + - pool_type (:obj:`str`): The type of pooling to use; 'avg' for Average Pooling. + - drop_rate (:obj:`float`): The dropout rate. + - use_conv (:obj:`bool`): Whether to use convolution or not. + """ super(ClassifierHead, self).__init__() self.drop_rate = drop_rate self.global_pool, num_pooled_features = _create_pool(in_chs, num_classes, pool_type, use_conv=use_conv) self.fc = _create_fc(num_pooled_features, num_classes, use_conv=use_conv) self.flatten = Flatten(1) if use_conv and pool_type else nn.Identity() - def forward(self, x): + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Overview: + Forward pass of the ClassifierHead. + Argument: + - x (:obj:`torch.Tensor`): Input tensor. + Returns: + - (:obj:`torch.Tensor`): Output tensor after classification. + """ x = self.global_pool(x) if self.drop_rate: x = F.dropout(x, p=float(self.drop_rate), training=self.training) @@ -106,36 +250,85 @@ def forward(self, x): return x -def create_attn(layer, plane): +def create_attn(layer: nn.Module, plane: int) -> None: + """ + Overview: + Create an attention mechanism. + Arguments: + - layer (:obj:`nn.Module`): The layer where the attention is to be applied. + - plane (:obj:`int`): The plane on which the attention is to be applied. + Returns: + - None + """ return None -def get_padding(kernel_size, stride, dilation=1): +def get_padding(kernel_size: int, stride: int, dilation: int = 1) -> int: + """ + Overview: + Compute the padding based on the kernel size, stride and dilation. + Arguments: + - kernel_size (:obj:`int`): The size of the kernel. + - stride (:obj:`int`): The stride of the convolution. + - dilation (:obj:`int`): The dilation factor. + Returns: + - padding (:obj:`int`): The computed padding. + """ padding = ((stride - 1) + dilation * (kernel_size - 1)) // 2 return padding class BasicBlock(nn.Module): + """ + Overview: + The basic building block for models like ResNet. This class extends pytorch's Module class. + It represents a standard block of layers including two convolutions, batch normalization, + an optional attention mechanism, and activation functions. + Interfaces: + ``__init__``, ``forward``, ``zero_init_last_bn`` + Properties: + - expansion (:obj:int): Specifies the expansion factor for the planes of the conv layers. + """ expansion = 1 def __init__( - self, - inplanes, - planes, - stride=1, - downsample=None, - cardinality=1, - base_width=64, - reduce_first=1, - dilation=1, - first_dilation=None, - act_layer=nn.ReLU, - norm_layer=nn.BatchNorm2d, - attn_layer=None, - aa_layer=None, - drop_block=None, - drop_path=None - ): + self, + inplanes: int, + planes: int, + stride: int = 1, + downsample: Callable = None, + cardinality: int = 1, + base_width: int = 64, + reduce_first: int = 1, + dilation: int = 1, + first_dilation: int = None, + act_layer: Callable = nn.ReLU, + norm_layer: Callable = nn.BatchNorm2d, + attn_layer: Callable = None, + aa_layer: Callable = None, + drop_block: Callable = None, + drop_path: Callable = None + ) -> None: + """ + Overview: + Initialize the BasicBlock with given parameters. + Arguments: + - inplanes (:obj:`int`): Number of input channels. + - planes (:obj:`int`): Number of output channels. + - stride (:obj:`int`): The stride of the convolutional layer. + - downsample (:obj:`Callable`): Function for downsampling the inputs. + - cardinality (:obj:`int`): Group size for grouped convolution. + - base_width (:obj:`int`): Base width of the convolutions. + - reduce_first (:obj:`int`): Reduction factor for first convolution of each block. + - dilation (:obj:`int`): Spacing between kernel points. + - first_dilation (:obj:`int`): First dilation value. + - act_layer (:obj:`Callable`): Function for activation layer. + - norm_layer (:obj:`Callable`): Function for normalization layer. + - attn_layer (:obj:`Callable`): Function for attention layer. + - aa_layer (:obj:`Callable`): Function for anti-aliasing layer. + - drop_block (:obj:`Callable`): Method for dropping block. + - drop_path (:obj:`Callable`): Method for dropping path. + """ super(BasicBlock, self).__init__() assert cardinality == 1, 'BasicBlock only supports cardinality of 1' @@ -170,10 +363,22 @@ def __init__( self.drop_block = drop_block self.drop_path = drop_path - def zero_init_last_bn(self): + def zero_init_last_bn(self) -> None: + """ + Overview: + Initialize the batch normalization layer with zeros. + """ nn.init.zeros_(self.bn2.weight) - def forward(self, x): + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Overview: + Defines the computation performed at every call. + Arguments: + - x (:obj:`torch.Tensor`): The input tensor. + Returns: + - output (:obj:`torch.Tensor`): The output tensor after passing through the BasicBlock. + """ shortcut = x x = self.conv1(x) @@ -204,26 +409,59 @@ def forward(self, x): class Bottleneck(nn.Module): + """ + Overview: + The Bottleneck class is a basic block used to build ResNet networks. It is a part of the PyTorch's + implementation of ResNet. This block is designed with several layers including a convolutional layer, + normalization layer, activation layer, attention layer, anti-aliasing layer, and a dropout layer. + Interfaces: + ``__init__``, ``forward``, ``zero_init_last_bn`` + Properties: + expansion, inplanes, planes, stride, downsample, cardinality, base_width, reduce_first, dilation, \ + first_dilation, act_layer, norm_layer, attn_layer, aa_layer, drop_block, drop_path + + """ expansion = 4 def __init__( - self, - inplanes, - planes, - stride=1, - downsample=None, - cardinality=1, - base_width=64, - reduce_first=1, - dilation=1, - first_dilation=None, - act_layer=nn.ReLU, - norm_layer=nn.BatchNorm2d, - attn_layer=None, - aa_layer=None, - drop_block=None, - drop_path=None - ): + self, + inplanes: int, + planes: int, + stride: int = 1, + downsample: Optional[nn.Module] = None, + cardinality: int = 1, + base_width: int = 64, + reduce_first: int = 1, + dilation: int = 1, + first_dilation: Optional[int] = None, + act_layer: Type[nn.Module] = nn.ReLU, + norm_layer: Type[nn.Module] = nn.BatchNorm2d, + attn_layer: Optional[Type[nn.Module]] = None, + aa_layer: Optional[Type[nn.Module]] = None, + drop_block: Callable = None, + drop_path: Callable = None + ) -> None: + """ + Overview: + Initialize the Bottleneck class with various parameters. + + Arguments: + - inplanes (:obj:`int`): The number of input planes. + - planes (:obj:`int`): The number of output planes. + - stride (:obj:`int`, optional): The stride size, defaults to 1. + - downsample (:obj:`nn.Module`, optional): The downsample method, defaults to None. + - cardinality (:obj:`int`, optional): The size of the group convolutions, defaults to 1. + - base_width (:obj:`int`, optional): The base width, defaults to 64. + - reduce_first (:obj:`int`, optional): The first reduction factor, defaults to 1. + - dilation (:obj:`int`, optional): The dilation factor, defaults to 1. + - first_dilation (:obj:`int`, optional): The first dilation factor, defaults to None. + - act_layer (:obj:`Type[nn.Module]`, optional): The activation layer type, defaults to nn.ReLU. + - norm_layer (:obj:`Type[nn.Module]`, optional): The normalization layer type, defaults to nn.BatchNorm2d. + - attn_layer (:obj:`Type[nn.Module]`, optional): The attention layer type, defaults to None. + - aa_layer (:obj:`Type[nn.Module]`, optional): The anti-aliasing layer type, defaults to None. + - drop_block (:obj:`Callable`): The dropout block, defaults to None. + - drop_path (:obj:`Callable`): The drop path, defaults to None. + """ super(Bottleneck, self).__init__() width = int(math.floor(planes * (base_width / 64)) * cardinality) @@ -262,10 +500,22 @@ def __init__( self.drop_block = drop_block self.drop_path = drop_path - def zero_init_last_bn(self): + def zero_init_last_bn(self) -> None: + """ + Overview: + Initialize the last batch normalization layer with zero. + """ nn.init.zeros_(self.bn3.weight) - def forward(self, x): + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Overview: + Defines the computation performed at every call. + Arguments: + - x (:obj:`Tensor`): The input tensor. + Returns: + - x (:obj:`Tensor`): The output tensor resulting from the computation. + """ shortcut = x x = self.conv1(x) @@ -301,7 +551,29 @@ def forward(self, x): return x -def downsample_conv(in_channels, out_channels, kernel_size, stride=1, dilation=1, first_dilation=None, norm_layer=None): +def downsample_conv( + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + dilation: int = 1, + first_dilation: int = None, + norm_layer: Type[nn.Module] = None +) -> nn.Sequential: + """ + Overview: + Create a sequential module for downsampling that includes a convolution layer and a normalization layer. + Arguments: + - in_channels (:obj:`int`): The number of input channels. + - out_channels (:obj:`int`): The number of output channels. + - kernel_size (:obj:`int`): The size of the kernel. + - stride (:obj:`int`, optional): The stride size, defaults to 1. + - dilation (:obj:`int`, optional): The dilation factor, defaults to 1. + - first_dilation (:obj:`int`, optional): The first dilation factor, defaults to None. + - norm_layer (:obj:`Type[nn.Module]`, optional): The normalization layer type, defaults to nn.BatchNorm2d. + Returns: + - nn.Sequential: A sequence of layers performing downsampling through convolution. + """ norm_layer = norm_layer or nn.BatchNorm2d kernel_size = 1 if stride == 1 and dilation == 1 else kernel_size first_dilation = (first_dilation or dilation) if kernel_size > 1 else 1 @@ -317,7 +589,30 @@ def downsample_conv(in_channels, out_channels, kernel_size, stride=1, dilation=1 ) -def downsample_avg(in_channels, out_channels, kernel_size, stride=1, dilation=1, first_dilation=None, norm_layer=None): +def downsample_avg( + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + dilation: int = 1, + first_dilation: int = None, + norm_layer: Type[nn.Module] = None +) -> nn.Sequential: + """ + Overview: + Create a sequential module for downsampling that includes an average pooling layer, a convolution layer, + and a normalization layer. + Arguments: + - in_channels (:obj:`int`): The number of input channels. + - out_channels (:obj:`int`): The number of output channels. + - kernel_size (:obj:`int`): The size of the kernel. + - stride (:obj:`int`, optional): The stride size, defaults to 1. + - dilation (:obj:`int`, optional): The dilation factor, defaults to 1. + - first_dilation (:obj:`int`, optional): The first dilation factor, defaults to None. + - norm_layer (:obj:`Type[nn.Module]`, optional): The normalization layer type, defaults to nn.BatchNorm2d. + Returns: + - nn.Sequential: A sequence of layers performing downsampling through average pooling. + """ norm_layer = norm_layer or nn.BatchNorm2d avg_stride = stride if dilation == 1 else 1 if stride == 1 and dilation == 1: @@ -333,24 +628,51 @@ def downsample_avg(in_channels, out_channels, kernel_size, stride=1, dilation=1, ) -def drop_blocks(drop_block_rate=0.): +def drop_blocks(drop_block_rate: float = 0.) -> List[None]: + """ + Overview: + Generate a list of None values based on the drop block rate. + Arguments: + - drop_block_rate (:obj:`float`, optional): The drop block rate, defaults to 0. + Returns: + - List[None]: A list of None values. + """ assert drop_block_rate == 0., drop_block_rate return [None for _ in range(4)] def make_blocks( - block_fn, - channels, - block_repeats, - inplanes, - reduce_first=1, - output_stride=32, - down_kernel_size=1, - avg_down=False, - drop_block_rate=0., - drop_path_rate=0., - **kwargs -): + block_fn: Type[nn.Module], + channels: List[int], + block_repeats: List[int], + inplanes: int, + reduce_first: int = 1, + output_stride: int = 32, + down_kernel_size: int = 1, + avg_down: bool = False, + drop_block_rate: float = 0., + drop_path_rate: float = 0., + **kwargs +) -> Tuple[List[Tuple[str, nn.Module]], List[Dict[str, Union[int, str]]]]: + """ + Overview: + Create a list of blocks for the network, with each block having a given number of repeats. Also, create a + feature info list that contains information about the output of each block. + Arguments: + - block_fn (:obj:`Type[nn.Module]`): The type of block to use. + - channels (:obj:`List[int]`): The list of output channels for each block. + - block_repeats (:obj:`List[int]`): The list of number of repeats for each block. + - inplanes (:obj:`int`): The number of input planes. + - reduce_first (:obj:`int`, optional): The first reduction factor, defaults to 1. + - output_stride (:obj:`int`, optional): The total stride of the network, defaults to 32. + - down_kernel_size (:obj:`int`, optional): The size of the downsample kernel, defaults to 1. + - avg_down (:obj:`bool`, optional): Whether to use average pooling for downsampling, defaults to False. + - drop_block_rate (:obj:`float`, optional): The drop block rate, defaults to 0. + - drop_path_rate (:obj:`float`, optional): The drop path rate, defaults to 0. + Returns: + - Tuple[List[Tuple[str, nn.Module]], List[Dict[str, Union[int, str]]]]: \ + A tuple that includes a list of blocks for the network and a feature info list. + """ stages = [] feature_info = [] net_num_blocks = sum(block_repeats) @@ -401,100 +723,70 @@ def make_blocks( class ResNet(nn.Module): - """ResNet / ResNeXt / SE-ResNeXt / SE-Net - - This class implements all variants of ResNet, ResNeXt, SE-ResNeXt, and SENet that - * have > 1 stride in the 3x3 conv layer of bottleneck - * have conv-bn-act ordering - - This ResNet impl supports a number of stem and downsample options based on the v1c, v1d, v1e, and v1s - variants included in the MXNet Gluon ResNetV1b model. The C and D variants are also discussed in the - 'Bag of Tricks' paper: https://arxiv.org/pdf/1812.01187. The B variant is equivalent to torchvision default. - - ResNet variants (the same modifications can be used in SE/ResNeXt models as well): - * normal, b - 7x7 stem, stem_width = 64, same as torchvision ResNet, NVIDIA ResNet 'v1.5', Gluon v1b - * c - 3 layer deep 3x3 stem, stem_width = 32 (32, 32, 64) - * d - 3 layer deep 3x3 stem, stem_width = 32 (32, 32, 64), average pool in downsample - * e - 3 layer deep 3x3 stem, stem_width = 64 (64, 64, 128), average pool in downsample - * s - 3 layer deep 3x3 stem, stem_width = 64 (64, 64, 128) - * t - 3 layer deep 3x3 stem, stem width = 32 (24, 48, 64), average pool in downsample - * tn - 3 layer deep 3x3 stem, stem width = 32 (24, 32, 64), average pool in downsample - - ResNeXt - * normal - 7x7 stem, stem_width = 64, standard cardinality and base widths - * same c,d, e, s variants as ResNet can be enabled - - SE-ResNeXt - * normal - 7x7 stem, stem_width = 64 - * same c, d, e, s variants as ResNet can be enabled - - SENet-154 - 3 layer deep 3x3 stem (same as v1c-v1s), stem_width = 64, cardinality=64, - reduction by 2 on width of first bottleneck convolution, 3x3 downsample convs after first block - - Parameters - ---------- - block : Block - Class for the residual block. Options are BasicBlockGl, BottleneckGl. - layers : list of int - Numbers of layers in each block - num_classes : int, default 1000 - Number of classification classes. - in_chans : int, default 3 - Number of input (color) channels. - cardinality : int, default 1 - Number of convolution groups for 3x3 conv in Bottleneck. - base_width : int, default 64 - Factor determining bottleneck channels. `planes * base_width / 64 * cardinality` - stem_width : int, default 64 - Number of channels in stem convolutions - stem_type : str, default '' - The type of stem: - * '', default - a single 7x7 conv with a width of stem_width - * 'deep' - three 3x3 convolution layers of widths stem_width, stem_width, stem_width * 2 - * 'deep_tiered' - three 3x3 conv layers of widths stem_width//4 * 3, stem_width, stem_width * 2 - block_reduce_first: int, default 1 - Reduction factor for first convolution output width of residual blocks, - 1 for all archs except senets, where 2 - down_kernel_size: int, default 1 - Kernel size of residual block downsampling path, 1x1 for most archs, 3x3 for senets - avg_down : bool, default False - Whether to use average pooling for projection skip connection between stages/downsample. - output_stride : int, default 32 - Set the output stride of the network, 32, 16, or 8. Typically used in segmentation. - act_layer : nn.Module, activation layer - norm_layer : nn.Module, normalization layer - aa_layer : nn.Module, anti-aliasing layer - drop_rate : float, default 0. - Dropout probability before classifier, for training - global_pool : str, default 'avg' - Global pooling type. One of 'avg', 'max', 'avgmax', 'catavgmax' + """ + Overview: + Implements ResNet, ResNeXt, SE-ResNeXt, and SENet models. This implementation supports various modifications + based on the v1c, v1d, v1e, and v1s variants included in the MXNet Gluon ResNetV1b model. For more details + about the variants and options, please refer to the 'Bag of Tricks' paper: https://arxiv.org/pdf/1812.01187. + Interfaces: + ``__init__``, ``forward``, ``zero_init_last_bn``, ``get_classifier`` """ def __init__( - self, - block, - layers, - num_classes=1000, - in_chans=3, - cardinality=1, - base_width=64, - stem_width=64, - stem_type='', - replace_stem_pool=False, - output_stride=32, - block_reduce_first=1, - down_kernel_size=1, - avg_down=False, - act_layer=nn.ReLU, - norm_layer=nn.BatchNorm2d, - aa_layer=None, - drop_rate=0.0, - drop_path_rate=0., - drop_block_rate=0., - global_pool='avg', - zero_init_last_bn=True, - block_args=None - ): + self, + block: nn.Module, + layers: List[int], + num_classes: int = 1000, + in_chans: int = 3, + cardinality: int = 1, + base_width: int = 64, + stem_width: int = 64, + stem_type: str = '', + replace_stem_pool: bool = False, + output_stride: int = 32, + block_reduce_first: int = 1, + down_kernel_size: int = 1, + avg_down: bool = False, + act_layer: nn.Module = nn.ReLU, + norm_layer: nn.Module = nn.BatchNorm2d, + aa_layer: Optional[nn.Module] = None, + drop_rate: float = 0.0, + drop_path_rate: float = 0.0, + drop_block_rate: float = 0.0, + global_pool: str = 'avg', + zero_init_last_bn: bool = True, + block_args: Optional[dict] = None + ) -> None: + """ + Overview: + Initialize the ResNet model with given block, layers and other configuration options. + Arguments: + - block (:obj:`nn.Module`): Class for the residual block. + - layers (:obj:`List[int]`): Numbers of layers in each block. + - num_classes (:obj:`int`, optional): Number of classification classes. Default is 1000. + - in_chans (:obj:`int`, optional): Number of input (color) channels. Default is 3. + - cardinality (:obj:`int`, optional): Number of convolution groups for 3x3 conv in Bottleneck. Default is 1. + - base_width (:obj:`int`, optional): Factor determining bottleneck channels. Default is 64. + - stem_width (:obj:`int`, optional): Number of channels in stem convolutions. Default is 64. + - stem_type (:obj:`str`, optional): The type of stem. Default is ''. + - replace_stem_pool (:obj:`bool`, optional): Whether to replace stem pooling. Default is False. + - output_stride (:obj:`int`, optional): Output stride of the network. Default is 32. + - block_reduce_first (:obj:`int`, optional): Reduction factor for first convolution output width of \ + residual blocks. Default is 1. + - down_kernel_size (:obj:`int`, optional): Kernel size of residual block downsampling path. Default is 1. + - avg_down (:obj:`bool`, optional): Whether to use average pooling for projection skip connection between + stages/downsample. Default is False. + - act_layer (:obj:`nn.Module`, optional): Activation layer. Default is nn.ReLU. + - norm_layer (:obj:`nn.Module`, optional): Normalization layer. Default is nn.BatchNorm2d. + - aa_layer (:obj:`Optional[nn.Module]`, optional): Anti-aliasing layer. Default is None. + - drop_rate (:obj:`float`, optional): Dropout probability before classifier, for training. Default is 0.0. + - drop_path_rate (:obj:`float`, optional): Drop path rate. Default is 0.0. + - drop_block_rate (:obj:`float`, optional): Drop block rate. Default is 0.0. + - global_pool (:obj:`str`, optional): Global pooling type. Default is 'avg'. + - zero_init_last_bn (:obj:`bool`, optional): Whether to initialize last batch normalization with zero. \ + Default is True. + - block_args (:obj:`Optional[dict]`, optional): Additional arguments for block. Default is None. + """ block_args = block_args or dict() assert output_stride in (8, 16, 32) self.num_classes = num_classes @@ -576,7 +868,14 @@ def __init__( self.init_weights(zero_init_last_bn=zero_init_last_bn) - def init_weights(self, zero_init_last_bn=True): + def init_weights(self, zero_init_last_bn: bool = True) -> None: + """ + Overview: + Initialize the weights in the model. + Arguments: + - zero_init_last_bn (:obj:`bool`, optional): Whether to initialize last batch normalization with zero. + Default is True. + """ for n, m in self.named_modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') @@ -588,14 +887,35 @@ def init_weights(self, zero_init_last_bn=True): if hasattr(m, 'zero_init_last_bn'): m.zero_init_last_bn() - def get_classifier(self): + def get_classifier(self) -> nn.Module: + """ + Overview: + Get the classifier module from the model. + Returns: + - classifier (:obj:`nn.Module`): The classifier module in the model. + """ return self.fc - def reset_classifier(self, num_classes, global_pool='avg'): + def reset_classifier(self, num_classes: int, global_pool: str = 'avg') -> None: + """ + Overview: + Reset the classifier with a new number of classes and pooling type. + Arguments: + - num_classes (:obj:`int`): New number of classification classes. + - global_pool (:obj:`str`, optional): New global pooling type. Default is 'avg'. + """ self.num_classes = num_classes self.global_pool, self.fc = create_classifier(self.num_features, self.num_classes, pool_type=global_pool) - def forward_features(self, x): + def forward_features(self, x: torch.Tensor) -> torch.Tensor: + """ + Overview: + Forward pass through the feature layers of the model. + Arguments: + - x (:obj:`torch.Tensor`): The input tensor. + Returns: + - x (:obj:`torch.Tensor`): The output tensor after passing through feature layers. + """ x = self.conv1(x) x = self.bn1(x) x = self.act1(x) @@ -607,7 +927,15 @@ def forward_features(self, x): x = self.layer4(x) return x - def forward(self, x): + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Overview: + Full forward pass through the model. + Arguments: + - x (:obj:`torch.Tensor`): The input tensor. + Returns: + - x (:obj:`torch.Tensor`): The output tensor after passing through the model. + """ x = self.forward_features(x) x = self.global_pool(x) x = x.view(x.shape[0], -1) @@ -617,5 +945,11 @@ def forward(self, x): return x -def resnet18(): +def resnet18() -> nn.Module: + """ + Overview: + Creates a ResNet18 model. + Returns: + - model (:obj:`nn.Module`): ResNet18 model. + """ return ResNet(block=BasicBlock, layers=[2, 2, 2, 2]) diff --git a/ding/torch_utils/network/rnn.py b/ding/torch_utils/network/rnn.py index 0f87ffbcb4..e24bd7e468 100644 --- a/ding/torch_utils/network/rnn.py +++ b/ding/torch_utils/network/rnn.py @@ -15,21 +15,25 @@ def is_sequence(data): """ Overview: - Judege whether input ``data`` is instance ``list`` or ``tuple``. + Determines if the input data is of type list or tuple. + Arguments: + - data: The input data to be checked. + Returns: + - boolean: True if the input is a list or a tuple, False otherwise. """ return isinstance(data, list) or isinstance(data, tuple) def sequence_mask(lengths: torch.Tensor, max_len: Optional[int] = None) -> torch.BoolTensor: - r""" + """ Overview: - create a mask for a batch sequences with different lengths + Generates a boolean mask for a batch of sequences with differing lengths. Arguments: - - lengths (:obj:`torch.Tensor`): lengths in each different sequences, shape could be (n, 1) or (n) - - max_len (:obj:`int`): the padding size, if max_len is None, the padding size is the \ - max length of sequences + - lengths (:obj:`torch.Tensor`): A tensor with the lengths of each sequence. Shape could be (n, 1) or (n). + - max_len (:obj:`int`, optional): The padding size. If max_len is None, the padding size is the max length of \ + sequences. Returns: - - masks (:obj:`torch.BoolTensor`): mask has the same device as lengths + - masks (:obj:`torch.BoolTensor`): A boolean mask tensor. The mask has the same device as lengths. """ if len(lengths.shape) == 1: lengths = lengths.unsqueeze(dim=1) @@ -42,24 +46,24 @@ def sequence_mask(lengths: torch.Tensor, max_len: Optional[int] = None) -> torch class LSTMForwardWrapper(object): - r""" + """ Overview: - A class which provides methods to use before and after `forward`, in order to wrap the LSTM `forward` method. + Class providing methods to use before and after the LSTM `forward` method. + Wraps the LSTM `forward` method. Interfaces: - _before_forward, _after_forward + ``_before_forward``, ``_after_forward`` """ def _before_forward(self, inputs: torch.Tensor, prev_state: Union[None, List[Dict]]) -> torch.Tensor: """ Overview: - Preprocess the inputs and previous states + Preprocesses the inputs and previous states before the LSTM `forward` method. Arguments: - - inputs (:obj:`torch.Tensor`): input vector of cell, tensor of size [seq_len, batch_size, input_size] - - prev_state (:obj:`Union[None, List[Dict]]`): None or tensor of size \ - [num_directions*num_layers, batch_size, hidden_size]. \ - If None then prv_state will be initialized to all zeros. + - inputs (:obj:`torch.Tensor`): Input vector of the LSTM cell. Shape: [seq_len, batch_size, input_size] + - prev_state (:obj:`Union[None, List[Dict]]`): Previous state tensor. Shape: [num_directions*num_layers, \ + batch_size, hidden_size]. If None, prv_state will be initialized to all zeros. Returns: - - prev_state (:obj:`torch.Tensor`): batch previous state in lstm + - prev_state (:obj:`torch.Tensor`): Preprocessed previous state for the LSTM batch. """ assert hasattr(self, 'num_layers') assert hasattr(self, 'hidden_size') @@ -103,14 +107,15 @@ def _before_forward(self, inputs: torch.Tensor, prev_state: Union[None, List[Dic def _after_forward(self, next_state: Tuple[torch.Tensor], list_next_state: bool = False) -> Union[List[Dict], Dict[str, torch.Tensor]]: - r""" + """ Overview: - Post-process the next_state, return list or tensor type next_states + Post-processes the next_state after the LSTM `forward` method. Arguments: - - next_state (:obj:`Tuple[torch.Tensor]`): Tuple which contains next state (h, c). - - list_next_state (:obj:`bool`): Whether to return next_state with list format, default set to False + - next_state (:obj:`Tuple[torch.Tensor]`): Tuple containing the next state (h, c). + - list_next_state (:obj:`bool`, optional): Determines the format of the returned next_state. \ + If True, returns next_state in list format. Default is False. Returns: - - next_state(:obj:`Union[List[Dict], Dict[str, torch.Tensor]]`): The formatted next_state. + - next_state(:obj:`Union[List[Dict], Dict[str, torch.Tensor]]`): The post-processed next_state. """ if list_next_state: h, c = next_state @@ -124,15 +129,15 @@ def _after_forward(self, class LSTM(nn.Module, LSTMForwardWrapper): - r""" + """ Overview: - Implimentation of LSTM cell with LN - Interface: - forward + Implementation of an LSTM cell with Layer Normalization (LN). + Interfaces: + ``__init__``, ``forward`` .. note:: - For beginners, you can refer to to learn the basics about lstm + For a primer on LSTM, refer to https://zhuanlan.zhihu.com/p/32085405. """ def __init__( @@ -143,15 +148,15 @@ def __init__( norm_type: Optional[str] = None, dropout: float = 0. ) -> None: - r""" + """ Overview: - Initializate the LSTM cell arguments and parameters + Initialize LSTM cell parameters. Arguments: - - input_size (:obj:`int`): size of the input vector - - hidden_size (:obj:`int`): size of the hidden state vector - - num_layers (:obj:`int`): number of lstm layers - - norm_type (:obj:`Optional[str]`): type of the normaliztion, (default: None) - - dropout (:obj:`float`): dropout rate, default to 0 + - input_size (:obj:`int`): Size of the input vector. + - hidden_size (:obj:`int`): Size of the hidden state vector. + - num_layers (:obj:`int`): Number of LSTM layers. + - norm_type (:obj:`Optional[str]`): Normalization type, default is None. + - dropout (:obj:`float`): Dropout rate, default is 0. """ super(LSTM, self).__init__() self.input_size = input_size @@ -173,6 +178,11 @@ def __init__( self._init() def _init(self): + """ + Overview: + Initialize the parameters of the LSTM cell. + """ + gain = math.sqrt(1. / self.hidden_size) for l in range(self.num_layers): torch.nn.init.uniform_(self.wx[l], -gain, gain) @@ -186,15 +196,15 @@ def forward(self, list_next_state: bool = True) -> Tuple[torch.Tensor, Union[torch.Tensor, list]]: """ Overview: - Take the previous state and the input and calculate the output and the nextstate + Compute output and next state given previous state and input. Arguments: - - inputs (:obj:`torch.Tensor`): input vector of cell, tensor of size [seq_len, batch_size, input_size] - - prev_state (:obj:`torch.Tensor`): None or tensor of size \ - [num_directions*num_layers, batch_size, hidden_size] - - list_next_state (:obj:`bool`): whether return next_state with list format, default set to False + - inputs (:obj:`torch.Tensor`): Input vector of cell, size [seq_len, batch_size, input_size]. + - prev_state (:obj:`torch.Tensor`): Previous state, \ + size [num_directions*num_layers, batch_size, hidden_size]. + - list_next_state (:obj:`bool`): Whether to return next_state in list format, default is True. Returns: - - x (:obj:`torch.Tensor`): output from lstm - - next_state (:obj:`Union[torch.Tensor, list]`): hidden state from lstm + - x (:obj:`torch.Tensor`): Output from LSTM. + - next_state (:obj:`Union[torch.Tensor, list]`): Hidden state from LSTM. """ seq_len, batch_size = inputs.shape[:2] prev_state = self._before_forward(inputs, prev_state) @@ -230,15 +240,12 @@ def forward(self, class PytorchLSTM(nn.LSTM, LSTMForwardWrapper): - r""" + """ Overview: - Wrap the PyTorch nn.LSTM, format the input and output - Interface: - forward - - .. note:: - - you can reference the + Wrapper class for PyTorch's nn.LSTM, formats the input and output. For more details on nn.LSTM, + refer to https://pytorch.org/docs/stable/generated/torch.nn.LSTM.html#torch.nn.LSTM + Interfaces: + ``forward`` """ def forward(self, @@ -247,16 +254,15 @@ def forward(self, list_next_state: bool = True) -> Tuple[torch.Tensor, Union[torch.Tensor, list]]: """ Overview: - Wrapped nn.LSTM.forward. + Executes nn.LSTM.forward with preprocessed input. Arguments: - - inputs (:obj:`torch.Tensor`): input vector of cell, tensor of size \ - [seq_len, batch_size, input_size] - - prev_state (:obj:`torch.Tensor`): None or tensor of size \ - [num_directions*num_layers, batch_size, hidden_size] - - list_next_state (:obj:`bool`): whether return next_state with list format, default set to False + - inputs (:obj:`torch.Tensor`): Input vector of cell, size [seq_len, batch_size, input_size]. + - prev_state (:obj:`torch.Tensor`): Previous state, size [num_directions*num_layers, batch_size, \ + hidden_size]. + - list_next_state (:obj:`bool`): Whether to return next_state in list format, default is True. Returns: - - output (:obj:`torch.Tensor`): output from lstm - - next_state (:obj:`Union[torch.Tensor, list]`): hidden state from lstm + - output (:obj:`torch.Tensor`): Output from LSTM. + - next_state (:obj:`Union[torch.Tensor, list]`): Hidden state from LSTM. """ prev_state = self._before_forward(inputs, prev_state) output, next_state = nn.LSTM.forward(self, inputs, prev_state) @@ -265,32 +271,48 @@ def forward(self, class GRU(nn.GRUCell, LSTMForwardWrapper): - r""" + """ Overview: - Wrap the torch.nn.GRUCell, format the input and output - Interface: - forward + This class extends the `torch.nn.GRUCell` and `LSTMForwardWrapper` classes, and formats inputs and outputs + accordingly. + Interfaces: + ``__init__``, ``forward`` + Properties: + hidden_size, num_layers .. note:: - you can reference the + For further details, refer to the official PyTorch documentation: + """ - def __init__(self, input_size, hidden_size, num_layers): + def __init__(self, input_size: int, hidden_size: int, num_layers: int) -> None: + """ + Overview: + Initialize the GRU class with input size, hidden size, and number of layers. + Arguments: + - input_size (:obj:`int`): The size of the input vector. + - hidden_size (:obj:`int`): The size of the hidden state vector. + - num_layers (:obj:`int`): The number of GRU layers. + """ super(GRU, self).__init__(input_size, hidden_size) self.hidden_size = hidden_size self.num_layers = num_layers - def forward(self, inputs, prev_state, list_next_state=True): - r""" + def forward(self, + inputs: torch.Tensor, + prev_state: Optional[torch.Tensor] = None, + list_next_state: bool = True) -> Tuple[torch.Tensor, Union[torch.Tensor, List]]: + """ Overview: - Wrapped nn.GRU.forward. + Wrap the `nn.GRU.forward` method. Arguments: - - inputs (:obj:`tensor`): input vector of cell, tensor of size [seq_len, batch_size, input_size] - - prev_state (:obj:`tensor`): None or tensor of size [num_directions*num_layers, batch_size, hidden_size] - - list_next_state (:obj:`bool`): whether return next_state with list format, default set to False + - inputs (:obj:`torch.Tensor`): Input vector of cell, tensor of size [seq_len, batch_size, input_size]. + - prev_state (:obj:`Optional[torch.Tensor]`): None or tensor of \ + size [num_directions*num_layers, batch_size, hidden_size]. + - list_next_state (:obj:`bool`): Whether to return next_state in list format (default is True). Returns: - - output (:obj:`tensor`): output from GRU - - next_state (:obj:`tensor` or :obj:`list`): hidden state from GRU + - output (:obj:`torch.Tensor`): Output from GRU. + - next_state (:obj:`torch.Tensor` or :obj:`list`): Hidden state from GRU. """ # for compatibility prev_state, _ = self._before_forward(inputs, prev_state) @@ -313,20 +335,20 @@ def get_lstm( seq_len: Optional[int] = None, batch_size: Optional[int] = None ) -> Union[LSTM, PytorchLSTM]: - r""" + """ Overview: - Build and return the corresponding LSTM cell + Build and return the corresponding LSTM cell based on the provided parameters. Arguments: - - lstm_type (:obj:`str`): version of rnn cell, now support ['normal', 'pytorch', 'hpc', 'gru'] - - input_size (:obj:`int`): size of the input vector - - hidden_size (:obj:`int`): size of the hidden state vector - - num_layers (:obj:`int`): number of lstm layers - - norm_type (:obj:`str`): type of the normaliztion, (default: None) - - dropout (:obj:float): dropout rate, default set to .0 - - seq_len (:obj:`Optional[int]`): seq len, default set to None - - batch_size (:obj:`Optional[int]`): batch_size len, default set to None + - lstm_type (:obj:`str`): Version of RNN cell. Supported options are ['normal', 'pytorch', 'hpc', 'gru']. + - input_size (:obj:`int`): Size of the input vector. + - hidden_size (:obj:`int`): Size of the hidden state vector. + - num_layers (:obj:`int`): Number of LSTM layers (default is 1). + - norm_type (:obj:`str`): Type of normalization (default is 'LN'). + - dropout (:obj:`float`): Dropout rate (default is 0.0). + - seq_len (:obj:`Optional[int]`): Sequence length (default is None). + - batch_size (:obj:`Optional[int]`): Batch size (default is None). Returns: - - lstm (:obj:`Union[LSTM, PytorchLSTM]`): the corresponding lstm cell + - lstm (:obj:`Union[LSTM, PytorchLSTM]`): The corresponding LSTM cell. """ assert lstm_type in ['normal', 'pytorch', 'hpc', 'gru'] if lstm_type == 'normal': diff --git a/ding/torch_utils/network/scatter_connection.py b/ding/torch_utils/network/scatter_connection.py index e0385fa240..d596f3aa1c 100644 --- a/ding/torch_utils/network/scatter_connection.py +++ b/ding/torch_utils/network/scatter_connection.py @@ -1,15 +1,19 @@ import torch import torch.nn as nn -from typing import Tuple +from typing import Tuple, List from ding.hpc_rl import hpc_wrapper -def shape_fn_scatter_connection(args, kwargs) -> list: - r""" +def shape_fn_scatter_connection(args, kwargs) -> List[int]: + """ Overview: - Return shape of scatter_connection for hpc + Return the shape of scatter_connection for HPC. + Arguments: + - args (:obj:`Tuple`): The arguments passed to the scatter_connection function. + - kwargs (:obj:`Dict`): The keyword arguments passed to the scatter_connection function. Returns: - - shape (:obj:`list`): List like [B, M, N, H, W, scatter_type] + - shape (:obj:`List[int]`): A list representing the shape of scatter_connection, \ + in the form of [B, M, N, H, W, scatter_type]. """ if len(args) <= 1: tmp = list(kwargs['x'].shape) @@ -24,20 +28,22 @@ def shape_fn_scatter_connection(args, kwargs) -> list: class ScatterConnection(nn.Module): - r""" + """ Overview: - Scatter feature to its corresponding location - In AlphaStar, each entity is embedded into a tensor, + Scatter feature to its corresponding location. In AlphaStar, each entity is embedded into a tensor, and these tensors are scattered into a feature map with map size. + Interfaces: + ``__init__``, ``forward``, ``xy_forward`` """ def __init__(self, scatter_type: str) -> None: - r""" + """ Overview: - Init class + Initialize the ScatterConnection object. Arguments: - - scatter_type (:obj:`str`): Supports ['add', 'cover']. If two entities have the same location, \ - scatter_type decides the first one should be covered or added to second one + - scatter_type (:obj:`str`): The scatter type, which decides the behavior when two entities have the \ + same location. It can be either 'add' or 'cover'. If 'add', the first one will be added to the \ + second one. If 'cover', the first one will be covered by the second one. """ super(ScatterConnection, self).__init__() self.scatter_type = scatter_type @@ -53,41 +59,63 @@ def __init__(self, scatter_type: str) -> None: def forward(self, x: torch.Tensor, spatial_size: Tuple[int, int], location: torch.Tensor) -> torch.Tensor: """ Overview: - scatter x into a spatial feature map + Scatter input tensor 'x' into a spatial feature map. Arguments: - - x (:obj:`tensor`): input tensor :math: `(B, M, N)` where `M` means the number of entity, `N` means \ - the dimension of entity attributes - - spatial_size (:obj:`tuple`): Tuple[H, W], the size of spatial feature x will be scattered into - - location (:obj:`tensor`): :math: `(B, M, 2)` torch.LongTensor, each location should be (y, x) + - x (:obj:`torch.Tensor`): The input tensor of shape `(B, M, N)`, where `B` is the batch size, `M` \ + is the number of entities, and `N` is the dimension of entity attributes. + - spatial_size (:obj:`Tuple[int, int]`): The size `(H, W)` of the spatial feature map into which 'x' \ + will be scattered, where `H` is the height and `W` is the width. + - location (:obj:`torch.Tensor`): The tensor of locations of shape `(B, M, 2)`. \ + Each location should be (y, x). Returns: - - output (:obj:`tensor`): :math: `(B, N, H, W)` where `H` and `W` are spatial_size, return the\ - scattered feature map - Shapes: - - Input: :math: `(B, M, N)` where `M` means the number of entity, `N` means \ - the dimension of entity attributes - - Size: Tuple type :math: `[H, W]` - - Location: :math: `(B, M, 2)` torch.LongTensor, each location should be (y, x) - - Output: :math: `(B, N, H, W)` where `H` and `W` are spatial_size - - .. note:: + - output (:obj:`torch.Tensor`): The scattered feature map of shape `(B, N, H, W)`. + Note: + When there are some overlapping in locations, 'cover' mode will result in the loss of information. + 'add' mode is used as a temporary substitute. + """ + device = x.device + B, M, N = x.shape + x = x.permute(0, 2, 1) + H, W = spatial_size + index = location[:, :, 1] + location[:, :, 0] * W + index = index.unsqueeze(dim=1).repeat(1, N, 1) + output = torch.zeros(size=(B, N, H, W), device=device).view(B, N, H * W) + if self.scatter_type == 'cover': + output.scatter_(dim=2, index=index, src=x) + elif self.scatter_type == 'add': + output.scatter_add_(dim=2, index=index, src=x) + output = output.view(B, N, H, W) + return output - When there are some overlapping in locations, ``cover`` mode will result in the loss of information, we - use the addition as temporal substitute. + def xy_forward( + self, x: torch.Tensor, spatial_size: Tuple[int, int], coord_x: torch.Tensor, coord_y + ) -> torch.Tensor: + """ + Overview: + Scatter input tensor 'x' into a spatial feature map using separate x and y coordinates. + Arguments: + - x (:obj:`torch.Tensor`): The input tensor of shape `(B, M, N)`, where `B` is the batch size, `M` \ + is the number of entities, and `N` is the dimension of entity attributes. + - spatial_size (:obj:`Tuple[int, int]`): The size `(H, W)` of the spatial feature map into which 'x' \ + will be scattered, where `H` is the height and `W` is the width. + - coord_x (:obj:`torch.Tensor`): The x-coordinates tensor of shape `(B, M)`. + - coord_y (:obj:`torch.Tensor`): The y-coordinates tensor of shape `(B, M)`. + Returns: + - output (:obj:`torch.Tensor`): The scattered feature map of shape `(B, N, H, W)`. + Note: + When there are some overlapping in locations, 'cover' mode will result in the loss of information. + 'add' mode is used as a temporary substitute. """ device = x.device B, M, N = x.shape + x = x.permute(0, 2, 1) H, W = spatial_size - index = location.view(-1, 2) - bias = torch.arange(B).mul_(H * W).unsqueeze(1).repeat(1, M).view(-1).to(device) - index = index[:, 0] * W + index[:, 1] - index += bias - index = index.repeat(N, 1) - x = x.view(-1, N).permute(1, 0) - output = torch.zeros(N, B * H * W, device=device) + index = (coord_x * W + coord_y).long() + index = index.unsqueeze(dim=1).repeat(1, N, 1) + output = torch.zeros(size=(B, N, H, W), device=device).view(B, N, H * W) if self.scatter_type == 'cover': - output.scatter_(dim=1, index=index, src=x) + output.scatter_(dim=2, index=index, src=x) elif self.scatter_type == 'add': - output.scatter_add_(dim=1, index=index, src=x) - output = output.reshape(N, B, H, W) - output = output.permute(1, 0, 2, 3).contiguous() + output.scatter_add_(dim=2, index=index, src=x) + output = output.view(B, N, H, W) return output diff --git a/ding/torch_utils/network/soft_argmax.py b/ding/torch_utils/network/soft_argmax.py index 458e3fb3b5..166d0bb8f6 100644 --- a/ding/torch_utils/network/soft_argmax.py +++ b/ding/torch_utils/network/soft_argmax.py @@ -4,49 +4,56 @@ class SoftArgmax(nn.Module): - r""" + """ Overview: - An nn.Module that computes SoftArgmax - Interface: - __init__, forward - - .. note: - - For more softargmax info, you can refere to - the wiki page or - the lecture + A neural network module that computes the SoftArgmax operation (essentially a 2-dimensional spatial softmax), + which is often used for location regression tasks. It converts a feature map (such as a heatmap) into precise + coordinate locations. + Interfaces: + ``__init__``, ``forward`` + .. note:: + For more information on SoftArgmax, you can refer to + and the paper . """ def __init__(self): - r""" + """ Overview: - Initialize the SoftArgmax module + Initialize the SoftArgmax module. """ super(SoftArgmax, self).__init__() def forward(self, x: torch.Tensor) -> torch.Tensor: - r""" + """ Overview: - Soft-argmax for location regression + Perform the forward pass of the SoftArgmax operation. Arguments: - - x (:obj:`torch.Tensor`): predict heat map + - x (:obj:`torch.Tensor`): The input tensor, typically a heatmap representing predicted locations. Returns: - - location (:obj:`torch.Tensor`): predict location + - location (:obj:`torch.Tensor`): The predicted coordinates as a result of the SoftArgmax operation. Shapes: - - x: :math:`(B, C, H, W)`, while B is the batch size, C is number of channels, \ - H and W stands for height and width - - location: :math:`(B, 2)`, while B is the batch size + - x: :math:`(B, C, H, W)`, where `B` is the batch size, `C` is the number of channels, \ + and `H` and `W` represent height and width respectively. + - location: :math:`(B, 2)`, where `B` is the batch size and 2 represents the coordinates (height, width). """ + # Unpack the dimensions of the input tensor B, C, H, W = x.shape device, dtype = x.device, x.dtype - # 1 channel - assert (x.shape[1] == 1) + # Ensure the input tensor has a single channel + assert C == 1, "Input tensor should have only one channel" + # Create a meshgrid for the height (h_kernel) and width (w_kernel) h_kernel = torch.arange(0, H, device=device).to(dtype) h_kernel = h_kernel.view(1, 1, H, 1).repeat(1, 1, 1, W) + w_kernel = torch.arange(0, W, device=device).to(dtype) w_kernel = w_kernel.view(1, 1, 1, W).repeat(1, 1, H, 1) + + # Apply the softmax function across the spatial dimensions (height and width) x = F.softmax(x.view(B, C, -1), dim=-1).view(B, C, H, W) - h = (x * h_kernel).sum(dim=[1, 2, 3]) - w = (x * w_kernel).sum(dim=[1, 2, 3]) + # Compute the expected values for height and width by multiplying the probability map by the meshgrids + h = (x * h_kernel).sum(dim=[1, 2, 3]) # Sum over the channel, height, and width dimensions + w = (x * w_kernel).sum(dim=[1, 2, 3]) # Sum over the channel, height, and width dimensions + + # Stack the height and width coordinates along a new dimension to form the final output tensor return torch.stack([h, w], dim=1) diff --git a/ding/torch_utils/network/tests/test_diffusion.py b/ding/torch_utils/network/tests/test_diffusion.py new file mode 100644 index 0000000000..f9794a726b --- /dev/null +++ b/ding/torch_utils/network/tests/test_diffusion.py @@ -0,0 +1,29 @@ +import pytest +import torch +from ding.torch_utils.network.diffusion import DiffusionUNet1d, TemporalValue + +batch_size = 2 +transition_dim = 10 +dim = 8 +dim_mults = [1, 2, 4] +horizon = 4 + + +@pytest.mark.unittest +class TestDiffusionNet: + + def test_DiffusionNet1d(self): + diffusion = DiffusionUNet1d(transition_dim, dim, dim_mults) + input = torch.rand(batch_size, horizon, transition_dim) + t = torch.randint(0, horizon, (batch_size, )).long() + cond = {t: torch.randn(batch_size, 2) for t in range(horizon)} + output = diffusion(input, cond, time=t) + assert output.shape == (batch_size, horizon, transition_dim) + + def test_TemporalValue(self): + value = TemporalValue(horizon, transition_dim, dim, dim_mults=dim_mults) + input = torch.rand(batch_size, horizon, transition_dim) + t = torch.randint(0, horizon, (batch_size, )).long() + cond = {t: torch.randn(batch_size, 2) for t in range(horizon)} + output = value(input, cond, time=t) + assert output.shape == (batch_size, 1) diff --git a/ding/torch_utils/network/tests/test_dreamer.py b/ding/torch_utils/network/tests/test_dreamer.py new file mode 100644 index 0000000000..accfb1d8c5 --- /dev/null +++ b/ding/torch_utils/network/tests/test_dreamer.py @@ -0,0 +1,73 @@ +import pytest +from easydict import EasyDict +import torch +from torch import distributions as torchd +from itertools import product +from ding.torch_utils.network.dreamer import DenseHead, SampleDist, OneHotDist, TwoHotDistSymlog, \ + SymlogDist, ContDist, Bernoulli, UnnormalizedHuber, weight_init, uniform_weight_init + +# arguments +shape = [255, (255, ), ()] +# to do +# dist = ['normal', 'huber', 'binary', 'twohot_symlog'] +dist = ['twohot_symlog'] +args = list(product(*[shape, dist])) + + +@pytest.mark.unittest +@pytest.mark.parametrize('shape, dist', args) +def test_DenseHead(shape, dist): + in_dim, layer_num, units, B, time = 1536, 2, 512, 16, 64 + head = DenseHead(in_dim, shape, layer_num, units, dist=dist) + x = torch.randn(B, time, in_dim) + a = torch.randn(B, time, 1) + y = head(x) + assert y.mode().shape == (B, time, 1) + assert y.log_prob(a).shape == (B, time) + + +B, time = 16, 64 +mean = torch.randn(B, time, 255) +std = 1.0 +a = torch.randn(B, time, 1) # or torch.randn(B, time, 255) +sample_shape = torch.Size([]) + + +@pytest.mark.unittest +def test_ContDist(): + dist_origin = torchd.normal.Normal(mean, std) + dist = torchd.independent.Independent(dist_origin, 1) + dist_new = ContDist(dist) + assert dist_new.mode().shape == (B, time, 255) + assert dist_new.log_prob(a).shape == (B, time) + assert dist_origin.log_prob(a).shape == (B, time, 255) + assert dist_new.sample().shape == (B, time, 255) + + +@pytest.mark.unittest +def test_UnnormalizedHuber(): + dist_origin = UnnormalizedHuber(mean, std) + dist = torchd.independent.Independent(dist_origin, 1) + dist_new = ContDist(dist) + assert dist_new.mode().shape == (B, time, 255) + assert dist_new.log_prob(a).shape == (B, time) + assert dist_origin.log_prob(a).shape == (B, time, 255) + assert dist_new.sample().shape == (B, time, 255) + + +@pytest.mark.unittest +def test_Bernoulli(): + dist_origin = torchd.bernoulli.Bernoulli(logits=mean) + dist = torchd.independent.Independent(dist_origin, 1) + dist_new = Bernoulli(dist) + assert dist_new.mode().shape == (B, time, 255) + assert dist_new.log_prob(a).shape == (B, time, 255) + # to do + # assert dist_new.sample().shape == (B, time, 255) + + +@pytest.mark.unittest +def test_TwoHotDistSymlog(): + dist = TwoHotDistSymlog(logits=mean) + assert dist.mode().shape == (B, time, 1) + assert dist.log_prob(a).shape == (B, time) diff --git a/ding/torch_utils/network/tests/test_merge.py b/ding/torch_utils/network/tests/test_merge.py new file mode 100644 index 0000000000..58f2ed34e4 --- /dev/null +++ b/ding/torch_utils/network/tests/test_merge.py @@ -0,0 +1,46 @@ +import pytest +import torch +from ding.torch_utils.network import GatingType, SumMerge, VectorMerge + + +@pytest.mark.unittest +def test_SumMerge(): + input_shape = (3, 5) + input = [torch.rand(input_shape).requires_grad_(True) for i in range(4)] + sum_merge = SumMerge() + + output = sum_merge(input) + assert output.shape == (3, 5) + loss = output.mean() + loss.backward() + assert isinstance(input[0].grad, torch.Tensor) + + +@pytest.mark.unittest +def test_VectorMerge(): + input_sizes = {'in1': 3, 'in2': 16, 'in3': 27} + output_size = 512 + input_dict = {} + for k, v in input_sizes.items(): + input_dict[k] = torch.rand((64, v)).requires_grad_(True) + + vector_merge = VectorMerge(input_sizes, output_size, GatingType.NONE) + output = vector_merge(input_dict) + assert output.shape == (64, output_size) + loss = output.mean() + loss.backward() + assert isinstance(input_dict['in1'].grad, torch.Tensor) + + vector_merge = VectorMerge(input_sizes, output_size, GatingType.GLOBAL) + output = vector_merge(input_dict) + assert output.shape == (64, output_size) + loss = output.mean() + loss.backward() + assert isinstance(input_dict['in1'].grad, torch.Tensor) + + vector_merge = VectorMerge(input_sizes, output_size, GatingType.POINTWISE) + output = vector_merge(input_dict) + assert output.shape == (64, output_size) + loss = output.mean() + loss.backward() + assert isinstance(input_dict['in1'].grad, torch.Tensor) diff --git a/ding/torch_utils/network/tests/test_nn_module.py b/ding/torch_utils/network/tests/test_nn_module.py index 43eb49db54..ac44e15cba 100644 --- a/ding/torch_utils/network/tests/test_nn_module.py +++ b/ding/torch_utils/network/tests/test_nn_module.py @@ -1,11 +1,15 @@ -import torch import pytest -from ding.torch_utils import build_activation, build_normalization -from ding.torch_utils.network.nn_module import conv1d_block, conv2d_block, fc_block, deconv2d_block, ChannelShuffle, \ - one_hot, NearestUpsample, BilinearUpsample, binary_encode, weight_init_, NaiveFlatten, normed_linear, normed_conv2d +import torch +from torch.testing import assert_allclose + +from ding.torch_utils import build_activation +from ding.torch_utils.network.nn_module import MLP, conv1d_block, conv2d_block, fc_block, deconv2d_block, \ + ChannelShuffle, one_hot, NearestUpsample, BilinearUpsample, binary_encode, weight_init_, NaiveFlatten, \ + normed_linear, normed_conv2d, NoiseLinearLayer batch_size = 2 in_channels = 2 +hidden_channels = 3 out_channels = 3 H = 2 W = 3 @@ -41,6 +45,50 @@ def test_weight_init(self): with pytest.raises(KeyError): weight_init_(weight, 'xxx') + def test_mlp(self): + layer_num = 3 + input_tensor = torch.rand(batch_size, in_channels).requires_grad_(True) + + for output_activation in [True, False]: + for output_norm in [True, False]: + for activation in [torch.nn.ReLU(), torch.nn.LeakyReLU(), torch.nn.Tanh(), None]: + for norm_type in ["LN", "BN", None]: + # Test case 1: MLP without last linear layer initialized to 0. + model = MLP( + in_channels, + hidden_channels, + out_channels, + layer_num, + activation=activation, + norm_type=norm_type, + output_activation=output_activation, + output_norm=output_norm + ) + output_tensor = self.run_model(input_tensor, model) + assert output_tensor.shape == (batch_size, out_channels) + + # Test case 2: MLP with last linear layer initialized to 0. + model = MLP( + in_channels, + hidden_channels, + out_channels, + layer_num, + activation=activation, + norm_type=norm_type, + output_activation=output_activation, + output_norm=output_norm, + last_linear_layer_init_zero=True + ) + output_tensor = self.run_model(input_tensor, model) + assert output_tensor.shape == (batch_size, out_channels) + last_linear_layer = None + for layer in reversed(model): + if isinstance(layer, torch.nn.Linear): + last_linear_layer = layer + break + assert_allclose(last_linear_layer.weight, torch.zeros_like(last_linear_layer.weight)) + assert_allclose(last_linear_layer.bias, torch.zeros_like(last_linear_layer.bias)) + def test_conv1d_block(self): length = 2 input = torch.rand(batch_size, in_channels, length).requires_grad_(True) @@ -190,3 +238,27 @@ def test_flatten(self): model3 = NaiveFlatten(1, 3) output3 = model2(inputs) assert output1.shape == (4, 3 * 8 * 8) + + def test_noise_linear_layer(self): + input = torch.rand(batch_size, in_channels).requires_grad_(True) + layer = NoiseLinearLayer(in_channels, out_channels, sigma0=0.5) + # No noise by default + output = self.run_model(input, layer) + assert output.shape == (batch_size, out_channels) + # Enable noise + layer.enable_noise = True + layer.reset_noise() + output_noise = self.run_model(input, layer) + assert output_noise.shape == (batch_size, out_channels) + # Check that outputs are different after resetting noise + with torch.no_grad(): + layer.reset_noise() + out1 = layer(input) + layer.reset_noise() + out2 = layer(input) + # The outputs should be different (very likely) + assert not torch.allclose(out1, out2) + # Check reset_parameters + layer.reset_parameters() + assert layer.weight_mu.shape == (out_channels, in_channels) + assert layer.bias_mu.shape == (out_channels, ) diff --git a/ding/torch_utils/network/tests/test_popart.py b/ding/torch_utils/network/tests/test_popart.py new file mode 100644 index 0000000000..8ea694f49d --- /dev/null +++ b/ding/torch_utils/network/tests/test_popart.py @@ -0,0 +1,36 @@ +import pytest +import torch +from ding.torch_utils import PopArt + +batch_size = 4 +input_features = 16 +output_features = 4 + + +@pytest.mark.unittest +class TestPopArt: + + def test_popart(self): + input = torch.rand((batch_size, input_features)).requires_grad_(True) + model = PopArt(input_features, output_features) + output = model(input) + loss = output['pred'].mean() + loss.backward() + assert isinstance(input.grad, torch.Tensor) + + # validate the shape of parameters and outputs + assert output['pred'].shape == (batch_size, output_features) + assert output['unnormalized_pred'].shape == (batch_size, output_features) + assert model.mu.shape == torch.Size([output_features]) + assert model.sigma.shape == torch.Size([output_features]) + assert model.v.shape == torch.Size([output_features]) + + # validate the normalization + assert torch.all(torch.abs(output['pred']) <= 1) + + model.update_parameters(torch.rand(batch_size, output_features)) + + # validate the non-empty of parameters + assert not torch.all(torch.isnan(model.mu)) + assert not torch.all(torch.isnan(model.sigma)) + assert not torch.all(torch.isnan(model.v)) diff --git a/ding/torch_utils/network/tests/test_res_block.py b/ding/torch_utils/network/tests/test_res_block.py index 2fd00fc6dc..26b0b946b3 100644 --- a/ding/torch_utils/network/tests/test_res_block.py +++ b/ding/torch_utils/network/tests/test_res_block.py @@ -7,7 +7,8 @@ H, W = 2, 3 activation = torch.nn.ReLU() norm_type = 'BN' -res_type = ['basic', 'bottleneck'] +res_type = ['basic', 'bottleneck', 'downsample'] +res_type_classic = ['basic', 'bottleneck'] @pytest.mark.unittest @@ -16,12 +17,14 @@ class TestResBlock: def test_res_blcok(self): input = torch.rand(batch_size, in_channels, 2, 3).requires_grad_(True) for r in res_type: - model = ResBlock(in_channels, activation, norm_type, r) - output = model(input) - loss = output.mean() - loss.backward() - assert output.shape == input.shape - assert isinstance(input.grad, torch.Tensor) + for norm_type in ['BN', 'LN', 'IN', 'GN', None]: + model = ResBlock(in_channels, activation, norm_type, r) + output = model(input) + loss = output.mean() + loss.backward() + if r in res_type_classic: + assert output.shape == input.shape + assert isinstance(input.grad, torch.Tensor) def test_res_fc_block(self): input = torch.rand(batch_size, in_channels).requires_grad_(True) diff --git a/ding/torch_utils/network/tests/test_scatter.py b/ding/torch_utils/network/tests/test_scatter.py index cfac0d883c..a79e107376 100644 --- a/ding/torch_utils/network/tests/test_scatter.py +++ b/ding/torch_utils/network/tests/test_scatter.py @@ -29,3 +29,17 @@ def test_naive(self): loss = output.mean() loss.backward() assert isinstance(input.grad, torch.Tensor) + + def test_xy_forward(self): + for scatter_type in ['add', 'cover']: + model = ScatterConnection(scatter_type) + B, M, N = 10, 20, 3 + spatial_size = (13, 17) + input = torch.randn(size=(B, M, N)).requires_grad_(True) + coord_x = torch.randint(low=0, high=13, size=(B, M)) + coord_y = torch.randint(low=0, high=17, size=(B, M)) + output = model.xy_forward(input, spatial_size, coord_x, coord_y) + loss = output.mean() + loss.backward() + assert isinstance(input.grad, torch.Tensor) + assert output.shape == (B, N, *spatial_size) diff --git a/ding/torch_utils/network/transformer.py b/ding/torch_utils/network/transformer.py index e707134a3f..14943e76ba 100644 --- a/ding/torch_utils/network/transformer.py +++ b/ding/torch_utils/network/transformer.py @@ -8,23 +8,23 @@ class Attention(nn.Module): - r""" + """ Overview: - For each entry embedding, compute individual attention across all entries, add them up to get output attention + For each entry embedding, compute individual attention across all entries, add them up to get output attention. Interfaces: - split, forward + ``__init__``, ``split``, ``forward`` """ def __init__(self, input_dim: int, head_dim: int, output_dim: int, head_num: int, dropout: nn.Module) -> None: - r""" + """ Overview: - Init attention + Initialize the Attention module with the provided dimensions and dropout layer. Arguments: - - input_dim (:obj:`int`): dimension of input - - head_dim (:obj:`int`): dimension of each head - - output_dim (:obj:`int`): dimension of output - - head_num (:obj:`int`): head num for multihead attention - - dropout (:obj:`nn.Module`): dropout layer + - input_dim (:obj:`int`): The dimension of the input. + - head_dim (:obj:`int`): The dimension of each head in the multi-head attention mechanism. + - output_dim (:obj:`int`): The dimension of the output. + - head_num (:obj:`int`): The number of heads in the multi-head attention mechanism. + - dropout (:obj:`nn.Module`): The dropout layer used in the attention mechanism. """ super(Attention, self).__init__() self.head_num = head_num @@ -34,14 +34,14 @@ def __init__(self, input_dim: int, head_dim: int, output_dim: int, head_num: int self.project = fc_block(head_dim * head_num, output_dim) def split(self, x: torch.Tensor, T: bool = False) -> List[torch.Tensor]: - r""" + """ Overview: - Split input to get multihead queries, keys, values + Split the input to get multi-head queries, keys, and values. Arguments: - - x (:obj:`torch.Tensor`): query or key or value - - T (:obj:`bool`): whether to transpose output + - x (:obj:`torch.Tensor`): The tensor to be split, which could be a query, key, or value. + - T (:obj:`bool`, optional): If True, transpose the output tensors. Defaults to False. Returns: - - x (:obj:`List[torch.Tensor]`): list of output tensors for each head + - x (:obj:`List[torch.Tensor]`): A list of output tensors for each head. """ B, N = x.shape[:2] x = x.view(B, N, self.head_num, self.head_dim) @@ -51,14 +51,15 @@ def split(self, x: torch.Tensor, T: bool = False) -> List[torch.Tensor]: return x def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor: - r""" + """ Overview: - Compute attention + Compute the attention from the input tensor. Arguments: - - x (:obj:`torch.Tensor`): input tensor - - mask (:obj:`Optional[torch.Tensor]`): mask out invalid entries + - x (:obj:`torch.Tensor`): The input tensor for the forward computation. + - mask (:obj:`Optional[torch.Tensor]`, optional): Optional mask to exclude invalid entries. + Defaults to None. Returns: - - attention (:obj:`torch.Tensor`): attention tensor + - attention (:obj:`torch.Tensor`): The computed attention tensor. """ assert (len(x.shape) == 3) B, N = x.shape[:2] @@ -82,27 +83,29 @@ def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch class TransformerLayer(nn.Module): - r""" + """ Overview: - In transformer layer, first computes entries's attention and applies a feedforward layer + In transformer layer, first computes entries's attention and applies a feedforward layer. + Interfaces: + ``__init__``, ``forward`` """ def __init__( self, input_dim: int, head_dim: int, hidden_dim: int, output_dim: int, head_num: int, mlp_num: int, dropout: nn.Module, activation: nn.Module ) -> None: - r""" + """ Overview: - Init transformer layer + Initialize the TransformerLayer with the provided dimensions, dropout layer, and activation function. Arguments: - - input_dim (:obj:`int`): dimension of input - - head_dim (:obj:`int`): dimension of each head - - hidden_dim (:obj:`int`): dimension of hidden layer in mlp - - output_dim (:obj:`int`): dimension of output - - head_num (:obj:`int`): number of heads for multihead attention - - mlp_num (:obj:`int`): number of mlp layers - - dropout (:obj:`nn.Module`): dropout layer - - activation (:obj:`nn.Module`): activation function + - input_dim (:obj:`int`): The dimension of the input. + - head_dim (:obj:`int`): The dimension of each head in the multi-head attention mechanism. + - hidden_dim (:obj:`int`): The dimension of the hidden layer in the MLP (Multi-Layer Perceptron). + - output_dim (:obj:`int`): The dimension of the output. + - head_num (:obj:`int`): The number of heads in the multi-head attention mechanism. + - mlp_num (:obj:`int`): The number of layers in the MLP. + - dropout (:obj:`nn.Module`): The dropout layer used in the attention mechanism. + - activation (:obj:`nn.Module`): The activation function used in the MLP. """ super(TransformerLayer, self).__init__() self.attention = Attention(input_dim, head_dim, output_dim, head_num, dropout) @@ -121,29 +124,33 @@ def __init__( def forward(self, inputs: Tuple[torch.Tensor, torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]: """ Overview: - Transformer layer forward + Compute the forward pass through the Transformer layer. Arguments: - - inputs (:obj:`Tuple[torch.Tensor, torch.Tensor]`): x and mask + - inputs (:obj:`Tuple[torch.Tensor, torch.Tensor]`): A tuple containing the input tensor `x` and + the mask tensor. Returns: - - output (:obj:`Tuple[torch.Tensor, torch.Tensor]`): predict value and mask + - output (:obj:`Tuple[torch.Tensor, torch.Tensor]`): A tuple containing the predicted value tensor and + the mask tensor. """ x, mask = inputs a = self.dropout(self.attention(x, mask)) x = self.layernorm1(x + a) m = self.dropout(self.mlp(x)) x = self.layernorm2(x + m) - return (x, mask) + return x, mask class Transformer(nn.Module): - ''' + """ Overview: - Transformer implementation + Implementation of the Transformer model. .. note:: + For more details, refer to "Attention is All You Need": http://arxiv.org/abs/1706.03762. - For details refer to Attention is all you need: http://arxiv.org/abs/1706.03762 - ''' + Interfaces: + ``__init__``, ``forward`` + """ def __init__( self, @@ -157,19 +164,20 @@ def __init__( dropout_ratio: float = 0., activation: nn.Module = nn.ReLU(), ): - r""" + """ Overview: - Init transformer + Initialize the Transformer with the provided dimensions, dropout layer, activation function, + and layer numbers. Arguments: - - input_dim (:obj:`int`): dimension of input - - head_dim (:obj:`int`): dimension of each head - - hidden_dim (:obj:`int`): dimension of hidden layer in mlp - - output_dim (:obj:`int`): dimension of output - - head_num (:obj:`int`): number of heads for multihead attention - - mlp_num (:obj:`int`): number of mlp layers - - layer_num (:obj:`int`): number of transformer layers - - dropout_ratio (:obj:`float`): dropout ratio - - activation (:obj:`nn.Module`): activation function + - input_dim (:obj:`int`): The dimension of the input. + - head_dim (:obj:`int`): The dimension of each head in the multi-head attention mechanism. + - hidden_dim (:obj:`int`): The dimension of the hidden layer in the MLP (Multi-Layer Perceptron). + - output_dim (:obj:`int`): The dimension of the output. + - head_num (:obj:`int`): The number of heads in the multi-head attention mechanism. + - mlp_num (:obj:`int`): The number of layers in the MLP. + - layer_num (:obj:`int`): The number of Transformer layers. + - dropout_ratio (:obj:`float`): The dropout ratio for the dropout layer. + - activation (:obj:`nn.Module`): The activation function used in the MLP. """ super(Transformer, self).__init__() self.embedding = fc_block(input_dim, output_dim, activation=activation) @@ -184,16 +192,17 @@ def __init__( self.main = nn.Sequential(*layers) def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor: - r""" + """ Overview: - Transformer forward + Perform the forward pass through the Transformer. Arguments: - - x (:obj:`torch.Tensor`): input tensor. Shape (B, N, C), B is batch size, \ - N is number of entries, C is feature dimension - - mask (:obj:`Optional[torch.Tensor]`): bool tensor, can be used to mask out invalid entries in attention. \ - Shape (B, N), B is batch size, N is number of entries + - x (:obj:`torch.Tensor`): The input tensor, with shape `(B, N, C)`, where `B` is batch size, \ + `N` is the number of entries, and `C` is the feature dimension. + - mask (:obj:`Optional[torch.Tensor]`, optional): The mask tensor (bool), used to mask out invalid \ + entries in attention. It has shape `(B, N)`, where `B` is batch size and `N` is number of \ + entries. Defaults to None. Returns: - - x (:obj:`torch.Tensor`): transformer output + - x (:obj:`torch.Tensor`): The output tensor from the Transformer. """ if mask is not None: mask = mask.unsqueeze(dim=1).repeat(1, mask.shape[1], 1).unsqueeze(dim=1) @@ -204,12 +213,25 @@ def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch class ScaledDotProductAttention(nn.Module): - ''' + """ Overview: - Implementation of dot product attentionn with scaling. - ''' + Implementation of Scaled Dot Product Attention, a key component of Transformer models. + This class performs the dot product of the query, key and value tensors, scales it with the square root of the + dimension of the key vector (d_k) and applies dropout for regularization. + Interfaces: + ``__init__``, ``forward`` + """ def __init__(self, d_k: int, dropout: float = 0.0) -> None: + """ + Overview: + Initialize the ScaledDotProductAttention module with the dimension of the key vector and the dropout rate. + Arguments: + - d_k (:obj:`int`): The dimension of the key vector. This will be used to scale the dot product of the \ + query and key. + - dropout (:obj:`float`, optional): The dropout rate to be applied after the softmax operation. \ + Defaults to 0.0. + """ super(ScaledDotProductAttention, self).__init__() self.d_k = d_k self.dropout = nn.Dropout(dropout) @@ -221,6 +243,18 @@ def forward( v: torch.Tensor, mask: Optional[torch.Tensor] = None ) -> torch.Tensor: + """ + Overview: + Perform the Scaled Dot Product Attention operation on the query, key and value tensors. + Arguments: + - q (:obj:`torch.Tensor`): The query tensor. + - k (:obj:`torch.Tensor`): The key tensor. + - v (:obj:`torch.Tensor`): The value tensor. + - mask (:obj:`Optional[torch.Tensor]`): An optional mask tensor to be applied on the attention scores. + Defaults to None. + Returns: + - output (:obj:`torch.Tensor`): The output tensor after the attention operation. + """ attn = torch.matmul(q / (self.d_k ** 0.5), k.transpose(2, 3)) if mask is not None: # inplace modification for reasonable softmax diff --git a/ding/torch_utils/optimizer_helper.py b/ding/torch_utils/optimizer_helper.py index cbd5b3bedd..ea3f7b0a73 100644 --- a/ding/torch_utils/optimizer_helper.py +++ b/ding/torch_utils/optimizer_helper.py @@ -1,8 +1,7 @@ import torch import math from torch.nn.utils import clip_grad_norm_, clip_grad_value_ -from torch._six import inf -from typing import Union, Iterable, Tuple, Callable +from typing import Union, Iterable, Tuple, Callable, List import torch.nn as nn import torch.nn.functional as F import torch.optim as optim @@ -11,9 +10,11 @@ import copy import random +inf = math.inf + def calculate_grad_norm(model: torch.nn.Module, norm_type=2) -> float: - r""" + """ Overview: calculate grad norm of the parameters whose grad norms are not None in the model. Arguments: @@ -37,7 +38,7 @@ def calculate_grad_norm(model: torch.nn.Module, norm_type=2) -> float: def calculate_grad_norm_without_bias_two_norm(model: torch.nn.Module) -> float: - r""" + """ Overview: calculate grad norm of the parameters whose grad norms are not None in the model. Arguments: @@ -53,6 +54,14 @@ def calculate_grad_norm_without_bias_two_norm(model: torch.nn.Module) -> float: def grad_ignore_norm(parameters, max_norm, norm_type=2): + """ + Overview: + Clip the gradient norm of an iterable of parameters. + Arguments: + - parameters (:obj:`Iterable`): an iterable of torch.Tensor + - max_norm (:obj:`float`): the max norm of the gradients + - norm_type (:obj:`float`): 2.0 means use norm2 to clip + """ if isinstance(parameters, torch.Tensor): parameters = [parameters] parameters = list(filter(lambda p: p.grad is not None, parameters)) @@ -74,6 +83,13 @@ def grad_ignore_norm(parameters, max_norm, norm_type=2): def grad_ignore_value(parameters, clip_value): + """ + Overview: + Clip the gradient value of an iterable of parameters. + Arguments: + - parameters (:obj:`Iterable`): an iterable of torch.Tensor + - clip_value (:obj:`float`): the value to start clipping + """ if isinstance(parameters, torch.Tensor): parameters = [parameters] clip_value = float(clip_value) @@ -89,11 +105,11 @@ def grad_ignore_value(parameters, clip_value): class Adam(torch.optim.Adam): - r""" + """ Overview: Rewrited Adam optimizer to support more features. - Interface: - __init__, step + Interfaces: + ``__init__``, ``step``, ``_state_init``, ``get_grad`` """ def __init__( @@ -117,9 +133,9 @@ def __init__( ignore_norm_type: float = 2.0, ignore_momentum_timestep: int = 100, ): - r""" + """ Overview: - init medthod of refactored Adam class + init method of refactored Adam class Arguments: - params (:obj:`iterable`): – an iterable of torch.Tensor s or dict s. \ Specifies what Tensors should be optimized @@ -188,15 +204,23 @@ def __init__( ) def _state_init(self, p, amsgrad): + """ + Overview: + Initialize the state of the optimizer + Arguments: + - p (:obj:`torch.Tensor`): the parameter to be optimized + - amsgrad (:obj:`bool`): whether to use the AMSGrad variant of this algorithm from the paper\ + On the Convergence of Adam and Beyond + """ state = self.state[p] state['thre_exp_avg_sq'] = torch.zeros_like(p.data, device=p.data.device) # others if torch.__version__ < "1.12.0": state['step'] = 0 - #TODO - #wait torch upgrad to 1.4, 1.3.1 didn't support memory format state['step'] = 0 + # TODO + # wait torch upgrad to 1.4, 1.3.1 didn't support memory format state['step'] = 0 else: - state['step'] = torch.zeros((1, ), dtype=torch.float, device=p.device) \ + state['step'] = torch.zeros((1,), dtype=torch.float, device=p.device) \ if self.defaults['capturable'] else torch.tensor(0.) state['exp_avg'] = torch.zeros_like(p.data) @@ -207,7 +231,7 @@ def _state_init(self, p, amsgrad): state['max_exp_avg_sq'] = torch.zeros_like(p.data) def step(self, closure: Union[Callable, None] = None): - r""" + """ Overview: Performs a single optimization step Arguments: @@ -235,7 +259,7 @@ def step(self, closure: Union[Callable, None] = None): if len(state) == 0: self._state_init(p, group['amsgrad']) grad = p.grad.data - #should we use same beta group? + # should we use same beta group? beta1, beta2 = group['betas'] bias_correction2 = 1 - beta2 ** state['step'] state['thre_exp_avg_sq'].mul_(beta2).addcmul_(1 - beta2, grad, grad) @@ -259,7 +283,7 @@ def step(self, closure: Union[Callable, None] = None): if len(state) == 0: self._state_init(p, group['amsgrad']) grad = p.grad.data - #should we use same beta group? + # should we use same beta group? beta1, beta2 = group['betas'] bias_correction2 = 1 - beta2 ** state['step'] state['thre_exp_avg_sq'].mul_(beta2).addcmul_(1 - beta2, grad, grad) @@ -267,7 +291,7 @@ def step(self, closure: Union[Callable, None] = None): param_norm = grad.norm(self._clip_norm_type) total_norm += param_norm.item() ** self._clip_norm_type - #sum momentum_norm + # sum momentum_norm momentum = ((state['thre_exp_avg_sq'].sqrt() / math.sqrt(bias_correction2)) * self._clip_coef).norm(self._clip_norm_type) total_momentum_norm += momentum.item() ** self._clip_norm_type @@ -294,7 +318,7 @@ def step(self, closure: Union[Callable, None] = None): if len(state) == 0: self._state_init(p, group['amsgrad']) grad = p.grad.data - #should we use same beta group? + # should we use same beta group? beta1, beta2 = group['betas'] bias_correction2 = 1 - beta2 ** state['step'] state['thre_exp_avg_sq'].mul_(beta2).addcmul_(1 - beta2, grad, grad) @@ -326,7 +350,7 @@ def step(self, closure: Union[Callable, None] = None): if len(state) == 0: self._state_init(p, group['amsgrad']) grad = p.grad.data - #should we use same beta group? + # should we use same beta group? beta1, beta2 = group['betas'] bias_correction2 = 1 - beta2 ** state['step'] state['thre_exp_avg_sq'].mul_(beta2).addcmul_(1 - beta2, grad, grad) @@ -334,7 +358,7 @@ def step(self, closure: Union[Callable, None] = None): param_norm = grad.norm(self._ignore_norm_type) total_norm += param_norm.item() ** self._ignore_norm_type - #sum momentum_norm + # sum momentum_norm momentum = ((state['thre_exp_avg_sq'].sqrt() / math.sqrt(bias_correction2)) * self._ignore_coef).norm(self._ignore_norm_type) total_momentum_norm += momentum.item() ** self._ignore_norm_type @@ -348,7 +372,7 @@ def step(self, closure: Union[Callable, None] = None): for p in group['params']: p.grad.zero_() - #Adam optim type + # Adam optim type if self._optim_type == 'adamw': for group in self.param_groups: for p in group['params']: @@ -372,8 +396,8 @@ class RMSprop(torch.optim.RMSprop): r""" Overview: Rewrited RMSprop optimizer to support more features. - Interface: - __init__, step + Interfaces: + ``__init__``, ``step``, ``_state_init``, ``get_grad`` """ def __init__( @@ -397,9 +421,9 @@ def __init__( ignore_norm_type: float = 2.0, ignore_momentum_timestep: int = 100, ): - r""" + """ Overview: - init medthod of refactored Adam class + init method of refactored Adam class Arguments: - params (:obj:`iterable`): – an iterable of torch.Tensor s or dict s. \ Specifies what Tensors should be optimized @@ -454,8 +478,24 @@ def __init__( ) def _state_init(self, p, momentum, centered): + """ + Overview: + Initialize the state of the optimizer + Arguments: + - p (:obj:`torch.Tensor`): the parameter to be optimized + - momentum (:obj:`float`): the momentum coefficient + - centered (:obj:`bool`): if True, compute the centered RMSprop, \ + the gradient is normalized by an estimation of its variance + """ + state = self.state[p] - state['step'] = 0 + if torch.__version__ < "1.12.0": + state['step'] = 0 + # TODO + # wait torch upgrad to 1.4, 1.3.1 didn't support memory format state['step'] = 0 + else: + state['step'] = torch.zeros((1,), dtype=torch.float, device=p.device) \ + if ('capturable' in self.defaults and self.defaults['capturable']) else torch.tensor(0.) state['thre_square_avg'] = torch.zeros_like(p.data, device=p.data.device) state['square_avg'] = torch.zeros_like(p.data, device=p.data.device) if momentum: @@ -464,7 +504,7 @@ def _state_init(self, p, momentum, centered): state['grad_avg'] = torch.zeros_like(p.data, device=p.data.device) def step(self, closure: Union[Callable, None] = None): - r""" + """ Overview: Performs a single optimization step Arguments: @@ -517,7 +557,7 @@ def step(self, closure: Union[Callable, None] = None): param_norm = grad.norm(self._clip_norm_type) total_norm += param_norm.item() ** self._clip_norm_type - #sum momentum_norm + # sum momentum_norm momentum = (state['thre_square_avg'].sqrt() * self._clip_coef).norm(self._clip_norm_type) total_momentum_norm += momentum.item() ** self._clip_norm_type step = min(step, state['step']) @@ -578,7 +618,7 @@ def step(self, closure: Union[Callable, None] = None): param_norm = grad.norm(self._ignore_norm_type) total_norm += param_norm.item() ** self._ignore_norm_type - #sum momentum_norm + # sum momentum_norm momentum = (state['thre_square_avg'].sqrt() * self._ignore_coef).norm(self._ignore_norm_type) total_momentum_norm += momentum.item() ** self._ignore_norm_type step = min(step, state['step']) @@ -594,6 +634,11 @@ def step(self, closure: Union[Callable, None] = None): return super().step(closure=closure) def get_grad(self) -> float: + """ + Overview: + calculate grad norm of the parameters whose grad norms are not None in the model. + """ + total_norm = 0. params = [t for group in self.param_groups for t in group['params'] if t.requires_grad and t.grad is not None] for p in params: @@ -607,35 +652,55 @@ class PCGrad(): Overview: PCGrad optimizer to support multi-task. you can view the paper in the following link https://arxiv.org/pdf/2001.06782.pdf + Interfaces: + ``__init__``, ``zero_grad``, ``step``, ``pc_backward`` + Properties: + - optimizer (:obj:`torch.optim`): the optimizer to be used """ def __init__(self, optimizer, reduction='mean'): + """ + Overview: + Initialization of PCGrad optimizer + Arguments: + - optimizer (:obj:`torch.optim`): the optimizer to be used + - reduction (:obj:`str`): the reduction method, support ['mean', 'sum'] + """ + self._optim, self._reduction = optimizer, reduction @property def optimizer(self): + """ + Overview: + get the optimizer + """ + return self._optim def zero_grad(self): - ''' - clear the gradient of the parameters - ''' + """ + Overview: + clear the gradient of the parameters + """ return self._optim.zero_grad(set_to_none=True) def step(self): - ''' - update the parameters with the gradient - ''' + """ + Overview: + update the parameters with the gradient + """ return self._optim.step() def pc_backward(self, objectives): - ''' - calculate the gradient of the parameters + """ + Overview: + calculate the gradient of the parameters Arguments: - objectives: a list of objectives - ''' + """ grads, shapes, has_grads = self._pack_grad(objectives) pc_grad = self._project_conflicting(grads, has_grads) @@ -644,6 +709,15 @@ def pc_backward(self, objectives): return def _project_conflicting(self, grads, has_grads, shapes=None): + """ + Overview: + project the conflicting gradient to the orthogonal space + Arguments: + - grads (:obj:`list`): a list of the gradient of the parameters + - has_grads (:obj:`list`): a list of mask represent whether the parameter has gradient + - shapes (:obj:`list`): a list of the shape of the parameters + """ + shared = torch.stack(has_grads).prod(0).bool() pc_grad, num_task = copy.deepcopy(grads), len(grads) for g_i in pc_grad: @@ -664,9 +738,12 @@ def _project_conflicting(self, grads, has_grads, shapes=None): return merged_grad def _set_grad(self, grads): - ''' - set the modified gradients to the network - ''' + """ + Overview: + set the modified gradients to the network + Arguments: + - grads (:obj:`list`): a list of the gradient of the parameters + """ idx = 0 for group in self._optim.param_groups: @@ -677,13 +754,16 @@ def _set_grad(self, grads): return def _pack_grad(self, objectives): - ''' - pack the gradient of the parameters of the network for each objective + """ + Overview: + pack the gradient of the parameters of the network for each objective + Arguments: + - objectives: a list of objectives Returns: - grad: a list of the gradient of the parameters - shape: a list of the shape of the parameters - has_grad: a list of mask represent whether the parameter has gradient - ''' + """ grads, shapes, has_grads = [], [], [] for obj in objectives: @@ -696,6 +776,14 @@ def _pack_grad(self, objectives): return grads, shapes, has_grads def _unflatten_grad(self, grads, shapes): + """ + Overview: + unflatten the gradient of the parameters of the network + Arguments: + - grads (:obj:`list`): a list of the gradient of the parameters + - shapes (:obj:`list`): a list of the shape of the parameters + """ + unflatten_grad, idx = [], 0 for shape in shapes: length = np.prod(shape) @@ -704,17 +792,26 @@ def _unflatten_grad(self, grads, shapes): return unflatten_grad def _flatten_grad(self, grads, shapes): + """ + Overview: + flatten the gradient of the parameters of the network + Arguments: + - grads (:obj:`list`): a list of the gradient of the parameters + - shapes (:obj:`list`): a list of the shape of the parameters + """ + flatten_grad = torch.cat([g.flatten() for g in grads]) return flatten_grad def _retrieve_grad(self): - ''' - get the gradient of the parameters of the network with specific objective + """ + Overview: + get the gradient of the parameters of the network with specific objective Returns: - grad: a list of the gradient of the parameters - shape: a list of the shape of the parameters - has_grad: a list of mask represent whether the parameter has gradient - ''' + """ grad, shape, has_grad = [], [], [] for group in self._optim.param_groups: @@ -730,3 +827,58 @@ def _retrieve_grad(self): grad.append(p.grad.clone()) has_grad.append(torch.ones_like(p).to(p.device)) return grad, shape, has_grad + + +def configure_weight_decay(model: nn.Module, weight_decay: float) -> List: + """ + Overview: + Separating out all parameters of the model into two buckets: those that will experience + weight decay for regularization and those that won't (biases, and layer-norm or embedding weights). + Arguments: + - model (:obj:`nn.Module`): The given PyTorch model. + - weight_decay (:obj:`float`): Weight decay value for optimizer. + Returns: + - optim groups (:obj:`List`): The parameter groups to be set in the latter optimizer. + """ + decay = set() + no_decay = set() + whitelist_weight_modules = (torch.nn.Linear, ) + blacklist_weight_modules = (torch.nn.LayerNorm, torch.nn.Embedding) + for mn, m in model.named_modules(): + for pn, p in m.named_parameters(): + fpn = '%s.%s' % (mn, pn) if mn else pn # full param name + # Because named_modules and named_parameters are recursive + # we will see the same tensors p many times. But doing it this way + # allows us to know which parent module any tensor p belongs to. + if pn.endswith('bias'): + # all biases will not be decayed + no_decay.add(fpn) + elif pn.endswith('weight') and isinstance(m, whitelist_weight_modules): + # weights of whitelist modules will be weight decayed + decay.add(fpn) + elif pn.endswith('weight') and isinstance(m, blacklist_weight_modules): + # weights of blacklist modules will NOT be weight decayed + no_decay.add(fpn) + else: + decay.add(fpn) + + decay = decay - no_decay + # validate that we considered every parameter + param_dict = {pn: p for pn, p in model.named_parameters()} + union_params = decay | no_decay + assert len( + param_dict.keys() - union_params) == 0, "parameters %s were not separated into either decay/no_decay set!" \ + % (str(param_dict.keys() - union_params),) + + optim_groups = [ + { + "params": [param_dict[pn] for pn in sorted(list(decay))], + "weight_decay": weight_decay + }, + { + "params": [param_dict[pn] for pn in sorted(list(no_decay))], + "weight_decay": 0.0 + }, + ] + + return optim_groups diff --git a/ding/torch_utils/parameter.py b/ding/torch_utils/parameter.py new file mode 100644 index 0000000000..08da7feb76 --- /dev/null +++ b/ding/torch_utils/parameter.py @@ -0,0 +1,89 @@ +from typing import Optional +import torch +from torch import nn +from torch.distributions.transforms import TanhTransform + + +class NonegativeParameter(nn.Module): + """ + Overview: + This module will output a non-negative parameter during the forward process. + Interfaces: + ``__init__``, ``forward``, ``set_data``. + """ + + def __init__(self, data: Optional[torch.Tensor] = None, requires_grad: bool = True, delta: float = 1e-8): + """ + Overview: + Initialize the NonegativeParameter object using the given arguments. + Arguments: + - data (:obj:`Optional[torch.Tensor]`): The initial value of generated parameter. If set to ``None``, the \ + default value is 0. + - requires_grad (:obj:`bool`): Whether this parameter requires grad. + - delta (:obj:`Any`): The delta of log function. + """ + super().__init__() + if data is None: + data = torch.zeros(1) + self.log_data = nn.Parameter(torch.log(data + delta), requires_grad=requires_grad) + + def forward(self) -> torch.Tensor: + """ + Overview: + Output the non-negative parameter during the forward process. + Returns: + parameter (:obj:`torch.Tensor`): The generated parameter. + """ + return torch.exp(self.log_data) + + def set_data(self, data: torch.Tensor) -> None: + """ + Overview： + Set the value of the non-negative parameter. + Arguments: + data (:obj:`torch.Tensor`): The new value of the non-negative parameter. + """ + self.log_data = nn.Parameter(torch.log(data + 1e-8), requires_grad=self.log_data.requires_grad) + + +class TanhParameter(nn.Module): + """ + Overview: + This module will output a tanh parameter during the forward process. + Interfaces: + ``__init__``, ``forward``, ``set_data``. + """ + + def __init__(self, data: Optional[torch.Tensor] = None, requires_grad: bool = True): + """ + Overview: + Initialize the TanhParameter object using the given arguments. + Arguments: + - data (:obj:`Optional[torch.Tensor]`): The initial value of generated parameter. If set to ``None``, the \ + default value is 1. + - requires_grad (:obj:`bool`): Whether this parameter requires grad. + """ + super().__init__() + if data is None: + data = torch.zeros(1) + self.transform = TanhTransform(cache_size=1) + + self.data_inv = nn.Parameter(self.transform.inv(data), requires_grad=requires_grad) + + def forward(self) -> torch.Tensor: + """ + Overview: + Output the tanh parameter during the forward process. + Returns: + parameter (:obj:`torch.Tensor`): The generated parameter. + """ + return self.transform(self.data_inv) + + def set_data(self, data: torch.Tensor) -> None: + """ + Overview: + Set the value of the tanh parameter. + Arguments: + data (:obj:`torch.Tensor`): The new value of the tanh parameter. + """ + self.data_inv = nn.Parameter(self.transform.inv(data), requires_grad=self.data_inv.requires_grad) diff --git a/ding/torch_utils/reshape_helper.py b/ding/torch_utils/reshape_helper.py index e9e3f60676..49a3e3b8d2 100644 --- a/ding/torch_utils/reshape_helper.py +++ b/ding/torch_utils/reshape_helper.py @@ -4,7 +4,7 @@ def fold_batch(x: Tensor, nonbatch_ndims: int = 1) -> Tuple[Tensor, Size]: - r""" + """ Overview: :math:`(T, B, X) \leftarrow (T*B, X)`\ Fold the first (ndim - nonbatch_ndims) dimensions of a tensor as batch dimension.\ @@ -39,7 +39,7 @@ def fold_batch(x: Tensor, nonbatch_ndims: int = 1) -> Tuple[Tensor, Size]: def unfold_batch(x: Tensor, batch_dims: Union[Size, Tuple]) -> Tensor: - r""" + """ Overview: Unfold the batch dimension of a tensor. @@ -62,7 +62,7 @@ def unfold_batch(x: Tensor, batch_dims: Union[Size, Tuple]) -> Tensor: def unsqueeze_repeat(x: Tensor, repeat_times: int, unsqueeze_dim: int = 0) -> Tensor: - r""" + """ Overview: Squeeze the tensor on `unsqueeze_dim` and then repeat in this dimension for `repeat_times` times.\ This is useful for preproprocessing the input to an model ensemble. diff --git a/ding/torch_utils/tests/test_backend_helper.py b/ding/torch_utils/tests/test_backend_helper.py new file mode 100644 index 0000000000..f598884693 --- /dev/null +++ b/ding/torch_utils/tests/test_backend_helper.py @@ -0,0 +1,21 @@ +import pytest +import torch + +from ding.torch_utils.backend_helper import enable_tf32 + + +@pytest.mark.cudatest +class TestBackendHelper: + + def test_tf32(self): + r""" + Overview: + Test the tf32. + """ + enable_tf32() + net = torch.nn.Linear(3, 4) + x = torch.randn(1, 3) + y = torch.sum(net(x)) + net.zero_grad() + y.backward() + assert net.weight.grad is not None diff --git a/ding/torch_utils/tests/test_data_helper.py b/ding/torch_utils/tests/test_data_helper.py index 218ce59ba7..5b61d96dd9 100644 --- a/ding/torch_utils/tests/test_data_helper.py +++ b/ding/torch_utils/tests/test_data_helper.py @@ -4,9 +4,10 @@ import torch import torch.nn as nn from torch.utils.data import DataLoader +import treetensor.torch as ttorch from ding.torch_utils import CudaFetcher, to_device, to_dtype, to_tensor, to_ndarray, to_list, \ - tensor_to_list, same_shape, build_log_buffer, get_tensor_data + tensor_to_list, same_shape, build_log_buffer, get_tensor_data, to_item from ding.utils import EasyTimer @@ -113,6 +114,36 @@ def test_to_list(self, setup_data_dict): with pytest.raises(TypeError): to_ndarray(EasyTimer()) + def test_to_item(self): + data = { + 'tensor': torch.randn(1), + 'list': [True, False, torch.randn(1)], + 'tuple': (4, 5, 6), + 'bool': True, + 'int': 10, + 'float': 10., + 'array': np.random.randn(1), + 'str': "asdf", + 'none': None, + } + assert not np.isscalar(data['tensor']) + assert not np.isscalar(data['array']) + assert not np.isscalar(data['list'][-1]) + new_data = to_item(data) + assert np.isscalar(new_data['tensor']) + assert np.isscalar(new_data['array']) + assert np.isscalar(new_data['list'][-1]) + + data = ttorch.randn({'a': 1}) + new_data = to_item(data) + assert np.isscalar(new_data.a) + + with pytest.raises((ValueError, RuntimeError)): + to_item({'a': torch.randn(4), 'b': torch.rand(1)}, ignore_error=False) + output = to_item({'a': torch.randn(4), 'b': torch.rand(1)}, ignore_error=True) + assert 'a' not in output + assert 'b' in output + def test_same_shape(self): tlist = [torch.randn(3, 5) for i in range(5)] assert same_shape(tlist) diff --git a/ding/torch_utils/tests/test_feature_merge.py b/ding/torch_utils/tests/test_feature_merge.py new file mode 100644 index 0000000000..41cfc57f5c --- /dev/null +++ b/ding/torch_utils/tests/test_feature_merge.py @@ -0,0 +1,131 @@ +import pytest +import torch +from ding.torch_utils.network.merge import TorchBilinearCustomized, TorchBilinear, BilinearGeneral, FiLM + + +@pytest.mark.unittest +def test_torch_bilinear_customized(): + batch_size = 10 + in1_features = 20 + in2_features = 30 + out_features = 40 + bilinear_customized = TorchBilinearCustomized(in1_features, in2_features, out_features) + x = torch.randn(batch_size, in1_features) + z = torch.randn(batch_size, in2_features) + out = bilinear_customized(x, z) + assert out.shape == (batch_size, out_features), "Output shape does not match expected shape." + + +@pytest.mark.unittest +def test_torch_bilinear(): + batch_size = 10 + in1_features = 20 + in2_features = 30 + out_features = 40 + torch_bilinear = TorchBilinear(in1_features, in2_features, out_features) + x = torch.randn(batch_size, in1_features) + z = torch.randn(batch_size, in2_features) + out = torch_bilinear(x, z) + assert out.shape == (batch_size, out_features), "Output shape does not match expected shape." + + +@pytest.mark.unittest +def test_bilinear_consistency(): + batch_size = 10 + in1_features = 20 + in2_features = 30 + out_features = 40 + + # Initialize weights and biases with set values + weight = torch.randn(out_features, in1_features, in2_features) + bias = torch.randn(out_features) + + # Create and initialize TorchBilinearCustomized and TorchBilinear models + bilinear_customized = TorchBilinearCustomized(in1_features, in2_features, out_features) + bilinear_customized.weight.data = weight.clone() + bilinear_customized.bias.data = bias.clone() + + torch_bilinear = TorchBilinear(in1_features, in2_features, out_features) + torch_bilinear.weight.data = weight.clone() + torch_bilinear.bias.data = bias.clone() + + # Provide same input to both models + x = torch.randn(batch_size, in1_features) + z = torch.randn(batch_size, in2_features) + + # Compute outputs + out_bilinear_customized = bilinear_customized(x, z) + out_torch_bilinear = torch_bilinear(x, z) + + # Compute the mean squared error between outputs + mse = torch.mean((out_bilinear_customized - out_torch_bilinear) ** 2) + + print(f"Mean Squared Error between outputs: {mse.item()}") + + # Check if outputs are the same + # assert torch.allclose(out_bilinear_customized, out_torch_bilinear), + # "Outputs of TorchBilinearCustomized and TorchBilinear are not the same." + + +def test_bilinear_general(): + """ + Overview: + Test for the `BilinearGeneral` class. + """ + # Define the input dimensions and batch size + in1_features = 20 + in2_features = 30 + out_features = 40 + batch_size = 10 + + # Create a BilinearGeneral instance + bilinear_general = BilinearGeneral(in1_features, in2_features, out_features) + + # Create random inputs + input1 = torch.randn(batch_size, in1_features) + input2 = torch.randn(batch_size, in2_features) + + # Perform forward pass + output = bilinear_general(input1, input2) + + # Check output shape + assert output.shape == (batch_size, out_features), "Output shape does not match expected shape." + + # Check parameter shapes + assert bilinear_general.W.shape == ( + out_features, in1_features, in2_features + ), "Weight W shape does not match expected shape." + assert bilinear_general.U.shape == (out_features, in2_features), "Weight U shape does not match expected shape." + assert bilinear_general.V.shape == (out_features, in1_features), "Weight V shape does not match expected shape." + assert bilinear_general.b.shape == (out_features, ), "Bias shape does not match expected shape." + + # Check parameter types + assert isinstance(bilinear_general.W, torch.nn.Parameter), "Weight W is not an instance of torch.nn.Parameter." + assert isinstance(bilinear_general.U, torch.nn.Parameter), "Weight U is not an instance of torch.nn.Parameter." + assert isinstance(bilinear_general.V, torch.nn.Parameter), "Weight V is not an instance of torch.nn.Parameter." + assert isinstance(bilinear_general.b, torch.nn.Parameter), "Bias is not an instance of torch.nn.Parameter." + + +@pytest.mark.unittest +def test_film_forward(): + # Set the feature and context dimensions + feature_dim = 128 + context_dim = 256 + + # Initialize the FiLM layer + film_layer = FiLM(feature_dim, context_dim) + + # Create random feature and context vectors + feature = torch.randn((32, feature_dim)) # batch size is 32 + context = torch.randn((32, context_dim)) # batch size is 32 + + # Forward propagation + conditioned_feature = film_layer(feature, context) + + # Check the output shape + assert conditioned_feature.shape == feature.shape, \ + f'Expected output shape {feature.shape}, but got {conditioned_feature.shape}' + + # Check that the output is different from the input + assert not torch.all(torch.eq(feature, conditioned_feature)), \ + 'The output feature is the same as the input feature' diff --git a/ding/torch_utils/tests/test_lr_scheduler.py b/ding/torch_utils/tests/test_lr_scheduler.py new file mode 100644 index 0000000000..4ba52d9e1f --- /dev/null +++ b/ding/torch_utils/tests/test_lr_scheduler.py @@ -0,0 +1,20 @@ +import pytest +import torch +from torch.optim import Adam + +from ding.torch_utils.lr_scheduler import cos_lr_scheduler + + +@pytest.mark.unittest +class TestLRSchedulerHelper: + + def test_cos_lr_scheduler(self): + r""" + Overview: + Test the cos lr scheduler. + """ + net = torch.nn.Linear(3, 4) + opt = Adam(net.parameters(), lr=1e-2) + scheduler = cos_lr_scheduler(opt, learning_rate=1e-2, min_lr=6e-5) + scheduler.step(101) + assert opt.param_groups[0]['lr'] == 6e-5 diff --git a/ding/torch_utils/tests/test_model_helper.py b/ding/torch_utils/tests/test_model_helper.py new file mode 100644 index 0000000000..e2dd72e54e --- /dev/null +++ b/ding/torch_utils/tests/test_model_helper.py @@ -0,0 +1,19 @@ +import pytest +import torch + +from ding.torch_utils.model_helper import get_num_params + + +@pytest.mark.unittest +class TestModelHelper: + + def test_model_helper(self): + r""" + Overview: + Test the model helper. + """ + net = torch.nn.Linear(3, 4, bias=False) + assert get_num_params(net) == 12 + + net = torch.nn.Conv2d(3, 3, kernel_size=3, bias=False) + assert get_num_params(net) == 81 diff --git a/ding/torch_utils/tests/test_optimizer.py b/ding/torch_utils/tests/test_optimizer.py index 6f985d506f..389346fe54 100644 --- a/ding/torch_utils/tests/test_optimizer.py +++ b/ding/torch_utils/tests/test_optimizer.py @@ -2,7 +2,7 @@ import torch.nn as nn import torch.optim as optim from ding.torch_utils.optimizer_helper import Adam, RMSprop, calculate_grad_norm, \ - calculate_grad_norm_without_bias_two_norm, PCGrad + calculate_grad_norm_without_bias_two_norm, PCGrad, configure_weight_decay import pytest import time @@ -177,3 +177,21 @@ def naive_test(self): pc_adam.pc_backward([loss1, loss2]) for p in net.parameters(): assert isinstance(p, torch.Tensor) + + +@pytest.mark.unittest +class TestWeightDecay: + + def test_wd(self): + net = nn.Sequential(nn.Linear(3, 4), nn.LayerNorm(4)) + x = torch.randn(1, 3) + group_params = configure_weight_decay(model=net, weight_decay=1e-4) + assert group_params[0]['weight_decay'] == 1e-4 + assert group_params[1]['weight_decay'] == 0 + assert len(group_params[0]['params']) == 1 + assert len(group_params[1]['params']) == 3 + opt = Adam(group_params, lr=1e-2) + opt.zero_grad() + y = torch.sum(net(x)) + y.backward() + opt.step() diff --git a/ding/torch_utils/tests/test_parameter.py b/ding/torch_utils/tests/test_parameter.py new file mode 100644 index 0000000000..3467c65829 --- /dev/null +++ b/ding/torch_utils/tests/test_parameter.py @@ -0,0 +1,25 @@ +import unittest +import pytest +import torch +from ding.torch_utils.parameter import NonegativeParameter, TanhParameter + + +@pytest.mark.unittest +def test_nonegative_parameter(): + nonegative_parameter = NonegativeParameter(torch.tensor([2.0, 3.0])) + assert torch.sum(torch.abs(nonegative_parameter() - torch.tensor([2.0, 3.0]))) == 0 + nonegative_parameter.set_data(torch.tensor(1)) + assert nonegative_parameter() == 1 + + +@pytest.mark.unittest +def test_tanh_parameter(): + tanh_parameter = TanhParameter(torch.tensor([0.5, -0.2])) + assert torch.isclose(tanh_parameter() - torch.tensor([0.5, -0.2]), torch.zeros(2), atol=1e-6).all() + tanh_parameter.set_data(torch.tensor(0.3)) + assert tanh_parameter() == 0.3 + + +if __name__ == "__main__": + test_nonegative_parameter() + test_tanh_parameter() diff --git a/ding/utils/__init__.py b/ding/utils/__init__.py index 5e262b2c38..3623cab824 100644 --- a/ding/utils/__init__.py +++ b/ding/utils/__init__.py @@ -1,10 +1,11 @@ import ding from .collection_helper import iter_mapping -from .compression_helper import get_data_compressor, get_data_decompressor +from .compression_helper import get_data_compressor, get_data_decompressor, CloudPickleWrapper from .default_helper import override, dicts_to_lists, lists_to_dicts, squeeze, default_get, error_wrapper, list_split, \ LimitedSpaceContainer, deep_merge_dicts, set_pkg_seed, flatten_dict, one_time_warning, split_data_generator, \ RunningMeanStd, make_key_as_identifier, remove_illegal_item from .design_helper import SingletonMetaclass +from .dict_helper import convert_easy_dict_to_dict from .file_helper import read_file, save_file, remove_file from .import_helper import try_import_ceph, try_import_mc, try_import_link, import_module, try_import_redis, \ try_import_rediscluster @@ -27,12 +28,16 @@ from .system_helper import get_ip, get_pid, get_task_uid, PropagatingThread, find_free_port from .time_helper import build_time_helper, EasyTimer, WatchDog from .type_helper import SequenceType -from .render_helper import render, fps +from .render_helper import render, fps, get_env_fps, render_env from .fast_copy import fastcopy +from .bfs_helper import get_vi_sequence +from .normalizer_helper import DatasetNormalizer +from .memory_helper import SimpleMemoryProfiler -if ding.enable_linklink: +if ding.enable_linklink: # False as default from .linklink_dist_helper import get_rank, get_world_size, dist_mode, dist_init, dist_finalize, \ allreduce, broadcast, DistContext, allreduce_async, synchronize else: from .pytorch_ddp_dist_helper import get_rank, get_world_size, dist_mode, dist_init, dist_finalize, \ - allreduce, broadcast, DistContext, allreduce_async, synchronize + allreduce, allreduce_with_indicator, broadcast, DDPContext, allreduce_async, synchronize, reduce_data, \ + broadcast_object_list, to_ddp_config, allreduce_data diff --git a/ding/utils/autolog/data.py b/ding/utils/autolog/data.py index 9b960e4fa0..e611b97f43 100644 --- a/ding/utils/autolog/data.py +++ b/ding/utils/autolog/data.py @@ -10,8 +10,24 @@ class RangedData(metaclass=ABCMeta): + """ + Overview: + A data structure that can store data for a period of time. + Interfaces: + ``__init__``, ``append``, ``extend``, ``current``, ``history``, ``expire``, ``__bool__``, ``_get_time``. + Properties: + - expire (:obj:`float`): The expire time. + """ def __init__(self, expire: float, use_pickle: bool = False): + """ + Overview: + Initialize the RangedData object. + Arguments: + - expire (:obj:`float`): The expire time of the data. + - use_pickle (:obj:`bool`): Whether to use pickle to serialize the data. + """ + self.__expire = expire self.__use_pickle = use_pickle self.__check_expire() @@ -25,6 +41,11 @@ def __init__(self, expire: float, use_pickle: bool = False): self.__lock = Lock() def __check_expire(self): + """ + Overview: + Check the expire time. + """ + if isinstance(self.__expire, (int, float)): if self.__expire <= 0: raise ValueError( @@ -36,6 +57,13 @@ def __check_expire(self): ) def __registry_data_item(self, data: _Tp) -> int: + """ + Overview: + Registry the data item. + Arguments: + - data (:obj:`_Tp`): The data item. + """ + with self.__data_lock: self.__data_max_id += 1 if self.__use_pickle: @@ -46,6 +74,13 @@ def __registry_data_item(self, data: _Tp) -> int: return self.__data_max_id def __get_data_item(self, data_id: int) -> _Tp: + """ + Overview: + Get the data item. + Arguments: + - data_id (:obj:`int`): The data id. + """ + with self.__data_lock: if self.__use_pickle: return pickle.loads(self.__data_items[data_id]) @@ -53,10 +88,24 @@ def __get_data_item(self, data_id: int) -> _Tp: return self.__data_items[data_id] def __remove_data_item(self, data_id: int): + """ + Overview: + Remove the data item. + Arguments: + - data_id (:obj:`int`): The data id. + """ + with self.__data_lock: del self.__data_items[data_id] def __check_time(self, time_: float): + """ + Overview: + Check the time. + Arguments: + - time_ (:obj:`float`): The time. + """ + if self.__queue: _time, _ = self.__queue[-1] if time_ < _time: @@ -67,9 +116,22 @@ def __check_time(self, time_: float): ) def __append_item(self, time_: float, data: _Tp): + """ + Overview: + Append the data item. + Arguments: + - time_ (:obj:`float`): The time. + - data (:obj:`_Tp`): The data item. + """ + self.__queue.append((time_, self.__registry_data_item(data))) def __flush_history(self): + """ + Overview: + Flush the history data. + """ + _time = self._get_time() _limit_time = _time - self.__expire while self.__queue: @@ -85,11 +147,21 @@ def __flush_history(self): self.__last_item = (_head_time, _head_id) def __append(self, time_: float, data: _Tp): + """ + Overview: + Append the data. + """ + self.__check_time(time_) self.__append_item(time_, data) self.__flush_history() def __current(self): + """ + Overview: + Get the current data. + """ + if self.__queue: _tail_time, _tail_id = self.__queue.pop() self.__queue.append((_tail_time, _tail_id)) @@ -101,6 +173,11 @@ def __current(self): raise ValueError("This range is empty.") def __history_yield(self): + """ + Overview: + Yield the history data. + """ + _time = self._get_time() _limit_time = _time - self.__expire _latest_time, _latest_id = None, None @@ -117,9 +194,19 @@ def __history_yield(self): yield _time, self.__get_data_item(_latest_id) def __history(self): + """ + Overview: + Get the history data. + """ + return list(self.__history_yield()) def append(self, data: _Tp): + """ + Overview: + Append the data. + """ + with self.__lock: self.__flush_history() _time = self._get_time() @@ -127,6 +214,11 @@ def append(self, data: _Tp): return self def extend(self, iter_: Iterable[_Tp]): + """ + Overview: + Extend the data. + """ + with self.__lock: self.__flush_history() _time = self._get_time() @@ -135,40 +227,92 @@ def extend(self, iter_: Iterable[_Tp]): return self def current(self) -> _Tp: + """ + Overview: + Get the current data. + """ + with self.__lock: self.__flush_history() return self.__current() def history(self) -> List[Tuple[Union[int, float], _Tp]]: + """ + Overview: + Get the history data. + """ + with self.__lock: self.__flush_history() return self.__history() @property def expire(self) -> float: + """ + Overview: + Get the expire time. + """ + with self.__lock: self.__flush_history() return self.__expire def __bool__(self): + """ + Overview: + Check whether the range is empty. + """ + with self.__lock: self.__flush_history() return not not (self.__queue or self.__last_item) @abstractmethod def _get_time(self) -> float: + """ + Overview: + Get the current time. + """ + raise NotImplementedError class TimeRangedData(RangedData): + """ + Overview: + A data structure that can store data for a period of time. + Interfaces: + ``__init__``, ``_get_time``, ``append``, ``extend``, ``current``, ``history``, ``expire``, ``__bool__``. + Properties: + - time (:obj:`BaseTime`): The time. + - expire (:obj:`float`): The expire time. + """ def __init__(self, time_: BaseTime, expire: float): + """ + Overview: + Initialize the TimeRangedData object. + Arguments: + - time_ (:obj:`BaseTime`): The time. + - expire (:obj:`float`): The expire time. + """ + RangedData.__init__(self, expire) self.__time = time_ def _get_time(self) -> float: + """ + Overview: + Get the current time. + """ + return self.__time.time() @property def time(self): + """ + Overview: + Get the time. + """ + return self.__time diff --git a/ding/utils/autolog/model.py b/ding/utils/autolog/model.py index 041c710caf..5c58bb6544 100644 --- a/ding/utils/autolog/model.py +++ b/ding/utils/autolog/model.py @@ -11,6 +11,12 @@ class _LoggedModelMeta(ABCMeta): + """ + Overview: + Metaclass of LoggedModel, used to find all LoggedValue properties and register them. + Interfaces: + ``__init__`` + """ def __init__(cls, name: str, bases: tuple, namespace: dict): @@ -75,14 +81,24 @@ class LoggedModel(metaclass=_LoggedModelMeta): >>> print(ll.range_values['value'](TimeMode.ABSOLUTE)) # use absolute time >>> print(ll.avg['value']()) # average value of last 10 secs - Interface: - __init__, fixed_time, current_time, freeze, unfreeze, register_attribute_value, __getattr__ + Interfaces: + ``__init__``, ``time``, ``expire``, ``fixed_time``, ``current_time``, ``freeze``, ``unfreeze``, \ + ``register_attribute_value``, ``__getattr__``, ``get_property_attribute`` Property: - time, expire + - time (:obj:`BaseTime`): The time. + - expire (:obj:`float`): The expire time. """ def __init__(self, time_: _TimeObjectType, expire: _TimeType): + """ + Overview: + Initialize the LoggedModel object using the given arguments. + Arguments: + - time_ (:obj:`BaseTime`): The time. + - expire (:obj:`float`): The expire time. + """ + self.__time = time_ self.__time_proxy = TimeProxy(self.__time, frozen=False) self.__init_time = self.__time_proxy.time() @@ -96,12 +112,29 @@ def __init__(self, time_: _TimeObjectType, expire: _TimeType): @property def __properties(self) -> List[str]: + """ + Overview: + Get all property names. + """ + return getattr(self, _LOGGED_MODEL__PROPERTIES) def __get_property_ranged_data(self, name: str) -> TimeRangedData: + """ + Overview: + Get ranged data of one property. + Arguments: + - name (:obj:`str`): The property name. + """ + return getattr(self, _LOGGED_MODEL__PROPERTY_ATTR_PREFIX + name) def __init_properties(self): + """ + Overview: + Initialize all properties. + """ + for name in self.__properties: setattr( self, _LOGGED_MODEL__PROPERTY_ATTR_PREFIX + name, @@ -109,6 +142,12 @@ def __init_properties(self): ) def __get_range_values_func(self, name: str): + """ + Overview: + Get range_values function of one property. + Arguments: + - name (:obj:`str`): The property name. + """ def _func(mode: TimeMode = TimeMode.RELATIVE_LIFECYCLE): _current_time = self.__time_proxy.time() @@ -130,6 +169,11 @@ def _func(mode: TimeMode = TimeMode.RELATIVE_LIFECYCLE): return _func def __register_default_funcs(self): + """ + Overview: + Register default functions. + """ + for name in self.__properties: self.register_attribute_value('range_values', name, self.__get_range_values_func(name)) @@ -196,6 +240,10 @@ def register_attribute_value(self, attribute_name: str, property_name: str, valu """ Overview: Register a new attribute for one of the values. Example can be found in overview of class. + Arguments: + - attribute_name (:obj:`str`): name of attribute + - property_name (:obj:`str`): name of property + - value (:obj:`Any`): value of attribute """ self.__methods[attribute_name] = self.__methods.get(attribute_name, {}) self.__methods[attribute_name][property_name] = value @@ -210,7 +258,7 @@ def __getattr__(self, attribute_name: str) -> Any: Overview: Support all methods registered. - Args: + Arguments: attribute_name (str): name of attribute Return: diff --git a/ding/utils/autolog/time_ctl.py b/ding/utils/autolog/time_ctl.py index e350dd92e9..637d27e568 100644 --- a/ding/utils/autolog/time_ctl.py +++ b/ding/utils/autolog/time_ctl.py @@ -9,6 +9,8 @@ class BaseTime(metaclass=ABCMeta): """ Overview: Abstract time interface + Interfaces: + ``time`` """ @abstractmethod @@ -27,7 +29,8 @@ class NaturalTime(BaseTime): """ Overview: Natural time object - + Interfaces: + ``__init__``, ``time`` Example: >>> from ding.utils.autolog.time_ctl import NaturalTime >>> time_ = NaturalTime() @@ -62,7 +65,8 @@ class TickTime(BaseTime): """ Overview: Tick time object - + Interfaces: + ``__init__``, ``step``, ``time`` Example: >>> from ding.utils.autolog.time_ctl import TickTime >>> time_ = TickTime() @@ -73,8 +77,8 @@ def __init__(self, init: int = 0): Overview: Constructor of TickTime - Args: - init (int, optional): init tick time, default is 1 + Arguments: + - init (:obj:`int`): initial time, default is 0 """ self.__tick_time = init @@ -83,11 +87,11 @@ def step(self, delta: int = 1) -> int: Overview Step the time forward for this TickTime - Args: - delta (int, optional): steps to step forward, default is 1 + Arguments: + - delta (:obj:`int`): steps to step forward, default is 1 Returns: - int: new time after stepping + - time (:obj:`int`): new time after stepping Example: >>> from ding.utils.autolog.time_ctl import TickTime @@ -128,8 +132,9 @@ class TimeProxy(BaseTime): Overview: Proxy of time object, it can freeze time, sometimes useful when reproducing. This object is thread-safe, and also freeze and unfreeze operation is strictly ordered. - - Example: + Interfaces: + ``__init__``, ``freeze``, ``unfreeze``, ``time``, ``current_time`` + Examples: >>> from ding.utils.autolog.time_ctl import TickTime, TimeProxy >>> tick_time_ = TickTime() >>> time_ = TimeProxy(tick_time_) @@ -150,10 +155,10 @@ def __init__(self, time_: BaseTime, frozen: bool = False, lock_type: LockContext Overview: Constructor for Time proxy - Args: - time_ (BaseTime): another time object it based on - frozen (bool, optional): this object will be frozen immediately if true, otherwise not, default is False - lock_type (LockContextType, optional): type of the lock, default is THREAD_LOCK + Arguments: + - time_ (:obj:`BaseTime`): another time object it based on + - frozen (:obj:`bool`): this object will be frozen immediately if true, otherwise not, default is False + - lock_type (:obj:`LockContextType`): type of the lock, default is THREAD_LOCK """ self.__time = time_ self.__current_time = self.__time.time() diff --git a/ding/utils/autolog/value.py b/ding/utils/autolog/value.py index 07b7609b33..98510a036a 100644 --- a/ding/utils/autolog/value.py +++ b/ding/utils/autolog/value.py @@ -11,22 +11,61 @@ class LoggedValue: This class's instances will be associated with their owner LoggedModel instance, all the LoggedValue of one LoggedModel will shared the only one time object (defined in time_ctl), so that timeline can be managed properly. + Interfaces: + ``__init__``, ``__get__``, ``__set__`` + Properties: + - __property_name (:obj:`str`): The name of the property. """ def __init__(self, type_: Type[_ValueType] = object): + """ + Overview: + Initialize the LoggedValue object. + Interfaces: + ``__init__`` + """ + self.__type = type_ @property def __property_name(self): + """ + Overview: + Get the name of the property. + """ + return getattr(self, _LOGGED_VALUE__PROPERTY_NAME) def __get_ranged_data(self, instance) -> TimeRangedData: + """ + Overview: + Get the ranged data. + Interfaces: + ``__get_ranged_data`` + """ + return getattr(instance, _LOGGED_MODEL__PROPERTY_ATTR_PREFIX + self.__property_name) def __get__(self, instance, owner): + """ + Overview: + Get the value. + Arguments: + - instance (:obj:`LoggedModel`): The owner LoggedModel instance. + - owner (:obj:`type`): The owner class. + """ + return self.__get_ranged_data(instance).current() def __set__(self, instance, value: _ValueType): + """ + Overview: + Set the value. + Arguments: + - instance (:obj:`LoggedModel`): The owner LoggedModel instance. + - value (:obj:`_ValueType`): The value to set. + """ + if isinstance(value, self.__type): return self.__get_ranged_data(instance).append(value) else: diff --git a/ding/utils/bfs_helper.py b/ding/utils/bfs_helper.py new file mode 100644 index 0000000000..948bba5cf9 --- /dev/null +++ b/ding/utils/bfs_helper.py @@ -0,0 +1,70 @@ +import numpy as np +import torch +from gym import Env +from typing import Tuple, List + + +def get_vi_sequence(env: Env, observation: np.ndarray) -> Tuple[np.ndarray, List]: + """ + Overview: + Given an instance of the maze environment and the current observation, using Broad-First-Search (BFS) \ + algorithm to plan an optimal path and record the result. + Arguments: + - env (:obj:`Env`): The instance of the maze environment. + - observation (:obj:`np.ndarray`): The current observation. + Returns: + - output (:obj:`Tuple[np.ndarray, List]`): The BFS result. ``output[0]`` contains the BFS map after each \ + iteration and ``output[1]`` contains the optimal actions before reaching the finishing point. + """ + xy = np.where(observation[Ellipsis, -1] == 1) + start_x, start_y = xy[0][0], xy[1][0] + target_location = env.target_location + nav_map = env.nav_map + current_points = [target_location] + chosen_actions = {target_location: 0} + visited_points = {target_location: True} + vi_sequence = [] + + vi_map = np.full((env.size, env.size), fill_value=env.n_action, dtype=np.int32) + + found_start = False + while current_points and not found_start: + next_points = [] + for point_x, point_y in current_points: + for (action, (next_point_x, next_point_y)) in [(0, (point_x - 1, point_y)), (1, (point_x, point_y - 1)), + (2, (point_x + 1, point_y)), (3, (point_x, point_y + 1))]: + + if (next_point_x, next_point_y) in visited_points: + continue + + if not (0 <= next_point_x < len(nav_map) and 0 <= next_point_y < len(nav_map[next_point_x])): + continue + + if nav_map[next_point_x][next_point_y] == 'x': + continue + + next_points.append((next_point_x, next_point_y)) + visited_points[(next_point_x, next_point_y)] = True + chosen_actions[(next_point_x, next_point_y)] = action + vi_map[next_point_x, next_point_y] = action + + if next_point_x == start_x and next_point_y == start_y: + found_start = True + vi_sequence.append(vi_map.copy()) + current_points = next_points + track_back = [] + if found_start: + cur_x, cur_y = start_x, start_y + while cur_x != target_location[0] or cur_y != target_location[1]: + act = vi_sequence[-1][cur_x, cur_y] + track_back.append((torch.FloatTensor(env.process_states([cur_x, cur_y], env.get_maze_map())), act)) + if act == 0: + cur_x += 1 + elif act == 1: + cur_y += 1 + elif act == 2: + cur_x -= 1 + elif act == 3: + cur_y -= 1 + + return np.array(vi_sequence), track_back diff --git a/ding/utils/collection_helper.py b/ding/utils/collection_helper.py index ea8fa715a8..7c8caed6b4 100644 --- a/ding/utils/collection_helper.py +++ b/ding/utils/collection_helper.py @@ -5,7 +5,7 @@ def iter_mapping(iter_: Iterable[_IterType], mapping: Callable[[_IterType], _IterTargetType]): - r""" + """ Overview: Map a list of iterable elements to input iteration callable Arguments: diff --git a/ding/utils/compression_helper.py b/ding/utils/compression_helper.py index 1d29e96309..f7ff9dfe3e 100644 --- a/ding/utils/compression_helper.py +++ b/ding/utils/compression_helper.py @@ -1,47 +1,108 @@ +from typing import Any, ByteString, Callable import pickle +import cloudpickle import zlib -import lz4.block import numpy as np -def dummy_compressor(data): - r""" +class CloudPickleWrapper: + """ Overview: - Return input data. + CloudPickleWrapper can be able to pickle more python object(e.g: an object with lambda expression). + Interfaces: + ``__init__``, ``__getstate__``, ``__setstate__``. + """ + + def __init__(self, data: Any) -> None: + """ + Overview: + Initialize the CloudPickleWrapper using the given arguments. + Arguments: + - data (:obj:`Any`): The object to be dumped. + """ + self.data = data + + def __getstate__(self) -> bytes: + """ + Overview: + Get the state of the CloudPickleWrapper. + Returns: + - data (:obj:`bytes`): The dumped byte-like result. + """ + + return cloudpickle.dumps(self.data) + + def __setstate__(self, data: bytes) -> None: + """ + Overview: + Set the state of the CloudPickleWrapper. + Arguments: + - data (:obj:`bytes`): The dumped byte-like result. + """ + + if isinstance(data, (tuple, list, np.ndarray)): # pickle is faster + self.data = pickle.loads(data) + else: + self.data = cloudpickle.loads(data) + + +def dummy_compressor(data: Any) -> Any: + """ + Overview: + Return the raw input data. + Arguments: + - data (:obj:`Any`): The input data of the compressor. + Returns: + - output (:obj:`Any`): This compressor will exactly return the input data. """ return data -def zlib_data_compressor(data): - r""" +def zlib_data_compressor(data: Any) -> bytes: + """ Overview: Takes the input compressed data and return the compressed original data (zlib compressor) in binary format. + Arguments: + - data (:obj:`Any`): The input data of the compressor. + Returns: + - output (:obj:`bytes`): The compressed byte-like result. Examples: >>> zlib_data_compressor("Hello") - b'x\x9ck`\x99\xca\xc9\x00\x01=\xac\x1e\xa999\xf9S\xf4\x00%L\x04j' """ return zlib.compress(pickle.dumps(data)) -def lz4_data_compressor(data): - r""" +def lz4_data_compressor(data: Any) -> bytes: + """ Overview: Return the compressed original data (lz4 compressor).The compressor outputs in binary format. + Arguments: + - data (:obj:`Any`): The input data of the compressor. + Returns: + - output (:obj:`bytes`): The compressed byte-like result. Examples: >>> lz4.block.compress(pickle.dumps("Hello")) b'\x14\x00\x00\x00R\x80\x04\x95\t\x00\x01\x00\x90\x8c\x05Hello\x94.' """ + try: + import lz4.block + except ImportError: + from ditk import logging + import sys + logging.warning("Please install lz4 first, such as `pip3 install lz4`") + sys.exit(1) return lz4.block.compress(pickle.dumps(data)) -def jpeg_data_compressor(data): +def jpeg_data_compressor(data: np.ndarray) -> bytes: """ Overview: - To reduce memory usage, we can choose to store the jpeg strings of image - instead of the numpy array in the buffer. - This function encodes the observation numpy arr to the jpeg strings. + To reduce memory usage, we can choose to store the jpeg strings of image instead of the numpy array in \ + the buffer. This function encodes the observation numpy arr to the jpeg strings. Arguments: - data (:obj:`np.array`): the observation numpy arr. + - data (:obj:`np.array`): the observation numpy arr. + Returns: + - img_str (:obj:`bytes`): The compressed byte-like result. """ try: import cv2 @@ -64,13 +125,13 @@ def jpeg_data_compressor(data): def get_data_compressor(name: str): - r""" + """ Overview: - Get the data compressor according to the input name + Get the data compressor according to the input name. Arguments: - name(:obj:`str`): Name of the compressor, support ``['lz4', 'zlib', 'jpeg', 'none']`` Return: - - (:obj:`Callable`): Corresponding data_compressor, taking input data returning compressed data. + - compressor (:obj:`Callable`): Corresponding data_compressor, taking input data returning compressed data. Example: >>> compress_fn = get_data_compressor('lz4') >>> compressed_data = compressed(input_data) @@ -78,40 +139,60 @@ def get_data_compressor(name: str): return _COMPRESSORS_MAP[name] -def dummy_decompressor(data): +def dummy_decompressor(data: Any) -> Any: """ Overview: - Return input data. + Return the input data. + Arguments: + - data (:obj:`Any`): The input data of the decompressor. + Returns: + - output (:obj:`bytes`): The decompressed result, which is exactly the input. """ return data -def lz4_data_decompressor(compressed_data): - r""" +def lz4_data_decompressor(compressed_data: bytes) -> Any: + """ Overview: Return the decompressed original data (lz4 compressor). + Arguments: + - data (:obj:`bytes`): The input data of the decompressor. + Returns: + - output (:obj:`Any`): The decompressed object. """ + try: + import lz4.block + except ImportError: + from ditk import logging + import sys + logging.warning("Please install lz4 first, such as `pip3 install lz4`") + sys.exit(1) return pickle.loads(lz4.block.decompress(compressed_data)) -def zlib_data_decompressor(compressed_data): - r""" +def zlib_data_decompressor(compressed_data: bytes) -> Any: + """ Overview: Return the decompressed original data (zlib compressor). + Arguments: + - data (:obj:`bytes`): The input data of the decompressor. + Returns: + - output (:obj:`Any`): The decompressed object. """ return pickle.loads(zlib.decompress(compressed_data)) -def jpeg_data_decompressor(compressed_data, gray_scale=False): +def jpeg_data_decompressor(compressed_data: bytes, gray_scale=False) -> np.ndarray: """ Overview: - To reduce memory usage, we can choose to store the jpeg strings of image - instead of the numpy array in the buffer. - This function decodes the observation numpy arr from the jpeg strings. + To reduce memory usage, we can choose to store the jpeg strings of image instead of the numpy array in the \ + buffer. This function decodes the observation numpy arr from the jpeg strings. Arguments: - compressed_data (:obj:`str`): the jpeg strings. - gray_scale (:obj:`bool`): if the observation is gray, ``gray_scale=True``, + - compressed_data (:obj:`bytes`): The jpeg strings. + - gray_scale (:obj:`bool`): If the observation is gray, ``gray_scale=True``, if the observation is RGB, ``gray_scale=False``. + Returns: + - arr (:obj:`np.ndarray`): The decompressed numpy array. """ try: import cv2 @@ -138,10 +219,10 @@ def jpeg_data_decompressor(compressed_data, gray_scale=False): } -def get_data_decompressor(name: str): - r""" +def get_data_decompressor(name: str) -> Callable: + """ Overview: - Get the data decompressor according to the input name + Get the data decompressor according to the input name. Arguments: - name(:obj:`str`): Name of the decompressor, support ``['lz4', 'zlib', 'none']`` @@ -150,7 +231,7 @@ def get_data_decompressor(name: str): For all the decompressors, the input of a bytes-like object is required. Returns: - - (:obj:`Callable`): Corresponding data_decompressor. + - decompressor (:obj:`Callable`): Corresponding data decompressor. Examples: >>> decompress_fn = get_data_decompressor('lz4') >>> origin_data = compressed(compressed_data) diff --git a/ding/utils/data/__init__.py b/ding/utils/data/__init__.py index 6544dd8f67..0f6b950c15 100644 --- a/ding/utils/data/__init__.py +++ b/ding/utils/data/__init__.py @@ -2,3 +2,5 @@ from .dataloader import AsyncDataLoader from .dataset import NaiveRLDataset, D4RLDataset, HDF5Dataset, BCODataset, \ create_dataset, hdf5_save, offline_data_save_type +from .rlhf_online_dataset import OnlineRLDataset +from .rlhf_offline_dataset import OfflineRLDataset diff --git a/ding/utils/data/base_dataloader.py b/ding/utils/data/base_dataloader.py index da93d8f24d..d19bd9fcde 100644 --- a/ding/utils/data/base_dataloader.py +++ b/ding/utils/data/base_dataloader.py @@ -4,7 +4,10 @@ def example_get_data_fn() -> Any: """ - Note: staticmethod or static function, all the operation is on CPU + Overview: + Get data from file or other middleware + .. note:: + staticmethod or static function, all the operation is on CPU """ # 1. read data from file or other middleware # 2. data post-processing(e.g.: normalization, to tensor) @@ -13,12 +16,20 @@ def example_get_data_fn() -> Any: class IDataLoader: + """ + Overview: + Base class of data loader + Interfaces: + ``__init__``, ``__next__``, ``__iter__``, ``_get_data``, ``close`` + """ def __next__(self, batch_size: Optional[int] = None) -> torch.Tensor: """ + Overview: + Get one batch data Arguments: - batch_size: sometimes, batch_size is specified by each iteration, if batch_size is None, - use default batch_size value + - batch_size (:obj:`Optional[int]`): sometimes, batch_size is specified by each iteration, \ + if batch_size is None, use default batch_size value """ # get one batch train data if batch_size is None: @@ -27,11 +38,29 @@ def __next__(self, batch_size: Optional[int] = None) -> torch.Tensor: return self._collate_fn(data) def __iter__(self) -> Iterable: + """ + Overview: + Get data iterator + """ + return self def _get_data(self, batch_size: Optional[int] = None) -> List[torch.Tensor]: + """ + Overview: + Get one batch data + Arguments: + - batch_size (:obj:`Optional[int]`): sometimes, batch_size is specified by each iteration, \ + if batch_size is None, use default batch_size value + """ + raise NotImplementedError def close(self) -> None: + """ + Overview: + Close data loader + """ + # release resource pass diff --git a/ding/utils/data/collate_fn.py b/ding/utils/data/collate_fn.py index 47eda3d31c..5397a9c450 100644 --- a/ding/utils/data/collate_fn.py +++ b/ding/utils/data/collate_fn.py @@ -4,11 +4,11 @@ import torch import treetensor.torch as ttorch import re -from torch._six import string_classes import collections.abc as container_abcs from ding.compatibility import torch_ge_131 int_classes = int +string_classes = (str, bytes) np_str_obj_array_pattern = re.compile(r'[SaUO]') default_collate_err_msg_format = ( @@ -17,7 +17,49 @@ ) -def ttorch_collate(x, json=False): +def ttorch_collate(x, json: bool = False, cat_1dim: bool = True): + """ + Overview: + Collates a list of tensors or nested dictionaries of tensors into a single tensor or nested \ + dictionary of tensors. + + Arguments: + - x : The input list of tensors or nested dictionaries of tensors. + - json (:obj:`bool`): If True, converts the output to JSON format. Defaults to False. + - cat_1dim (:obj:`bool`): If True, concatenates tensors with shape (B, 1) along the last dimension. \ + Defaults to True. + + Returns: + The collated output tensor or nested dictionary of tensors. + + Examples: + >>> # case 1: Collate a list of tensors + >>> tensors = [torch.tensor([1, 2, 3]), torch.tensor([4, 5, 6]), torch.tensor([7, 8, 9])] + >>> collated = ttorch_collate(tensors) + collated = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) + >>> # case 2: Collate a nested dictionary of tensors + >>> nested_dict = { + 'a': torch.tensor([1, 2, 3]), + 'b': torch.tensor([4, 5, 6]), + 'c': torch.tensor([7, 8, 9]) + } + >>> collated = ttorch_collate(nested_dict) + collated = { + 'a': torch.tensor([1, 2, 3]), + 'b': torch.tensor([4, 5, 6]), + 'c': torch.tensor([7, 8, 9]) + } + >>> # case 3: Collate a list of nested dictionaries of tensors + >>> nested_dicts = [ + {'a': torch.tensor([1, 2, 3]), 'b': torch.tensor([4, 5, 6])}, + {'a': torch.tensor([7, 8, 9]), 'b': torch.tensor([10, 11, 12])} + ] + >>> collated = ttorch_collate(nested_dicts) + collated = { + 'a': torch.tensor([[1, 2, 3], [7, 8, 9]]), + 'b': torch.tensor([[4, 5, 6], [10, 11, 12]]) + } + """ def inplace_fn(t): for k in t.keys(): @@ -28,7 +70,8 @@ def inplace_fn(t): inplace_fn(t[k]) x = ttorch.stack(x) - inplace_fn(x) + if cat_1dim: + inplace_fn(x) if json: x = x.json() return x @@ -40,6 +83,17 @@ def default_collate(batch: Sequence, """ Overview: Put each data field into a tensor with outer dimension batch size. + + Arguments: + - batch (:obj:`Sequence`): A data sequence, whose length is batch size, whose element is one piece of data. + - cat_1dim (:obj:`bool`): Whether to concatenate tensors with shape (B, 1) to (B), defaults to True. + - ignore_prefix (:obj:`list`): A list of prefixes to ignore when collating dictionaries, \ + defaults to ['collate_ignore']. + + Returns: + - ret (:obj:`Union[torch.Tensor, Mapping, Sequence]`): the collated data, with batch size into each data \ + field. The return dtype depends on the original element dtype, can be [torch.Tensor, Mapping, Sequence]. + Example: >>> # a list with B tensors shaped (m, n) -->> a tensor shaped (B, m, n) >>> a = [torch.zeros(2,3) for _ in range(4)] @@ -59,17 +113,13 @@ def default_collate(batch: Sequence, >>> b = default_collate(a) >>> print(b[2].shape, b[3].shape) torch.Size([4, 2, 3]) torch.Size([4, 3, 4]) - Arguments: - - batch (:obj:`Sequence`): a data sequence, whose length is batch size, whose element is one piece of data - Returns: - - ret (:obj:`Union[torch.Tensor, Mapping, Sequence]`): the collated data, with batch size into each data field.\ - the return dtype depends on the original element dtype, can be [torch.Tensor, Mapping, Sequence]. """ - elem = batch[0] - elem_type = type(elem) if isinstance(batch, ttorch.Tensor): return batch.json() + + elem = batch[0] + elem_type = type(elem) if isinstance(elem, torch.Tensor): out = None if torch_ge_131() and torch.utils.data.get_worker_info() is not None: @@ -85,7 +135,7 @@ def default_collate(batch: Sequence, else: return torch.stack(batch, 0, out=out) elif isinstance(elem, ttorch.Tensor): - return ttorch_collate(batch, json=True) + return ttorch_collate(batch, json=True, cat_1dim=cat_1dim) elif elem_type.__module__ == 'numpy' and elem_type.__name__ != 'str_' \ and elem_type.__name__ != 'string_': if elem_type.__name__ == 'ndarray': @@ -122,17 +172,27 @@ def default_collate(batch: Sequence, def timestep_collate(batch: List[Dict[str, Any]]) -> Dict[str, Union[torch.Tensor, list]]: """ Overview: - Put each timestepped data field into a tensor with outer dimension batch size using ``default_collate``. - For short, this process can be represented by: - [len=B, ele={dict_key: [len=T, ele=Tensor(any_dims)]}] -> {dict_key: Tensor([T, B, any_dims])} + Collates a batch of timestepped data fields into tensors with the outer dimension being the batch size. \ + Each timestepped data field is represented as a tensor with shape [T, B, any_dims], where T is the length \ + of the sequence, B is the batch size, and any_dims represents the shape of the tensor at each timestep. + Arguments: - - batch (:obj:`List[Dict[str, Any]]`): a list of dicts with length B, each element is {some_key: some_seq} \ - ('prev_state' should be a key in the dict); \ - some_seq is a sequence with length T, each element is a torch.Tensor with any shape. + - batch(:obj:`List[Dict[str, Any]]`): A list of dictionaries with length B, where each dictionary represents \ + a timestepped data field. Each dictionary contains a key-value pair, where the key is the name of the \ + data field and the value is a sequence of torch.Tensor objects with any shape. + Returns: - - ret (:obj:`Dict[str, Union[torch.Tensor, list]]`): the collated data, with timestep and batch size \ - into each data field. By using ``default_collate``, timestep would come to the first dim. \ - So the final shape is :math:`(T, B, dim1, dim2, ...)` + - ret(:obj:`Dict[str, Union[torch.Tensor, list]]`): The collated data, with the timestep and batch size \ + incorporated into each data field. The shape of each data field is [T, B, dim1, dim2, ...]. + + Examples: + >>> batch = [ + {'data0': [torch.tensor([1, 2, 3]), torch.tensor([4, 5, 6])]}, + {'data1': [torch.tensor([7, 8, 9]), torch.tensor([10, 11, 12])]} + ] + >>> collated_data = timestep_collate(batch) + >>> print(collated_data['data'].shape) + torch.Size([2, 2, 3]) """ def stack(data): @@ -166,14 +226,37 @@ def stack(data): def diff_shape_collate(batch: Sequence) -> Union[torch.Tensor, Mapping, Sequence]: """ Overview: - Similar to ``default_collate``, put each data field into a tensor with outer dimension batch size. - The main difference is that, ``diff_shape_collate`` allows tensors in the batch have `None`, - which is quite common StarCraft observation. + Collates a batch of data with different shapes. + This function is similar to `default_collate`, but it allows tensors in the batch to have `None` values, \ + which is common in StarCraft observations. + Arguments: - - batch (:obj:`Sequence`): a data sequence, whose length is batch size, whose element is one piece of data + - batch (:obj:`Sequence`): A sequence of data, where each element is a piece of data. + Returns: - - ret (:obj:`Union[torch.Tensor, Mapping, Sequence]`): the collated data, with batch size into each data field.\ - the return dtype depends on the original element dtype, can be [torch.Tensor, Mapping, Sequence]. + - ret (:obj:`Union[torch.Tensor, Mapping, Sequence]`): The collated data, with the batch size applied \ + to each data field. The return type depends on the original element type and can be a torch.Tensor, \ + Mapping, or Sequence. + + Examples: + >>> # a list with B tensors shaped (m, n) -->> a tensor shaped (B, m, n) + >>> a = [torch.zeros(2,3) for _ in range(4)] + >>> diff_shape_collate(a).shape + torch.Size([4, 2, 3]) + >>> + >>> # a list with B lists, each list contains m elements -->> a list of m tensors, each with shape (B, ) + >>> a = [[0 for __ in range(3)] for _ in range(4)] + >>> diff_shape_collate(a) + [tensor([0, 0, 0, 0]), tensor([0, 0, 0, 0]), tensor([0, 0, 0, 0])] + >>> + >>> # a list with B dicts, whose values are tensors shaped :math:`(m, n)` -->> + >>> # a dict whose values are tensors with shape :math:`(B, m, n)` + >>> a = [{i: torch.zeros(i,i+1) for i in range(2, 4)} for _ in range(4)] + >>> print(a[0][2].shape, a[0][3].shape) + torch.Size([2, 3]) torch.Size([3, 4]) + >>> b = diff_shape_collate(a) + >>> print(b[2].shape, b[3].shape) + torch.Size([4, 2, 3]) torch.Size([4, 3, 4]) """ elem = batch[0] elem_type = type(elem) @@ -213,19 +296,39 @@ def default_decollate( ) -> List[Any]: """ Overview: - Drag out batch_size collated data's batch size to decollate it, - which is the reverse operation of ``default_collate``. + Drag out batch_size collated data's batch size to decollate it, which is the reverse operation of \ + ``default_collate``. + Arguments: - - batch (:obj:`Union[torch.Tensor, Sequence, Mapping]`): can refer to the Returns of ``default_collate`` - - ignore(:obj:`List[str]`): a list of names to be ignored, only function if input ``batch`` is a dict. \ - If key is in this list, its value would stay the same with no decollation. + - batch (:obj:`Union[torch.Tensor, Sequence, Mapping]`): The collated data batch. It can be a tensor, \ + sequence, or mapping. + - ignore(:obj:`List[str]`): A list of names to be ignored. Only applicable if the input ``batch`` is a \ + dictionary. If a key is in this list, its value will remain the same without decollation. Defaults to \ + ['prev_state', 'prev_actor_state', 'prev_critic_state']. + Returns: - - ret (:obj:`List[Any]`): a list with B elements. + - ret (:obj:`List[Any]`): A list with B elements, where B is the batch size. + + Examples: + >>> batch = { + 'a': [ + [1, 2, 3], + [4, 5, 6] + ], + 'b': [ + [7, 8, 9], + [10, 11, 12] + ]} + >>> default_decollate(batch) + { + 0: {'a': [1, 2, 3], 'b': [7, 8, 9]}, + 1: {'a': [4, 5, 6], 'b': [10, 11, 12]}, + } """ if isinstance(batch, torch.Tensor): batch = torch.split(batch, 1, dim=0) - # squeeze if original batch's shape is like (B, dim1, dim2, ...); - # otherwise directly return the list. + # Squeeze if the original batch's shape is like (B, dim1, dim2, ...); + # otherwise, directly return the list. if len(batch[0].shape) > 1: batch = [elem.squeeze(0) for elem in batch] return list(batch) @@ -235,5 +338,7 @@ def default_decollate( tmp = {k: v if k in ignore else default_decollate(v) for k, v in batch.items()} B = len(list(tmp.values())[0]) return [{k: tmp[k][i] for k in tmp.keys()} for i in range(B)] + elif isinstance(batch, torch.distributions.Distribution): # For compatibility + return [None for _ in range(batch.batch_shape[0])] - raise TypeError("not support batch type: {}".format(type(batch))) + raise TypeError("Not supported batch type: {}".format(type(batch))) diff --git a/ding/utils/data/dataloader.py b/ding/utils/data/dataloader.py index cb81925fd3..2c5f76a293 100644 --- a/ding/utils/data/dataloader.py +++ b/ding/utils/data/dataloader.py @@ -13,11 +13,12 @@ class AsyncDataLoader(IDataLoader): - r""" + """ Overview: An asynchronous dataloader. - Interface: - __init__, __iter__, __next__, close + Interfaces: + ``__init__``, ``__iter__``, ``__next__``, ``_get_data``, ``_async_loop``, ``_worker_loop``, ``_cuda_loop``, \ + ``_get_data``, ``close`` """ def __init__( @@ -100,7 +101,7 @@ def __init__( if self.batch_size != self.chunk_size: # job_result {batch_id: result_list} is used to store processed result in temporal. self.job_result = self.manager.dict() - self.job_result_lock = LockContext(type_=LockContextType.PROCESS_LOCK) + self.job_result_lock = LockContext(lock_type=LockContextType.PROCESS_LOCK) self.job_queue = self.mp_context.Queue(maxsize=queue_maxsize) self.worker = [ self.mp_context.Process( @@ -186,6 +187,7 @@ def _async_loop(self, p: tm.multiprocessing.connection, c: tm.multiprocessing.co - p (:obj:`tm.multiprocessing.connection`): Parent connection. - c (:obj:`tm.multiprocessing.connection`): Child connection. """ + torch.set_num_threads(1) p.close() # Close unused p, only use c while not self.end_flag: if self.num_workers > 1: @@ -337,6 +339,10 @@ def __next__(self) -> Any: self.async_train_queue.join_thread() def __del__(self) -> None: + """ + Overview: + Delete this dataloader. + """ self.close() def close(self) -> None: diff --git a/ding/utils/data/dataset.py b/ding/utils/data/dataset.py old mode 100644 new mode 100755 index a2dec4c2ea..bd24e2eaf8 --- a/ding/utils/data/dataset.py +++ b/ding/utils/data/dataset.py @@ -1,20 +1,26 @@ from typing import List, Dict, Tuple -import pickle -import torch -import numpy as np from ditk import logging from copy import deepcopy -from dataclasses import dataclass - from easydict import EasyDict from torch.utils.data import Dataset +from dataclasses import dataclass + +import pickle +import easydict +import torch +import numpy as np -from ding.utils import DATASET_REGISTRY, import_module +from ding.utils.bfs_helper import get_vi_sequence +from ding.utils import DATASET_REGISTRY, import_module, DatasetNormalizer from ding.rl_utils import discount_cumsum @dataclass class DatasetStatistics: + """ + Overview: + Dataset statistics. + """ mean: np.ndarray # obs std: np.ndarray # obs action_bounds: np.ndarray @@ -22,8 +28,21 @@ class DatasetStatistics: @DATASET_REGISTRY.register('naive') class NaiveRLDataset(Dataset): + """ + Overview: + Naive RL dataset, which is used for offline RL algorithms. + Interfaces: + ``__init__``, ``__len__``, ``__getitem__`` + """ def __init__(self, cfg) -> None: + """ + Overview: + Initialization method. + Arguments: + - cfg (:obj:`dict`): Config dict. + """ + assert type(cfg) in [str, EasyDict], "invalid cfg type: {}".format(type(cfg)) if isinstance(cfg, EasyDict): self._data_path = cfg.policy.collect.data_path @@ -33,21 +52,51 @@ def __init__(self, cfg) -> None: self._data: List[Dict[str, torch.Tensor]] = pickle.load(f) def __len__(self) -> int: + """ + Overview: + Get the length of the dataset. + """ + return len(self._data) def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + """ + Overview: + Get the item of the dataset. + """ + return self._data[idx] @DATASET_REGISTRY.register('d4rl') class D4RLDataset(Dataset): + """ + Overview: + D4RL dataset, which is used for offline RL algorithms. + Interfaces: + ``__init__``, ``__len__``, ``__getitem__`` + Properties: + - mean (:obj:`np.ndarray`): Mean of the dataset. + - std (:obj:`np.ndarray`): Std of the dataset. + - action_bounds (:obj:`np.ndarray`): Action bounds of the dataset. + - statistics (:obj:`dict`): Statistics of the dataset. + """ def __init__(self, cfg: dict) -> None: + """ + Overview: + Initialization method. + Arguments: + - cfg (:obj:`dict`): Config dict. + """ + import gym try: import d4rl # register d4rl enviroments with open ai gym except ImportError: + import sys logging.warning("not found d4rl env, please install it, refer to https://github.com/rail-berkeley/d4rl") + sys.exit(1) # Init parameters data_path = cfg.policy.collect.get('data_path', None) @@ -65,16 +114,69 @@ def __init__(self, cfg: dict) -> None: except (KeyError, AttributeError): # do not normalize pass + if hasattr(cfg.env, "reward_norm"): + if cfg.env.reward_norm == "normalize": + dataset['rewards'] = (dataset['rewards'] - dataset['rewards'].mean()) / dataset['rewards'].std() + elif cfg.env.reward_norm == "iql_antmaze": + dataset['rewards'] = dataset['rewards'] - 1.0 + elif cfg.env.reward_norm == "iql_locomotion": + + def return_range(dataset, max_episode_steps): + returns, lengths = [], [] + ep_ret, ep_len = 0.0, 0 + for r, d in zip(dataset["rewards"], dataset["terminals"]): + ep_ret += float(r) + ep_len += 1 + if d or ep_len == max_episode_steps: + returns.append(ep_ret) + lengths.append(ep_len) + ep_ret, ep_len = 0.0, 0 + # returns.append(ep_ret) # incomplete trajectory + lengths.append(ep_len) # but still keep track of number of steps + assert sum(lengths) == len(dataset["rewards"]) + return min(returns), max(returns) + + min_ret, max_ret = return_range(dataset, 1000) + dataset['rewards'] /= max_ret - min_ret + dataset['rewards'] *= 1000 + elif cfg.env.reward_norm == "cql_antmaze": + dataset['rewards'] = (dataset['rewards'] - 0.5) * 4.0 + elif cfg.env.reward_norm == "antmaze": + dataset['rewards'] = (dataset['rewards'] - 0.25) * 2.0 + else: + raise NotImplementedError + self._data = [] self._load_d4rl(dataset) + @property + def data(self) -> List: + return self._data + def __len__(self) -> int: + """ + Overview: + Get the length of the dataset. + """ + return len(self._data) def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + """ + Overview: + Get the item of the dataset. + """ + return self._data[idx] def _load_d4rl(self, dataset: Dict[str, np.ndarray]) -> None: + """ + Overview: + Load the d4rl dataset. + Arguments: + - dataset (:obj:`Dict[str, np.ndarray]`): The d4rl dataset. + """ + for i in range(len(dataset['observations'])): trans_data = {} trans_data['obs'] = torch.from_numpy(dataset['observations'][i]) @@ -85,6 +187,15 @@ def _load_d4rl(self, dataset: Dict[str, np.ndarray]) -> None: self._data.append(trans_data) def _cal_statistics(self, dataset, env, eps=1e-3, add_action_buffer=True): + """ + Overview: + Calculate the statistics of the dataset. + Arguments: + - dataset (:obj:`Dict[str, np.ndarray]`): The d4rl dataset. + - env (:obj:`gym.Env`): The environment. + - eps (:obj:`float`): Epsilon. + """ + self._mean = dataset['observations'].mean(0) self._std = dataset['observations'].std(0) + eps action_max = dataset['actions'].max(0) @@ -96,36 +207,89 @@ def _cal_statistics(self, dataset, env, eps=1e-3, add_action_buffer=True): self._action_bounds = np.stack([action_min, action_max], axis=0) def _normalize_states(self, dataset): + """ + Overview: + Normalize the states. + Arguments: + - dataset (:obj:`Dict[str, np.ndarray]`): The d4rl dataset. + """ + dataset['observations'] = (dataset['observations'] - self._mean) / self._std dataset['next_observations'] = (dataset['next_observations'] - self._mean) / self._std return dataset @property def mean(self): + """ + Overview: + Get the mean of the dataset. + """ + return self._mean @property def std(self): + """ + Overview: + Get the std of the dataset. + """ + return self._std @property def action_bounds(self) -> np.ndarray: + """ + Overview: + Get the action bounds of the dataset. + """ + return self._action_bounds @property def statistics(self) -> dict: + """ + Overview: + Get the statistics of the dataset. + """ + return DatasetStatistics(mean=self.mean, std=self.std, action_bounds=self.action_bounds) @DATASET_REGISTRY.register('hdf5') class HDF5Dataset(Dataset): + """ + Overview: + HDF5 dataset is saved in hdf5 format, which is used for offline RL algorithms. + The hdf5 format is a common format for storing large numerical arrays in Python. + For more details, please refer to https://support.hdfgroup.org/HDF5/. + Interfaces: + ``__init__``, ``__len__``, ``__getitem__`` + Properties: + - mean (:obj:`np.ndarray`): Mean of the dataset. + - std (:obj:`np.ndarray`): Std of the dataset. + - action_bounds (:obj:`np.ndarray`): Action bounds of the dataset. + - statistics (:obj:`dict`): Statistics of the dataset. + """ def __init__(self, cfg: dict) -> None: + """ + Overview: + Initialization method. + Arguments: + - cfg (:obj:`dict`): Config dict. + """ + try: import h5py except ImportError: - logging.warning("not found h5py package, please install it trough 'pip install h5py' ") + import sys + logging.warning("not found h5py package, please install it trough `pip install h5py ") + sys.exit(1) data_path = cfg.policy.collect.get('data_path', None) + if 'dataset' in cfg: + self.context_len = cfg.dataset.context_len + else: + self.context_len = 0 data = h5py.File(data_path, 'r') self._load_data(data) self._cal_statistics() @@ -137,18 +301,57 @@ def __init__(self, cfg: dict) -> None: pass def __len__(self) -> int: - return len(self._data['obs']) + """ + Overview: + Get the length of the dataset. + """ + + return len(self._data['obs']) - self.context_len def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: - return {k: self._data[k][idx] for k in self._data.keys()} + """ + Overview: + Get the item of the dataset. + Arguments: + - idx (:obj:`int`): The index of the dataset. + """ + + if self.context_len == 0: # for other offline RL algorithms + return {k: self._data[k][idx] for k in self._data.keys()} + else: # for decision transformer + block_size = self.context_len + done_idx = idx + block_size + idx = done_idx - block_size + states = torch.as_tensor( + np.array(self._data['obs'][idx:done_idx]), dtype=torch.float32 + ).view(block_size, -1) + actions = torch.as_tensor(self._data['action'][idx:done_idx], dtype=torch.long) + rtgs = torch.as_tensor(self._data['reward'][idx:done_idx, 0], dtype=torch.float32) + timesteps = torch.as_tensor(range(idx, done_idx), dtype=torch.int64) + traj_mask = torch.ones(self.context_len, dtype=torch.long) + return timesteps, states, actions, rtgs, traj_mask def _load_data(self, dataset: Dict[str, np.ndarray]) -> None: + """ + Overview: + Load the dataset. + Arguments: + - dataset (:obj:`Dict[str, np.ndarray]`): The dataset. + """ + self._data = {} for k in dataset.keys(): logging.info(f'Load {k} data.') self._data[k] = dataset[k][:] - def _cal_statistics(self, eps=1e-3): + def _cal_statistics(self, eps: float = 1e-3): + """ + Overview: + Calculate the statistics of the dataset. + Arguments: + - eps (:obj:`float`): Epsilon. + """ + self._mean = self._data['obs'].mean(0) self._std = self._data['obs'].std(0) + eps action_max = self._data['action'].max(0) @@ -159,28 +362,59 @@ def _cal_statistics(self, eps=1e-3): self._action_bounds = np.stack([action_min, action_max], axis=0) def _normalize_states(self): + """ + Overview: + Normalize the states. + """ + self._data['obs'] = (self._data['obs'] - self._mean) / self._std self._data['next_obs'] = (self._data['next_obs'] - self._mean) / self._std @property def mean(self): + """ + Overview: + Get the mean of the dataset. + """ + return self._mean @property def std(self): + """ + Overview: + Get the std of the dataset. + """ + return self._std @property def action_bounds(self) -> np.ndarray: + """ + Overview: + Get the action bounds of the dataset. + """ + return self._action_bounds @property def statistics(self) -> dict: + """ + Overview: + Get the statistics of the dataset. + """ + return DatasetStatistics(mean=self.mean, std=self.std, action_bounds=self.action_bounds) @DATASET_REGISTRY.register('d4rl_trajectory') class D4RLTrajectoryDataset(Dataset): + """ + Overview: + D4RL trajectory dataset, which is used for offline RL algorithms. + Interfaces: + ``__init__``, ``__len__``, ``__getitem__`` + """ # from infos.py from official d4rl github repo REF_MIN_SCORE = { @@ -322,7 +556,344 @@ class D4RLTrajectoryDataset(Dataset): }, } + def __init__(self, cfg: dict) -> None: + """ + Overview: + Initialization method. + Arguments: + - cfg (:obj:`dict`): Config dict. + """ + + dataset_path = cfg.dataset.data_dir_prefix + rtg_scale = cfg.dataset.rtg_scale + self.context_len = cfg.dataset.context_len + self.env_type = cfg.dataset.env_type + + if 'hdf5' in dataset_path: # for mujoco env + try: + import h5py + import collections + except ImportError: + import sys + logging.warning("not found h5py package, please install it trough `pip install h5py ") + sys.exit(1) + dataset = h5py.File(dataset_path, 'r') + + N = dataset['rewards'].shape[0] + data_ = collections.defaultdict(list) + + use_timeouts = False + if 'timeouts' in dataset: + use_timeouts = True + + episode_step = 0 + paths = [] + for i in range(N): + done_bool = bool(dataset['terminals'][i]) + if use_timeouts: + final_timestep = dataset['timeouts'][i] + else: + final_timestep = (episode_step == 1000 - 1) + for k in ['observations', 'actions', 'rewards', 'terminals']: + data_[k].append(dataset[k][i]) + if done_bool or final_timestep: + episode_step = 0 + episode_data = {} + for k in data_: + episode_data[k] = np.array(data_[k]) + paths.append(episode_data) + data_ = collections.defaultdict(list) + episode_step += 1 + + self.trajectories = paths + + # calculate state mean and variance and returns_to_go for all traj + states = [] + for traj in self.trajectories: + traj_len = traj['observations'].shape[0] + states.append(traj['observations']) + # calculate returns to go and rescale them + traj['returns_to_go'] = discount_cumsum(traj['rewards'], 1.0) / rtg_scale + + # used for input normalization + states = np.concatenate(states, axis=0) + self.state_mean, self.state_std = np.mean(states, axis=0), np.std(states, axis=0) + 1e-6 + + # normalize states + for traj in self.trajectories: + traj['observations'] = (traj['observations'] - self.state_mean) / self.state_std + + elif 'pkl' in dataset_path: + if 'dqn' in dataset_path: + # load dataset + with open(dataset_path, 'rb') as f: + self.trajectories = pickle.load(f) + + if isinstance(self.trajectories[0], list): + # for our collected dataset, e.g. cartpole/lunarlander case + trajectories_tmp = [] + + original_keys = ['obs', 'next_obs', 'action', 'reward'] + keys = ['observations', 'next_observations', 'actions', 'rewards'] + trajectories_tmp = [ + { + key: np.stack( + [ + self.trajectories[eps_index][transition_index][o_key] + for transition_index in range(len(self.trajectories[eps_index])) + ], + axis=0 + ) + for key, o_key in zip(keys, original_keys) + } for eps_index in range(len(self.trajectories)) + ] + self.trajectories = trajectories_tmp + + states = [] + for traj in self.trajectories: + # traj_len = traj['observations'].shape[0] + states.append(traj['observations']) + # calculate returns to go and rescale them + traj['returns_to_go'] = discount_cumsum(traj['rewards'], 1.0) / rtg_scale + + # used for input normalization + states = np.concatenate(states, axis=0) + self.state_mean, self.state_std = np.mean(states, axis=0), np.std(states, axis=0) + 1e-6 + + # normalize states + for traj in self.trajectories: + traj['observations'] = (traj['observations'] - self.state_mean) / self.state_std + else: + # load dataset + with open(dataset_path, 'rb') as f: + self.trajectories = pickle.load(f) + + states = [] + for traj in self.trajectories: + states.append(traj['observations']) + # calculate returns to go and rescale them + traj['returns_to_go'] = discount_cumsum(traj['rewards'], 1.0) / rtg_scale + + # used for input normalization + states = np.concatenate(states, axis=0) + self.state_mean, self.state_std = np.mean(states, axis=0), np.std(states, axis=0) + 1e-6 + + # normalize states + for traj in self.trajectories: + traj['observations'] = (traj['observations'] - self.state_mean) / self.state_std + else: + # -- load data from memory (make more efficient) + obss = [] + actions = [] + returns = [0] + done_idxs = [] + stepwise_returns = [] + + transitions_per_buffer = np.zeros(50, dtype=int) + num_trajectories = 0 + while len(obss) < cfg.dataset.num_steps: + buffer_num = np.random.choice(np.arange(50 - cfg.dataset.num_buffers, 50), 1)[0] + i = transitions_per_buffer[buffer_num] + frb = FixedReplayBuffer( + data_dir=cfg.dataset.data_dir_prefix + '/1/replay_logs', + replay_suffix=buffer_num, + observation_shape=(84, 84), + stack_size=4, + update_horizon=1, + gamma=0.99, + observation_dtype=np.uint8, + batch_size=32, + replay_capacity=100000 + ) + if frb._loaded_buffers: + done = False + curr_num_transitions = len(obss) + trajectories_to_load = cfg.dataset.trajectories_per_buffer + while not done: + states, ac, ret, next_states, next_action, next_reward, terminal, indices = \ + frb.sample_transition_batch(batch_size=1, indices=[i]) + states = states.transpose((0, 3, 1, 2))[0] # (1, 84, 84, 4) --> (4, 84, 84) + obss.append(states) + actions.append(ac[0]) + stepwise_returns.append(ret[0]) + if terminal[0]: + done_idxs.append(len(obss)) + returns.append(0) + if trajectories_to_load == 0: + done = True + else: + trajectories_to_load -= 1 + returns[-1] += ret[0] + i += 1 + if i >= 100000: + obss = obss[:curr_num_transitions] + actions = actions[:curr_num_transitions] + stepwise_returns = stepwise_returns[:curr_num_transitions] + returns[-1] = 0 + i = transitions_per_buffer[buffer_num] + done = True + num_trajectories += (cfg.dataset.trajectories_per_buffer - trajectories_to_load) + transitions_per_buffer[buffer_num] = i + + actions = np.array(actions) + returns = np.array(returns) + stepwise_returns = np.array(stepwise_returns) + done_idxs = np.array(done_idxs) + + # -- create reward-to-go dataset + start_index = 0 + rtg = np.zeros_like(stepwise_returns) + for i in done_idxs: + i = int(i) + curr_traj_returns = stepwise_returns[start_index:i] + for j in range(i - 1, start_index - 1, -1): # start from i-1 + rtg_j = curr_traj_returns[j - start_index:i - start_index] + rtg[j] = sum(rtg_j) + start_index = i + + # -- create timestep dataset + start_index = 0 + timesteps = np.zeros(len(actions) + 1, dtype=int) + for i in done_idxs: + i = int(i) + timesteps[start_index:i + 1] = np.arange(i + 1 - start_index) + start_index = i + 1 + + self.obss = obss + self.actions = actions + self.done_idxs = done_idxs + self.rtgs = rtg + self.timesteps = timesteps + # return obss, actions, returns, done_idxs, rtg, timesteps + + def get_max_timestep(self) -> int: + """ + Overview: + Get the max timestep of the dataset. + """ + + return max(self.timesteps) + + def get_state_stats(self) -> Tuple[np.ndarray, np.ndarray]: + """ + Overview: + Get the state mean and std of the dataset. + """ + + return deepcopy(self.state_mean), deepcopy(self.state_std) + + def get_d4rl_dataset_stats(self, env_d4rl_name: str) -> Dict[str, list]: + """ + Overview: + Get the d4rl dataset stats. + Arguments: + - env_d4rl_name (:obj:`str`): The d4rl env name. + """ + + return self.D4RL_DATASET_STATS[env_d4rl_name] + + def __len__(self) -> int: + """ + Overview: + Get the length of the dataset. + """ + + if self.env_type != 'atari': + return len(self.trajectories) + else: + return len(self.obss) - self.context_len + + def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Overview: + Get the item of the dataset. + Arguments: + - idx (:obj:`int`): The index of the dataset. + """ + + if self.env_type != 'atari': + traj = self.trajectories[idx] + traj_len = traj['observations'].shape[0] + + if traj_len > self.context_len: + # sample random index to slice trajectory + si = np.random.randint(0, traj_len - self.context_len) + + states = torch.from_numpy(traj['observations'][si:si + self.context_len]) + actions = torch.from_numpy(traj['actions'][si:si + self.context_len]) + returns_to_go = torch.from_numpy(traj['returns_to_go'][si:si + self.context_len]) + timesteps = torch.arange(start=si, end=si + self.context_len, step=1) + + # all ones since no padding + traj_mask = torch.ones(self.context_len, dtype=torch.long) + + else: + padding_len = self.context_len - traj_len + + # padding with zeros + states = torch.from_numpy(traj['observations']) + states = torch.cat( + [states, torch.zeros(([padding_len] + list(states.shape[1:])), dtype=states.dtype)], dim=0 + ) + + actions = torch.from_numpy(traj['actions']) + actions = torch.cat( + [actions, torch.zeros(([padding_len] + list(actions.shape[1:])), dtype=actions.dtype)], dim=0 + ) + + returns_to_go = torch.from_numpy(traj['returns_to_go']) + returns_to_go = torch.cat( + [ + returns_to_go, + torch.zeros(([padding_len] + list(returns_to_go.shape[1:])), dtype=returns_to_go.dtype) + ], + dim=0 + ) + + timesteps = torch.arange(start=0, end=self.context_len, step=1) + + traj_mask = torch.cat( + [torch.ones(traj_len, dtype=torch.long), + torch.zeros(padding_len, dtype=torch.long)], dim=0 + ) + return timesteps, states, actions, returns_to_go, traj_mask + else: # mean cost less than 0.001s + block_size = self.context_len + done_idx = idx + block_size + for i in self.done_idxs: + if i > idx: # first done_idx greater than idx + done_idx = min(int(i), done_idx) + break + idx = done_idx - block_size + states = torch.as_tensor( + np.array(self.obss[idx:done_idx]), dtype=torch.float32 + ).view(block_size, -1) # (block_size, 4*84*84) + states = states / 255. + actions = torch.as_tensor(self.actions[idx:done_idx], dtype=torch.long).unsqueeze(1) # (block_size, 1) + rtgs = torch.as_tensor(self.rtgs[idx:done_idx], dtype=torch.float32).unsqueeze(1) + timesteps = torch.as_tensor(self.timesteps[idx:idx + 1], dtype=torch.int64).unsqueeze(1) + traj_mask = torch.ones(self.context_len, dtype=torch.long) + return timesteps, states, actions, rtgs, traj_mask + + +@DATASET_REGISTRY.register('d4rl_diffuser') +class D4RLDiffuserDataset(Dataset): + """ + Overview: + D4RL diffuser dataset, which is used for offline RL algorithms. + Interfaces: + ``__init__``, ``__len__``, ``__getitem__`` + """ + def __init__(self, dataset_path: str, context_len: int, rtg_scale: float) -> None: + """ + Overview: + Initialization method of D4RLDiffuserDataset. + Arguments: + - dataset_path (:obj:`str`): The dataset path. + - context_len (:obj:`int`): The length of the context. + - rtg_scale (:obj:`float`): The scale of the returns to go. + """ self.context_len = context_len @@ -341,7 +912,7 @@ def __init__(self, dataset_path: str, context_len: int, rtg_scale: float) -> Non { key: np.stack( [ - self.trjectories[eps_index][transition_index][o_key] + self.trajectories[eps_index][transition_index][o_key] for transition_index in range(len(self.trajectories[eps_index])) ], axis=0 @@ -365,94 +936,573 @@ def __init__(self, dataset_path: str, context_len: int, rtg_scale: float) -> Non for traj in self.trajectories: traj['observations'] = (traj['observations'] - self.state_mean) / self.state_std - def get_state_stats(self) -> Tuple[np.ndarray, np.ndarray]: - return deepcopy(self.state_mean), deepcopy(self.state_std) - - def get_d4rl_dataset_stats(self, env_d4rl_name: str) -> Dict[str, list]: - return self.D4RL_DATASET_STATS[env_d4rl_name] - def __len__(self) -> int: - return len(self.trajectories) - - def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - traj = self.trajectories[idx] - traj_len = traj['observations'].shape[0] +class FixedReplayBuffer(object): + """ + Overview: + Object composed of a list of OutofGraphReplayBuffers. + Interfaces: + ``__init__``, ``get_transition_elements``, ``sample_transition_batch`` + """ + + def __init__(self, data_dir: str, replay_suffix: int, *args, **kwargs): + """ + Overview: + Initialize the FixedReplayBuffer class. + Arguments: + - data_dir (:obj:`str`): Log directory from which to load the replay buffer. + - replay_suffix (:obj:`int`): If not None, then only load the replay buffer \ + corresponding to the specific suffix in data directory. + - args (:obj:`list`): Arbitrary extra arguments. + - kwargs (:obj:`dict`): Arbitrary keyword arguments. + """ + + self._args = args + self._kwargs = kwargs + self._data_dir = data_dir + self._loaded_buffers = False + self.add_count = np.array(0) + self._replay_suffix = replay_suffix + if not self._loaded_buffers: + if replay_suffix is not None: + assert replay_suffix >= 0, 'Please pass a non-negative replay suffix' + self.load_single_buffer(replay_suffix) + else: + pass + # self._load_replay_buffers(num_buffers=50) + + def load_single_buffer(self, suffix): + """ + Overview: + Load a single replay buffer. + Arguments: + - suffix (:obj:`int`): The suffix of the replay buffer. + """ + + replay_buffer = self._load_buffer(suffix) + if replay_buffer is not None: + self._replay_buffers = [replay_buffer] + self.add_count = replay_buffer.add_count + self._num_replay_buffers = 1 + self._loaded_buffers = True + + def _load_buffer(self, suffix): + """ + Overview: + Loads a OutOfGraphReplayBuffer replay buffer. + Arguments: + - suffix (:obj:`int`): The suffix of the replay buffer. + """ - if traj_len >= self.context_len: - # sample random index to slice trajectory - si = np.random.randint(0, traj_len - self.context_len) - - states = torch.from_numpy(traj['observations'][si:si + self.context_len]) - actions = torch.from_numpy(traj['actions'][si:si + self.context_len]) - returns_to_go = torch.from_numpy(traj['returns_to_go'][si:si + self.context_len]) - timesteps = torch.arange(start=si, end=si + self.context_len, step=1) - - # all ones since no padding - traj_mask = torch.ones(self.context_len, dtype=torch.long) + try: + from dopamine.replay_memory import circular_replay_buffer + STORE_FILENAME_PREFIX = circular_replay_buffer.STORE_FILENAME_PREFIX + # pytype: disable=attribute-error + replay_buffer = circular_replay_buffer.OutOfGraphReplayBuffer(*self._args, **self._kwargs) + replay_buffer.load(self._data_dir, suffix) + # pytype: enable=attribute-error + return replay_buffer + # except tf.errors.NotFoundError: + except: + raise ('can not load') + + def get_transition_elements(self): + """ + Overview: + Returns the transition elements. + """ + + return self._replay_buffers[0].get_transition_elements() + + def sample_transition_batch(self, batch_size=None, indices=None): + """ + Overview: + Returns a batch of transitions (including any extra contents). + Arguments: + - batch_size (:obj:`int`): The batch size. + - indices (:obj:`list`): The indices of the batch. + """ + + buffer_index = np.random.randint(self._num_replay_buffers) + return self._replay_buffers[buffer_index].sample_transition_batch(batch_size=batch_size, indices=indices) + + +class PCDataset(Dataset): + """ + Overview: + Dataset for Procedure Cloning. + Interfaces: + ``__init__``, ``__len__``, ``__getitem__`` + """ + + def __init__(self, all_data): + """ + Overview: + Initialization method of PCDataset. + Arguments: + - all_data (:obj:`tuple`): The tuple of all data. + """ + + self._data = all_data + + def __getitem__(self, item): + """ + Overview: + Get the item of the dataset. + Arguments: + - item (:obj:`int`): The index of the dataset. + """ + + return {'obs': self._data[0][item], 'bfs_in': self._data[1][item], 'bfs_out': self._data[2][item]} + def __len__(self): + """ + Overview: + Get the length of the dataset. + """ + + return self._data[0].shape[0] + + +def load_bfs_datasets(train_seeds=1, test_seeds=5): + """ + Overview: + Load BFS datasets. + Arguments: + - train_seeds (:obj:`int`): The number of train seeds. + - test_seeds (:obj:`int`): The number of test seeds. + """ + + from dizoo.maze.envs import Maze + + def load_env(seed): + ccc = easydict.EasyDict({'size': 16}) + e = Maze(ccc) + e.seed(seed) + e.reset() + return e + + envs = [load_env(i) for i in range(train_seeds + test_seeds)] + + observations_train = [] + observations_test = [] + bfs_input_maps_train = [] + bfs_input_maps_test = [] + bfs_output_maps_train = [] + bfs_output_maps_test = [] + for idx, env in enumerate(envs): + if idx < train_seeds: + observations = observations_train + bfs_input_maps = bfs_input_maps_train + bfs_output_maps = bfs_output_maps_train else: - padding_len = self.context_len - traj_len - - # padding with zeros - states = torch.from_numpy(traj['observations']) - states = torch.cat( - [states, torch.zeros(([padding_len] + list(states.shape[1:])), dtype=states.dtype)], dim=0 - ) - - actions = torch.from_numpy(traj['actions']) - actions = torch.cat( - [actions, torch.zeros(([padding_len] + list(actions.shape[1:])), dtype=actions.dtype)], dim=0 - ) - - returns_to_go = torch.from_numpy(traj['returns_to_go']) - returns_to_go = torch.cat( - [ - returns_to_go, - torch.zeros(([padding_len] + list(returns_to_go.shape[1:])), dtype=returns_to_go.dtype) - ], - dim=0 - ) - - timesteps = torch.arange(start=0, end=self.context_len, step=1) - - traj_mask = torch.cat( - [torch.ones(traj_len, dtype=torch.long), - torch.zeros(padding_len, dtype=torch.long)], dim=0 - ) - - return timesteps, states, actions, returns_to_go, traj_mask + observations = observations_test + bfs_input_maps = bfs_input_maps_test + bfs_output_maps = bfs_output_maps_test + + start_obs = env.process_states(env._get_obs(), env.get_maze_map()) + _, track_back = get_vi_sequence(env, start_obs) + env_observations = torch.stack([track_back[i][0] for i in range(len(track_back))], dim=0) + + for i in range(env_observations.shape[0]): + bfs_sequence, _ = get_vi_sequence(env, env_observations[i].numpy().astype(np.int32)) # [L, W, W] + bfs_input_map = env.n_action * np.ones([env.size, env.size], dtype=np.long) + + for j in range(bfs_sequence.shape[0]): + bfs_input_maps.append(torch.from_numpy(bfs_input_map)) + bfs_output_maps.append(torch.from_numpy(bfs_sequence[j])) + observations.append(env_observations[i]) + bfs_input_map = bfs_sequence[j] + + train_data = PCDataset( + ( + torch.stack(observations_train, dim=0), + torch.stack(bfs_input_maps_train, dim=0), + torch.stack(bfs_output_maps_train, dim=0), + ) + ) + test_data = PCDataset( + ( + torch.stack(observations_test, dim=0), + torch.stack(bfs_input_maps_test, dim=0), + torch.stack(bfs_output_maps_test, dim=0), + ) + ) + + return train_data, test_data @DATASET_REGISTRY.register('bco') class BCODataset(Dataset): + """ + Overview: + Dataset for Behavioral Cloning from Observation. + Interfaces: + ``__init__``, ``__len__``, ``__getitem__`` + Properties: + - obs (:obj:`np.ndarray`): The observation array. + - action (:obj:`np.ndarray`): The action array. + """ def __init__(self, data=None): + """ + Overview: + Initialization method of BCODataset. + Arguments: + - data (:obj:`dict`): The data dict. + """ + if data is None: raise ValueError('Dataset can not be empty!') else: self._data = data def __len__(self): + """ + Overview: + Get the length of the dataset. + """ + return len(self._data['obs']) def __getitem__(self, idx): + """ + Overview: + Get the item of the dataset. + Arguments: + - idx (:obj:`int`): The index of the dataset. + """ + return {k: self._data[k][idx] for k in self._data.keys()} @property def obs(self): + """ + Overview: + Get the observation array. + """ + return self._data['obs'] @property def action(self): + """ + Overview: + Get the action array. + """ + return self._data['action'] +@DATASET_REGISTRY.register('diffuser_traj') +class SequenceDataset(torch.utils.data.Dataset): + """ + Overview: + Dataset for diffuser. + Interfaces: + ``__init__``, ``__len__``, ``__getitem__`` + """ + + def __init__(self, cfg): + """ + Overview: + Initialization method of SequenceDataset. + Arguments: + - cfg (:obj:`dict`): The config dict. + """ + + import gym + + env_id = cfg.env.env_id + data_path = cfg.policy.collect.get('data_path', None) + env = gym.make(env_id) + + dataset = env.get_dataset() + + self.returns_scale = cfg.env.returns_scale + self.horizon = cfg.env.horizon + self.max_path_length = cfg.env.max_path_length + self.discount = cfg.policy.learn.discount_factor + self.discounts = self.discount ** np.arange(self.max_path_length)[:, None] + self.use_padding = cfg.env.use_padding + self.include_returns = cfg.env.include_returns + self.env_id = cfg.env.env_id + itr = self.sequence_dataset(env, dataset) + self.n_episodes = 0 + + fields = {} + for k in dataset.keys(): + if 'metadata' in k: + continue + fields[k] = [] + fields['path_lengths'] = [] + + for i, episode in enumerate(itr): + path_length = len(episode['observations']) + assert path_length <= self.max_path_length + fields['path_lengths'].append(path_length) + for key, val in episode.items(): + if key not in fields: + fields[key] = [] + if val.ndim < 2: + val = np.expand_dims(val, axis=-1) + shape = (self.max_path_length, val.shape[-1]) + arr = np.zeros(shape, dtype=np.float32) + arr[:path_length] = val + fields[key].append(arr) + if episode['terminals'].any() and cfg.env.termination_penalty and 'timeouts' in episode: + assert not episode['timeouts'].any(), 'Penalized a timeout episode for early termination' + fields['rewards'][-1][path_length - 1] += cfg.env.termination_penalty + self.n_episodes += 1 + + for k in fields.keys(): + fields[k] = np.array(fields[k]) + + self.normalizer = DatasetNormalizer(fields, cfg.policy.normalizer, path_lengths=fields['path_lengths']) + self.indices = self.make_indices(fields['path_lengths'], self.horizon) + + self.observation_dim = cfg.env.obs_dim + self.action_dim = cfg.env.action_dim + self.fields = fields + self.normalize() + self.normed = False + if cfg.env.normed: + self.vmin, self.vmax = self._get_bounds() + self.normed = True + + # shapes = {key: val.shape for key, val in self.fields.items()} + # print(f'[ datasets/mujoco ] Dataset fields: {shapes}') + + def sequence_dataset(self, env, dataset=None): + """ + Overview: + Sequence the dataset. + Arguments: + - env (:obj:`gym.Env`): The gym env. + """ + + import collections + N = dataset['rewards'].shape[0] + if 'maze2d' in env.spec.id: + dataset = self.maze2d_set_terminals(env, dataset) + data_ = collections.defaultdict(list) + + # The newer version of the dataset adds an explicit + # timeouts field. Keep old method for backwards compatability. + use_timeouts = 'timeouts' in dataset + + episode_step = 0 + for i in range(N): + done_bool = bool(dataset['terminals'][i]) + if use_timeouts: + final_timestep = dataset['timeouts'][i] + else: + final_timestep = (episode_step == env._max_episode_steps - 1) + + for k in dataset: + if 'metadata' in k: + continue + data_[k].append(dataset[k][i]) + + if done_bool or final_timestep: + episode_step = 0 + episode_data = {} + for k in data_: + episode_data[k] = np.array(data_[k]) + if 'maze2d' in env.spec.id: + episode_data = self.process_maze2d_episode(episode_data) + yield episode_data + data_ = collections.defaultdict(list) + + episode_step += 1 + + def maze2d_set_terminals(self, env, dataset): + """ + Overview: + Set the terminals for maze2d. + Arguments: + - env (:obj:`gym.Env`): The gym env. + - dataset (:obj:`dict`): The dataset dict. + """ + + goal = env.get_target() + threshold = 0.5 + + xy = dataset['observations'][:, :2] + distances = np.linalg.norm(xy - goal, axis=-1) + at_goal = distances < threshold + timeouts = np.zeros_like(dataset['timeouts']) + + # timeout at time t iff + # at goal at time t and + # not at goal at time t + 1 + timeouts[:-1] = at_goal[:-1] * ~at_goal[1:] + + timeout_steps = np.where(timeouts)[0] + path_lengths = timeout_steps[1:] - timeout_steps[:-1] + + print( + f'[ utils/preprocessing ] Segmented {env.spec.id} | {len(path_lengths)} paths | ' + f'min length: {path_lengths.min()} | max length: {path_lengths.max()}' + ) + + dataset['timeouts'] = timeouts + return dataset + + def process_maze2d_episode(self, episode): + """ + Overview: + Process the maze2d episode, adds in `next_observations` field to episode. + Arguments: + - episode (:obj:`dict`): The episode dict. + """ + + assert 'next_observations' not in episode + length = len(episode['observations']) + next_observations = episode['observations'][1:].copy() + for key, val in episode.items(): + episode[key] = val[:-1] + episode['next_observations'] = next_observations + return episode + + def normalize(self, keys=['observations', 'actions']): + """ + Overview: + Normalize the dataset, normalize fields that will be predicted by the diffusion model + Arguments: + - keys (:obj:`list`): The list of keys. + """ + + for key in keys: + array = self.fields[key].reshape(self.n_episodes * self.max_path_length, -1) + normed = self.normalizer.normalize(array, key) + self.fields[f'normed_{key}'] = normed.reshape(self.n_episodes, self.max_path_length, -1) + + def make_indices(self, path_lengths, horizon): + """ + Overview: + Make indices for sampling from dataset. Each index maps to a datapoint. + Arguments: + - path_lengths (:obj:`np.ndarray`): The path length array. + - horizon (:obj:`int`): The horizon. + """ + + indices = [] + for i, path_length in enumerate(path_lengths): + max_start = min(path_length - 1, self.max_path_length - horizon) + if not self.use_padding: + max_start = min(max_start, path_length - horizon) + for start in range(max_start): + end = start + horizon + indices.append((i, start, end)) + indices = np.array(indices) + return indices + + def get_conditions(self, observations): + """ + Overview: + Get the conditions on current observation for planning. + Arguments: + - observations (:obj:`np.ndarray`): The observation array. + """ + + if 'maze2d' in self.env_id: + return {'condition_id': [0, self.horizon - 1], 'condition_val': [observations[0], observations[-1]]} + else: + return {'condition_id': [0], 'condition_val': [observations[0]]} + + def __len__(self): + """ + Overview: + Get the length of the dataset. + """ + + return len(self.indices) + + def _get_bounds(self): + """ + Overview: + Get the bounds of the dataset. + """ + + print('[ datasets/sequence ] Getting value dataset bounds...', end=' ', flush=True) + vmin = np.inf + vmax = -np.inf + for i in range(len(self.indices)): + value = self.__getitem__(i)['returns'].item() + vmin = min(value, vmin) + vmax = max(value, vmax) + print('✓') + return vmin, vmax + + def normalize_value(self, value): + """ + Overview: + Normalize the value. + Arguments: + - value (:obj:`np.ndarray`): The value array. + """ + + # [0, 1] + normed = (value - self.vmin) / (self.vmax - self.vmin) + # [-1, 1] + normed = normed * 2 - 1 + return normed + + def __getitem__(self, idx, eps=1e-4): + """ + Overview: + Get the item of the dataset. + Arguments: + - idx (:obj:`int`): The index of the dataset. + - eps (:obj:`float`): The epsilon. + """ + + path_ind, start, end = self.indices[idx] + + observations = self.fields['normed_observations'][path_ind, start:end] + actions = self.fields['normed_actions'][path_ind, start:end] + done = self.fields['terminals'][path_ind, start:end] + + # conditions = self.get_conditions(observations) + trajectories = np.concatenate([actions, observations], axis=-1) + + if self.include_returns: + rewards = self.fields['rewards'][path_ind, start:] + discounts = self.discounts[:len(rewards)] + returns = (discounts * rewards).sum() + if self.normed: + returns = self.normalize_value(returns) + returns = np.array([returns / self.returns_scale], dtype=np.float32) + batch = { + 'trajectories': trajectories, + 'returns': returns, + 'done': done, + 'action': actions, + } + else: + batch = { + 'trajectories': trajectories, + 'done': done, + 'action': actions, + } + + batch.update(self.get_conditions(observations)) + return batch + + def hdf5_save(exp_data, expert_data_path): + """ + Overview: + Save the data to hdf5. + """ + try: import h5py except ImportError: + import sys logging.warning("not found h5py package, please install it trough 'pip install h5py' ") - import numpy as np + sys.exit(1) dataset = dataset = h5py.File('%s_demos.hdf5' % expert_data_path.replace('.pkl', ''), 'w') dataset.create_dataset('obs', data=np.array([d['obs'].numpy() for d in exp_data]), compression='gzip') dataset.create_dataset('action', data=np.array([d['action'].numpy() for d in exp_data]), compression='gzip') @@ -462,15 +1512,30 @@ def hdf5_save(exp_data, expert_data_path): def naive_save(exp_data, expert_data_path): + """ + Overview: + Save the data to pickle. + """ + with open(expert_data_path, 'wb') as f: pickle.dump(exp_data, f) def offline_data_save_type(exp_data, expert_data_path, data_type='naive'): + """ + Overview: + Save the offline data. + """ + globals()[data_type + '_save'](exp_data, expert_data_path) def create_dataset(cfg, **kwargs) -> Dataset: + """ + Overview: + Create dataset. + """ + cfg = EasyDict(cfg) import_module(cfg.get('import_names', [])) return DATASET_REGISTRY.build(cfg.policy.collect.data_type, cfg=cfg, **kwargs) diff --git a/ding/utils/data/rlhf_offline_dataset.py b/ding/utils/data/rlhf_offline_dataset.py new file mode 100644 index 0000000000..2b95010f1d --- /dev/null +++ b/ding/utils/data/rlhf_offline_dataset.py @@ -0,0 +1,277 @@ +from typing import Iterable, Dict, List, Union, Any, Callable +from functools import partial +from tqdm import tqdm +from torch.utils.data import Dataset +from torch.distributed import get_rank +from transformers import AutoTokenizer +import torch +import torch.nn.functional as F + + +def zero_pad_sequences(sequences: List[torch.Tensor], side: str = "left", value: int = 0) -> torch.Tensor: + """ + Overview: + Pad sequences with zeros to create a batch tensor of uniform length. + Arguments: + - sequences (List[torch.Tensor]): A list of PyTorch tensors to be padded. + - side (str): The side to pad ('left' or 'right'), default is 'left'. + - value (int): The padding value to use, default is 0. + Returns: + - padded_sequences (torch.Tensor): A padded tensor of shape [batch_size, max_sequence_length]. + """ + assert side in ("left", "right"), side + max_len = max(seq.size(-1) for seq in sequences) + padded_sequences = [] + for seq in sequences: + pad_len = max_len - seq.size(-1) + padding = (pad_len, 0) if side == "left" else (0, pad_len) + padded_sequences.append(F.pad(seq, padding, value=value)) + return torch.stack(padded_sequences, dim=0) + + +class OfflineRLDataset(Dataset): + """ + Overview: + PyTorch Dataset for OfflineRL LLM training like KTO and DPO. + This dataset supports pure text input, as well as image, video, audio, etc. + """ + + def __init__( + self, + dataset: Iterable[Dict], + tokenizer: AutoTokenizer, + max_length: int, + input_key: str = "input", + extra_input_keys: List[str] = [], + output_key: str = "output", + label_key: str = "label", + apply_chat_template: bool = False, + tokenizer_chat_template: str = None, + input_template: str = None, + num_processors: int = 8, + parallel_load: bool = True + ) -> None: + """ + Overview: + Initialize the OfflineRLDataset. + Arguments: + - dataset (Iterable[Dict]): The iterable dataset object to be used, such as list or huggingface dataset. + - tokenizer (AutoTokenizer): The tokenizer to be used. + - max_length (int): The maximum length of the input. + - input_key (str): The key of the input, default is "input". + - extra_input_keys (List[str]): The extra input keys, such as "image", "video", "audio", etc. + - output_key (str): The key of the output, default is "output". + - label_key (str): The key of the label, default is "label". + - apply_chat_template (bool): Whether to apply the chat template, default is False. + - tokenizer_chat_template (str): The chat template to be used. + - input_template (str): The input template to be used. + - num_processors (int): The number of processors to be used, default is 8. + - parallel_load (bool): Whether to parallel load the dataset in the `__init__` method, default is True. + Parallel loading is usually used for huggingface dataset. + """ + super().__init__() + self.tokenizer = tokenizer + self.max_length = max_length + self.extra_input_keys = extra_input_keys + + if apply_chat_template: + apply_chat_template = self.tokenizer.apply_chat_template + if tokenizer_chat_template: + self.tokenizer.chat_template = tokenizer_chat_template + + # Parallel loading datasets + if parallel_load: + preprocess_data_fn = partial( + self._preprocess_data, + input_template=input_template, + input_key=input_key, + extra_input_keys=extra_input_keys, + output_key=output_key, + label_key=label_key, + apply_chat_template=apply_chat_template + ) + processed_dataset = dataset.map( + preprocess_data_fn, remove_columns=dataset.column_names, num_proc=num_processors + ) + # preprocess function may return None, so filter out the None + processed_dataset = processed_dataset.filter(lambda x: x["prompt"] is not None) + + self.prompts = processed_dataset["prompt"] + self.responses = processed_dataset["response"] + self.labels = processed_dataset["label"] + self.prompt_ids_lens = processed_dataset["prompt_ids_len"] + for key in extra_input_keys: + setattr(self, key, processed_dataset[key]) + else: + self.prompts = [] + self.responses = [] + self.labels = [] + self.prompt_ids_lens = [] + for key in extra_input_keys: + setattr(self, key, []) + for data in tqdm(dataset, desc="Preprocessing data", disable=not get_rank() == 0): + processed_data = self._preprocess_data(data) + if processed_data["prompt"] is not None: + self.prompts.append(processed_data["prompt"]) + self.responses.append(processed_data["response"]) + self.labels.append(processed_data["label"]) + self.prompt_ids_lens.append(processed_data["prompt_ids_len"]) + for key in extra_input_keys: + getattr(self, key).append(processed_data[key]) + + def _preprocess_data( + self, + data: Dict[str, Any], + input_template: str = None, + input_key: str = "input", + extra_input_keys: List[str] = [], + output_key: str = "output", + label_key: str = "label", + apply_chat_template: Union[bool, Callable] = False, + ) -> Dict[str, Any]: + """ + Overview: + Preprocess the data and return the processed data. + Arguments: + - data (Dict[str, Any]): The data to be processed. + - input_template (str): The input template to be used. + - input_key (str): The key of the input, default is "input". + - extra_input_keys (List[str]): The extra input keys, such as "image", "video", "audio", etc. + - output_key (str): The key of the output, default is "output". + - label_key (str): The key of the label, default is "label". + - apply_chat_template (Union[bool, Callable]): Controls chat template application. If True, uses the \ + tokenizer's default template. If a Callable is provided, uses that function to apply the template \ + (typically tokenizer.apply_chat_template). + Returns: + - processed_data (Dict[str, Any]): The processed data. + """ + label = data[label_key] + if extra_input_keys: + extra_inputs = {key: data[key] for key in extra_input_keys} + else: + extra_inputs = {} + + if apply_chat_template: + if output_key: + prompt = apply_chat_template(data[input_key], tokenize=False, add_generation_prompt=True) + response = apply_chat_template(data[input_key] + data[output_key], tokenize=False)[len(prompt):] + else: + prompt = apply_chat_template(data[input_key][:-1], tokenize=False, add_generation_prompt=True) + response = apply_chat_template(data[input_key], tokenize=False)[len(prompt):] + else: + prompt = data[input_key] + response = data[output_key] + if input_template: + prompt = input_template.format(prompt) + + prompt_token = self.tokenizer( + prompt, + max_length=self.max_length, + # use the batch max length (in `collate_fn`) to pad rather than the global max length + padding=False, + truncation=True, + return_tensors="pt", + # add special tokens for the prompt in `collate_fn` + add_special_tokens=False, + ) + prompt_ids_len = prompt_token["attention_mask"].int().sum().item() + + # filter the sample whose length is greater than max_length (2 for answer length) + if prompt_ids_len >= self.max_length - 2: + prompt = None + + return { + "prompt": prompt, + "response": response, + "label": label, + "prompt_ids_len": prompt_ids_len, + **extra_inputs + } + + def __len__(self) -> int: + """ + Overview: + Get the length of the dataset. + Returns: + - length (int): The length of the dataset. + """ + return len(self.prompts) + + def __getitem__(self, idx: int) -> Dict[str, Union[torch.Tensor, int]]: + """ + Overview: + Get the item at the given index. + Arguments: + - idx (int): The index of the item to get. + Returns: + - item (Dict[str, Union[torch.Tensor, int]]): The item at the given index. + """ + # extra inputs: usually image, video, audio, etc. + if self.extra_input_keys: + extra_inputs = {key: getattr(self, key)[idx] for key in self.extra_input_keys} + else: + extra_inputs = {} + return { + "prompt": self.prompts[idx], + "response": self.responses[idx], + "label": self.labels[idx], + "prompt_ids_len": self.prompt_ids_lens[idx], + **extra_inputs + } + + def collate_fn(self, item_list: List[Dict[str, Union[torch.Tensor, int]]]): + """ + Overview: + Collate the items into a batch, which is used to create a batch for training. + Arguments: + - item_list (List[Dict[str, Union[torch.Tensor, int]]]): The list of items to be collated. + Returns: + - collated_items (Dict[str, Union[torch.Tensor, int]]): The collated items. + """ + + def tokenizer(prompt: str, response: str): + text = (prompt + response).rstrip("\n") + if not text.endswith(self.tokenizer.eos_token): + text += " " + self.tokenizer.eos_token + inputs = self.tokenizer( + text, + max_length=self.max_length, + padding=False, + truncation=True, + return_tensors="pt", + add_special_tokens=False, + ) + + inputs["input_ids"][0][-1] = self.tokenizer.eos_token_id + inputs["attention_mask"][0][-1] = True + return inputs["input_ids"], inputs["attention_mask"] + + # tot_extra_inputs: Dict[str, List[torch.Tensor]] + tot_ids, tot_masks, tot_labels, prompt_ids_lens, tot_extra_inputs = [], [], [], [], {} + for item in item_list: + input_ids, attention_mask = tokenizer(item["prompt"], item["response"]) + tot_ids.append(input_ids) + tot_masks.append(attention_mask) + tot_labels.append(item["label"]) + prompt_ids_lens.append(item["prompt_ids_len"]) + for key in self.extra_input_keys: + if key not in tot_extra_inputs: + tot_extra_inputs[key] = [] + tot_extra_inputs[key].append(item[key]) + + # add unmatched y'| x (used to estimate the KL divergence between policy and reference) + for idx in range(len(item_list)): + next_idx = (idx + 1) % len(item_list) + input_ids, attention_mask = tokenizer(item_list[idx]["prompt"], item_list[next_idx]["response"]) + tot_ids.append(input_ids) + tot_masks.append(attention_mask) + tot_labels.append(-1) + prompt_ids_lens.append(item_list[idx]["prompt_ids_len"]) + for key in self.extra_input_keys: + if key not in tot_extra_inputs: + tot_extra_inputs[key] = [] + tot_extra_inputs[key].append(item_list[idx][key]) + + input_ids = zero_pad_sequences(tot_ids, side="right", value=self.tokenizer.pad_token_id) + attention_mask = zero_pad_sequences(tot_masks, side="right") + return input_ids, attention_mask, torch.LongTensor(tot_labels), prompt_ids_lens, tot_extra_inputs diff --git a/ding/utils/data/rlhf_online_dataset.py b/ding/utils/data/rlhf_online_dataset.py new file mode 100644 index 0000000000..d307f09a32 --- /dev/null +++ b/ding/utils/data/rlhf_online_dataset.py @@ -0,0 +1,99 @@ +from typing import Any, Dict, Union, Callable, Iterable +from tqdm import tqdm +from torch.utils.data import Dataset +from torch.distributed import get_rank +from transformers import AutoTokenizer + + +class OnlineRLDataset(Dataset): + """ + Overview: + PyTorch Dataset for OnlineRL LLM training like PPO. + This dataset only supports pure text input now. + """ + + def __init__( + self, + dataset: Iterable[Dict], + tokenizer: AutoTokenizer, + input_key: str = "input", + apply_chat_template: bool = False, + input_template: str = None, + ) -> None: + """ + Overview: + Initialize the OnlineRLDataset. + Arguments: + - dataset (torch.utils.data.Dataset): The dataset to preprocess. + - tokenizer (AutoTokenizer): The tokenizer to preprocess the data. + - input_key (str): The key of the input data, default is "input". + - apply_chat_template (bool): Whether to apply the chat template, default is False. + - input_template (str): The template to format the data. + """ + super().__init__() + self.tokenizer = tokenizer + self.input_template = input_template + + if apply_chat_template: + apply_chat_template = self.tokenizer.apply_chat_template + + self.prompts = [] + try: + rank = get_rank() + except ValueError: # not initialized yet, which is the case in unit test + rank = 0 + for data in tqdm(dataset, desc="Preprocessing data", disable=not rank == 0): + prompt = self._preprocess_data(data, input_template, input_key, apply_chat_template) + self.prompts.append(prompt) + + def __len__(self) -> int: + """ + Overview: + Get the length of the dataset. + Returns: + - length (int): The length of the dataset. + """ + return len(self.prompts) + + def __getitem__(self, idx: int) -> str: + """ + Overview: + Get the item at the given index. + Arguments: + - idx (int): The index of the item to get. + Returns: + - item (str): The item at the given index. + """ + return self.prompts[idx] + + def _preprocess_data( + self, + data: Dict[str, Any], + input_template: str = None, + input_key: str = "input", + apply_chat_template: Union[bool, Callable] = False, + ) -> str: + """ + Overview: + Preprocess the data to get the formatted prompt. + Arguments: + - data (Dict[str, Any]): The data to preprocess. + - input_template (str): The template to format the data. + - input_key (str): The key of the input data. + - apply_chat_template (Union[bool, Callable]): Controls chat template application. If True, uses the \ + tokenizer's default template. If a Callable is provided, uses that function to apply the template \ + (typically tokenizer.apply_chat_template). + Returns: + - prompt (str): The formatted prompt. + """ + if apply_chat_template: + chat = data[input_key] + if isinstance(chat, str): + chat = [{"role": "user", "content": chat}] + assert isinstance(chat, list) and all(isinstance(t, dict) for t in chat), "chat must be a list of dict" + prompt = apply_chat_template(chat, tokenize=False, add_generation_prompt=True) + else: + prompt = data[input_key] + if input_template: + prompt = input_template.format(prompt) + return prompt diff --git a/ding/utils/data/structure/__init__.py b/ding/utils/data/structure/__init__.py index 9e8011f9d4..3cc58828a6 100644 --- a/ding/utils/data/structure/__init__.py +++ b/ding/utils/data/structure/__init__.py @@ -1 +1,2 @@ from .cache import Cache +from .lifo_deque import LifoDeque diff --git a/ding/utils/data/structure/cache.py b/ding/utils/data/structure/cache.py index a7235220cc..faed797e3e 100644 --- a/ding/utils/data/structure/cache.py +++ b/ding/utils/data/structure/cache.py @@ -7,17 +7,17 @@ class Cache: - r""" + """ Overview: Data cache for reducing concurrent pressure, with timeout and full queue eject mechanism - Interface: - __init__, push_data, get_cached_data_iter, run, close + Interfaces: + ``__init__``, ``push_data``, ``get_cached_data_iter``, ``run``, ``close`` Property: remain_data_count """ def __init__(self, maxlen: int, timeout: float, monitor_interval: float = 1.0, _debug: bool = False) -> None: - r""" + """ Overview: Initialize the cache object. Arguments: @@ -34,13 +34,13 @@ def __init__(self, maxlen: int, timeout: float, monitor_interval: float = 1.0, _ # two separate receive and send queue for reducing interaction frequency and interference self.receive_queue = Queue(maxlen) self.send_queue = Queue(maxlen) - self.receive_lock = LockContext(type_=LockContextType.THREAD_LOCK) + self.receive_lock = LockContext(lock_type=LockContextType.THREAD_LOCK) self._timeout_thread = Thread(target=self._timeout_monitor) # the bool flag for gracefully shutting down the timeout monitor thread self._timeout_thread_flag = True def push_data(self, data: Any) -> None: - r""" + """ Overview: Push data into receive queue, if the receive queue is full(after push), then push all the data in receive queue into send queue. @@ -60,7 +60,7 @@ def push_data(self, data: Any) -> None: self.send_queue.put(self.receive_queue.get()[0]) def get_cached_data_iter(self) -> 'callable_iterator': # noqa - r""" + """ Overview: Get the iterator of the send queue. Once a data is pushed into send queue, it can be accessed by this iterator. 'STOP' is the end flag of this iterator. @@ -70,7 +70,7 @@ def get_cached_data_iter(self) -> 'callable_iterator': # noqa return iter(self.send_queue.get, 'STOP') def _timeout_monitor(self) -> None: - r""" + """ Overview: The workflow of the timeout monitor thread. """ @@ -88,7 +88,7 @@ def _timeout_monitor(self) -> None: break def _warn_if_timeout(self) -> bool: - r""" + """ Overview: Return whether is timeout. Returns @@ -107,14 +107,14 @@ def _warn_if_timeout(self) -> bool: return False def run(self) -> None: - r""" + """ Overview: Launch the cache internal thread, e.g. timeout monitor thread. """ self._timeout_thread.start() def close(self) -> None: - r""" + """ Overview: Shut down the cache internal thread and send the end flag to send queue's iterator. """ @@ -122,7 +122,7 @@ def close(self) -> None: self.send_queue.put('STOP') def dprint(self, s: str) -> None: - r""" + """ Overview: In debug mode, print debug str. Arguments: @@ -133,7 +133,7 @@ def dprint(self, s: str) -> None: @property def remain_data_count(self) -> int: - r""" + """ Overview: Return receive queue's remain data count Returns: diff --git a/ding/utils/data/structure/lifo_deque.py b/ding/utils/data/structure/lifo_deque.py new file mode 100644 index 0000000000..b18c4a0608 --- /dev/null +++ b/ding/utils/data/structure/lifo_deque.py @@ -0,0 +1,15 @@ +from queue import LifoQueue +from collections import deque + + +class LifoDeque(LifoQueue): + """ + Overview: + Like LifoQueue, but automatically replaces the oldest data when the queue is full. + Interfaces: + ``_init``, ``_put``, ``_get`` + """ + + def _init(self, maxsize): + self.maxsize = maxsize + 1 + self.queue = deque(maxlen=maxsize) diff --git a/ding/utils/data/tests/test_rlhf_offline_dataset.py b/ding/utils/data/tests/test_rlhf_offline_dataset.py new file mode 100644 index 0000000000..62167c90d3 --- /dev/null +++ b/ding/utils/data/tests/test_rlhf_offline_dataset.py @@ -0,0 +1,113 @@ +import pytest +from datasets import load_dataset, concatenate_datasets +from ding.utils.data import OfflineRLDataset +from transformers import AutoTokenizer + +IMG_CONTEXT_TOKEN = '' +IMG_START_TOKEN = '' +IMG_END_TOKEN = '' +IMG_CONTEXT_NUM = 10 # user-defined number of image patches in the context + + +@pytest.fixture +def dataset(): + # Load a sample dataset + hf_dataset = load_dataset("MMInstruction/VL-RewardBench", split='test') + # split pair data into two separate datasets + hf_dataset_1 = hf_dataset.map( + lambda x: { + "query": f"{IMG_START_TOKEN}{IMG_CONTEXT_TOKEN * IMG_CONTEXT_NUM}{IMG_END_TOKEN}\n{x['query']}", + "image": x["image"], + "response": x["response"][0], + "human_ranking": x["human_ranking"][0] + } + ) + hf_dataset_2 = hf_dataset.map( + lambda x: { + "query": f"{IMG_START_TOKEN}{IMG_CONTEXT_TOKEN * IMG_CONTEXT_NUM}{IMG_END_TOKEN}\n{x['query']}", + "image": x["image"], + "response": x["response"][1], + "human_ranking": x["human_ranking"][1] + } + ) + # combine two datasets + hf_dataset = concatenate_datasets([hf_dataset_1, hf_dataset_2]) + # shuffle the dataset + hf_dataset = hf_dataset.shuffle(seed=42) + return hf_dataset + + +@pytest.fixture +def tokenizer(): + # Load a tokenizer + return AutoTokenizer.from_pretrained("OpenGVLab/InternVL2_5-4B") + + +@pytest.mark.unittest +def test_offline_rl_dataset_initialization(dataset, tokenizer): + # Test the initialization of the OfflineRLDataset + offline_dataset = OfflineRLDataset( + dataset=dataset, + tokenizer=tokenizer, + max_length=1024, + input_key="query", + extra_input_keys=["image"], + output_key="response", + label_key="human_ranking" + ) + assert len(offline_dataset) == len(dataset) + offline_dataset = OfflineRLDataset( + dataset=dataset, + tokenizer=tokenizer, + max_length=256, + input_key="query", + extra_input_keys=["image"], + output_key="response", + label_key="human_ranking" + ) + # lower max_length will filter out some samples + assert len(offline_dataset) < len(dataset) + + +@pytest.mark.unittest +def test_offline_rl_dataset_item_retrieval(dataset, tokenizer): + # Test retrieving an item from the OfflineRLDataset + offline_dataset = OfflineRLDataset( + dataset=dataset, + tokenizer=tokenizer, + max_length=256, + input_key="query", + extra_input_keys=["image"], + output_key="response", + label_key="human_ranking" + ) + item = offline_dataset[0] + assert "prompt" in item + assert "response" in item + assert "label" in item + assert "prompt_ids_len" in item + assert "image" in item + print(item) + + +@pytest.mark.unittest +def test_offline_rl_dataset_collate_fn(dataset, tokenizer): + # Test the collate function of the OfflineRLDataset + offline_dataset = OfflineRLDataset( + dataset=dataset, + tokenizer=tokenizer, + max_length=256, + input_key="query", + output_key="response", + label_key="human_ranking" + ) + B = 10 + item_list = [offline_dataset[i] for i in range(B)] + input_ids, attention_mask, labels, prompt_ids_lens, extra_inputs = offline_dataset.collate_fn(item_list) + assert input_ids.size(0) == len(item_list) * 2 # because of the unmatched y'| x + assert attention_mask.size(0) == len(item_list) * 2 + assert labels.size(0) == len(item_list) * 2 + assert len(prompt_ids_lens) == len(item_list) * 2 + for key in offline_dataset.extra_input_keys: + assert key in extra_inputs + assert extra_inputs[key].size(0) == len(item_list) * 2 diff --git a/ding/utils/data/tests/test_rlhf_online_dataset.py b/ding/utils/data/tests/test_rlhf_online_dataset.py new file mode 100644 index 0000000000..cba9e7947c --- /dev/null +++ b/ding/utils/data/tests/test_rlhf_online_dataset.py @@ -0,0 +1,39 @@ +import pytest +from datasets import load_dataset +from transformers import AutoTokenizer +from ding.utils.data import OnlineRLDataset + + +@pytest.fixture +def dataset(): + # Load the dataset + hf_dataset = load_dataset("cat-searcher/minif2f-lean4")['validation'] + print(hf_dataset) + return hf_dataset + + +@pytest.fixture +def tokenizer(): + return AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-7B") + + +@pytest.mark.unittest +def test_onlinerl_dataset_initialization(dataset, tokenizer): + # Initialize OnlineRLDataset + online_rl_dataset = OnlineRLDataset( + dataset=dataset, tokenizer=tokenizer, input_key="formal_statement", apply_chat_template=True + ) + # Check if the dataset is initialized correctly + assert len(online_rl_dataset) == len(dataset) + + +@pytest.mark.unittest +def test_onlinerl_dataset_getitem(dataset, tokenizer): + # Initialize OnlineRLDataset + online_rl_dataset = OnlineRLDataset( + dataset=dataset, tokenizer=tokenizer, input_key="formal_statement", apply_chat_template=True + ) + # Check if __getitem__ returns the expected formatted prompt + item = online_rl_dataset[0] + print(item) + assert isinstance(item, str) diff --git a/ding/utils/default_helper.py b/ding/utils/default_helper.py index 108c08d58c..1881ca6cc0 100644 --- a/ding/utils/default_helper.py +++ b/ding/utils/default_helper.py @@ -5,13 +5,44 @@ from functools import lru_cache # in python3.9, we can change to cache import numpy as np import torch +import treetensor.torch as ttorch + + +def get_shape0(data: Union[List, Dict, torch.Tensor, ttorch.Tensor]) -> int: + """ + Overview: + Get shape[0] of data's torch tensor or treetensor + Arguments: + - data (:obj:`Union[List,Dict,torch.Tensor,ttorch.Tensor]`): data to be analysed + Returns: + - shape[0] (:obj:`int`): first dimension length of data, usually the batchsize. + """ + if isinstance(data, list) or isinstance(data, tuple): + return get_shape0(data[0]) + elif isinstance(data, dict): + for k, v in data.items(): + return get_shape0(v) + elif isinstance(data, torch.Tensor): + return data.shape[0] + elif isinstance(data, ttorch.Tensor): + + def fn(t): + item = list(t.values())[0] + if np.isscalar(item[0]): + return item[0] + else: + return fn(item) + + return fn(data.shape) + else: + raise TypeError("Error in getting shape0, not support type: {}".format(data)) def lists_to_dicts( data: Union[List[Union[dict, NamedTuple]], Tuple[Union[dict, NamedTuple]]], recursive: bool = False, ) -> Union[Mapping[object, object], NamedTuple]: - r""" + """ Overview: Transform a list of dicts to a dict of lists. Arguments: @@ -46,7 +77,7 @@ def lists_to_dicts( def dicts_to_lists(data: Mapping[object, List[object]]) -> List[Mapping[object, object]]: - r""" + """ Overview: Transform a dict of lists to a list of dicts. @@ -90,6 +121,8 @@ def squeeze(data: object) -> object: """ Overview: Squeeze data from tuple, list or dict to single object + Arguments: + - data (:obj:`object`): data to be squeezed Example: >>> a = (4, ) >>> a = squeeze(a) @@ -117,7 +150,7 @@ def default_get( default_fn: Optional[Callable] = None, judge_fn: Optional[Callable] = None ) -> Any: - r""" + """ Overview: Getting the value by input, checks generically on the inputs with \ at least ``data`` and ``name``. If ``name`` exists in ``data``, \ @@ -149,7 +182,7 @@ def default_get( def list_split(data: list, step: int) -> List[list]: - r""" + """ Overview: Split list of data by step. Arguments: @@ -179,7 +212,7 @@ def list_split(data: list, step: int) -> List[list]: def error_wrapper(fn, default_ret, warning_msg=""): - r""" + """ Overview: wrap the function, so that any Exception in the function will be catched and return the default_ret Arguments: @@ -208,10 +241,10 @@ def wrapper(*args, **kwargs): class LimitedSpaceContainer: - r""" + """ Overview: A space simulator. - Interface: + Interfaces: ``__init__``, ``get_residual_space``, ``release_space`` """ @@ -407,21 +440,49 @@ def set_pkg_seed(seed: int, use_cuda: bool = True) -> None: @lru_cache() def one_time_warning(warning_msg: str) -> None: + """ + Overview: + Print warning message only once. + Arguments: + - warning_msg (:obj:`str`): Warning message. + """ + logging.warning(warning_msg) def split_fn(data, indices, start, end): + """ + Overview: + Split data by indices + Arguments: + - data (:obj:`Union[List, Dict, torch.Tensor, ttorch.Tensor]`): data to be analysed + - indices (:obj:`np.ndarray`): indices to split + - start (:obj:`int`): start index + - end (:obj:`int`): end index + """ + if data is None: return None elif isinstance(data, list): return [split_fn(d, indices, start, end) for d in data] elif isinstance(data, dict): return {k1: split_fn(v1, indices, start, end) for k1, v1 in data.items()} + elif isinstance(data, str): + return data else: return data[indices[start:end]] def split_data_generator(data: dict, split_size: int, shuffle: bool = True) -> dict: + """ + Overview: + Split data into batches + Arguments: + - data (:obj:`dict`): data to be analysed + - split_size (:obj:`int`): split size + - shuffle (:obj:`bool`): whether shuffle + """ + assert isinstance(data, dict), type(data) length = [] for k, v in data.items(): @@ -430,7 +491,12 @@ def split_data_generator(data: dict, split_size: int, shuffle: bool = True) -> d elif k in ['prev_state', 'prev_actor_state', 'prev_critic_state']: length.append(len(v)) elif isinstance(v, list) or isinstance(v, tuple): - length.append(len(v[0])) + if isinstance(v[0], str): + # some buffer data contains useless string infos, such as 'buffer_id', + # which should not be split, so we just skip it + continue + else: + length.append(get_shape0(v[0])) elif isinstance(v, dict): length.append(len(v[list(v.keys())[0]])) else: @@ -439,7 +505,7 @@ def split_data_generator(data: dict, split_size: int, shuffle: bool = True) -> d # assert len(set(length)) == 1, "data values must have the same length: {}".format(length) # if continuous action, data['logit'] is list of length 2 length = length[0] - assert split_size >= 1 and split_size <= length + assert split_size >= 1 if shuffle: indices = np.random.permutation(length) else: @@ -455,7 +521,7 @@ class RunningMeanStd(object): """ Overview: Wrapper to update new variable, new mean, and new count - Interface: + Interfaces: ``__init__``, ``update``, ``reset``, ``new_shape`` Properties: - ``mean``, ``std``, ``_epsilon``, ``_shape``, ``_mean``, ``_var``, ``_count`` diff --git a/ding/utils/deprecation.py b/ding/utils/deprecation.py new file mode 100644 index 0000000000..b71c14977f --- /dev/null +++ b/ding/utils/deprecation.py @@ -0,0 +1,52 @@ +import functools +import textwrap +import warnings +from typing import Optional + + +def deprecated(since: str, removed_in: str, up_to: Optional[str] = None): + """ + Overview: + Decorate a function to signify its deprecation. + Arguments: + - since (:obj:`str`): the version when the function was first deprecated. + - removed_in (:obj:`str`): the version when the function will be removed. + - up_to (:obj:`Optional[str]`): the new API users should use. + Returns: + - decorator (:obj:`Callable`): decorated function. + Examples: + >>> from ding.utils.deprecation import deprecated + >>> @deprecated('0.4.1', '0.5.1') + >>> def hello(): + >>> print('hello') + """ + + def decorator(func): + existing_docstring = func.__doc__ or "" + + deprecated_doc = f'.. deprecated:: {since}\n Deprecated and will be removed in version {removed_in}' + + if up_to is not None: + deprecated_doc += f', please use `{up_to}` instead.' + else: + deprecated_doc += '.' + + func.__doc__ = deprecated_doc + "\n\n" + textwrap.dedent(existing_docstring) + + @functools.wraps(func) + def wrapper(*args, **kwargs): + warning_msg = ( + f'API `{func.__module__}.{func.__name__}` is deprecated since version {since} ' + f'and will be removed in version {removed_in}' + ) + if up_to is not None: + warning_msg += f", please use `{up_to}` instead." + else: + warning_msg += "." + + warnings.warn(warning_msg, category=FutureWarning, stacklevel=2) + return func(*args, **kwargs) + + return wrapper + + return decorator diff --git a/ding/utils/design_helper.py b/ding/utils/design_helper.py index 0850728d1a..24805218ff 100644 --- a/ding/utils/design_helper.py +++ b/ding/utils/design_helper.py @@ -4,15 +4,20 @@ # ABCMeta is a subclass of type, extending ABCMeta makes this metaclass is compatible with some classes # which extends ABC class SingletonMetaclass(ABCMeta): - r""" + """ Overview: Returns the given type instance in input class - Interface: + Interfaces: ``__call__`` """ instances = {} def __call__(cls: type, *args, **kwargs) -> object: + """ + Overview: + Returns the given type instance in input class + """ + if cls not in SingletonMetaclass.instances: SingletonMetaclass.instances[cls] = super(SingletonMetaclass, cls).__call__(*args, **kwargs) cls.instance = SingletonMetaclass.instances[cls] diff --git a/ding/utils/dict_helper.py b/ding/utils/dict_helper.py new file mode 100644 index 0000000000..e6a61799e6 --- /dev/null +++ b/ding/utils/dict_helper.py @@ -0,0 +1,13 @@ +from easydict import EasyDict + + +def convert_easy_dict_to_dict(easy_dict: EasyDict) -> dict: + """ + Overview: + Convert an EasyDict object to a dict object recursively. + Arguments: + - easy_dict (:obj:`EasyDict`): The EasyDict object to be converted. + Returns: + - dict: The converted dict object. + """ + return {k: convert_easy_dict_to_dict(v) if isinstance(v, EasyDict) else v for k, v in easy_dict.items()} diff --git a/ding/utils/fake_linklink.py b/ding/utils/fake_linklink.py index a8faf69e0a..5998030b36 100644 --- a/ding/utils/fake_linklink.py +++ b/ding/utils/fake_linklink.py @@ -2,16 +2,30 @@ class FakeClass: + """ + Overview: + Fake class. + """ def __init__(self, *args, **kwargs): pass class FakeNN: + """ + Overview: + Fake nn class. + """ + SyncBatchNorm2d = FakeClass class FakeLink: + """ + Overview: + Fake link class. + """ + nn = FakeNN() syncbnVarMode_t = namedtuple("syncbnVarMode_t", "L2")(L2=None) allreduceOp_t = namedtuple("allreduceOp_t", ['Sum', 'Max']) diff --git a/ding/utils/fast_copy.py b/ding/utils/fast_copy.py index 0f1d53fde6..cf4185ecbd 100644 --- a/ding/utils/fast_copy.py +++ b/ding/utils/fast_copy.py @@ -5,13 +5,21 @@ class _FastCopy: """ - The idea of this class comes from this article - https://newbedev.com/what-is-a-fast-pythonic-way-to-deepcopy-just-data-from-a-python-dict-or-list. - We use recursive calls to copy each object that needs to be copied, which will be 5x faster - than copy.deepcopy. + Overview: + The idea of this class comes from this article \ + https://newbedev.com/what-is-a-fast-pythonic-way-to-deepcopy-just-data-from-a-python-dict-or-list. + We use recursive calls to copy each object that needs to be copied, which will be 5x faster \ + than copy.deepcopy. + Interfaces: + ``__init__``, ``_copy_list``, ``_copy_dict``, ``_copy_tensor``, ``_copy_ndarray``, ``copy``. """ def __init__(self): + """ + Overview: + Initialize the _FastCopy object. + """ + dispatch = {} dispatch[list] = self._copy_list dispatch[dict] = self._copy_dict @@ -20,6 +28,13 @@ def __init__(self): self.dispatch = dispatch def _copy_list(self, l: List) -> dict: + """ + Overview: + Copy the list. + Arguments: + - l (:obj:`List`): The list to be copied. + """ + ret = l.copy() for idx, item in enumerate(ret): cp = self.dispatch.get(type(item)) @@ -28,6 +43,13 @@ def _copy_list(self, l: List) -> dict: return ret def _copy_dict(self, d: dict) -> dict: + """ + Overview: + Copy the dict. + Arguments: + - d (:obj:`dict`): The dict to be copied. + """ + ret = d.copy() for key, value in ret.items(): cp = self.dispatch.get(type(value)) @@ -37,12 +59,33 @@ def _copy_dict(self, d: dict) -> dict: return ret def _copy_tensor(self, t: torch.Tensor) -> torch.Tensor: + """ + Overview: + Copy the tensor. + Arguments: + - t (:obj:`torch.Tensor`): The tensor to be copied. + """ + return t.clone() def _copy_ndarray(self, a: np.ndarray) -> np.ndarray: + """ + Overview: + Copy the ndarray. + Arguments: + - a (:obj:`np.ndarray`): The ndarray to be copied. + """ + return np.copy(a) def copy(self, sth: Any) -> Any: + """ + Overview: + Copy the object. + Arguments: + - sth (:obj:`Any`): The object to be copied. + """ + cp = self.dispatch.get(type(sth)) if cp is None: return sth diff --git a/ding/utils/file_helper.py b/ding/utils/file_helper.py index 94a70d0925..b14c42de79 100644 --- a/ding/utils/file_helper.py +++ b/ding/utils/file_helper.py @@ -266,7 +266,7 @@ def save_file_rediscluster(path, data): def read_file(path: str, fs_type: Union[None, str] = None, use_lock: bool = False) -> object: - r""" + """ Overview: Read file from path Arguments: @@ -296,7 +296,7 @@ def read_file(path: str, fs_type: Union[None, str] = None, use_lock: bool = Fals def save_file(path: str, data: object, fs_type: Union[None, str] = None, use_lock: bool = False) -> None: - r""" + """ Overview: Save data to file of path Arguments: @@ -327,7 +327,7 @@ def save_file(path: str, data: object, fs_type: Union[None, str] = None, use_loc def remove_file(path: str, fs_type: Union[None, str] = None) -> None: - r""" + """ Overview: Remove file Arguments: diff --git a/ding/utils/k8s_helper.py b/ding/utils/k8s_helper.py index 9a3572a729..e30bba497d 100644 --- a/ding/utils/k8s_helper.py +++ b/ding/utils/k8s_helper.py @@ -20,14 +20,15 @@ def get_operator_server_kwargs(cfg: EasyDict) -> dict: - r''' + """ Overview: Get kwarg dict from config file Arguments: - cfg (:obj:`EasyDict`) System config Returns: - result (:obj:`dict`) Containing ``api_version``, ``namespace``, ``name``, ``port``, ``host``. - ''' + """ + namespace = os.environ.get('KUBERNETES_POD_NAMESPACE', DEFAULT_NAMESPACE) name = os.environ.get('KUBERNETES_POD_NAME', DEFAULT_POD_NAME) url = cfg.get('system_addr', None) or os.environ.get('KUBERNETES_SERVER_URL', None) @@ -49,10 +50,24 @@ def get_operator_server_kwargs(cfg: EasyDict) -> dict: def exist_operator_server() -> bool: + """ + Overview: + Check if the 'KUBERNETES_SERVER_URL' environment variable exists. + """ + return 'KUBERNETES_SERVER_URL' in os.environ def pod_exec_command(kubeconfig: str, name: str, namespace: str, cmd: str) -> Tuple[int, str]: + """ + Overview: + Execute command in pod + Arguments: + - kubeconfig (:obj:`str`) The path of kubeconfig file + - name (:obj:`str`) The name of pod + - namespace (:obj:`str`) The namespace of pod + """ + try: from kubernetes import config from kubernetes.client import CoreV1Api @@ -102,10 +117,20 @@ class K8sType(Enum): class K8sLauncher(object): """ - Overview: object to manage the K8s cluster + Overview: + object to manage the K8s cluster + Interfaces: + ``__init__``, ``_load``, ``create_cluster``, ``_check_k3d_tools``, ``delete_cluster``, ``preload_images`` """ def __init__(self, config_path: str) -> None: + """ + Overview: + Initialize the K8sLauncher object. + Arguments: + - config_path (:obj:`str`): The path of the config file. + """ + self.name = None self.servers = 1 self.agents = 0 @@ -116,6 +141,13 @@ def __init__(self, config_path: str) -> None: self._check_k3d_tools() def _load(self, config_path: str) -> None: + """ + Overview: + Load the config file. + Arguments: + - config_path (:obj:`str`): The path of the config file. + """ + with open(config_path, 'r') as f: data = yaml.safe_load(f) self.name = data.get('name') if data.get('name') else self.name @@ -140,6 +172,11 @@ def _load(self, config_path: str) -> None: self._images = data.get('preload_images') def _check_k3d_tools(self) -> None: + """ + Overview: + Check if the k3d tools exist. + """ + if self.type != K8sType.K3s: return args = ['which', 'k3d'] @@ -151,6 +188,11 @@ def _check_k3d_tools(self) -> None: ) def create_cluster(self) -> None: + """ + Overview: + Create the k8s cluster. + """ + print('Creating k8s cluster...') if self.type != K8sType.K3s: return @@ -168,6 +210,11 @@ def create_cluster(self) -> None: self.preload_images(self._images) def delete_cluster(self) -> None: + """ + Overview: + Delete the k8s cluster. + """ + print('Deleting k8s cluster...') if self.type != K8sType.K3s: return @@ -180,6 +227,11 @@ def delete_cluster(self) -> None: raise RuntimeError(f'Failed to delete cluster {self.name}: {err_str}') def preload_images(self, images: list) -> None: + """ + Overview: + Preload images. + """ + if self.type != K8sType.K3s or len(images) == 0: return args = ['k3d', 'image', 'import', f'--cluster={self.name}'] diff --git a/ding/utils/linklink_dist_helper.py b/ding/utils/linklink_dist_helper.py index a886c3cbff..adfb26fa83 100644 --- a/ding/utils/linklink_dist_helper.py +++ b/ding/utils/linklink_dist_helper.py @@ -20,7 +20,7 @@ def is_fake_link(): def get_rank() -> int: - r""" + """ Overview: Get the rank of ``linklink`` model, return 0 if use ``FakeLink``. @@ -33,7 +33,7 @@ def get_rank() -> int: def get_world_size() -> int: - r""" + """ Overview: Get the ``world_size`` of ``linklink model``, return 0 if use ``FakeLink``. @@ -46,7 +46,7 @@ def get_world_size() -> int: def broadcast(value: torch.Tensor, rank: int) -> None: - r""" + """ Overview: Use ``linklink.broadcast`` and raise error when using ``FakeLink`` Arguments: @@ -59,7 +59,7 @@ def broadcast(value: torch.Tensor, rank: int) -> None: def allreduce(data: torch.Tensor, op: str = 'sum') -> None: - r""" + """ Overview: Call ``linklink.allreduce`` on the data Arguments: @@ -79,7 +79,7 @@ def allreduce(data: torch.Tensor, op: str = 'sum') -> None: def allreduce_async(data: torch.Tensor, op: str = 'sum') -> None: - r""" + """ Overview: Call ``linklink.allreduce_async`` on the data Arguments: @@ -99,7 +99,7 @@ def allreduce_async(data: torch.Tensor, op: str = 'sum') -> None: def get_group(group_size: int) -> List: - r""" + """ Overview: Get the group segmentation of ``group_size`` each group Arguments: @@ -114,9 +114,11 @@ def get_group(group_size: int) -> List: def dist_mode(func: Callable) -> Callable: - r""" + """ Overview: Wrap the function so that in can init and finalize automatically before each call + Arguments: + - func (:obj:`Callable`): the function to wrap """ def wrapper(*args, **kwargs): @@ -128,7 +130,7 @@ def wrapper(*args, **kwargs): def dist_init(method: str = 'slurm', device_id: int = 0) -> Tuple[int, int]: - r""" + """ Overview: Init the distribution Arguments: @@ -152,7 +154,7 @@ def dist_init(method: str = 'slurm', device_id: int = 0) -> Tuple[int, int]: def dist_finalize() -> None: - r""" + """ Overview: Finalize ``linklink``, see ``linklink.finalize()`` """ @@ -160,25 +162,54 @@ def dist_finalize() -> None: class DistContext: + """ + Overview: + A context manager for ``linklink`` distribution + Interfaces: + ``__init__``, ``__enter__``, ``__exit__`` + """ def __init__(self) -> None: + """ + Overview: + Initialize the ``DistContext`` + """ + pass def __enter__(self) -> None: + """ + Overview: + Initialize ``linklink`` distribution + """ + dist_init() def __exit__(self, *args, **kwargs) -> Any: + """ + Overview: + Finalize ``linklink`` distribution + Arugments: + - args (:obj:`Tuple`): The arguments passed to the ``__exit__`` function. + - kwargs (:obj:`Dict`): The keyword arguments passed to the ``__exit__`` function. + """ + dist_finalize() def simple_group_split(world_size: int, rank: int, num_groups: int) -> List: - r""" + """ Overview: Split the group according to ``worldsize``, ``rank`` and ``num_groups`` + Arguments: + - world_size (:obj:`int`): The world size + - rank (:obj:`int`): The rank + - num_groups (:obj:`int`): The number of groups .. note:: With faulty input, raise ``array split does not result in an equal division`` """ + groups = [] rank_list = np.split(np.arange(world_size), num_groups) rank_list = [list(map(int, x)) for x in rank_list] @@ -189,4 +220,9 @@ def simple_group_split(world_size: int, rank: int, num_groups: int) -> List: def synchronize(): + """ + Overview: + Synchronize the process + """ + get_link().synchronize() diff --git a/ding/utils/loader/base.py b/ding/utils/loader/base.py index cd6f9ec390..fd55bc8b62 100644 --- a/ding/utils/loader/base.py +++ b/ding/utils/loader/base.py @@ -6,6 +6,13 @@ def _to_exception(exception) -> Callable[[Any], Exception]: + """ + Overview: + Convert exception to callable exception. + Arguments: + - exception (:obj:`Exception`): The exception to be converted. + """ + if hasattr(exception, '__call__'): return exception elif isinstance(exception, Exception): @@ -21,6 +28,13 @@ def _to_exception(exception) -> Callable[[Any], Exception]: def _to_loader(value) -> 'ILoaderClass': + """ + Overview: + Convert value to loader. + Arguments: + - value (:obj:`Any`): The value to be converted. + """ + if isinstance(value, ILoaderClass): return value elif isinstance(value, tuple): @@ -78,6 +92,11 @@ def _load(self, value_): def _reset_exception(loader, eg: Callable[[Any, Exception], Exception]): + """ + Overview: + Reset exception of loader. + """ + loader = Loader(loader) def _load(value): @@ -90,15 +109,42 @@ def _load(value): class ILoaderClass: + """ + Overview: + Base class of loader. + Interfaces: + ``__init__``, ``_load``, ``load``, ``check``, ``__call__``, ``__and__``, ``__or__``, ``__rshift__`` + """ @abstractmethod def _load(self, value: _ValueType) -> _ValueType: + """ + Overview: + Load the value. + Arguments: + - value (:obj:`_ValueType`): The value to be loaded. + """ + raise NotImplementedError def __load(self, value: _ValueType) -> _ValueType: + """ + Overview: + Load the value. + Arguments: + - value (:obj:`_ValueType`): The value to be loaded. + """ + return self._load(value) def __check(self, value: _ValueType) -> bool: + """ + Overview: + Check whether the value is valid. + Arguments: + - value (:obj:`_ValueType`): The value to be checked. + """ + try: self._load(value) except CAPTURE_EXCEPTIONS: @@ -107,15 +153,42 @@ def __check(self, value: _ValueType) -> bool: return True def load(self, value: _ValueType) -> _ValueType: + """ + Overview: + Load the value. + Arguments: + - value (:obj:`_ValueType`): The value to be loaded. + """ + return self.__load(value) def check(self, value: _ValueType) -> bool: + """ + Overview: + Check whether the value is valid. + Arguments: + - value (:obj:`_ValueType`): The value to be checked. + """ + return self.__check(value) def __call__(self, value: _ValueType) -> _ValueType: + """ + Overview: + Load the value. + Arguments: + - value (:obj:`_ValueType`): The value to be loaded. + """ + return self.__load(value) def __and__(self, other) -> 'ILoaderClass': + """ + Overview: + Combine two loaders. + Arguments: + - other (:obj:`ILoaderClass`): The other loader. + """ def _load(value: _ValueType) -> _ValueType: self.load(value) @@ -124,9 +197,22 @@ def _load(value: _ValueType) -> _ValueType: return Loader(_load) def __rand__(self, other) -> 'ILoaderClass': + """ + Overview: + Combine two loaders. + Arguments: + - other (:obj:`ILoaderClass`): The other loader. + """ + return Loader(other) & self def __or__(self, other) -> 'ILoaderClass': + """ + Overview: + Combine two loaders. + Arguments: + - other (:obj:`ILoaderClass`): The other loader. + """ def _load(value: _ValueType) -> _ValueType: try: @@ -137,9 +223,22 @@ def _load(value: _ValueType) -> _ValueType: return Loader(_load) def __ror__(self, other) -> 'ILoaderClass': + """ + Overview: + Combine two loaders. + Arguments: + - other (:obj:`ILoaderClass`): The other loader. + """ + return Loader(other) | self def __rshift__(self, other) -> 'ILoaderClass': + """ + Overview: + Combine two loaders. + Arguments: + - other (:obj:`ILoaderClass`): The other loader. + """ def _load(value: _ValueType) -> _ValueType: _return_value = self.load(value) @@ -148,4 +247,11 @@ def _load(value: _ValueType) -> _ValueType: return Loader(_load) def __rrshift__(self, other) -> 'ILoaderClass': + """ + Overview: + Combine two loaders. + Arguments: + - other (:obj:`ILoaderClass`): The other loader. + """ + return Loader(other) >> self diff --git a/ding/utils/loader/collection.py b/ding/utils/loader/collection.py index cbac490df4..770e6c6c64 100644 --- a/ding/utils/loader/collection.py +++ b/ding/utils/loader/collection.py @@ -9,8 +9,23 @@ class CollectionError(CompositeStructureError): + """ + Overview: + Collection error. + Interfaces: + ``__init__``, ``errors`` + Properties: + ``errors`` + """ def __init__(self, errors: COLLECTION_ERRORS): + """ + Overview: + Initialize the CollectionError. + Arguments: + - errors (:obj:`COLLECTION_ERRORS`): The errors. + """ + self.__errors = list(errors or []) CompositeStructureError.__init__( self, '{count} error(s) found in collection.'.format(count=repr(list(self.__errors))) @@ -18,10 +33,23 @@ def __init__(self, errors: COLLECTION_ERRORS): @property def errors(self) -> COLLECTION_ERRORS: + """ + Overview: + Get the errors. + """ + return self.__errors def collection(loader, type_back: bool = True) -> ILoaderClass: + """ + Overview: + Create a collection loader. + Arguments: + - loader (:obj:`ILoaderClass`): The loader. + - type_back (:obj:`bool`): Whether to convert the type back. + """ + loader = Loader(loader) def _load(value): @@ -47,6 +75,13 @@ def _load(value): def tuple_(*loaders) -> ILoaderClass: + """ + Overview: + Create a tuple loader. + Arguments: + - loaders (:obj:`tuple`): The loaders. + """ + loaders = [Loader(loader) for loader in loaders] def _load(value: tuple): @@ -56,6 +91,13 @@ def _load(value: tuple): def length(min_length: Optional[int] = None, max_length: Optional[int] = None) -> ILoaderClass: + """ + Overview: + Create a length loader. + Arguments: + - min_length (:obj:`int`): The minimum length. + - max_length (:obj:`int`): The maximum length. + """ def _load(value): _length = len(value) @@ -74,10 +116,23 @@ def _load(value): def length_is(length_: int) -> ILoaderClass: + """ + Overview: + Create a length loader. + Arguments: + - length_ (:obj:`int`): The length. + """ + return length(min_length=length_, max_length=length_) def contains(content) -> ILoaderClass: + """ + Overview: + Create a contains loader. + Arguments: + - content (:obj:`Any`): The content. + """ def _load(value): if content not in value: @@ -89,6 +144,13 @@ def _load(value): def cofilter(checker: Callable[[Any], bool], type_back: bool = True) -> ILoaderClass: + """ + Overview: + Create a cofilter loader. + Arguments: + - checker (:obj:`Callable[[Any], bool]`): The checker. + - type_back (:obj:`bool`): Whether to convert the type back. + """ def _load(value): _result = [item for item in value if checker(item)] @@ -100,6 +162,12 @@ def _load(value): def tpselector(*indices) -> ILoaderClass: + """ + Overview: + Create a tuple selector loader. + Arguments: + - indices (:obj:`tuple`): The indices. + """ def _load(value: tuple): return tuple([value[index] for index in indices]) diff --git a/ding/utils/loader/dict.py b/ding/utils/loader/dict.py index 9bda68a46a..a14d3ff9f8 100644 --- a/ding/utils/loader/dict.py +++ b/ding/utils/loader/dict.py @@ -7,16 +7,43 @@ class DictError(CompositeStructureError): + """ + Overview: + Dict error. + Interfaces: + ``__init__``, ``errors`` + Properties: + ``errors`` + """ def __init__(self, errors: DICT_ERRORS): + """ + Overview: + Initialize the DictError. + Arguments: + - errors (:obj:`DICT_ERRORS`): The errors. + """ + self.__error = errors @property def errors(self) -> DICT_ERRORS: + """ + Overview: + Get the errors. + """ + return self.__error def dict_(**kwargs) -> ILoaderClass: + """ + Overview: + Create a dict loader. + Arguments: + - kwargs (:obj:`Mapping[str, ILoaderClass]`): The loaders. + """ + kwargs = [(k, Loader(v)) for k, v in kwargs.items()] def _load(value): diff --git a/ding/utils/loader/exception.py b/ding/utils/loader/exception.py index 96f2b53ad5..9358f1c85e 100644 --- a/ding/utils/loader/exception.py +++ b/ding/utils/loader/exception.py @@ -7,8 +7,21 @@ class CompositeStructureError(ValueError, metaclass=ABCMeta): + """ + Overview: + Composite structure error. + Interfaces: + ``__init__``, ``errors`` + Properties: + ``errors`` + """ @property @abstractmethod def errors(self) -> ERROR_ITEMS: + """ + Overview: + Get the errors. + """ + raise NotImplementedError diff --git a/ding/utils/loader/mapping.py b/ding/utils/loader/mapping.py index 0bb9e4e85a..c3993c2366 100644 --- a/ding/utils/loader/mapping.py +++ b/ding/utils/loader/mapping.py @@ -10,23 +10,61 @@ class MappingError(CompositeStructureError): + """ + Overview: + Mapping error. + Interfaces: + ``__init__``, ``errors`` + """ def __init__(self, key_errors: MAPPING_ERRORS, value_errors: MAPPING_ERRORS): + """ + Overview: + Initialize the MappingError. + Arguments: + - key_errors (:obj:`MAPPING_ERRORS`): The key errors. + - value_errors (:obj:`MAPPING_ERRORS`): The value errors. + """ + self.__key_errors = list(key_errors or []) self.__value_errors = list(value_errors or []) self.__errors = self.__key_errors + self.__value_errors def key_errors(self) -> MAPPING_ERRORS: + """ + Overview: + Get the key errors. + """ + return self.__key_errors def value_errors(self) -> MAPPING_ERRORS: + """ + Overview: + Get the value errors. + """ + return self.__value_errors def errors(self) -> MAPPING_ERRORS: + """ + Overview: + Get the errors. + """ + return self.__errors def mapping(key_loader, value_loader, type_back: bool = True) -> ILoaderClass: + """ + Overview: + Create a mapping loader. + Arguments: + - key_loader (:obj:`ILoaderClass`): The key loader. + - value_loader (:obj:`ILoaderClass`): The value loader. + - type_back (:obj:`bool`): Whether to convert the type back. + """ + key_loader = Loader(key_loader) value_loader = Loader(value_loader) @@ -67,6 +105,13 @@ def _load(value): def mpfilter(check: Callable[[Any, Any], bool], type_back: bool = True) -> ILoaderClass: + """ + Overview: + Create a mapping filter loader. + Arguments: + - check (:obj:`Callable[[Any, Any], bool]`): The check function. + - type_back (:obj:`bool`): Whether to convert the type back. + """ def _load(value): _result = {key_: value_ for key_, value_ in value.items() if check(key_, value_)} @@ -79,14 +124,29 @@ def _load(value): def mpkeys() -> ILoaderClass: + """ + Overview: + Create a mapping keys loader. + """ + return method('items') & method('keys') & Loader(lambda v: set(v.keys())) def mpvalues() -> ILoaderClass: + """ + Overview: + Create a mapping values loader. + """ + return method('items') & method('values') & Loader(lambda v: set(v.values())) def mpitems() -> ILoaderClass: + """ + Overview: + Create a mapping items loader. + """ + return method('items') & Loader(lambda v: set([(key, value) for key, value in v.items()])) @@ -94,10 +154,25 @@ def mpitems() -> ILoaderClass: def item(key) -> ILoaderClass: + """ + Overview: + Create a item loader. + Arguments: + - key (:obj:`Any`): The key. + """ + return _INDEX_PRECHECK & Loader( (lambda v: key in v.keys(), lambda v: v[key], KeyError('key {key} not found'.format(key=repr(key)))) ) def item_or(key, default) -> ILoaderClass: + """ + Overview: + Create a item or loader. + Arguments: + - key (:obj:`Any`): The key. + - default (:obj:`Any`): The default value. + """ + return _INDEX_PRECHECK & (item(key) | raw(default)) diff --git a/ding/utils/loader/norm.py b/ding/utils/loader/norm.py index ea99692a93..af142ed4e6 100644 --- a/ding/utils/loader/norm.py +++ b/ding/utils/loader/norm.py @@ -7,6 +7,12 @@ def _callable_to_norm(func: Callable[[Any], Any]) -> 'INormClass': + """ + Overview: + Convert callable to norm. + Arguments: + - func (:obj:`Callable[[Any], Any]`): The callable to be converted. + """ class _Norm(INormClass): @@ -17,6 +23,13 @@ def _call(self, value): def norm(value) -> 'INormClass': + """ + Overview: + Convert value to norm. + Arguments: + - value (:obj:`Any`): The value to be converted. + """ + if isinstance(value, INormClass): return value elif isinstance(value, ILoaderClass): @@ -26,6 +39,12 @@ def norm(value) -> 'INormClass': def normfunc(func): + """ + Overview: + Convert function to norm function. + Arguments: + - func (:obj:`Callable[[Any], Any]`): The function to be converted. + """ @wraps(func) def _new_func(*args_norm, **kwargs_norm): @@ -47,14 +66,37 @@ def _callable(v): def _unary(a: 'INormClass', func: UNARY_FUNC) -> 'INormClass': + """ + Overview: + Create a unary norm. + Arguments: + - a (:obj:`INormClass`): The norm. + - func (:obj:`UNARY_FUNC`): The function. + """ + return _callable_to_norm(lambda v: func(a(v))) def _binary(a: 'INormClass', b: 'INormClass', func: BINARY_FUNC) -> 'INormClass': + """ + Overview: + Create a binary norm. + Arguments: + - a (:obj:`INormClass`): The first norm. + - b (:obj:`INormClass`): The second norm. + - func (:obj:`BINARY_FUNC`): The function. + """ return _callable_to_norm(lambda v: func(a(v), b(v))) def _binary_reducing(func: BINARY_FUNC, zero): + """ + Overview: + Create a binary reducing norm. + Arguments: + - func (:obj:`BINARY_FUNC`): The function. + - zero (:obj:`Any`): The zero value. + """ @wraps(func) def _new_func(*args) -> 'INormClass': @@ -67,118 +109,382 @@ def _new_func(*args) -> 'INormClass': class INormClass: + """ + Overview: + The norm class. + Interfaces: + ``__call__``, ``__add__``, ``__radd__``, ``__sub__``, ``__rsub__``, ``__mul__``, ``__rmul__``, ``__matmul__``, + ``__rmatmul__``, ``__truediv__``, ``__rtruediv__``, ``__floordiv__``, ``__rfloordiv__``, ``__mod__``, + ``__rmod__``, ``__pow__``, ``__rpow__``, ``__lshift__``, ``__rlshift__``, ``__rshift__``, ``__rrshift__``, + ``__and__``, ``__rand__``, ``__or__``, ``__ror__``, ``__xor__``, ``__rxor__``, ``__invert__``, ``__pos__``, + ``__neg__``, ``__eq__``, ``__ne__``, ``__lt__``, ``__le__``, ``__gt__``, ``__ge__`` + """ @abstractmethod def _call(self, value): + """ + Overview: + Call the norm. + Arguments: + - value (:obj:`Any`): The value to be normalized. + """ + raise NotImplementedError def __call__(self, value): + """ + Overview: + Call the norm. + Arguments: + - value (:obj:`Any`): The value to be normalized. + """ + return self._call(value) def __add__(self, other): + """ + Overview: + Add the norm. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return _binary(self, norm(other), operator.__add__) def __radd__(self, other): + """ + Overview: + Add the norm. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return norm(other) + self def __sub__(self, other): + """ + Overview: + Subtract the norm. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return _binary(self, norm(other), operator.__sub__) def __rsub__(self, other): + """ + Overview: + Subtract the norm. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return norm(other) - self def __mul__(self, other): + """ + Overview: + Multiply the norm. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return _binary(self, norm(other), operator.__mul__) def __rmul__(self, other): + """ + Overview: + Multiply the norm. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return norm(other) * self def __matmul__(self, other): + """ + Overview: + Matrix multiply the norm. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return _binary(self, norm(other), operator.__matmul__) def __rmatmul__(self, other): + """ + Overview: + Matrix multiply the norm. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return norm(other) @ self def __truediv__(self, other): + """ + Overview: + Divide the norm. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return _binary(self, norm(other), operator.__truediv__) def __rtruediv__(self, other): + """ + Overview: + Divide the norm. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return norm(other) / self def __floordiv__(self, other): + """ + Overview: + Floor divide the norm. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return _binary(self, norm(other), operator.__floordiv__) def __rfloordiv__(self, other): + """ + Overview: + Floor divide the norm. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return norm(other) // self def __mod__(self, other): + """ + Overview: + Mod the norm. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return _binary(self, norm(other), operator.__mod__) def __rmod__(self, other): + """ + Overview: + Mod the norm. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return norm(other) % self def __pow__(self, power, modulo=None): + """ + Overview: + Power the norm. + Arguments: + - power (:obj:`Any`): The power. + - modulo (:obj:`Any`): The modulo. + """ + return _binary(self, norm(power), operator.__pow__) def __rpow__(self, other): + """ + Overview: + Power the norm. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return norm(other) ** self def __lshift__(self, other): + """ + Overview: + Lshift the norm. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return _binary(self, norm(other), operator.__lshift__) def __rlshift__(self, other): + """ + Overview: + Lshift the norm. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return norm(other) << self def __rshift__(self, other): + """ + Overview: + Rshift the norm. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return _binary(self, norm(other), operator.__rshift__) def __rrshift__(self, other): + """ + Overview: + Rshift the norm. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return norm(other) >> self def __and__(self, other): + """ + Overview: + And operation the norm. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return _binary(self, norm(other), operator.__and__) def __rand__(self, other): + """ + Overview: + And operation the norm. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return norm(other) & self def __or__(self, other): + """ + Overview: + Or operation the norm. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return _binary(self, norm(other), operator.__or__) def __ror__(self, other): + """ + Overview: + Or operation the norm. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return norm(other) | self def __xor__(self, other): + """ + Overview: + Xor operation the norm. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return _binary(self, norm(other), operator.__xor__) def __rxor__(self, other): + """ + Overview: + Xor operation the norm. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return norm(other) ^ self def __invert__(self): + """ + Overview: + Invert the norm. + """ + return _unary(self, operator.__invert__) def __pos__(self): + """ + Overview: + Positive the norm. + """ + return _unary(self, operator.__pos__) def __neg__(self): + """ + Overview: + Negative the norm. + """ + return _unary(self, operator.__neg__) # Attention: DO NOT USE LINKING COMPARE OPERATORS, IT WILL CAUSE ERROR. def __eq__(self, other): + """ + Overview: + Compare the norm if they are equal. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return _binary(self, norm(other), operator.__eq__) def __ne__(self, other): + """ + Overview: + Compare the norm if they are not equal. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return _binary(self, norm(other), operator.__ne__) def __lt__(self, other): + """ + Overview: + Compare the norm if it is less than the other norm. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return _binary(self, norm(other), operator.__lt__) def __le__(self, other): + """ + Overview: + Compare the norm if it is less than or equal to the other norm. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return _binary(self, norm(other), operator.__le__) def __gt__(self, other): + """ + Overview: + Compare the norm if it is greater than the other norm. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return _binary(self, norm(other), operator.__gt__) def __ge__(self, other): + """ + Overview: + Compare the norm if it is greater than or equal to the other norm. + Arguments: + - other (:obj:`Any`): The other norm. + """ + return _binary(self, norm(other), operator.__ge__) @@ -204,6 +510,14 @@ def __ge__(self, other): @normfunc def lcmp(first, *items): + """ + Overview: + Compare the items. + Arguments: + - first (:obj:`Any`): The first item. + - items (:obj:`Any`): The other items. + """ + if len(items) % 2 == 1: raise ValueError('Count of items should be odd number but {number} found.'.format(number=len(items) + 1)) diff --git a/ding/utils/loader/number.py b/ding/utils/loader/number.py index 61b4e97225..a9fdc59a16 100644 --- a/ding/utils/loader/number.py +++ b/ding/utils/loader/number.py @@ -10,6 +10,15 @@ def numeric(int_ok: bool = True, float_ok: bool = True, inf_ok: bool = True) -> ILoaderClass: + """ + Overview: + Create a numeric loader. + Arguments: + - int_ok (:obj:`bool`): Whether int is allowed. + - float_ok (:obj:`bool`): Whether float is allowed. + - inf_ok (:obj:`bool`): Whether inf is allowed. + """ + if not int_ok and not float_ok: raise ValueError('Either int or float should be allowed.') @@ -42,6 +51,17 @@ def interval( right_ok: bool = True, eps=0.0 ) -> ILoaderClass: + """ + Overview: + Create a interval loader. + Arguments: + - left (:obj:`Optional[NUMBER_TYPING]`): The left bound. + - right (:obj:`Optional[NUMBER_TYPING]`): The right bound. + - left_ok (:obj:`bool`): Whether left bound is allowed. + - right_ok (:obj:`bool`): Whether right bound is allowed. + - eps (:obj:`float`): The epsilon. + """ + if left is None: left = -math.inf if right is None: @@ -93,70 +113,170 @@ def _load(value) -> NUMBER_TYPING: def is_negative() -> ILoaderClass: + """ + Overview: + Create a negative loader. + """ + return Loader((lambda x: x < 0, lambda x: ValueError('negative required but {value} found'.format(value=repr(x))))) def is_positive() -> ILoaderClass: + """ + Overview: + Create a positive loader. + """ + return Loader((lambda x: x > 0, lambda x: ValueError('positive required but {value} found'.format(value=repr(x))))) def non_negative() -> ILoaderClass: + """ + Overview: + Create a non-negative loader. + """ + return Loader( (lambda x: x >= 0, lambda x: ValueError('non-negative required but {value} found'.format(value=repr(x)))) ) def non_positive() -> ILoaderClass: + """ + Overview: + Create a non-positive loader. + """ + return Loader( (lambda x: x <= 0, lambda x: ValueError('non-positive required but {value} found'.format(value=repr(x)))) ) def negative() -> ILoaderClass: + """ + Overview: + Create a negative loader. + """ + return Loader(lambda x: -x) def positive() -> ILoaderClass: + """ + Overview: + Create a positive loader. + """ + return Loader(lambda x: +x) def _math_binary(func: Callable[[Any, Any], Any], attachment) -> ILoaderClass: + """ + Overview: + Create a math binary loader. + Arguments: + - func (:obj:`Callable[[Any, Any], Any]`): The function. + - attachment (:obj:`Any`): The attachment. + """ + return Loader(lambda x: func(x, Loader(attachment)(x))) def plus(addend) -> ILoaderClass: + """ + Overview: + Create a plus loader. + Arguments: + - addend (:obj:`Any`): The addend. + """ + return _math_binary(lambda x, y: x + y, addend) def minus(subtrahend) -> ILoaderClass: + """ + Overview: + Create a minus loader. + Arguments: + - subtrahend (:obj:`Any`): The subtrahend. + """ + return _math_binary(lambda x, y: x - y, subtrahend) def minus_with(minuend) -> ILoaderClass: + """ + Overview: + Create a minus loader. + Arguments: + - minuend (:obj:`Any`): The minuend. + """ + return _math_binary(lambda x, y: y - x, minuend) def multi(multiplier) -> ILoaderClass: + """ + Overview: + Create a multi loader. + Arguments: + - multiplier (:obj:`Any`): The multiplier. + """ + return _math_binary(lambda x, y: x * y, multiplier) def divide(divisor) -> ILoaderClass: + """ + Overview: + Create a divide loader. + Arguments: + - divisor (:obj:`Any`): The divisor. + """ + return _math_binary(lambda x, y: x / y, divisor) def divide_with(dividend) -> ILoaderClass: + """ + Overview: + Create a divide loader. + Arguments: + - dividend (:obj:`Any`): The dividend. + """ + return _math_binary(lambda x, y: y / x, dividend) def power(index) -> ILoaderClass: + """ + Overview: + Create a power loader. + Arguments: + - index (:obj:`Any`): The index. + """ + return _math_binary(lambda x, y: x ** y, index) def power_with(base) -> ILoaderClass: + """ + Overview: + Create a power loader. + Arguments: + - base (:obj:`Any`): The base. + """ + return _math_binary(lambda x, y: y ** x, base) def msum(*items) -> ILoaderClass: + """ + Overview: + Create a sum loader. + Arguments: + - items (:obj:`tuple`): The items. + """ def _load(value): return sum([item(value) for item in items]) @@ -165,6 +285,12 @@ def _load(value): def mmulti(*items) -> ILoaderClass: + """ + Overview: + Create a multi loader. + Arguments: + - items (:obj:`tuple`): The items. + """ def _load(value): _result = 1 @@ -186,6 +312,15 @@ def _load(value): def _msinglecmp(first, op, second) -> ILoaderClass: + """ + Overview: + Create a single compare loader. + Arguments: + - first (:obj:`Any`): The first item. + - op (:obj:`str`): The operator. + - second (:obj:`Any`): The second item. + """ + first = Loader(first) second = Loader(second) @@ -206,6 +341,14 @@ def _msinglecmp(first, op, second) -> ILoaderClass: def mcmp(first, *items) -> ILoaderClass: + """ + Overview: + Create a multi compare loader. + Arguments: + - first (:obj:`Any`): The first item. + - items (:obj:`tuple`): The items. + """ + if len(items) % 2 == 1: raise ValueError('Count of items should be odd number but {number} found.'.format(number=len(items) + 1)) diff --git a/ding/utils/loader/string.py b/ding/utils/loader/string.py index 4e7c9a8da3..16d4827cc4 100644 --- a/ding/utils/loader/string.py +++ b/ding/utils/loader/string.py @@ -9,6 +9,13 @@ def enum(*items, case_sensitive: bool = True) -> ILoaderClass: + """ + Overview: + Create an enum loader. + Arguments: + - items (:obj:`Iterable[str]`): The items. + - case_sensitive (:obj:`bool`): Whether case sensitive. + """ def _case_sensitive(func: STRING_PROCESSOR) -> STRING_PROCESSOR: if case_sensitive: @@ -38,6 +45,13 @@ def _load(value: str): def _to_regexp(regexp) -> Pattern: + """ + Overview: + Convert regexp to re.Pattern. + Arguments: + - regexp (:obj:`Union[str, re.Pattern]`): The regexp. + """ + if isinstance(regexp, Pattern): return regexp elif isinstance(regexp, str): @@ -49,6 +63,13 @@ def _to_regexp(regexp) -> Pattern: def rematch(regexp: Union[str, Pattern]) -> ILoaderClass: + """ + Overview: + Create a rematch loader. + Arguments: + - regexp (:obj:`Union[str, re.Pattern]`): The regexp. + """ + regexp = _to_regexp(regexp) def _load(value: str): @@ -66,6 +87,14 @@ def _load(value: str): def regrep(regexp: Union[str, Pattern], group: int = 0) -> ILoaderClass: + """ + Overview: + Create a regrep loader. + Arguments: + - regexp (:obj:`Union[str, re.Pattern]`): The regexp. + - group (:obj:`int`): The group. + """ + regexp = _to_regexp(regexp) def _load(value: str): diff --git a/ding/utils/loader/types.py b/ding/utils/loader/types.py index b868b05032..6039395ca6 100644 --- a/ding/utils/loader/types.py +++ b/ding/utils/loader/types.py @@ -5,6 +5,13 @@ def is_type(type_: type) -> ILoaderClass: + """ + Overview: + Create a type loader. + Arguments: + - type_ (:obj:`type`): The type. + """ + if isinstance(type_, type): return Loader(type_) else: @@ -12,10 +19,22 @@ def is_type(type_: type) -> ILoaderClass: def to_type(type_: type) -> ILoaderClass: + """ + Overview: + Create a type loader. + Arguments: + - type_ (:obj:`type`): The type. + """ + return Loader(lambda v: type_(v)) def is_callable() -> ILoaderClass: + """ + Overview: + Create a callable loader. + """ + return _reset_exception( check_only(prop('__call__')), lambda v, e: TypeError('callable expected but {func} not found'.format(func=repr('__call__'))) @@ -23,6 +42,13 @@ def is_callable() -> ILoaderClass: def prop(attr_name: str) -> ILoaderClass: + """ + Overview: + Create a attribute loader. + Arguments: + - attr_name (:obj:`str`): The attribute name. + """ + return Loader( ( lambda v: hasattr(v, attr_name), lambda v: getattr(v, attr_name), @@ -32,6 +58,13 @@ def prop(attr_name: str) -> ILoaderClass: def method(method_name: str) -> ILoaderClass: + """ + Overview: + Create a method loader. + Arguments: + - method_name (:obj:`str`): The method name. + """ + return _reset_exception( prop(method_name) >> is_callable(), lambda v, e: TypeError('type {type} not support function {func}'.format(type=repr(type(v).__name__), func=repr('__iter__'))) @@ -39,8 +72,24 @@ def method(method_name: str) -> ILoaderClass: def fcall(*args, **kwargs) -> ILoaderClass: + """ + Overview: + Create a function loader. + Arguments: + - args (:obj:`Tuple[Any]`): The args. + - kwargs (:obj:`Dict[str, Any]`): The kwargs. + """ + return Loader(lambda v: v(*args, **kwargs)) def fpartial(*args, **kwargs) -> ILoaderClass: + """ + Overview: + Create a partial function loader. + Arguments: + - args (:obj:`Tuple[Any]`): The args. + - kwargs (:obj:`Dict[str, Any]`): The kwargs. + """ + return Loader(lambda v: partial(v, *args, **kwargs)) diff --git a/ding/utils/loader/utils.py b/ding/utils/loader/utils.py index ac1998e26c..140bbf033d 100644 --- a/ding/utils/loader/utils.py +++ b/ding/utils/loader/utils.py @@ -2,20 +2,51 @@ def keep() -> ILoaderClass: + """ + Overview: + Create a keep loader. + """ + return Loader(lambda v: v) def raw(value) -> ILoaderClass: + """ + Overview: + Create a raw loader. + """ + return Loader(lambda v: value) def optional(loader) -> ILoaderClass: + """ + Overview: + Create a optional loader. + Arguments: + - loader (:obj:`ILoaderClass`): The loader. + """ + return Loader(loader) | None def check_only(loader) -> ILoaderClass: + """ + Overview: + Create a check only loader. + Arguments: + - loader (:obj:`ILoaderClass`): The loader. + """ + return Loader(loader) & keep() def check(loader) -> ILoaderClass: + """ + Overview: + Create a check loader. + Arguments: + - loader (:obj:`ILoaderClass`): The loader. + """ + return Loader(lambda x: Loader(loader).check(x)) diff --git a/ding/utils/lock_helper.py b/ding/utils/lock_helper.py index cb4a9c13b5..a7ecbb4166 100644 --- a/ding/utils/lock_helper.py +++ b/ding/utils/lock_helper.py @@ -4,7 +4,6 @@ import platform from enum import Enum, unique -from readerwriterlock import rwlock from pathlib import Path if platform.system().lower() != 'windows': import fcntl @@ -15,7 +14,8 @@ @unique class LockContextType(Enum): """ - Enum to express the type of the lock + Overview: + Enum to express the type of the lock. """ THREAD_LOCK = 1 PROCESS_LOCK = 2 @@ -33,37 +33,51 @@ class LockContext(object): Generate a LockContext in order to make sure the thread safety. Interfaces: - ``__init__``, ``__enter__``, ``__exit__`` + ``__init__``, ``__enter__``, ``__exit__``. Example: >>> with LockContext() as lock: >>> print("Do something here.") """ - def __init__(self, type_: LockContextType = LockContextType.THREAD_LOCK): - r""" + def __init__(self, lock_type: LockContextType = LockContextType.THREAD_LOCK): + """ Overview: - Init the lock according to given type + Init the lock according to the given type. + + Arguments: + - lock_type (:obj:`LockContextType`): The type of lock to be used. Defaults to LockContextType.THREAD_LOCK. """ - self.lock = _LOCK_TYPE_MAPPING[type_]() + self.lock = _LOCK_TYPE_MAPPING[lock_type]() def acquire(self): + """ + Overview: + Acquires the lock. + """ self.lock.acquire() def release(self): + """ + Overview: + Releases the lock. + """ self.lock.release() def __enter__(self): """ Overview: - Entering the context and acquire lock + Enters the context and acquires the lock. """ self.lock.acquire() def __exit__(self, *args, **kwargs): """ Overview: - Quiting the context and release lock + Exits the context and releases the lock. + Arguments: + - args (:obj:`Tuple`): The arguments passed to the ``__exit__`` function. + - kwargs (:obj:`Dict`): The keyword arguments passed to the ``__exit__`` function. """ self.lock.release() @@ -72,16 +86,23 @@ def __exit__(self, *args, **kwargs): def get_rw_file_lock(name: str, op: str): - r''' + """ Overview: Get generated file lock with name and operator Arguments: - - name (:obj:`str`) Lock's name. - - op (:obj:`str`) Assigned operator, i.e. ``read`` or ``write``. + - name (:obj:`str`): Lock's name. + - op (:obj:`str`): Assigned operator, i.e. ``read`` or ``write``. Returns: - - (:obj:`RWLockFairD`) Generated rwlock - ''' + - (:obj:`RWLockFairD`): Generated rwlock + """ assert op in ['read', 'write'] + try: + from readerwriterlock import rwlock + except ImportError: + import sys + from ditk import logging + logging.warning("Please install readerwriterlock first, such as `pip3 install readerwriterlock`.") + sys.exit(1) if name not in rw_lock_mapping: rw_lock_mapping[name] = rwlock.RWLockFairD() lock = rw_lock_mapping[name] @@ -92,22 +113,63 @@ def get_rw_file_lock(name: str, op: str): class FcntlContext: + """ + Overview: + A context manager that acquires an exclusive lock on a file using fcntl. \ + This is useful for preventing multiple processes from running the same code. + + Interfaces: + ``__init__``, ``__enter__``, ``__exit__``. + + Example: + >>> lock_path = "/path/to/lock/file" + >>> with FcntlContext(lock_path) as lock: + >>> # Perform operations while the lock is held + + """ def __init__(self, lock_path: str) -> None: + """ + Overview: + Initialize the LockHelper object. + + Arguments: + - lock_path (:obj:`str`): The path to the lock file. + """ self.lock_path = lock_path self.f = None def __enter__(self) -> None: + """ + Overview: + Acquires the lock and opens the lock file in write mode. \ + If the lock file does not exist, it is created. + """ assert self.f is None, self.lock_path self.f = open(self.lock_path, 'w') fcntl.flock(self.f.fileno(), fcntl.LOCK_EX) def __exit__(self, *args, **kwargs) -> None: + """ + Overview: + Closes the file and releases any resources used by the lock_helper object. + Arguments: + - args (:obj:`Tuple`): The arguments passed to the ``__exit__`` function. + - kwargs (:obj:`Dict`): The keyword arguments passed to the ``__exit__`` function. + """ self.f.close() self.f = None -def get_file_lock(name: str, op: str) -> None: +def get_file_lock(name: str, op: str) -> FcntlContext: + """ + Overview: + Acquires a file lock for the specified file. \ + + Arguments: + - name (:obj:`str`): The name of the file. + - op (:obj:`str`): The operation to perform on the file lock. + """ if fcntl is None: return get_rw_file_lock(name, op) else: diff --git a/ding/utils/log_helper.py b/ding/utils/log_helper.py index 9ba7e0cddf..0a966532ac 100644 --- a/ding/utils/log_helper.py +++ b/ding/utils/log_helper.py @@ -19,7 +19,7 @@ def build_logger( need_text: bool = True, text_level: Union[int, str] = logging.INFO ) -> Tuple[Optional[logging.Logger], Optional['SummaryWriter']]: # noqa - r""" + """ Overview: Build text logger and tensorboard logger. Arguments: @@ -34,13 +34,22 @@ def build_logger( """ if name is None: name = 'default' - logger = LoggerFactory.create_logger(path, name=name) if need_text else None + logger = LoggerFactory.create_logger(path, name=name, level=text_level) if need_text else None tb_name = name + '_tb_logger' tb_logger = TBLoggerFactory.create_logger(os.path.join(path, tb_name)) if need_tb else None return logger, tb_logger class TBLoggerFactory(object): + """ + Overview: + TBLoggerFactory is a factory class for ``SummaryWriter``. + Interfaces: + ``create_logger`` + Properties: + - ``tb_loggers`` (:obj:`Dict[str, SummaryWriter]`): A dict that stores ``SummaryWriter`` instances. + """ + tb_loggers = {} @classmethod @@ -53,10 +62,16 @@ def create_logger(cls: type, logdir: str) -> DistributedWriter: class LoggerFactory(object): + """ + Overview: + LoggerFactory is a factory class for ``logging.Logger``. + Interfaces: + ``create_logger``, ``get_tabulate_vars``, ``get_tabulate_vars_hor`` + """ @classmethod def create_logger(cls, path: str, name: str = 'default', level: Union[int, str] = logging.INFO) -> logging.Logger: - r""" + """ Overview: Create logger using logging Arguments: @@ -80,7 +95,7 @@ def create_logger(cls, path: str, name: str = 'default', level: Union[int, str] @staticmethod def get_tabulate_vars(variables: Dict[str, Any]) -> str: - r""" + """ Overview: Get the text description in tabular form of all vars Arguments: @@ -97,6 +112,13 @@ def get_tabulate_vars(variables: Dict[str, Any]) -> str: @staticmethod def get_tabulate_vars_hor(variables: Dict[str, Any]) -> str: + """ + Overview: + Get the text description in tabular form of all vars + Arguments: + - variables (:obj:`List[str]`): Names of the vars to query. + """ + column_to_divide = 5 # which includes the header "Name & Value" datak = [] @@ -131,7 +153,7 @@ def get_tabulate_vars_hor(variables: Dict[str, Any]) -> str: def pretty_print(result: dict, direct_print: bool = True) -> str: - r""" + """ Overview: Print a dict ``result`` in a pretty way Arguments: diff --git a/ding/utils/log_writer_helper.py b/ding/utils/log_writer_helper.py index 7efbc32416..0f8a1c5115 100644 --- a/ding/utils/log_writer_helper.py +++ b/ding/utils/log_writer_helper.py @@ -16,10 +16,22 @@ class DistributedWriter(SummaryWriter): A simple subclass of SummaryWriter that supports writing to one process in multi-process mode. The best way is to use it in conjunction with the ``router`` to take advantage of the message \ and event components of the router (see ``writer.plugin``). + Interfaces: + ``get_instance``, ``plugin``, ``initialize``, ``__del__`` """ root = None def __init__(self, *args, **kwargs): + """ + Overview: + Initialize the DistributedWriter object. + Arguments: + - args (:obj:`Tuple`): The arguments passed to the ``__init__`` function of the parent class, \ + SummaryWriter. + - kwargs (:obj:`Dict`): The keyword arguments passed to the ``__init__`` function of the parent class, \ + SummaryWriter. + """ + self._default_writer_to_disk = kwargs.get("write_to_disk") if "write_to_disk" in kwargs else True # We need to write data to files lazily, so we should not use file writer in __init__, # On the contrary, we will initialize the file writer when the user calls the @@ -37,6 +49,11 @@ def get_instance(cls, *args, **kwargs) -> "DistributedWriter": Overview: Get instance and set the root level instance on the first called. If args and kwargs is none, this method will return root instance. + Arguments: + - args (:obj:`Tuple`): The arguments passed to the ``__init__`` function of the parent class, \ + SummaryWriter. + - kwargs (:obj:`Dict`): The keyword arguments passed to the ``__init__`` function of the parent class, \ + SummaryWriter. """ if args or kwargs: ins = cls(*args, **kwargs) @@ -52,6 +69,9 @@ def plugin(self, router: "Parallel", is_writer: bool = False) -> "DistributedWri Plugin ``router``, so when using this writer with active router, it will automatically send requests\ to the main writer instead of writing it to the disk. So we can collect data from multiple processes\ and write them into one file. + Arguments: + - router (:obj:`Parallel`): The router to be plugged in. + - is_writer (:obj:`bool`): Whether this writer is the main writer. Examples: >>> DistributedWriter().plugin(router, is_writer=True) """ @@ -66,20 +86,44 @@ def plugin(self, router: "Parallel", is_writer: bool = False) -> "DistributedWri return self def _on_distributed_writer(self, fn_name: str, *args, **kwargs): + """ + Overview: + This method is called when the router receives a request to write data. + Arguments: + - fn_name (:obj:`str`): The name of the function to be called. + - args (:obj:`Tuple`): The arguments passed to the function to be called. + - kwargs (:obj:`Dict`): The keyword arguments passed to the function to be called. + """ + if self._is_writer: getattr(self, fn_name)(*args, **kwargs) def initialize(self): + """ + Overview: + Initialize the file writer. + """ self.close() self._write_to_disk = self._default_writer_to_disk self._get_file_writer() self._lazy_initialized = True def __del__(self): + """ + Overview: + Close the file writer. + """ self.close() def enable_parallel(fn_name, fn): + """ + Overview: + Decorator to enable parallel writing. + Arguments: + - fn_name (:obj:`str`): The name of the function to be called. + - fn (:obj:`Callable`): The function to be called. + """ def _parallel_fn(self: DistributedWriter, *args, **kwargs): if not self._lazy_initialized: diff --git a/ding/utils/memory_helper.py b/ding/utils/memory_helper.py new file mode 100644 index 0000000000..735a302a31 --- /dev/null +++ b/ding/utils/memory_helper.py @@ -0,0 +1,760 @@ +from typing import Any, Dict, List, Tuple +from collections import OrderedDict +from functools import partial, reduce + +import os +import time +import torch +try: + import pyecharts +except ImportError: + import logging + logging.warning("Please install pyecharts first, you can install it by running 'pip install pyecharts'") + pyecharts = None + +MegaByte = 1024 * 1024 + + +class SimpleMemState: + """ + Overview: + A class to represent the memory state of a model layer. + Properties: + ``layer_mem``, ``total_mem`` + Interfaces: + ``add``, ``delete``, ``update_total_memory``, ``find_layer_state``, ``dump``, ``to_json`` + """ + + def __init__(self, layer_name: str, layer_mem: int = 0) -> None: + """ + Overview: + Initialize the memory state of a model/tensors with the specific name. + Arguments: + - layer_name (:obj:`str`): The name of the layer. + - layer_mem (:obj:`int`, optional): The memory usage of the layer in bytes. Defaults to 0. + """ + self.layer_name = layer_name + + # Memory status of the current model layer. + self._layer_mem: int = layer_mem + # Total memory status of the model and sub-models, initialized with layer memory. + self._total_mem: int = self._layer_mem + # SimpleMemState of sub-models. + self.sub_model_stats = OrderedDict() + + @property + def layer_mem(self) -> int: + """ + Overview: + Get the memory usage of the layer. + + Returns: + - layer_mem (:obj:`int`): The memory usage of the layer in bytes. + """ + return self._layer_mem + + @layer_mem.setter + def layer_mem(self, new_layer_mem: int) -> None: + """ + Overview: + Set the memory usage of the layer and update the total memory. + Arguments: + - new_layer_mem (:obj:`int`): The new memory usage of the layer in bytes. + """ + diff = new_layer_mem - self._layer_mem + self._layer_mem = new_layer_mem + self._total_mem += diff + + @property + def total_mem(self) -> int: + """ + Overview: + Get the total memory usage of the model and sub-models. + + Returns: + - total_mem (:obj:`int`): The total memory usage of the model and sub-models in bytes. + """ + return self._total_mem + + def add(self, layer_name: str, layer_mem: int = 0, flush: bool = True) -> None: + """ + Overview: + Add a layer to the memory state. + Arguments: + - layer_name (:obj:`str`): The name of the layer. + - layer_mem (:obj:`int`, optional): The memory usage of the layer in bytes. Defaults to 0. + - flush (:obj:`Optional[bool]`): Whether to update the total memory usage. Defaults to True. + """ + path = layer_name.split(".") + + target = self.find_layer_state(path, create=True) + target.layer_mem = layer_mem + + if flush: + self.update_total_memory() + + def delete(self, layer_name: str, flush: bool = True) -> None: + """ + Overview: + Delete a layer from the memory state. + Arguments: + - layer_name (:obj:`str`): The name of the layer. + - flush (:obj:`Optional[bool]`): Whether to update the total memory usage. Defaults to True. + """ + path = layer_name.split(".") + assert len(path) >= 2, f"Only support deleting non-root layers, layer_name: {layer_name}" + + parent_path = path[0:-1] + layer = path[-1] + parent = self.find_layer_state(parent_path) + + if parent is not None and layer in parent.sub_model_stats: + del parent.sub_model_stats[layer] + + if flush: + self.update_total_memory() + + def update_total_memory(self) -> None: + """ + Overview: + Update the total memory usage of the model and sub-models. + """ + self._total_mem = self._layer_mem + + for stat in self.sub_model_stats.values(): + # Update sub-model status first. + stat.update_total_memory() + # Add sub-model total_mem to model total_mem. + self._total_mem += stat._total_mem + + def find_layer_state(self, path: Tuple[str], create: bool = False) -> "SimpleMemState": + """ + Overview: + Find the memory state of a layer. + Arguments: + - path (:obj:`Tuple[str]`): The path to the layer. + - create (:obj:`Optional[bool]`): Whether to create the layer if it doesn't exist. Defaults to False. + Returns: + - state (:obj:`SimpleMemState`): The memory state of the layer. + """ + current_node = self + + for _node in path: + if _node not in current_node.sub_model_stats: + if not create: + return None + # Create a layer node. + current_node.sub_model_stats[_node] = SimpleMemState(_node) + + current_node = current_node.sub_model_stats[_node] + + return current_node + + def dump(self, prefix: str = "") -> str: + """ + Overview: + Dump the memory state of the model and sub-models. + Arguments: + - prefix (:obj:`Optional[str]`): The prefix to add to the layer names. Defaults to "". + Returns: + - result (:obj:`str`): The memory state information. + """ + cur_prefix = prefix + "." + self.layer_name if prefix != "" else self.layer_name + res = f"layer: {cur_prefix}, layer_mem: {self.layer_mem / MegaByte:.2f} MB, total_mem: {self.total_mem / MegaByte:.2f} MB\n" # noqa + + for sub_layer in self.sub_model_stats.values(): + res += sub_layer.dump(cur_prefix) + + return res + + def to_json(self, base: int = 1024 * 1024) -> dict: + """ + Overview: + Convert the memory state to a JSON structure. + Arguments: + - base (:obj:`Optional[int]`): The base value to convert the memory usage to. Defaults to 1024 * 1024, \ + which converts the memory usage to MB. + Returns: + - result (:obj:`dict`): The JSON structure of the memory state. + """ + children = [child.to_json() for child in self.sub_model_stats.values()] + if len(children) == 0: + return {"name": self.layer_name, "value": self.layer_mem // base} + else: + return {"name": self.layer_name, "children": children} + + +class ActivationMemState: + """ + Overview: + A class to represent the memory state of activation tensors. + Properties: + ``total_mem`` + Interfaces: + ``add``, ``dump``, ``to_json`` + """ + + def __init__(self, num_chunks: int) -> None: + """ + Overview: + Initialize the memory state of activation tensors. + Arguments: + - num_chunks (:obj:`int`): The number of chunks, multiple chunks are used in some large-scale models. + """ + self._num_chunks = num_chunks + + self.inited: List[bool] = [False for _ in range(num_chunks)] + self.states: List[SimpleMemState] = [SimpleMemState(f"activations_{idx}") for idx in range(num_chunks)] + + @property + def total_mem(self) -> int: + """ + Overview: + Get the total memory usage of the activation tensors. + Returns: + - total_mem (:obj:`int`): The total memory usage of the activation tensors in bytes. + """ + return sum(state.total_mem for state in self.states) + + def dump(self, prefix: str = "") -> str: + """ + Overview: + Dump the memory state of the activation tensors. + Arguments: + - prefix (:obj:`Optional[str]`): The prefix to add to the layer names. Defaults to "". + Returns: + - result (:obj:`str`): The memory state information. + """ + return reduce(lambda x, y: x + y, [state.dump(prefix) for state in self.states]) + + def to_json(self, base: int = 1024 * 1024) -> List[dict]: + """ + Overview: + Convert the memory state to a JSON structure. + Arguments: + - base (:obj:`Optional[int]`): The base value to convert the memory usage to. Defaults to 1024 * 1024, \ + which converts the memory usage to MB. + Returns: + - result (:obj:`List[dict]`): The JSON structure of the memory state. + """ + return [state.to_json(base) for state in self.states] + + +def _unpack_naive_wrapper(model: torch.nn.Module) -> Tuple[torch.nn.Module, int]: + num_chunks = len(model) if isinstance(model, torch.nn.ModuleList) else 1 + + return model, num_chunks + + +class SimpleMemoryProfiler: + """ + Overview: + A memory profiler for a PyTorch neural network model. + Interfaces: + ``point``, ``step`` + """ + + def __init__( + self, + model: torch.nn.Module, + optimizer: torch.optim.Optimizer, + log_folder: str, + total_steps: int = 5, + ): + """ + Overview: + Initialize the memory profiler. + Arguments: + - model (:obj:`torch.nn.Module`): The model to profile. + - optimizer (:obj:`torch.optim.Optimizer`): The optimizer used for training the model. + - log_folder (:obj:`str`): The folder to write the memory state information to. + - total_steps (:obj:`Optional[int]`): The number of steps to trace. Defaults to 5. + """ + self._model, self._num_model_chunks = _unpack_naive_wrapper(model) + self._optimizer = optimizer + self._log_folder = log_folder + self._remaining_steps = total_steps + + self._stoped = False + self._record_start_time = time.time() + + # For activation memory state. + + self._activation_mem: int = 0 + self._activation_mem_max: int = 0 + self._activation_base_mems = ActivationMemState(self._num_model_chunks) + + # Check or create log folder + os.makedirs(self._log_folder, exist_ok=True) + + # Register activation memory tracking hooks + if self._num_model_chunks > 1: + for chunk_id in range(self._num_model_chunks): + self._register_activation_trace_hooks(chunk_id, self._model[chunk_id]) + else: + self._register_activation_trace_hooks(0, self._model) + + # Calculate static parameter cuda memory + self._param_mem_state = SimpleMemState("param_mem") + self._calc_tensor_memory(self._param_mem_state, self._model.named_parameters()) + # Calculate static grad cuda memory + self._grad_mem_state = SimpleMemState("grad_mem") + self._calc_tensor_memory(self._grad_mem_state, self._model.named_parameters(), True) + # Calculate static optimizer state cuda memory + self._os_params_mem_state = SimpleMemState("os_params_mem") + self._os_state_mem_state = SimpleMemState("os_state_mem") + self._calc_tensor_group_memory(self._os_params_mem_state, list(enumerate(self._optimizer.param_groups))) + + # Generate the first memory record + self.point(with_options="params,grads,os_params", create=True) + + def point(self, with_options: str = "", create: bool = False) -> None: + """ + Overview: + Record the memory state of the model and optimizer at current point. + Arguments: + - with_options (:obj:`Optional[str]`): The options to include in the memory state. Defaults to "". + - create (:obj:`Optional[bool]`): Whether to create a new memory record. Defaults to False. + """ + now = time.time() + file = f"{self._log_folder}/memory.log" + + if with_options == "all": + options = ["params", "grads", "os_params", "os_state", "activation_base"] + else: + options = with_options.split(",") + + total_mem = ( + self._param_mem_state.total_mem + self._grad_mem_state.total_mem + self._os_params_mem_state.total_mem + + self._os_state_mem_state.total_mem + self._activation_mem + ) / MegaByte + + # Generate summary information for memory state + summary_info = ( + f"total_memory: {total_mem:.2f} MB" + "\n" + + f"params_memory: {self._param_mem_state.total_mem / MegaByte:.2f} MB, " + + f"grads_memory: {self._grad_mem_state.total_mem / MegaByte:.2f} MB, " + + f"os_params_memory: {self._os_params_mem_state.total_mem / MegaByte:.2f} MB, " + + f"os_state_memory: {self._os_state_mem_state.total_mem / MegaByte:.2f} MB, " + + f"activation_memory: {self._activation_mem / MegaByte:.2f} MB" + ) + + # Generate layout information based on selected options + layout_info = "" + if "params" in options: + layout_info += "params_layout:\n" + self._param_mem_state.dump() + if "grads" in options: + layout_info += "grads_layout:\n" + self._grad_mem_state.dump() + if "os_params" in options: + layout_info += "os_params_layout:\n" + self._os_params_mem_state.dump() + if "os_state" in options: + layout_info += "os_state_layout:\n" + self._os_state_mem_state.dump() + if "activation_base" in options: + layout_info += "activation_base_layout:\n" + self._activation_base_mems.dump() + + # Write memory state information to log file + file_mode = "w" if create else "a" + with open(file, file_mode, encoding="utf-8") as writer: + writer.write( + "Memory State:\n" + f"time: {now - self._record_start_time}\n" + "---summary---\n" + summary_info + "\n" + ) + if layout_info != "": + writer.write("---Layout---\n" + layout_info) + writer.write("\n") + + def step(self) -> None: + """ + Overview: + Update the memory state of the optimizer state (e.g., momentum, learning rate) and record the memory state. + """ + if self._stoped: + return + + self._remaining_steps -= 1 + if self._remaining_steps == 0: + self._stoped = True + + # Update os state memory usage + self._os_state_mem_state = SimpleMemState("os_state_mem") + self._calc_tensor_group_memory(self._os_state_mem_state, list(self._optimizer.state_dict()["state"].items())) + + if not self._stoped: + # Do we need to print os_state_layout every time? Is it always constant? + self.point(with_options="os_state") + else: + # Dump memory layout + self.point(with_options="all") + # Generate sunburst charts + self._render_sunburst_chart(self._param_mem_state.to_json()["children"], "params_memory_sunburst") + self._render_sunburst_chart(self._grad_mem_state.to_json()["children"], "grads_memory_sunburst") + self._render_sunburst_chart( + [self._os_params_mem_state.to_json(), + self._os_state_mem_state.to_json()], + "os_memory_sunburst", + ) + self._render_sunburst_chart(self._activation_base_mems.to_json(), "activation_memory_sunburst") + # Generate summary sunburst chart + summary_sunburst_data = [ + { + "name": "params", + "value": self._param_mem_state.total_mem // MegaByte + }, + { + "name": "grads", + "value": self._grad_mem_state.total_mem // MegaByte + }, + { + "name": "os_params", + "value": self._os_params_mem_state.total_mem // MegaByte + }, + { + "name": "os_state", + "value": self._os_state_mem_state.total_mem // MegaByte + }, + { + "name": "activation", + "value": self._activation_mem_max // MegaByte + }, + ] + + self._render_sunburst_chart(summary_sunburst_data, "summary_sunburst") + + def _render_sunburst_chart(self, data: Any, name: str) -> None: + """ + Overview: + Render a sunburst chart for the memory state with pyecharts. + Arguments: + - data (:obj:`Any`): The data to render. + - name (:obj:`str`): The name of the chart. + """ + pyecharts.charts.Sunburst(init_opts=pyecharts.options.InitOpts(width="1000px", height="1000px")).add( + name, + data_pair=data, + highlight_policy="ancestor", + radius=[0, "95%"], + levels=[ + {}, + { + "r0": "10%", + "r": "35%", + "itemStyle": { + "borderWidth": 3 + }, + "label": { + "align": "left" + }, + }, + { + "r0": "35%", + "r": "55%", + "label": { + "align": "left" + } + }, + { + "r0": "55%", + "r": "70%", + "label": { + "align": "left" + } + }, + { + "r0": "70%", + "r": "80%", + "label": { + "align": "left" + } + }, + { + "r0": "80%", + "r": "90%", + "label": { + "align": "left" + } + }, + { + "r0": "90%", + "r": "92%", + "label": { + "position": "outside", + "padding": 3, + "silent": False + }, + "itemStyle": { + "borderWidth": 3 + }, + }, + ], + ).set_global_opts(title_opts=pyecharts.options.TitleOpts(title="CUDA Memory") + ).set_series_opts(label_opts=pyecharts.options.LabelOpts(formatter="{b}") + ).render(f"{self._log_folder}/{name}.html") + + def _inner_activation_trace_hook( + self, + chunk_id: int, + layer_name: str, + model: Any, + inputs: Any, + output: torch.Tensor, + ) -> None: + """ + Overview: + Hook function to trace the activation memory usage for a inner layer. + + .. note:: + For more details about hook mechanism, please refer to the PyTorch documentation. + + Arguments: + - chunk_id (:obj:`int`): The model chunk id. + - layer_name (:obj:`str`): The name of the layer. + - model (:obj:`Any`): The model to trace. + - inputs (:obj:`Any`): The inputs to the layer. + - output (:obj:`torch.Tensor`): The output tensor. + """ + del model, inputs + assert isinstance(output, torch.Tensor), f"Invalid output type: {type(output)}" + + if self._stoped or self._activation_base_mems.inited[chunk_id]: + return + + # Delay updating the total_mem of activation_base_mem here, it will be handled in the forward ending hook. + self._activation_base_mems.states[chunk_id].add( + layer_name, output.element_size() * output.nelement(), flush=False + ) + + def _activation_trace_hook_forward(self, chunk_id: int, model: Any, inputs: Any, output: Any) -> None: + """ + Overview: + Hook function to trace the activation memory usage for a forward pass. + + .. note:: + For more details about hook mechanism, please refer to the PyTorch documentation. + + Arguments: + - chunk_id (:obj:`int`): The model chunk id. + - model (:obj:`Any`): The model to trace. + - inputs (:obj:`Any`): The inputs to the model. + - output (:obj:`Any`): The output of the model. + """ + del model, inputs + + if self._stoped: + return + + # Check if the activation memory has been initialized + if self._activation_base_mems.inited[chunk_id] is False: + self._activation_base_mems.inited[chunk_id] = True + # Update the total memory of the activation base memory state + self._activation_base_mems.states[chunk_id].update_total_memory() + # Set with_options to "activation_base" to include activation_base_layout in the memory dump + with_options = "activation_base" + else: + with_options = "" + + # Accumulate activation memory usage for each forward pass + self._activation_mem += self._activation_base_mems.states[chunk_id].total_mem + if self._activation_mem > self._activation_mem_max: + self._activation_mem_max = self._activation_mem + + # Trigger a memory record + self.point(with_options) + + def _activation_tarce_hook_backward(self, chunk_id: int, model: Any, inputs: Any, grad_outputs: Any) -> None: + """ + Overview: + Hook function to trace the activation memory usage for a backward pass. + + .. note:: + For more details about hook mechanism, please refer to the PyTorch documentation. + + Arguments: + - chunk_id (:obj:`int`): The model chunk id. + - model (:obj:`Any`): The model to trace. + - inputs (:obj:`Any`): The inputs to the model. + - grad_outputs (:obj:`Any`): The gradients of the outputs. + """ + del model, inputs, grad_outputs + + if self._stoped: + return + + # Release activation memory usage for each backward pass + self._activation_mem -= self._activation_base_mems.states[chunk_id].total_mem + + # Trigger a memory record + self.point() + + def _register_activation_trace_hooks(self, chunk_id: int, model_chunk: torch.nn.Module) -> None: + """ + Overview: + Register activation trace hooks for the model and each submodule in the model. + Arguments: + - chunk_id (:obj:`int`): The model chunk id. + - model_chunk (:obj:`torch.nn.Module`): The model chunk to trace. + """ + + # Register inner activation trace hooks for each submodule in the model + for layer_name, sub_model in model_chunk.named_modules(): + # Register the hook + if len(sub_model._modules) != 0: + continue # TODO: in some special cases, we may need some additional configuration to correct + + sub_model.register_forward_hook(partial(self._inner_activation_trace_hook, chunk_id, layer_name)) + + # Register a forward hook for the main model to track activation memory usage + model_chunk.register_forward_hook(partial(self._activation_trace_hook_forward, chunk_id)) + # Register a backward hook for the main model to release activation memory usage + model_chunk.register_full_backward_hook(partial(self._activation_tarce_hook_backward, chunk_id)) + + def _calc_tensor_memory( + self, + root_stat: SimpleMemState, + named_tensors: Dict[str, torch.Tensor], + require_grad: bool = False + ) -> None: + """ + Overview: + Core function to calculate the memory usage of tensors and update the memory state. + Arguments: + - root_stat (:obj:`SimpleMemState`): The root memory state. + - named_tensors (:obj:`Dict[str, torch.Tensor]`): A dictionary containing the named tensors. + - require_grad (:obj:`Optional[bool]`): Whether to consider tensors with gradients. Defaults to False. + """ + for name, tensor in named_tensors: + if require_grad and not tensor.requires_grad: + continue + + layer_splits = name.split(sep=".") + layer_stat = root_stat.find_layer_state(layer_splits, create=True) + layer_stat.layer_mem = tensor.element_size() * tensor.nelement() + + root_stat.update_total_memory() + + def _calc_tensor_group_memory(self, root_stat: SimpleMemState, tensor_groups: List[Tuple[int, torch.Tensor]]): + """ + Overview: + Core function to calculate the memory usage of a group of tensors and update the memory state. + Arguments: + - root_stat (:obj:`SimpleMemState`): The root memory state. + - tensor_groups (:obj:`List[Tuple[int, torch.Tensor]]`): A list of tuples containing the tensor groups. + """ + + def _normalize_helper(named_tensors: Dict[str, Any]) -> List[Tuple[str, Any]]: + res = {} + + for name, tensors in named_tensors.items(): + if isinstance(tensors, torch.Tensor): + res[name] = tensors + elif isinstance(tensors, (list, tuple)): + for index, tensor in enumerate(tensors): + res[f"{name}.{index}"] = tensor + elif isinstance(tensors, dict): + for subname, tensor in tensors.items(): + res[f"{name}.{subname}"] = tensor + else: + raise TypeError(f"unsupported normalize value type: {type(tensors)}") + + return list(res.items()) + + def _value_check(tensor_or_tensors): + if torch.is_tensor(tensor_or_tensors): + return True + elif isinstance(tensor_or_tensors, (list, tuple)) and all(torch.is_tensor(x) for x in tensor_or_tensors): + return True + elif isinstance(tensor_or_tensors, dict) and all(torch.is_tensor(x) for x in tensor_or_tensors.values()): + return True + else: + return False + + # Calculate the memory usage of a group of tensors. + for idx, tensors in tensor_groups: + # Normalize the named tensors + named_tensors = {f"{idx}.{k}": v for k, v in tensors.items() if _value_check(v)} + named_tensors = _normalize_helper(named_tensors) + # Calculate the memory usage of the tensors and update the memory state + self._calc_tensor_memory(root_stat, named_tensors) + + +def get_current_device() -> torch.device: + """ + Overview: + Get the current PyTorch tensor device. + + Returns: + - device (:obj:`torch.device`): The current device. + """ + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +def multi_chunk_test(): + """ + Overview: + A test function to demonstrate the memory profiler for a model with multiple chunks. + """ + + class SimpleModel(torch.nn.Module): + + def __init__(self, skip_layer2: bool = False): + super().__init__() + self.layer1 = torch.nn.Linear(5120, 5120, True) + self.layer3 = torch.nn.Linear(5120, 5120, False) + + if skip_layer2: + self.layer2 = None + else: + self.layer2 = SimpleModel(skip_layer2=True) + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + output1 = self.layer1(inputs) + if self.layer2 is not None: + output2 = self.layer2(output1) + else: + output2 = output1 + output = self.layer3(output2) + + return output + + def _simple_schedule(_num_chunks, _model_chunks, _input) -> torch.Tensor: + if _num_chunks > 1: + _output = _input + for _model_chunk in _model_chunks: + _output = _model_chunk(_output) + else: + _output = _model_chunks(_input) + + return _output + + # num_chunks config + _num_chunks = 1 + + # init model and optimizer + if _num_chunks > 1: + _chunks = [SimpleModel(skip_layer2=idx % 2 == 0) for idx in range(_num_chunks)] + _model = torch.nn.ModuleList(_chunks).to(get_current_device()) + else: + _model: torch.nn.Module = SimpleModel().to(get_current_device()) + _optimizer = torch.optim.Adam(_model.parameters()) + + # init profiler + profiler = SimpleMemoryProfiler(_model, _optimizer, "./test_simple_memory_profiler_multi_chunk", total_steps=1) + + _optimizer.zero_grad() + + # inputs + x1 = torch.randn((128, 5120)).to(get_current_device()) + x2 = torch.randn((128, 5120)).to(get_current_device()) + # forward + out1 = _simple_schedule(_num_chunks, _model, x1) + out2 = _simple_schedule(_num_chunks, _model, x2) + # backward + out1.mean().backward() + out2.mean().backward() + + _optimizer.step() + + # Update the optimizer state memory usage and record the memory state + profiler.step() + + +if __name__ == "__main__": + multi_chunk_test() diff --git a/ding/utils/normalizer_helper.py b/ding/utils/normalizer_helper.py new file mode 100755 index 0000000000..0fc914f30e --- /dev/null +++ b/ding/utils/normalizer_helper.py @@ -0,0 +1,493 @@ +import numpy as np + + +class DatasetNormalizer: + """ + Overview: + The `DatasetNormalizer` class provides functionality to normalize and unnormalize data in a dataset. + It takes a dataset as input and applies a normalizer function to each key in the dataset. + + Interfaces: + ``__init__``, ``__repr__``, ``normalize``, ``unnormalize``. + """ + + def __init__(self, dataset: np.ndarray, normalizer: str, path_lengths: list = None): + """ + Overview: + Initialize the NormalizerHelper object. + + Arguments: + - dataset (:obj:`np.ndarray`): The dataset to be normalized. + - normalizer (:obj:`str`): The type of normalizer to be used. Can be a string representing the name of \ + the normalizer class. + - path_lengths (:obj:`list`): The length of the paths in the dataset. Defaults to None. + """ + dataset = flatten(dataset, path_lengths) + + self.observation_dim = dataset['observations'].shape[1] + self.action_dim = dataset['actions'].shape[1] + + if isinstance(normalizer, str): + normalizer = eval(normalizer) + + self.normalizers = {} + for key, val in dataset.items(): + try: + self.normalizers[key] = normalizer(val) + except: + print(f'[ utils/normalization ] Skipping {key} | {normalizer}') + # key: normalizer(val) + # for key, val in dataset.items() + + def __repr__(self) -> str: + """ + Overview: + Returns a string representation of the NormalizerHelper object. \ + The string representation includes the key-value pairs of the normalizers \ + stored in the NormalizerHelper object. + Returns: + - ret (:obj:`str`):A string representation of the NormalizerHelper object. + """ + string = '' + for key, normalizer in self.normalizers.items(): + string += f'{key}: {normalizer}]\n' + return string + + def normalize(self, x: np.ndarray, key: str) -> np.ndarray: + """ + Overview: + Normalize the input data using the specified key. + + Arguments: + - x (:obj:`np.ndarray`): The input data to be normalized. + - key (:obj`str`): The key to identify the normalizer. + + Returns: + - ret (:obj:`np.ndarray`): The normalized value of the input data. + """ + return self.normalizers[key].normalize(x) + + def unnormalize(self, x: np.ndarray, key: str) -> np.ndarray: + """ + Overview: + Unnormalizes the given value `x` using the specified `key`. + + Arguments: + - x (:obj:`np.ndarray`): The value to be unnormalized. + - key (:obj`str`): The key to identify the normalizer. + + Returns: + - ret (:obj:`np.ndarray`): The unnormalized value. + """ + return self.normalizers[key].unnormalize(x) + + +def flatten(dataset: dict, path_lengths: list) -> dict: + """ + Overview: + Flattens dataset of { key: [ n_episodes x max_path_length x dim ] } \ + to { key : [ (n_episodes * sum(path_lengths)) x dim ] } + + Arguments: + - dataset (:obj:`dict`): The dataset to be flattened. + - path_lengths (:obj:`list`): A list of path lengths for each episode. + + Returns: + - flattened (:obj:`dict`): The flattened dataset. + """ + flattened = {} + for key, xs in dataset.items(): + assert len(xs) == len(path_lengths) + if key == 'path_lengths': + continue + flattened[key] = np.concatenate([x[:length] for x, length in zip(xs, path_lengths)], axis=0) + return flattened + + +class Normalizer: + """ + Overview: + Parent class, subclass by defining the `normalize` and `unnormalize` methods + + Interfaces: + ``__init__``, ``__repr__``, ``normalize``, ``unnormalize``. + """ + + def __init__(self, X): + """ + Overview: + Initialize the Normalizer object. + Arguments: + - X (:obj:`np.ndarray`): The data to be normalized. + """ + + self.X = X.astype(np.float32) + self.mins = X.min(axis=0) + self.maxs = X.max(axis=0) + + def __repr__(self) -> str: + """ + Overview: + Returns a string representation of the Normalizer object. + Returns: + - ret (:obj:`str`): A string representation of the Normalizer object. + """ + + return ( + f"""[ Normalizer ] dim: {self.mins.size}\n -: """ + f"""{np.round(self.mins, 2)}\n +: {np.round(self.maxs, 2)}\n""" + ) + + def normalize(self, *args, **kwargs): + """ + Overview: + Normalize the input data. + Arguments: + - args (:obj:`list`): The arguments passed to the ``normalize`` function. + - kwargs (:obj:`dict`): The keyword arguments passed to the ``normalize`` function. + """ + + raise NotImplementedError() + + def unnormalize(self, *args, **kwargs): + """ + Overview: + Unnormalize the input data. + Arguments: + - args (:obj:`list`): The arguments passed to the ``unnormalize`` function. + - kwargs (:obj:`dict`): The keyword arguments passed to the ``unnormalize`` function. + """ + + raise NotImplementedError() + + +class GaussianNormalizer(Normalizer): + """ + Overview: + A class that normalizes data to zero mean and unit variance. + + Interfaces: + ``__init__``, ``__repr__``, ``normalize``, ``unnormalize``. + """ + + def __init__(self, *args, **kwargs): + """ + Overview: + Initialize the GaussianNormalizer object. + Arguments: + - args (:obj:`list`): The arguments passed to the ``__init__`` function of the parent class, \ + i.e., the Normalizer class. + - kwargs (:obj:`dict`): The keyword arguments passed to the ``__init__`` function of the parent class, \ + i.e., the Normalizer class. + """ + + super().__init__(*args, **kwargs) + self.means = self.X.mean(axis=0) + self.stds = self.X.std(axis=0) + self.z = 1 + + def __repr__(self) -> str: + """ + Overview: + Returns a string representation of the GaussianNormalizer object. + Returns: + - ret (:obj:`str`): A string representation of the GaussianNormalizer object. + """ + + return ( + f"""[ Normalizer ] dim: {self.mins.size}\n """ + f"""means: {np.round(self.means, 2)}\n """ + f"""stds: {np.round(self.z * self.stds, 2)}\n""" + ) + + def normalize(self, x: np.ndarray) -> np.ndarray: + """ + Overview: + Normalize the input data. + + Arguments: + - x (:obj:`np.ndarray`): The input data to be normalized. + + Returns: + - ret (:obj:`np.ndarray`): The normalized data. + """ + return (x - self.means) / self.stds + + def unnormalize(self, x: np.ndarray) -> np.ndarray: + """ + Overview: + Unnormalize the input data. + + Arguments: + - x (:obj:`np.ndarray`): The input data to be unnormalized. + + Returns: + - ret (:obj:`np.ndarray`): The unnormalized data. + """ + return x * self.stds + self.means + + +class CDFNormalizer(Normalizer): + """ + Overview: + A class that makes training data uniform (over each dimension) by transforming it with marginal CDFs. + + Interfaces: + ``__init__``, ``__repr__``, ``normalize``, ``unnormalize``. + """ + + def __init__(self, X): + """ + Overview: + Initialize the CDFNormalizer object. + Arguments: + - X (:obj:`np.ndarray`): The data to be normalized. + """ + + super().__init__(atleast_2d(X)) + self.dim = self.X.shape[1] + self.cdfs = [CDFNormalizer1d(self.X[:, i]) for i in range(self.dim)] + + def __repr__(self) -> str: + """ + Overview: + Returns a string representation of the CDFNormalizer object. + Returns: + - ret (:obj:`str`): A string representation of the CDFNormalizer object. + """ + + return f'[ CDFNormalizer ] dim: {self.mins.size}\n' + ' | '.join( + f'{i:3d}: {cdf}' for i, cdf in enumerate(self.cdfs) + ) + + def wrap(self, fn_name: str, x: np.ndarray) -> np.ndarray: + """ + Overview: + Wraps the given function name and applies it to the input data. + + Arguments: + - fn_name (:obj:`str`): The name of the function to be applied. + - x (:obj:`np.ndarray`): The input data. + + Returns: + - ret: The output of the function applied to the input data. + """ + shape = x.shape + # reshape to 2d + x = x.reshape(-1, self.dim) + out = np.zeros_like(x) + for i, cdf in enumerate(self.cdfs): + fn = getattr(cdf, fn_name) + out[:, i] = fn(x[:, i]) + return out.reshape(shape) + + def normalize(self, x: np.ndarray) -> np.ndarray: + """ + Overview: + Normalizes the input data. + + Arguments: + - x (:obj:`np.ndarray`): The input data. + + Returns: + - ret (:obj:`np.ndarray`): The normalized data. + """ + return self.wrap('normalize', x) + + def unnormalize(self, x: np.ndarray) -> np.ndarray: + """ + Overview: + Unnormalizes the input data. + + Arguments: + - x (:obj:`np.ndarray`): The input data. + + Returns: + - ret (:obj:`np.ndarray`):: The unnormalized data. + """ + return self.wrap('unnormalize', x) + + +class CDFNormalizer1d: + """ + Overview: + CDF normalizer for a single dimension. This class provides methods to normalize and unnormalize data \ + using the Cumulative Distribution Function (CDF) approach. + Interfaces: + ``__init__``, ``__repr__``, ``normalize``, ``unnormalize``. + """ + + def __init__(self, X: np.ndarray): + """ + Overview: + Initialize the CDFNormalizer1d object. + Arguments: + - X (:obj:`np.ndarray`): The data to be normalized. + """ + + import scipy.interpolate as interpolate + assert X.ndim == 1 + self.X = X.astype(np.float32) + if self.X.max() == self.X.min(): + self.constant = True + else: + self.constant = False + quantiles, cumprob = empirical_cdf(self.X) + self.fn = interpolate.interp1d(quantiles, cumprob) + self.inv = interpolate.interp1d(cumprob, quantiles) + + self.xmin, self.xmax = quantiles.min(), quantiles.max() + self.ymin, self.ymax = cumprob.min(), cumprob.max() + + def __repr__(self) -> str: + """ + Overview: + Returns a string representation of the CDFNormalizer1d object. + """ + + return (f'[{np.round(self.xmin, 2):.4f}, {np.round(self.xmax, 2):.4f}') + + def normalize(self, x: np.ndarray) -> np.ndarray: + """ + Overview: + Normalize the input data. + + Arguments: + - x (:obj:`np.ndarray`): The data to be normalized. + + Returns: + - ret (:obj:`np.ndarray`): The normalized data. + """ + if self.constant: + return x + + x = np.clip(x, self.xmin, self.xmax) + # [ 0, 1 ] + y = self.fn(x) + # [ -1, 1 ] + y = 2 * y - 1 + return y + + def unnormalize(self, x: np.ndarray, eps: float = 1e-4) -> np.ndarray: + """ + Overview: + Unnormalize the input data. + + Arguments: + - x (:obj:`np.ndarray`): The data to be unnormalized. + - eps (:obj:`float`): A small value used for numerical stability. Defaults to 1e-4. + + Returns: + - ret (:obj:`np.ndarray`): The unnormalized data. + """ + # [ -1, 1 ] --> [ 0, 1 ] + if self.constant: + return x + + x = (x + 1) / 2. + + if (x < self.ymin - eps).any() or (x > self.ymax + eps).any(): + print( + f"""[ dataset/normalization ] Warning: out of range in unnormalize: """ + f"""[{x.min()}, {x.max()}] | """ + f"""x : [{self.xmin}, {self.xmax}] | """ + f"""y: [{self.ymin}, {self.ymax}]""" + ) + + x = np.clip(x, self.ymin, self.ymax) + + y = self.inv(x) + return y + + +def empirical_cdf(sample: np.ndarray) -> (np.ndarray, np.ndarray): + """ + Overview: + Compute the empirical cumulative distribution function (CDF) of a given sample. + + Arguments: + - sample (:obj:`np.ndarray`): The input sample for which to compute the empirical CDF. + + Returns: + - quantiles (:obj:`np.ndarray`): The unique values in the sample. + - cumprob (:obj:`np.ndarray`): The cumulative probabilities corresponding to the quantiles. + + References: + - Stack Overflow: https://stackoverflow.com/a/33346366 + """ + + # find the unique values and their corresponding counts + quantiles, counts = np.unique(sample, return_counts=True) + + # take the cumulative sum of the counts and divide by the sample size to + # get the cumulative probabilities between 0 and 1 + cumprob = np.cumsum(counts).astype(np.double) / sample.size + + return quantiles, cumprob + + +def atleast_2d(x: np.ndarray) -> np.ndarray: + """ + Overview: + Ensure that the input array has at least two dimensions. + + Arguments: + - x (:obj:`np.ndarray`): The input array. + + Returns: + - ret (:obj:`np.ndarray`): The input array with at least two dimensions. + """ + if x.ndim < 2: + x = x[:, None] + return x + + +class LimitsNormalizer(Normalizer): + """ + Overview: + A class that normalizes and unnormalizes values within specified limits. \ + This class maps values within the range [xmin, xmax] to the range [-1, 1]. + + Interfaces: + ``__init__``, ``__repr__``, ``normalize``, ``unnormalize``. + """ + + def normalize(self, x: np.ndarray) -> np.ndarray: + """ + Overview: + Normalizes the input values. + + Argments: + - x (:obj:`np.ndarray`): The input values to be normalized. + + Returns: + - ret (:obj:`np.ndarray`): The normalized values. + + """ + # [ 0, 1 ] + x = (x - self.mins) / (self.maxs - self.mins) + # [ -1, 1 ] + x = 2 * x - 1 + return x + + def unnormalize(self, x: np.ndarray, eps: float = 1e-4) -> np.ndarray: + """ + Overview: + Unnormalizes the input values. + + Arguments: + - x (:obj:`np.ndarray`): The input values to be unnormalized. + - eps (:obj:`float`): A small value used for clipping. Defaults to 1e-4. + + Returns: + - ret (:obj:`np.ndarray`): The unnormalized values. + + """ + if x.max() > 1 + eps or x.min() < -1 - eps: + # print(f'[ datasets/mujoco ] Warning: sample out of range | ({x.min():.4f}, {x.max():.4f})') + x = np.clip(x, -1, 1) + + # [ -1, 1 ] --> [ 0, 1 ] + x = (x + 1) / 2. + + return x * (self.maxs - self.mins) + self.mins diff --git a/ding/utils/orchestrator_launcher.py b/ding/utils/orchestrator_launcher.py index 4b18195e72..69324ecc08 100644 --- a/ding/utils/orchestrator_launcher.py +++ b/ding/utils/orchestrator_launcher.py @@ -6,7 +6,10 @@ class OrchestratorLauncher(object): """ - Overview: object to manage di-orchestrator in existing k8s cluster + Overview: + Object to manage di-orchestrator in existing k8s cluster + Interfaces: + ``__init__``, ``create_orchestrator``, ``delete_orchestrator`` """ def __init__( @@ -18,6 +21,18 @@ def __init__( cert_manager_version: str = 'v1.3.1', cert_manager_registry: str = 'quay.io/jetstack' ) -> None: + """ + Overview: + Initialize the OrchestratorLauncher object. + Arguments: + - version (:obj:`str`): The version of di-orchestrator. + - name (:obj:`str`): The name of di-orchestrator. + - cluster (:obj:`K8sLauncher`): The k8s cluster to deploy di-orchestrator. + - registry (:obj:`str`): The docker registry to pull images. + - cert_manager_version (:obj:`str`): The version of cert-manager. + - cert_manager_registry (:obj:`str`): The docker registry to pull cert-manager images. + """ + self.name = name self.version = version self.cluster = cluster @@ -47,6 +62,11 @@ def __init__( self._check_kubectl_tools() def _check_kubectl_tools(self) -> None: + """ + Overview: + Check if kubectl tools is installed. + """ + args = ['which', 'kubectl'] proc = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE) out, _ = proc.communicate() @@ -56,6 +76,11 @@ def _check_kubectl_tools(self) -> None: ) def create_orchestrator(self) -> None: + """ + Overview: + Create di-orchestrator in k8s cluster. + """ + print('Creating orchestrator...') if self.cluster is not None: self.cluster.preload_images(self._images) @@ -69,6 +94,11 @@ def create_orchestrator(self) -> None: wait_to_be_ready(self._namespace, self._webhook) def delete_orchestrator(self) -> None: + """ + Overview: + Delete di-orchestrator in k8s cluster. + """ + print('Deleting orchestrator...') for item in [self.cert_manager, self.installer]: args = ['kubectl', 'delete', '-f', f'{item}'] @@ -81,6 +111,13 @@ def delete_orchestrator(self) -> None: def create_components_from_config(config: str) -> None: + """ + Overview: + Create components from config file. + Arguments: + - config (:obj:`str`): The config file. + """ + args = ['kubectl', 'create', '-f', f'{config}'] proc = subprocess.Popen(args, stderr=subprocess.PIPE) _, err = proc.communicate() @@ -93,6 +130,15 @@ def create_components_from_config(config: str) -> None: def wait_to_be_ready(namespace: str, component: str, timeout: int = 120) -> None: + """ + Overview: + Wait for the component to be ready. + Arguments: + - namespace (:obj:`str`): The namespace of the component. + - component (:obj:`str`): The name of the component. + - timeout (:obj:`int`): The timeout of waiting. + """ + try: from kubernetes import config, client, watch except ModuleNotFoundError: diff --git a/ding/utils/profiler_helper.py b/ding/utils/profiler_helper.py index 20e1024629..96c2a1a076 100644 --- a/ding/utils/profiler_helper.py +++ b/ding/utils/profiler_helper.py @@ -10,15 +10,42 @@ def register_profiler(write_profile, pr, folder_path): class Profiler: + """ + Overview: + A class for profiling code execution. It can be used as a context manager or a decorator. + + Interfaces: + ``__init__``, ``mkdir``, ``write_profile``, ``profile``. + """ def __init__(self): + """ + Overview: + Initialize the Profiler object. + """ + self.pr = cProfile.Profile() - def mkdir(self, directory): + def mkdir(self, directory: str): + """ + OverView: + Create a directory if it doesn't exist. + + Arguments: + - directory (:obj:`str`): The path of the directory to be created. + """ if not os.path.exists(directory): os.makedirs(directory) - def write_profile(self, pr, folder_path): + def write_profile(self, pr: cProfile.Profile, folder_path: str): + """ + OverView: + Write the profiling results to files. + + Arguments: + - pr (:obj:`cProfile.Profile`): The profiler object containing the profiling results. + - folder_path (:obj:`str`): The path of the folder where the profiling files will be saved. + """ pr.disable() s_tottime = io.StringIO() s_cumtime = io.StringIO() @@ -36,6 +63,14 @@ def write_profile(self, pr, folder_path): pr.dump_stats(folder_path + "/profile.prof") def profile(self, folder_path="./tmp"): + """ + OverView: + Enable profiling and save the results to files. + + Arguments: + - folder_path (:obj:`str`): The path of the folder where the profiling files will be saved. \ + Defaults to "./tmp". + """ self.mkdir(folder_path) self.pr.enable() register_profiler(self.write_profile, self.pr, folder_path) diff --git a/ding/utils/pytorch_ddp_dist_helper.py b/ding/utils/pytorch_ddp_dist_helper.py index 96847be357..fcea6c81e8 100644 --- a/ding/utils/pytorch_ddp_dist_helper.py +++ b/ding/utils/pytorch_ddp_dist_helper.py @@ -1,9 +1,11 @@ -import os -from typing import Callable, Tuple, List, Any +from typing import Callable, Tuple, List, Any, Union +from easydict import EasyDict +import os import numpy as np import torch import torch.distributed as dist +import datetime from .default_helper import error_wrapper @@ -11,7 +13,7 @@ def get_rank() -> int: - r""" + """ Overview: Get the rank of current process in total world_size """ @@ -20,7 +22,7 @@ def get_rank() -> int: def get_world_size() -> int: - r""" + """ Overview: Get the world_size(total process number in data parallel training) """ @@ -30,23 +32,105 @@ def get_world_size() -> int: broadcast = dist.broadcast allgather = dist.all_gather +broadcast_object_list = dist.broadcast_object_list def allreduce(x: torch.Tensor) -> None: + """ + Overview: + All reduce the tensor ``x`` in the world + Arguments: + - x (:obj:`torch.Tensor`): the tensor to be reduced + """ + dist.all_reduce(x) x.div_(get_world_size()) +def allreduce_with_indicator(grad: torch.Tensor, indicator: torch.Tensor) -> None: + """ + Overview: + Custom allreduce: Sum both the gradient and indicator tensors across all processes. + Then, if at least one process contributed (i.e., the summation of indicator > 0), + divide the gradient by the summed indicator. This ensures that if only a subset of + GPUs contributed a gradient, the averaging is performed based on the actual number + of contributors rather than the total number of GPUs. + Arguments: + - grad (torch.Tensor): Local gradient tensor to be reduced. + - indicator (torch.Tensor): A tensor flag (1 if the gradient is computed, 0 otherwise). + """ + # Allreduce (sum) the gradient and indicator + dist.all_reduce(grad) + dist.all_reduce(indicator) + + # Avoid division by zero. If indicator is close to 0 (extreme case), grad remains zeros. + if not torch.isclose(indicator, torch.tensor(0.0)): + grad.div_(indicator.item()) + + def allreduce_async(name: str, x: torch.Tensor) -> None: + """ + Overview: + All reduce the tensor ``x`` in the world asynchronously + Arguments: + - name (:obj:`str`): the name of the tensor + - x (:obj:`torch.Tensor`): the tensor to be reduced + """ + x.div_(get_world_size()) dist.all_reduce(x, async_op=True) +def reduce_data(x: Union[int, float, torch.Tensor], dst: int) -> Union[int, float, torch.Tensor]: + """ + Overview: + Reduce the tensor ``x`` to the destination process ``dst`` + Arguments: + - x (:obj:`Union[int, float, torch.Tensor]`): the tensor to be reduced + - dst (:obj:`int`): the destination process + """ + + if np.isscalar(x): + x_tensor = torch.as_tensor([x]).cuda() + dist.reduce(x_tensor, dst) + return x_tensor.item() + elif isinstance(x, torch.Tensor): + dist.reduce(x, dst) + return x + else: + raise TypeError("not supported type: {}".format(type(x))) + + +def allreduce_data(x: Union[int, float, torch.Tensor], op: str) -> Union[int, float, torch.Tensor]: + """ + Overview: + All reduce the tensor ``x`` in the world + Arguments: + - x (:obj:`Union[int, float, torch.Tensor]`): the tensor to be reduced + - op (:obj:`str`): the operation to perform on data, support ``['sum', 'avg']`` + """ + + assert op in ['sum', 'avg'], op + if np.isscalar(x): + x_tensor = torch.as_tensor([x]).cuda() + dist.all_reduce(x_tensor) + if op == 'avg': + x_tensor.div_(get_world_size()) + return x_tensor.item() + elif isinstance(x, torch.Tensor): + dist.all_reduce(x) + if op == 'avg': + x.div_(get_world_size()) + return x + else: + raise TypeError("not supported type: {}".format(type(x))) + + synchronize = torch.cuda.synchronize def get_group(group_size: int) -> List: - r""" + """ Overview: Get the group segmentation of ``group_size`` each group Arguments: @@ -61,9 +145,11 @@ def get_group(group_size: int) -> List: def dist_mode(func: Callable) -> Callable: - r""" + """ Overview: Wrap the function so that in can init and finalize automatically before each call + Arguments: + - func (:obj:`Callable`): the function to be wrapped """ def wrapper(*args, **kwargs): @@ -74,15 +160,27 @@ def wrapper(*args, **kwargs): return wrapper -def dist_init(backend: str = 'nccl', - addr: str = None, - port: str = None, - rank: int = None, - world_size: int = None) -> Tuple[int, int]: - r""" +def dist_init( + backend: str = 'nccl', + addr: str = None, + port: str = None, + rank: int = None, + world_size: int = None, + timeout: datetime.timedelta = datetime.timedelta(seconds=60000) +) -> Tuple[int, int]: + """ Overview: - Init the distributed training setting + Initialize the distributed training setting. + Arguments: + - backend (:obj:`str`): The backend of the distributed training, supports ``['nccl', 'gloo']``. + - addr (:obj:`str`): The address of the master node. + - port (:obj:`str`): The port of the master node. + - rank (:obj:`int`): The rank of the current process. + - world_size (:obj:`int`): The total number of processes. + - timeout (:obj:`datetime.timedelta`): The timeout for operations executed against the process group. \ + Default is 60000 seconds. """ + assert backend in ['nccl', 'gloo'], backend os.environ['MASTER_ADDR'] = addr or os.environ.get('MASTER_ADDR', "localhost") os.environ['MASTER_PORT'] = port or os.environ.get('MASTER_PORT', "10314") # hard-code @@ -100,7 +198,7 @@ def dist_init(backend: str = 'nccl', else: world_size = int(ntasks) - dist.init_process_group(backend=backend, rank=rank, world_size=world_size) + dist.init_process_group(backend=backend, rank=rank, world_size=world_size, timeout=timeout) num_gpus = torch.cuda.device_count() torch.cuda.set_device(rank % num_gpus) @@ -110,29 +208,56 @@ def dist_init(backend: str = 'nccl', def dist_finalize() -> None: - r""" + """ Overview: Finalize distributed training resources """ - dist.destroy_process_group() + # This operation usually hangs out so we ignore it temporally. + # dist.destroy_process_group() + pass -class DistContext: +class DDPContext: + """ + Overview: + A context manager for ``linklink`` distribution + Interfaces: + ``__init__``, ``__enter__``, ``__exit__`` + """ def __init__(self) -> None: + """ + Overview: + Initialize the ``DDPContext`` + """ + pass def __enter__(self) -> None: + """ + Overview: + Initialize ``linklink`` distribution + """ + dist_init() def __exit__(self, *args, **kwargs) -> Any: + """ + Overview: + Finalize ``linklink`` distribution + """ + dist_finalize() def simple_group_split(world_size: int, rank: int, num_groups: int) -> List: - r""" + """ Overview: Split the group according to ``worldsize``, ``rank`` and ``num_groups`` + Arguments: + - world_size (:obj:`int`): The world size + - rank (:obj:`int`): The rank + - num_groups (:obj:`int`): The number of groups .. note:: With faulty input, raise ``array split does not result in an equal division`` @@ -144,3 +269,23 @@ def simple_group_split(world_size: int, rank: int, num_groups: int) -> List: groups.append(dist.new_group(rank_list[i])) group_size = world_size // num_groups return groups[rank // group_size] + + +def to_ddp_config(cfg: EasyDict) -> EasyDict: + """ + Overview: + Convert the config to ddp config + Arguments: + - cfg (:obj:`EasyDict`): The config to be converted + """ + + w = get_world_size() + if 'batch_size' in cfg.policy: + cfg.policy.batch_size = int(np.ceil(cfg.policy.batch_size / w)) + if 'batch_size' in cfg.policy.learn: + cfg.policy.learn.batch_size = int(np.ceil(cfg.policy.learn.batch_size / w)) + if 'n_sample' in cfg.policy.collect: + cfg.policy.collect.n_sample = int(np.ceil(cfg.policy.collect.n_sample / w)) + if 'n_episode' in cfg.policy.collect: + cfg.policy.collect.n_episode = int(np.ceil(cfg.policy.collect.n_episode / w)) + return cfg diff --git a/ding/utils/registry.py b/ding/utils/registry.py index e0c98d29f7..44810391ed 100644 --- a/ding/utils/registry.py +++ b/ding/utils/registry.py @@ -10,28 +10,40 @@ class Registry(dict): """ - A helper class for managing registering modules, it extends a dictionary - and provides a register functions. - - Eg. creeting a registry: - some_registry = Registry({"default": default_module}) - - There're two ways of registering new modules: - 1): normal way is just calling register function: - def foo(): - ... - some_registry.register("foo_module", foo) - 2): used as decorator when declaring the module: - @some_registry.register("foo_module") - @some_registry.register("foo_modeul_nickname") - def foo(): - ... - - Access of module is just like using a dictionary, eg: - f = some_registry["foo_module"] + Overview: + A helper class for managing registering modules, it extends a dictionary + and provides a register functions. + Interfaces: + ``__init__``, ``register``, ``get``, ``build``, ``query``, ``query_details`` + Examples (creating): + >>> some_registry = Registry({"default": default_module}) + + Examples (registering: normal way): + >>> def foo(): + >>> ... + >>> some_registry.register("foo_module", foo) + + Examples (registering: decorator way): + >>> @some_registry.register("foo_module") + >>> @some_registry.register("foo_modeul_nickname") + >>> def foo(): + >>> ... + + Examples (accessing): + >>> f = some_registry["foo_module"] """ def __init__(self, *args, **kwargs) -> None: + """ + Overview: + Initialize the Registry object. + Arguments: + - args (:obj:`Tuple`): The arguments passed to the ``__init__`` function of the parent class, \ + dict. + - kwargs (:obj:`Dict`): The keyword arguments passed to the ``__init__`` function of the parent class, \ + dict. + """ + super(Registry, self).__init__(*args, **kwargs) self.__trace__ = dict() @@ -41,6 +53,15 @@ def register( module: Optional[Callable] = None, force_overwrite: bool = False ) -> Callable: + """ + Overview: + Register the module. + Arguments: + - module_name (:obj:`Optional[str]`): The name of the module. + - module (:obj:`Optional[Callable]`): The module to be registered. + - force_overwrite (:obj:`bool`): Whether to overwrite the module with the same name. + """ + if _DI_ENGINE_REG_TRACE_IS_ON: frame = inspect.stack()[1][0] info = inspect.getframeinfo(frame) @@ -69,14 +90,40 @@ def register_fn(fn: Callable) -> Callable: @staticmethod def _register_generic(module_dict: dict, module_name: str, module: Callable, force_overwrite: bool = False) -> None: + """ + Overview: + Register the module. + Arguments: + - module_dict (:obj:`dict`): The dict to store the module. + - module_name (:obj:`str`): The name of the module. + - module (:obj:`Callable`): The module to be registered. + - force_overwrite (:obj:`bool`): Whether to overwrite the module with the same name. + """ + if not force_overwrite: assert module_name not in module_dict, module_name module_dict[module_name] = module def get(self, module_name: str) -> Callable: + """ + Overview: + Get the module. + Arguments: + - module_name (:obj:`str`): The name of the module. + """ + return self[module_name] def build(self, obj_type: str, *obj_args, **obj_kwargs) -> object: + """ + Overview: + Build the object. + Arguments: + - obj_type (:obj:`str`): The type of the object. + - obj_args (:obj:`Tuple`): The arguments passed to the object. + - obj_kwargs (:obj:`Dict`): The keyword arguments passed to the object. + """ + try: build_fn = self[obj_type] return build_fn(*obj_args, **obj_kwargs) @@ -88,17 +135,30 @@ def build(self, obj_type: str, *obj_args, **obj_kwargs) -> object: global _innest_error if _innest_error: argspec = inspect.getfullargspec(build_fn) - message = 'for {}(alias={})'.format(build_fn, obj_type) - message += '\nExpected args are:{}'.format(argspec) - message += '\nGiven args are:{}/{}'.format(argspec, obj_kwargs.keys()) - message += '\nGiven args details are:{}/{}'.format(argspec, obj_kwargs) + message = 'Hint: for {}(alias={})'.format(build_fn, obj_type) + message += '\n\nExpected args are:\n {}\nGiven arguments keys are:\n{}\n'.format( + argspec, obj_kwargs.keys() + ) + print(message) _innest_error = False raise e def query(self) -> Iterable: + """ + Overview: + all registered module names. + """ + return self.keys() def query_details(self, aliases: Optional[Iterable] = None) -> OrderedDict: + """ + Overview: + Get the details of the registered modules. + Arguments: + - aliases (:obj:`Optional[Iterable]`): The aliases of the modules. + """ + assert _DI_ENGINE_REG_TRACE_IS_ON, "please exec 'export DIENGINEREGTRACE=ON' first" if aliases is None: aliases = self.keys() diff --git a/ding/utils/render_helper.py b/ding/utils/render_helper.py index 5535b290f5..11aed75941 100644 --- a/ding/utils/render_helper.py +++ b/ding/utils/render_helper.py @@ -5,8 +5,26 @@ from ding.envs import BaseEnv, BaseEnvManager +def render_env(env, render_mode: Optional[str] = 'rgb_array') -> "ndarray": + """ + Overview: + Render the environment's current frame. + Arguments: + - env (:obj:`gym.Env`): DI-engine env instance. + - render_mode (:obj:`str`): Render mode. + Returns: + - frame (:obj:`numpy.ndarray`): [H * W * C] + """ + if hasattr(env, 'sim'): + # mujoco: mujoco frame is unside-down by default + return env.sim.render(camera_name='track', height=128, width=128)[::-1] + else: + # other envs + return env.render(mode=render_mode) + + def render(env: "BaseEnv", render_mode: Optional[str] = 'rgb_array') -> "ndarray": - ''' + """ Overview: Render the environment's current frame. Arguments: @@ -14,37 +32,45 @@ def render(env: "BaseEnv", render_mode: Optional[str] = 'rgb_array') -> "ndarray - render_mode (:obj:`str`): Render mode. Returns: - frame (:obj:`numpy.ndarray`): [H * W * C] - ''' + """ gym_env = env._env - if hasattr(gym_env, 'sim'): - # mujoco: mujoco frame is unside-down by default - return gym_env.sim.render(camera_name='track', height=128, width=128)[::-1] + return render_env(gym_env, render_mode=render_mode) + + +def get_env_fps(env) -> "int": + """ + Overview: + Get the environment's fps. + Arguments: + - env (:obj:`gym.Env`): DI-engine env instance. + Returns: + - fps (:obj:`int`). + """ + + if hasattr(env, 'model'): + # mujoco + fps = 1 / env.model.opt.timestep + elif hasattr(env, 'env') and 'video.frames_per_second' in env.env.metadata.keys(): + # classic control + fps = env.env.metadata['video.frames_per_second'] else: - # other envs - return gym_env.render(mode=render_mode) + # atari and other envs + fps = 30 + return fps def fps(env_manager: "BaseEnvManager") -> "int": - ''' + """ Overview: Render the environment's fps. Arguments: - env (:obj:`BaseEnvManager`): DI-engine env manager instance. Returns: - fps (:obj:`int`). - ''' + """ try: # env_ref is a ding gym environment gym_env = env_manager.env_ref._env - if hasattr(gym_env, 'model'): - # mujoco - fps = 1 / gym_env.model.opt.timestep - elif hasattr(gym_env, 'env') and 'video.frames_per_second' in gym_env.env.metadata.keys(): - # classic control - fps = gym_env.env.metadata['video.frames_per_second'] - else: - # atari and other envs - fps = 30 - return fps + return get_env_fps(gym_env) except: return 30 diff --git a/ding/utils/scheduler_helper.py b/ding/utils/scheduler_helper.py index 9b1c2600a4..d37ce97c52 100644 --- a/ding/utils/scheduler_helper.py +++ b/ding/utils/scheduler_helper.py @@ -9,7 +9,7 @@ class Scheduler(object): For example, models often benefits from reducing entropy weight once the learning process stagnates. This scheduler reads a metrics quantity and if no improvement is seen for a 'patience' number of epochs, the corresponding parameter is increased or decreased, which decides on the 'schedule_mode'. - Args: + Arguments: - schedule_flag (:obj:`bool`): Indicates whether to use scheduler in training pipeline. Default: False - schedule_mode (:obj:`str`): One of 'reduce', 'add','multi','div'. The schecule_mode @@ -48,13 +48,13 @@ class Scheduler(object): ) def __init__(self, merged_scheduler_config: EasyDict) -> None: - ''' + """ Overview: Initialize the scheduler. - Args: + Arguments: - merged_scheduler_config (:obj:`EasyDict`): the scheduler config, which merges the user config and defaul config - ''' + """ schedule_mode = merged_scheduler_config.schedule_mode factor = merged_scheduler_config.factor @@ -100,7 +100,7 @@ def __init__(self, merged_scheduler_config: EasyDict) -> None: self.bad_epochs_num = 0 def step(self, metrics: float, param: float) -> float: - ''' + """ Overview: Decides whether to update the scheduled parameter Args: @@ -108,7 +108,7 @@ def step(self, metrics: float, param: float) -> float: - param (:obj:`float`): parameter need to be updated Returns: - step_param (:obj:`float`): parameter after one step - ''' + """ assert isinstance(metrics, float), 'The metrics should be converted to a float number' cur_metrics = metrics @@ -129,14 +129,14 @@ def step(self, metrics: float, param: float) -> float: return param def update_param(self, param: float) -> float: - ''' + """ Overview: update the scheduling parameter Args: - param (:obj:`float`): parameter need to be updated Returns: - updated param (:obj:`float`): parameter after updating - ''' + """ schedule_fn = { 'reduce': lambda x, y, z: max(x - y, z[0]), 'add': lambda x, y, z: min(x + y, z[1]), @@ -153,20 +153,20 @@ def update_param(self, param: float) -> float: @property def in_cooldown(self) -> bool: - ''' + """ Overview: Checks whether the scheduler is in cooldown peried. If in cooldown, the scheduler will ignore any bad epochs. - ''' + """ return self.cooldown_counter > 0 def is_better(self, cur: float) -> bool: - ''' + """ Overview: Checks whether the current metrics is better than last matric with respect to threshold. Args: - cur (:obj:`float`): current metrics - ''' + """ if self.last_metrics is None: return True diff --git a/ding/utils/segment_tree.py b/ding/utils/segment_tree.py index b92dbed742..5c87280ab4 100644 --- a/ding/utils/segment_tree.py +++ b/ding/utils/segment_tree.py @@ -9,6 +9,11 @@ @lru_cache() def njit(): + """ + Overview: + Decorator to compile a function using numba. + """ + try: if ding.enable_numba: import numba @@ -34,7 +39,7 @@ class SegmentTree: Overview: Segment tree data structure, implemented by the tree-like array. Only the leaf nodes are real value, non-leaf nodes are to do some operations on its left and right child. - Interface: + Interfaces: ``__init__``, ``reduce``, ``__setitem__``, ``__getitem__`` """ @@ -111,6 +116,11 @@ def __getitem__(self, idx: int) -> float: return self.value[idx + self.capacity] def _compile(self) -> None: + """ + Overview: + Compile the functions using numba. + """ + f64 = np.array([0, 1], dtype=np.float64) f32 = np.array([0, 1], dtype=np.float32) i64 = np.array([0, 1], dtype=np.int64) @@ -121,11 +131,19 @@ def _compile(self) -> None: class SumSegmentTree(SegmentTree): + """ + Overview: + Sum segment tree, which is inherited from ``SegmentTree``. Init by passing ``operation='sum'``. + Interfaces: + ``__init__``, ``find_prefixsum_idx`` + """ def __init__(self, capacity: int) -> None: """ Overview: Init sum segment tree by passing ``operation='sum'`` + Arguments: + - capacity (:obj:`int`): Capacity of the tree (the number of the leaf nodes). """ super(SumSegmentTree, self).__init__(capacity, operation='sum') @@ -148,17 +166,35 @@ def find_prefixsum_idx(self, prefixsum: float, trust_caller: bool = True) -> int class MinSegmentTree(SegmentTree): + """ + Overview: + Min segment tree, which is inherited from ``SegmentTree``. Init by passing ``operation='min'``. + Interfaces: + ``__init__`` + """ def __init__(self, capacity: int) -> None: """ Overview: - Init sum segment tree by passing ``operation='min'`` + Initialize sum segment tree by passing ``operation='min'`` + Arguments: + - capacity (:obj:`int`): Capacity of the tree (the number of the leaf nodes). """ super(MinSegmentTree, self).__init__(capacity, operation='min') @njit() def _setitem(tree: np.ndarray, idx: int, val: float, operation: str) -> None: + """ + Overview: + Set ``tree[idx] = val``; Then update the related nodes. + Arguments: + - tree (:obj:`np.ndarray`): The tree array. + - idx (:obj:`int`): The index of the leaf node. + - val (:obj:`float`): The value that will be assigned to ``leaf[idx]``. + - operation (:obj:`str`): The operation function to construct the tree, e.g. sum, max, min, etc. + """ + tree[idx] = val # Update from specified node to the root node while idx > 1: @@ -172,6 +208,18 @@ def _setitem(tree: np.ndarray, idx: int, val: float, operation: str) -> None: @njit() def _reduce(tree: np.ndarray, start: int, end: int, neutral_element: float, operation: str) -> float: + """ + Overview: + Reduce the tree in range ``[start, end)`` + Arguments: + - tree (:obj:`np.ndarray`): The tree array. + - start (:obj:`int`): Start index(relative index, the first leaf node is 0). + - end (:obj:`int`): End index(relative index). + - neutral_element (:obj:`float`): The value of the neutral element, which is used to init \ + all nodes value in the tree. + - operation (:obj:`str`): The operation function to construct the tree, e.g. sum, max, min, etc. + """ + # Nodes in 【start, end) will be aggregated result = neutral_element while start < end: @@ -197,6 +245,18 @@ def _reduce(tree: np.ndarray, start: int, end: int, neutral_element: float, oper @njit() def _find_prefixsum_idx(tree: np.ndarray, capacity: int, prefixsum: float, neutral_element: float) -> int: + """ + Overview: + Find the highest non-zero index i, sum_{j}leaf[j] <= ``prefixsum`` (where 0 <= j < i) + and sum_{j}leaf[j] > ``prefixsum`` (where 0 <= j < i+1) + Arguments: + - tree (:obj:`np.ndarray`): The tree array. + - capacity (:obj:`int`): Capacity of the tree (the number of the leaf nodes). + - prefixsum (:obj:`float`): The target prefixsum. + - neutral_element (:obj:`float`): The value of the neutral element, which is used to init \ + all nodes value in the tree. + """ + # The function is to find a non-leaf node's index which satisfies: # self.value[idx] > input prefixsum and self.value[idx + 1] <= input prefixsum # In other words, we can assume that there are intervals: [num_0, num_1), [num_1, num_2), ... [num_k, num_k+1), diff --git a/ding/utils/slurm_helper.py b/ding/utils/slurm_helper.py index 4a1c35f4d2..c03b3e9463 100644 --- a/ding/utils/slurm_helper.py +++ b/ding/utils/slurm_helper.py @@ -14,6 +14,11 @@ def get_ip() -> str: + """ + Overview: + Get the ip of the current node + """ + assert os.environ.get('SLURMD_NODENAME'), 'not found SLURMD_NODENAME env variable' # expecting nodename to be like: 'SH-IDC1-10-5-36-64' nodename = os.environ.get('SLURMD_NODENAME', '') @@ -22,9 +27,11 @@ def get_ip() -> str: def get_manager_node_ip(node_ip: Optional[str] = None) -> str: - r""" + """ Overview: Look up the manager node of the slurm cluster and return the node ip + Arguments: + - node_ip (:obj:`Optional[str]`): The ip of the current node """ if 'SLURM_JOB_ID' not in os.environ: from ditk import logging @@ -44,6 +51,11 @@ def get_manager_node_ip(node_ip: Optional[str] = None) -> str: # get all info of cluster def get_cls_info() -> Dict[str, list]: + """ + Overview: + Get the cluster info + """ + ret_dict = {} info = subprocess.getoutput('sinfo -Nh').split('\n') for line in info: @@ -61,6 +73,13 @@ def get_cls_info() -> Dict[str, list]: def node_to_partition(target_node: str) -> Tuple[str, str]: + """ + Overview: + Get the partition of the target node + Arguments: + - target_node (:obj:`str`): The target node + """ + info = subprocess.getoutput('sinfo -Nh').split('\n') for line in info: line = line.strip().split() @@ -73,10 +92,24 @@ def node_to_partition(target_node: str) -> Tuple[str, str]: def node_to_host(node: str) -> str: + """ + Overview: + Get the host of the node + Arguments: + - node (:obj:`str`): The node + """ + return '.'.join(node.split('-')[-4:]) def find_free_port_slurm(node: str) -> int: + """ + Overview: + Find a free port on the node + Arguments: + - node (:obj:`str`): The node + """ + partition = node_to_partition(node) if partition == 'spring_scheduler': comment = '--comment=spring-submit' diff --git a/ding/utils/system_helper.py b/ding/utils/system_helper.py index 118513d36b..915ef380e9 100644 --- a/ding/utils/system_helper.py +++ b/ding/utils/system_helper.py @@ -8,7 +8,7 @@ def get_ip() -> str: - r""" + """ Overview: Get the ``ip(host)`` of socket Returns: @@ -22,7 +22,7 @@ def get_ip() -> str: def get_pid() -> int: - r""" + """ Overview: ``os.getpid`` """ @@ -30,7 +30,7 @@ def get_pid() -> int: def get_task_uid() -> str: - r""" + """ Overview: Get the slurm ``job_id``, ``pid`` and ``uid`` """ @@ -41,7 +41,7 @@ class PropagatingThread(Thread): """ Overview: Subclass of Thread that propagates execution exception in the thread to the caller - Interface: + Interfaces: ``run``, ``join`` Examples: >>> def func(): @@ -52,6 +52,11 @@ class PropagatingThread(Thread): """ def run(self) -> None: + """ + Overview: + Run the thread + """ + self.exc = None try: self.ret = self._target(*self._args, **self._kwargs) @@ -59,6 +64,11 @@ def run(self) -> None: self.exc = e def join(self) -> Any: + """ + Overview: + Join the thread + """ + super(PropagatingThread, self).join() if self.exc: raise RuntimeError('Exception in thread({})'.format(id(self))) from self.exc @@ -66,9 +76,11 @@ def join(self) -> Any: def find_free_port(host: str) -> int: - r""" + """ Overview: Look up the free port list and return one + Arguments: + - host (:obj:`str`): The host """ with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s: s.bind(('', 0)) diff --git a/ding/utils/tests/test_bfs_helper.py b/ding/utils/tests/test_bfs_helper.py new file mode 100644 index 0000000000..7f095907a0 --- /dev/null +++ b/ding/utils/tests/test_bfs_helper.py @@ -0,0 +1,26 @@ +import easydict +import numpy +import pytest + +from ding.utils import get_vi_sequence +from dizoo.maze.envs.maze_env import Maze + + +@pytest.mark.unittest +class TestBFSHelper: + + def test_bfs(self): + + def load_env(seed): + ccc = easydict.EasyDict({'size': 16}) + e = Maze(ccc) + e.seed(seed) + e.reset() + return e + + env = load_env(314) + start_obs = env.process_states(env._get_obs(), env.get_maze_map()) + vi_sequence, track_back = get_vi_sequence(env, start_obs) + assert vi_sequence.shape[1:] == (16, 16) + assert track_back[0][0].shape == (16, 16, 3) + assert isinstance(track_back[0][1], numpy.int32) diff --git a/ding/utils/tests/test_default_helper.py b/ding/utils/tests/test_default_helper.py index 2540436d36..c48b1d05ae 100644 --- a/ding/utils/tests/test_default_helper.py +++ b/ding/utils/tests/test_default_helper.py @@ -3,15 +3,42 @@ import numpy as np import pytest import torch +import treetensor.torch as ttorch from ding.utils.default_helper import lists_to_dicts, dicts_to_lists, squeeze, default_get, override, error_wrapper, \ list_split, LimitedSpaceContainer, set_pkg_seed, deep_merge_dicts, deep_update, flatten_dict, RunningMeanStd, \ - one_time_warning, split_data_generator + one_time_warning, split_data_generator, get_shape0 @pytest.mark.unittest class TestDefaultHelper(): + def test_get_shape0(self): + a = { + 'a': { + 'b': torch.randn(4, 3) + }, + 'c': { + 'd': torch.randn(4) + }, + } + b = [a, a] + c = (a, a) + d = { + 'a': { + 'b': ["a", "b", "c", "d"] + }, + 'c': { + 'd': torch.randn(4) + }, + } + a = ttorch.as_tensor(a) + assert get_shape0(a) == 4 + assert get_shape0(b) == 4 + assert get_shape0(c) == 4 + with pytest.raises(Exception) as e_info: + assert get_shape0(d) == 4 + def test_lists_to_dicts(self): set_pkg_seed(12) with pytest.raises(ValueError): diff --git a/ding/utils/tests/test_deprecation.py b/ding/utils/tests/test_deprecation.py new file mode 100644 index 0000000000..c5304825a2 --- /dev/null +++ b/ding/utils/tests/test_deprecation.py @@ -0,0 +1,28 @@ +import pytest +import warnings +from ding.utils.deprecation import deprecated + + +@pytest.mark.unittest +def test_deprecated(): + + @deprecated('0.4.1', '0.5.1') + def deprecated_func1(): + pass + + @deprecated('0.4.1', '0.5.1', 'deprecated_func3') + def deprecated_func2(): + pass + + with warnings.catch_warnings(record=True) as w: + deprecated_func1() + assert ( + 'API `test_deprecation.deprecated_func1` is deprecated ' + 'since version 0.4.1 and will be removed in version 0.5.1.' + ) == str(w[-1].message) + deprecated_func2() + assert ( + 'API `test_deprecation.deprecated_func2` is deprecated ' + 'since version 0.4.1 and will be removed in version 0.5.1, ' + 'please use `deprecated_func3` instead.' + ) == str(w[-1].message) diff --git a/ding/utils/tests/test_k8s_launcher.py b/ding/utils/tests/test_k8s_launcher.py index 30478e6b25..6145c17d1b 100644 --- a/ding/utils/tests/test_k8s_launcher.py +++ b/ding/utils/tests/test_k8s_launcher.py @@ -9,7 +9,7 @@ except ImportError: _test_mark = pytest.mark.ignore else: - _test_mark = pytest.mark.unittest + _test_mark = pytest.mark.envtest @_test_mark diff --git a/ding/utils/tests/test_log_helper.py b/ding/utils/tests/test_log_helper.py index 439015ac0c..21b78015f3 100644 --- a/ding/utils/tests/test_log_helper.py +++ b/ding/utils/tests/test_log_helper.py @@ -1,6 +1,7 @@ import random import pytest from easydict import EasyDict +from ditk import logging from ding.utils.log_helper import build_logger, pretty_print from ding.utils.file_helper import remove_file @@ -38,16 +39,16 @@ def test_pretty_print(self): pretty_print(cfg) def test_logger(self): - logger, tb_logger = build_logger(cfg.common.save_path, name="fake_test", need_tb=True) - vars = {'aa': 3.0, 'bb': 4, 'cc': 3e4} + logger, tb_logger = build_logger(cfg.common.save_path, name="fake_test", need_tb=True, text_level=logging.DEBUG) + variables = {'aa': 3.0, 'bb': 4, 'cc': 3e4} # text logger logger.info("I'm an info") logger.debug("I'm a bug") logger.error("I'm an error") - logger.info(logger.get_tabulate_vars(vars)) + logger.info(logger.get_tabulate_vars(variables)) # tensorboard logger for i in range(10): - new_vars = {k: v * (i + random.random()) for k, v in vars.items()} + new_vars = {k: v * (i + random.random()) for k, v in variables.items()} for k, v in new_vars.items(): tb_logger.add_scalar(k, v, i) remove_file(cfg.common.save_path) diff --git a/ding/utils/tests/test_memory_helper.py b/ding/utils/tests/test_memory_helper.py new file mode 100644 index 0000000000..010eaee40c --- /dev/null +++ b/ding/utils/tests/test_memory_helper.py @@ -0,0 +1,39 @@ +import pytest +import os +import time +import torch +from ding.utils import SimpleMemoryProfiler +from ding.utils.memory_helper import multi_chunk_test +from ding.model import DRQN + + +@pytest.mark.unittest +def test_memory_profiler(): + multi_chunk_test() + assert os.path.exists('test_simple_memory_profiler_multi_chunk') + assert os.path.exists('test_simple_memory_profiler_multi_chunk/memory.log') + assert os.path.exists('test_simple_memory_profiler_multi_chunk/summary_sunburst.html') + + model = DRQN( + obs_shape=(4, 84, 84), + action_shape=6, + encoder_hidden_size_list=[128, 256, 512], + head_layer_num=3, + norm_type='LN', + ) + # print(model) + optimizer = torch.optim.Adam(model.parameters(), lr=0.001) + profiler = SimpleMemoryProfiler(model, optimizer, 'test_simple_memory_profiler_drqn', total_steps=1) + + x = torch.randn(3, 8, 4, 84, 84) + inputs = {'obs': x, 'prev_state': None} + y = model(inputs)['logit'] + + optimizer.zero_grad() + y.mean().backward() + optimizer.step() + + profiler.step() + + time.sleep(0.3) + os.popen('rm -rf test_simple_memory_profiler*') diff --git a/ding/utils/tests/test_normalizer_helper.py b/ding/utils/tests/test_normalizer_helper.py new file mode 100755 index 0000000000..d3339a00b4 --- /dev/null +++ b/ding/utils/tests/test_normalizer_helper.py @@ -0,0 +1,38 @@ +import easydict +import numpy +import pytest + +from ding.utils.normalizer_helper import DatasetNormalizer + + +# TODO(nyz): fix unittest bugs +@pytest.mark.tmp +class TestNormalizerHelper: + + def test_normalizer(self): + x = numpy.random.randn(10) + mean = x.mean() + std = x.std() + mins = x.min() + maxs = x.max() + normalizer = DatasetNormalizer({'test': x}, 'GaussianNormalizer', 10) + test = numpy.random.randn(1) + normal_test = normalizer.normalize(test, 'test') + unnormal_test = normalizer.unnormalize(normal_test, 'test') + assert unnormal_test == test + assert normal_test == (test - mean) / std + + normalizer = DatasetNormalizer({'test': x}, 'LimitsNormalizer', 10) + test = numpy.random.randn(1) + normal_test1 = (test - mins) / (maxs - mins) + normal_test1 = 2 * normal_test1 - 1 + normal_test = normalizer.normalize(test, 'test') + unnormal_test = normalizer.unnormalize(normal_test, 'test') + assert unnormal_test == test + assert normal_test == normal_test1 + + normalizer = DatasetNormalizer({'test': x}, 'CDFNormalizer', 10) + test = numpy.random.randn(1) + normal_test = normalizer.normalize(test, 'test') + unnormal_test = normalizer.unnormalize(normal_test, 'test') + assert unnormal_test == test diff --git a/ding/utils/time_helper.py b/ding/utils/time_helper.py index f723939aa5..06498e4546 100644 --- a/ding/utils/time_helper.py +++ b/ding/utils/time_helper.py @@ -9,7 +9,7 @@ def build_time_helper(cfg: EasyDict = None, wrapper_type: str = None) -> Callable[[], 'TimeWrapper']: - r""" + """ Overview: Build the timehelper @@ -46,11 +46,11 @@ def build_time_helper(cfg: EasyDict = None, wrapper_type: str = None) -> Callabl class EasyTimer: - r""" + """ Overview: A decent timer wrapper that can be used easily. - Interface: + Interfaces: ``__init__``, ``__enter__``, ``__exit__`` Example: @@ -61,7 +61,7 @@ class EasyTimer: """ def __init__(self, cuda=True): - r""" + """ Overview: Init class EasyTimer @@ -76,7 +76,7 @@ def __init__(self, cuda=True): self.value = 0.0 def __enter__(self): - r""" + """ Overview: Enter timer, start timing """ @@ -84,7 +84,7 @@ def __enter__(self): self._timer.start_time() def __exit__(self, *args): - r""" + """ Overview: Exit timer, stop timing """ @@ -92,29 +92,29 @@ def __exit__(self, *args): class TimeWrapperTime(TimeWrapper): - r""" + """ Overview: A class method that inherit from ``TimeWrapper`` class - Interface: + Interfaces: ``start_time``, ``end_time`` """ # overwrite @classmethod def start_time(cls): - r""" + """ Overview: - Implement and overide the ``start_time`` method in ``TimeWrapper`` class + Implement and override the ``start_time`` method in ``TimeWrapper`` class """ cls.start = time.time() # overwrite @classmethod def end_time(cls): - r""" + """ Overview: - Implement and overide the end_time method in ``TimeWrapper`` class + Implement and override the end_time method in ``TimeWrapper`` class Returns: - time(:obj:`float`): The time between ``start_time`` and end_time @@ -134,7 +134,7 @@ class WatchDog(object): .. note:: If it is not reset before exceeding this value, ``TimeourError`` raised. - Interface: + Interfaces: ``start``, ``stop`` Examples: @@ -146,11 +146,18 @@ class WatchDog(object): """ def __init__(self, timeout: int = 1): + """ + Overview: + Initialize watchdog with ``timeout`` value. + Arguments: + - timeout (:obj:`int`): Timeout value of the ``watchdog [seconds]``. + """ + self._timeout = timeout + 1 self._failed = False def start(self): - r""" + """ Overview: Start watchdog. """ @@ -159,10 +166,18 @@ def start(self): @staticmethod def _event(signum: Any, frame: Any): + """ + Overview: + Event handler for watchdog. + Arguments: + - signum (:obj:`Any`): Signal number. + - frame (:obj:`Any`): Current stack frame. + """ + raise TimeoutError() def stop(self): - r""" + """ Overview: Stop watchdog with ``alarm(0)``, ``SIGALRM``, and ``SIG_DFL`` signals. """ diff --git a/ding/utils/time_helper_base.py b/ding/utils/time_helper_base.py index c77aa1c8d4..86f58d0fe8 100644 --- a/ding/utils/time_helper_base.py +++ b/ding/utils/time_helper_base.py @@ -1,19 +1,19 @@ class TimeWrapper(object): - r""" + """ Overview: Abstract class method that defines ``TimeWrapper`` class - Interface: + Interfaces: ``wrapper``, ``start_time``, ``end_time`` """ @classmethod def wrapper(cls, fn): - r""" + """ Overview: Classmethod wrapper, wrap a function and automatically return its running time - - - fn (:obj:`function`): The function to be wrap and timed + Arguments: + - fn (:obj:`function`): The function to be wrap and timed """ def time_func(*args, **kwargs): @@ -26,7 +26,7 @@ def time_func(*args, **kwargs): @classmethod def start_time(cls): - r""" + """ Overview: Abstract classmethod, start timing """ @@ -34,7 +34,7 @@ def start_time(cls): @classmethod def end_time(cls): - r""" + """ Overview: Abstract classmethod, stop timing """ diff --git a/ding/utils/time_helper_cuda.py b/ding/utils/time_helper_cuda.py index da691bb6b3..51ea5e925a 100644 --- a/ding/utils/time_helper_cuda.py +++ b/ding/utils/time_helper_cuda.py @@ -4,7 +4,7 @@ def get_cuda_time_wrapper() -> Callable[[], 'TimeWrapper']: - r""" + """ Overview: Return the ``TimeWrapperCuda`` class, this wrapper aims to ensure compatibility in no cuda device @@ -18,7 +18,7 @@ def get_cuda_time_wrapper() -> Callable[[], 'TimeWrapper']: # TODO find a way to autodoc the class within method class TimeWrapperCuda(TimeWrapper): - r""" + """ Overview: A class method that inherit from ``TimeWrapper`` class @@ -26,7 +26,7 @@ class TimeWrapperCuda(TimeWrapper): Must use torch.cuda.synchronize(), reference: \ - Interface: + Interfaces: ``start_time``, ``end_time`` """ # cls variable is initialized on loading this class @@ -36,7 +36,7 @@ class TimeWrapperCuda(TimeWrapper): # overwrite @classmethod def start_time(cls): - r""" + """ Overview: Implement and overide the ``start_time`` method in ``TimeWrapper`` class """ @@ -46,7 +46,7 @@ def start_time(cls): # overwrite @classmethod def end_time(cls): - r""" + """ Overview: Implement and overide the end_time method in ``TimeWrapper`` class Returns: diff --git a/ding/worker/collector/base_serial_collector.py b/ding/worker/collector/base_serial_collector.py index 75f4bbe8be..caa88410af 100644 --- a/ding/worker/collector/base_serial_collector.py +++ b/ding/worker/collector/base_serial_collector.py @@ -197,28 +197,33 @@ def append(self, data: Any) -> None: super().append(data) -def to_tensor_transitions(data: List[Dict[str, Any]]) -> List[Dict[str, Any]]: +def to_tensor_transitions(data: List[Dict[str, Any]], shallow_copy_next_obs: bool = True) -> List[Dict[str, Any]]: """ Overview: - transitions collected data to tensor. + Transform ths original transition return from env to tensor format. Argument: - - data (:obj:`List[Dict[str, Any]]`): the data that will be transited to tensor. + - data (:obj:`List[Dict[str, Any]]`): The data that will be transformed to tensor. + - shallow_copy_next_obs (:obj:`bool`): Whether to shallow copy next_obs. Default: True. Return: - - data (:obj:`List[Dict[str, Any]]`): the data that can be transited to tensor. + - data (:obj:`List[Dict[str, Any]]`): The transformed tensor-like data. .. tip:: In order to save memory, If there are next_obs in the passed data, we do special \ - treatment on next_obs so that the next_obs of each state in the data fragment is \ - the next state's obs and the next_obs of the last state is its own next_obs, \ - and we make transform_scalar is False. + treatment on next_obs so that the next_obs of each state in the data fragment is \ + the next state's obs and the next_obs of the last state is its own next_obsself. \ + Besides, we set transform_scalar to False to avoid the extra ``.item()`` operation. """ if 'next_obs' not in data[0]: return to_tensor(data, transform_scalar=False) else: - # for save memory of next_obs - data = to_tensor(data, transform_scalar=False) - data = to_tensor(data, ignore_keys=['next_obs'], transform_scalar=False) - for i in range(len(data) - 1): - data[i]['next_obs'] = data[i + 1]['obs'] - data[-1]['next_obs'] = to_tensor(data[-1]['next_obs'], transform_scalar=False) + # to_tensor will assign the separate memory to next_obs, if shallow_copy_next_obs is True, + # we can add ignore_keys to avoid this data copy for saving memory of next_obs. + if shallow_copy_next_obs: + data = to_tensor(data, ignore_keys=['next_obs'], transform_scalar=False) + for i in range(len(data) - 1): + data[i]['next_obs'] = data[i + 1]['obs'] + data[-1]['next_obs'] = to_tensor(data[-1]['next_obs'], transform_scalar=False) + return data + else: + data = to_tensor(data, transform_scalar=False) return data diff --git a/ding/worker/collector/base_serial_evaluator.py b/ding/worker/collector/base_serial_evaluator.py index 9c0517b30e..f5a31319ec 100644 --- a/ding/worker/collector/base_serial_evaluator.py +++ b/ding/worker/collector/base_serial_evaluator.py @@ -83,7 +83,7 @@ class VectorEvalMonitor(object): any thing, it is likely that we will get more short episodes than long episodes. As a result, \ our average reward will have a bias and may not be accurate. we use VectorEvalMonitor to solve the problem. Interfaces: - __init__, is_finished, update_info, update_reward, get_episode_reward, get_latest_reward, get_current_episode,\ + __init__, is_finished, update_info, update_reward, get_episode_return, get_latest_reward, get_current_episode,\ get_episode_info """ @@ -156,7 +156,7 @@ def get_video(self): """ videos = sum([list(v) for v in self._video.values()], []) videos = [np.transpose(np.stack(video, 0), [0, 3, 1, 2]) for video in videos] - sortarg = np.argsort(self.get_episode_reward()) + sortarg = np.argsort(self.get_episode_return()) # worst, median(s), best if len(sortarg) == 1: idxs = [sortarg[0]] @@ -179,10 +179,10 @@ def get_video(self): assert len(videos.shape) == 5, 'Need [N, T, C, H, W] input tensor for video logging!' return videos - def get_episode_reward(self) -> list: + def get_episode_return(self) -> list: """ Overview: - Get the total reward of one episode. + Sum up all reward and get the total return of one episode. """ return sum([list(v) for v in self._reward.values()], []) # sum(iterable, start) @@ -205,7 +205,7 @@ def get_current_episode(self) -> int: def get_episode_info(self) -> dict: """ Overview: - Get all episode information, such as total reward of one episode. + Get all episode information, such as total return of one episode. """ if len(self._info[0]) == 0: return None diff --git a/ding/worker/collector/battle_episode_serial_collector.py b/ding/worker/collector/battle_episode_serial_collector.py index df556f5503..156ab986d6 100644 --- a/ding/worker/collector/battle_episode_serial_collector.py +++ b/ding/worker/collector/battle_episode_serial_collector.py @@ -162,10 +162,20 @@ def envstep(self) -> int: Overview: Print the total envstep count. Return: - - envstep (:obj:`int`): the total envstep count + - envstep (:obj:`int`): The total envstep count. """ return self._total_envstep_count + @envstep.setter + def envstep(self, value: int) -> None: + """ + Overview: + Set the total envstep count. + Arguments: + - value (:obj:`int`): The total envstep count. + """ + self._total_envstep_count = value + def close(self) -> None: """ Overview: @@ -260,7 +270,9 @@ def collect(self, self._traj_buffer[env_id][policy_id].append(transition) # prepare data if timestep.done: - transitions = to_tensor_transitions(self._traj_buffer[env_id][policy_id]) + transitions = to_tensor_transitions( + self._traj_buffer[env_id][policy_id], not self._deepcopy_obs + ) if self._cfg.get_train_sample: train_sample = self._policy[policy_id].get_train_sample(transitions) return_data[policy_id].extend(train_sample) @@ -274,8 +286,8 @@ def collect(self, if timestep.done: self._total_episode_count += 1 info = { - 'reward0': timestep.info[0]['final_eval_reward'], - 'reward1': timestep.info[1]['final_eval_reward'], + 'reward0': timestep.info[0]['eval_episode_return'], + 'reward1': timestep.info[1]['eval_episode_return'], 'time': self._env_info[env_id]['time'], 'step': self._env_info[env_id]['step'], } @@ -306,8 +318,8 @@ def _output_log(self, train_iter: int) -> None: episode_count = len(self._episode_info) envstep_count = sum([d['step'] for d in self._episode_info]) duration = sum([d['time'] for d in self._episode_info]) - episode_reward0 = [d['reward0'] for d in self._episode_info] - episode_reward1 = [d['reward1'] for d in self._episode_info] + episode_return0 = [d['reward0'] for d in self._episode_info] + episode_return1 = [d['reward1'] for d in self._episode_info] self._total_duration += duration info = { 'episode_count': episode_count, @@ -316,14 +328,14 @@ def _output_log(self, train_iter: int) -> None: 'avg_envstep_per_sec': envstep_count / duration, 'avg_episode_per_sec': episode_count / duration, 'collect_time': duration, - 'reward0_mean': np.mean(episode_reward0), - 'reward0_std': np.std(episode_reward0), - 'reward0_max': np.max(episode_reward0), - 'reward0_min': np.min(episode_reward0), - 'reward1_mean': np.mean(episode_reward1), - 'reward1_std': np.std(episode_reward1), - 'reward1_max': np.max(episode_reward1), - 'reward1_min': np.min(episode_reward1), + 'reward0_mean': np.mean(episode_return0), + 'reward0_std': np.std(episode_return0), + 'reward0_max': np.max(episode_return0), + 'reward0_min': np.min(episode_return0), + 'reward1_mean': np.mean(episode_return1), + 'reward1_std': np.std(episode_return1), + 'reward1_max': np.max(episode_return1), + 'reward1_min': np.min(episode_return1), 'total_envstep_count': self._total_envstep_count, 'total_episode_count': self._total_episode_count, 'total_duration': self._total_duration, diff --git a/ding/worker/collector/battle_interaction_serial_evaluator.py b/ding/worker/collector/battle_interaction_serial_evaluator.py index f631359de6..9d4700ca57 100644 --- a/ding/worker/collector/battle_interaction_serial_evaluator.py +++ b/ding/worker/collector/battle_interaction_serial_evaluator.py @@ -9,7 +9,7 @@ from ding.utils import build_logger, EasyTimer, deep_merge_dicts, lists_to_dicts, dicts_to_lists, \ SERIAL_EVALUATOR_REGISTRY from ding.envs import BaseEnvManager -from ding.torch_utils import to_tensor, to_ndarray, tensor_to_list +from ding.torch_utils import to_tensor, to_ndarray, tensor_to_list, to_item from .base_serial_evaluator import ISerialEvaluator, VectorEvalMonitor @@ -132,7 +132,7 @@ def reset(self, _policy: Optional[List[namedtuple]] = None, _env: Optional[BaseE self.reset_env(_env) if _policy is not None: self.reset_policy(_policy) - self._max_eval_reward = float("-inf") + self._max_episode_return = float("-inf") self._last_eval_iter = 0 self._end_flag = False @@ -192,6 +192,7 @@ def eval( assert n_episode is not None, "please indicate eval n_episode" envstep_count = 0 info = {} + # TODO replace return_info with episode_info (validated by the league demo case) return_info = [[] for _ in range(self._policy_num)] eval_monitor = VectorEvalMonitor(self._env.env_num, n_episode) self._env.reset() @@ -219,7 +220,7 @@ def eval( for p in self._policy: p.reset([env_id]) # policy0 is regarded as main policy default - reward = t.info[0]['final_eval_reward'] + reward = t.info[0]['eval_episode_return'] if 'episode_info' in t.info[0]: eval_monitor.update_info(env_id, t.info[0]['episode_info']) eval_monitor.update_reward(env_id, reward) @@ -232,7 +233,7 @@ def eval( ) envstep_count += 1 duration = self._timer.value - episode_reward = eval_monitor.get_episode_reward() + episode_return = eval_monitor.get_episode_return() info = { 'train_iter': train_iter, 'ckpt_name': 'iteration_{}.pth.tar'.format(train_iter), @@ -242,11 +243,11 @@ def eval( 'evaluate_time': duration, 'avg_envstep_per_sec': envstep_count / duration, 'avg_time_per_episode': n_episode / duration, - 'reward_mean': np.mean(episode_reward), - 'reward_std': np.std(episode_reward), - 'reward_max': np.max(episode_reward), - 'reward_min': np.min(episode_reward), - # 'each_reward': episode_reward, + 'reward_mean': np.mean(episode_return), + 'reward_std': np.std(episode_return), + 'reward_max': np.max(episode_return), + 'reward_min': np.min(episode_return), + # 'each_reward': episode_return, } episode_info = eval_monitor.get_episode_info() if episode_info is not None: @@ -260,16 +261,17 @@ def eval( continue self._tb_logger.add_scalar('{}_iter/'.format(self._instance_name) + k, v, train_iter) self._tb_logger.add_scalar('{}_step/'.format(self._instance_name) + k, v, envstep) - eval_reward = np.mean(episode_reward) - if eval_reward > self._max_eval_reward: + episode_return = np.mean(episode_return) + if episode_return > self._max_episode_return: if save_ckpt_fn: save_ckpt_fn('ckpt_best.pth.tar') - self._max_eval_reward = eval_reward - stop_flag = eval_reward >= self._stop_value and train_iter > 0 + self._max_episode_return = episode_return + stop_flag = episode_return >= self._stop_value and train_iter > 0 if stop_flag: self._logger.info( "[DI-engine serial pipeline] " + - "Current eval_reward: {} is greater than stop_value: {}".format(eval_reward, self._stop_value) + + "Current episode_return: {} is greater than stop_value: {}".format(episode_return, self._stop_value) + ", so your RL agent is converged, you can refer to 'log/evaluator/evaluator_logger.txt' for details." ) + return_info = to_item(return_info) return stop_flag, return_info diff --git a/ding/worker/collector/battle_sample_serial_collector.py b/ding/worker/collector/battle_sample_serial_collector.py index 672b6164f5..33f4df11c9 100644 --- a/ding/worker/collector/battle_sample_serial_collector.py +++ b/ding/worker/collector/battle_sample_serial_collector.py @@ -175,10 +175,20 @@ def envstep(self) -> int: Overview: Print the total envstep count. Return: - - envstep (:obj:`int`): the total envstep count + - envstep (:obj:`int`): The total envstep count. """ return self._total_envstep_count + @envstep.setter + def envstep(self, value: int) -> None: + """ + Overview: + Set the total envstep count. + Arguments: + - value (:obj:`int`): The total envstep count. + """ + self._total_envstep_count = value + def close(self) -> None: """ Overview: @@ -276,7 +286,9 @@ def collect( self._traj_buffer[env_id][policy_id].append(transition) # prepare data if timestep.done or len(self._traj_buffer[env_id][policy_id]) == self._traj_len: - transitions = to_tensor_transitions(self._traj_buffer[env_id][policy_id]) + transitions = to_tensor_transitions( + self._traj_buffer[env_id][policy_id], not self._deepcopy_obs + ) train_sample = self._policy[policy_id].get_train_sample(transitions) return_data[policy_id].extend(train_sample) self._total_train_sample_count += len(train_sample) @@ -295,7 +307,7 @@ def collect( 'train_sample': self._env_info[env_id]['train_sample'], } for i in range(self._policy_num): - info['reward{}'.format(i)] = timestep.info[i]['final_eval_reward'] + info['reward{}'.format(i)] = timestep.info[i]['eval_episode_return'] self._episode_info.append(info) for i, p in enumerate(self._policy): p.reset([env_id]) @@ -322,10 +334,10 @@ def _output_log(self, train_iter: int) -> None: episode_count = len(self._episode_info) envstep_count = sum([d['step'] for d in self._episode_info]) duration = sum([d['time'] for d in self._episode_info]) - episode_reward = [] + episode_return = [] for i in range(self._policy_num): - episode_reward_item = [d['reward{}'.format(i)] for d in self._episode_info] - episode_reward.append(episode_reward_item) + episode_return_item = [d['reward{}'.format(i)] for d in self._episode_info] + episode_return.append(episode_return_item) self._total_duration += duration info = { 'episode_count': episode_count, @@ -341,7 +353,7 @@ def _output_log(self, train_iter: int) -> None: for k, fn in {'mean': np.mean, 'std': np.std, 'max': np.max, 'min': np.min}.items(): for i in range(self._policy_num): # such as reward0_mean - info['reward{}_{}'.format(i, k)] = fn(episode_reward[i]) + info['reward{}_{}'.format(i, k)] = fn(episode_return[i]) self._episode_info.clear() self._logger.info("collect end:\n{}".format('\n'.join(['{}: {}'.format(k, v) for k, v in info.items()]))) for k, v in info.items(): diff --git a/ding/worker/collector/episode_serial_collector.py b/ding/worker/collector/episode_serial_collector.py index 9fcf733779..5147f5a456 100644 --- a/ding/worker/collector/episode_serial_collector.py +++ b/ding/worker/collector/episode_serial_collector.py @@ -21,7 +21,9 @@ class EpisodeSerialCollector(ISerialCollector): envstep """ - config = dict(deepcopy_obs=False, transform_obs=False, collect_print_freq=100, get_train_sample=False) + config = dict( + deepcopy_obs=False, transform_obs=False, collect_print_freq=100, get_train_sample=False, reward_shaping=False + ) def __init__( self, @@ -91,9 +93,10 @@ def reset_policy(self, _policy: Optional[namedtuple] = None) -> None: assert hasattr(self, '_env'), "please set env first" if _policy is not None: self._policy = _policy - self._default_n_episode = _policy.get_attribute('cfg').collect.get('n_episode', None) + self._policy_cfg = self._policy.get_attribute('cfg') + self._default_n_episode = _policy.get_attribute('n_episode') self._unroll_len = _policy.get_attribute('unroll_len') - self._on_policy = _policy.get_attribute('cfg').on_policy + self._on_policy = _policy.get_attribute('on_policy') self._traj_len = INF self._logger.debug( 'Set default n_episode mode(n_episode({}), env_num({}), traj_len({}))'.format( @@ -154,10 +157,20 @@ def envstep(self) -> int: Overview: Print the total envstep count. Return: - - envstep (:obj:`int`): the total envstep count + - envstep (:obj:`int`): The total envstep count. """ return self._total_envstep_count + @envstep.setter + def envstep(self, value: int) -> None: + """ + Overview: + Set the total envstep count. + Arguments: + - value (:obj:`int`): The total envstep count. + """ + self._total_envstep_count = value + def close(self) -> None: """ Overview: @@ -250,7 +263,9 @@ def collect(self, self._total_envstep_count += 1 # prepare data if timestep.done: - transitions = to_tensor_transitions(self._traj_buffer[env_id]) + transitions = to_tensor_transitions(self._traj_buffer[env_id], not self._deepcopy_obs) + if self._cfg.reward_shaping: + self._env.reward_shaping(env_id, transitions) if self._cfg.get_train_sample: train_sample = self._policy.get_train_sample(transitions) return_data.extend(train_sample) @@ -263,7 +278,7 @@ def collect(self, # If env is done, record episode info and reset if timestep.done: self._total_episode_count += 1 - reward = timestep.info['final_eval_reward'] + reward = timestep.info['eval_episode_return'] info = { 'reward': reward, 'time': self._env_info[env_id]['time'], @@ -293,7 +308,7 @@ def _output_log(self, train_iter: int) -> None: episode_count = len(self._episode_info) envstep_count = sum([d['step'] for d in self._episode_info]) duration = sum([d['time'] for d in self._episode_info]) - episode_reward = [d['reward'] for d in self._episode_info] + episode_return = [d['reward'] for d in self._episode_info] self._total_duration += duration info = { 'episode_count': episode_count, @@ -302,14 +317,14 @@ def _output_log(self, train_iter: int) -> None: 'avg_envstep_per_sec': envstep_count / duration, 'avg_episode_per_sec': episode_count / duration, 'collect_time': duration, - 'reward_mean': np.mean(episode_reward), - 'reward_std': np.std(episode_reward), - 'reward_max': np.max(episode_reward), - 'reward_min': np.min(episode_reward), + 'reward_mean': np.mean(episode_return), + 'reward_std': np.std(episode_return), + 'reward_max': np.max(episode_return), + 'reward_min': np.min(episode_return), 'total_envstep_count': self._total_envstep_count, 'total_episode_count': self._total_episode_count, 'total_duration': self._total_duration, - # 'each_reward': episode_reward, + # 'each_reward': episode_return, } self._episode_info.clear() self._logger.info("collect end:\n{}".format('\n'.join(['{}: {}'.format(k, v) for k, v in info.items()]))) diff --git a/ding/worker/collector/interaction_serial_evaluator.py b/ding/worker/collector/interaction_serial_evaluator.py index 3e4645260a..ea92243f0a 100644 --- a/ding/worker/collector/interaction_serial_evaluator.py +++ b/ding/worker/collector/interaction_serial_evaluator.py @@ -1,13 +1,12 @@ -from typing import Optional, Callable, Tuple +from typing import Optional, Callable, Tuple, Dict, List from collections import namedtuple import numpy as np import torch -import torch.distributed as dist from ding.envs import BaseEnvManager -from ding.torch_utils import to_tensor, to_ndarray +from ding.torch_utils import to_tensor, to_ndarray, to_item from ding.utils import build_logger, EasyTimer, SERIAL_EVALUATOR_REGISTRY -from ding.utils import get_world_size, get_rank +from ding.utils import get_world_size, get_rank, broadcast_object_list from .base_serial_evaluator import ISerialEvaluator, VectorEvalMonitor @@ -23,13 +22,15 @@ class InteractionSerialEvaluator(ISerialEvaluator): """ config = dict( - # Evaluate every "eval_freq" training iterations. + # (int) Evaluate every "eval_freq" training iterations. eval_freq=1000, render=dict( - # tensorboard video render is disabled by default + # Tensorboard video render is disabled by default. render_freq=-1, mode='train_iter', - ) + ), + # (str) File path for visualize environment information. + figure_path=None, ) def __init__( @@ -43,7 +44,7 @@ def __init__( ) -> None: """ Overview: - Init method. Load config and use ``self._cfg`` setting to build common serial evaluator components, + Init method. Load config and use ``self._cfg`` setting to build common serial evaluator components, \ e.g. logger helper, timer. Arguments: - cfg (:obj:`EasyDict`): Configuration EasyDict. @@ -65,10 +66,7 @@ def __init__( './{}/log/{}'.format(self._exp_name, self._instance_name), self._instance_name ) else: - self._logger, _ = build_logger( - './{}/log/{}'.format(self._exp_name, self._instance_name), self._instance_name, need_tb=False - ) - self._tb_logger = None + self._logger, self._tb_logger = None, None # for close elegantly self.reset(policy, env) self._timer = EasyTimer() @@ -110,6 +108,7 @@ def reset_policy(self, _policy: Optional[namedtuple] = None) -> None: assert hasattr(self, '_env'), "please set env first" if _policy is not None: self._policy = _policy + self._policy_cfg = self._policy.get_attribute('cfg') self._policy.reset() def reset(self, _policy: Optional[namedtuple] = None, _env: Optional[BaseEnvManager] = None) -> None: @@ -130,7 +129,10 @@ def reset(self, _policy: Optional[namedtuple] = None, _env: Optional[BaseEnvMana self.reset_env(_env) if _policy is not None: self.reset_policy(_policy) - self._max_eval_reward = float("-inf") + if self._policy_cfg.type == 'dreamer_command': + self._states = None + self._resets = np.array([False for i in range(self._env_num)]) + self._max_episode_return = float("-inf") self._last_eval_iter = -1 self._end_flag = False self._last_render_iter = -1 @@ -186,7 +188,8 @@ def eval( envstep: int = -1, n_episode: Optional[int] = None, force_render: bool = False, - ) -> Tuple[bool, dict]: + policy_kwargs: Optional[Dict] = {}, + ) -> Tuple[bool, Dict[str, List]]: ''' Overview: Evaluate policy and store the best policy based on whether it reaches the highest historical reward. @@ -197,16 +200,12 @@ def eval( - n_episode (:obj:`int`): Number of evaluation episodes. Returns: - stop_flag (:obj:`bool`): Whether this training program can be ended. - - return_info (:obj:`dict`): Current evaluation return information. + - episode_info (:obj:`Dict[str, List]`): Current evaluation episode information. ''' - if get_world_size() > 1: - # sum up envstep to rank0 - envstep_tensor = torch.tensor(envstep).cuda() - dist.reduce(envstep_tensor, dst=0) - envstep = envstep_tensor.item() - # evaluator only work on rank0 - stop_flag, return_info = False, [] + stop_flag = False + episode_info = None # Initialize to ensure it's defined in all ranks + if get_rank() == 0: if n_episode is None: n_episode = self._default_n_episode @@ -229,7 +228,14 @@ def eval( if render: eval_monitor.update_video(self._env.ready_imgs) - policy_output = self._policy.forward(obs) + if self._policy_cfg.type == 'dreamer_command': + policy_output = self._policy.forward( + obs, **policy_kwargs, reset=self._resets, state=self._states + ) + #self._states = {env_id: output['state'] for env_id, output in policy_output.items()} + self._states = [output['state'] for output in policy_output.values()] + else: + policy_output = self._policy.forward(obs, **policy_kwargs) actions = {i: a['action'] for i, a in policy_output.items()} actions = to_ndarray(actions) timesteps = self._env.step(actions) @@ -239,22 +245,27 @@ def eval( # If there is an abnormal timestep, reset all the related variables(including this env). self._policy.reset([env_id]) continue + if self._policy_cfg.type == 'dreamer_command': + self._resets[env_id] = t.done if t.done: # Env reset is done by env_manager automatically. + if 'figure_path' in self._cfg and self._cfg.figure_path is not None: + self._env.enable_save_figure(env_id, self._cfg.figure_path) self._policy.reset([env_id]) - reward = t.info['final_eval_reward'] + reward = t.info['eval_episode_return'] + saved_info = {'eval_episode_return': t.info['eval_episode_return']} if 'episode_info' in t.info: - eval_monitor.update_info(env_id, t.info['episode_info']) + saved_info.update(t.info['episode_info']) + eval_monitor.update_info(env_id, saved_info) eval_monitor.update_reward(env_id, reward) - return_info.append(t.info) self._logger.info( - "[EVALUATOR]env {} finish episode, final reward: {}, current episode: {}".format( + "[EVALUATOR]env {} finish episode, final reward: {:.4f}, current episode: {}".format( env_id, eval_monitor.get_latest_reward(env_id), eval_monitor.get_current_episode() ) ) envstep_count += 1 duration = self._timer.value - episode_reward = eval_monitor.get_episode_reward() + episode_return = eval_monitor.get_episode_return() info = { 'train_iter': train_iter, 'ckpt_name': 'iteration_{}.pth.tar'.format(train_iter), @@ -264,11 +275,11 @@ def eval( 'evaluate_time': duration, 'avg_envstep_per_sec': envstep_count / duration, 'avg_time_per_episode': n_episode / duration, - 'reward_mean': np.mean(episode_reward), - 'reward_std': np.std(episode_reward), - 'reward_max': np.max(episode_reward), - 'reward_min': np.min(episode_reward), - # 'each_reward': episode_reward, + 'reward_mean': np.mean(episode_return), + 'reward_std': np.std(episode_return), + 'reward_max': np.max(episode_return), + 'reward_min': np.min(episode_return), + # 'each_reward': episode_return, } episode_info = eval_monitor.get_episode_info() if episode_info is not None: @@ -290,23 +301,25 @@ def eval( from ding.utils import fps self._tb_logger.add_video(video_title, videos, render_iter, fps(self._env)) - eval_reward = np.mean(episode_reward) - if eval_reward > self._max_eval_reward: + episode_return = np.mean(episode_return) + if episode_return > self._max_episode_return: if save_ckpt_fn: save_ckpt_fn('ckpt_best.pth.tar') - self._max_eval_reward = eval_reward - stop_flag = eval_reward >= self._stop_value and train_iter > 0 + self._max_episode_return = episode_return + stop_flag = episode_return >= self._stop_value and train_iter > 0 if stop_flag: self._logger.info( - "[DI-engine serial pipeline] " + - "Current eval_reward: {} is greater than stop_value: {}".format(eval_reward, self._stop_value) + - ", so your RL agent is converged, you can refer to " + + "[DI-engine serial pipeline] " + "Current episode_return: {:.4f} is greater than stop_value: {}". + format(episode_return, self._stop_value) + ", so your RL agent is converged, you can refer to " + "'log/evaluator/evaluator_logger.txt' for details." ) if get_world_size() > 1: - objects = [stop_flag, return_info] - dist.broadcast_object_list(objects, src=0) - stop_flag, return_info = objects + objects = [stop_flag, episode_info] + broadcast_object_list(objects, src=0) + stop_flag, episode_info = objects + + # Ensure episode_info is converted to the correct format + episode_info = to_item(episode_info) if episode_info is not None else {} - return stop_flag, return_info + return stop_flag, episode_info diff --git a/ding/worker/collector/marine_parallel_collector.py b/ding/worker/collector/marine_parallel_collector.py index 6f1bd99ae9..c659c70398 100644 --- a/ding/worker/collector/marine_parallel_collector.py +++ b/ding/worker/collector/marine_parallel_collector.py @@ -218,7 +218,7 @@ def _process_timestep(self, timestep: Dict[int, namedtuple]) -> None: self._policy_output_pool.reset(env_id) for p in self._policy: p.reset([env_id]) - reward = t[0].info['final_eval_reward'] + reward = t[0].info['eval_episode_return'] # Only left player's reward will be recorded. left_reward = reward[0] if isinstance(left_reward, torch.Tensor): diff --git a/ding/worker/collector/metric_serial_evaluator.py b/ding/worker/collector/metric_serial_evaluator.py index 4313e745e2..a160e437fc 100644 --- a/ding/worker/collector/metric_serial_evaluator.py +++ b/ding/worker/collector/metric_serial_evaluator.py @@ -219,7 +219,7 @@ def eval( if stop_flag: self._logger.info( "[DI-engine serial pipeline] " + - "Current eval_reward: {} is greater than stop_value: {}".format(avg_eval_result, self._stop_value) + + "Current episode_return: {} is greater than stop_value: {}".format(avg_eval_result, self._stop_value) + ", so your RL agent is converged, you can refer to 'log/evaluator/evaluator_logger.txt' for details." ) return stop_flag, avg_eval_result diff --git a/ding/worker/collector/sample_serial_collector.py b/ding/worker/collector/sample_serial_collector.py index beb691967b..697a8d0bcf 100644 --- a/ding/worker/collector/sample_serial_collector.py +++ b/ding/worker/collector/sample_serial_collector.py @@ -6,7 +6,8 @@ import torch from ding.envs import BaseEnvManager -from ding.utils import build_logger, EasyTimer, SERIAL_COLLECTOR_REGISTRY, one_time_warning +from ding.utils import build_logger, EasyTimer, SERIAL_COLLECTOR_REGISTRY, one_time_warning, get_rank, get_world_size, \ + allreduce_data from ding.torch_utils import to_tensor, to_ndarray from .base_serial_collector import ISerialCollector, CachePool, TrajBuffer, INF, to_tensor_transitions @@ -47,21 +48,32 @@ def __init__( self._exp_name = exp_name self._instance_name = instance_name self._collect_print_freq = cfg.collect_print_freq - self._deepcopy_obs = cfg.deepcopy_obs # avoid shallow copy, e.g., ovelap of s_t and s_t+1 + self._deepcopy_obs = cfg.deepcopy_obs # whether to deepcopy each data self._transform_obs = cfg.transform_obs self._cfg = cfg self._timer = EasyTimer() self._end_flag = False + self._rank = get_rank() + self._world_size = get_world_size() - if tb_logger is not None: + if self._rank == 0: + if tb_logger is not None: + self._logger, _ = build_logger( + path='./{}/log/{}'.format(self._exp_name, self._instance_name), + name=self._instance_name, + need_tb=False + ) + self._tb_logger = tb_logger + else: + self._logger, self._tb_logger = build_logger( + path='./{}/log/{}'.format(self._exp_name, self._instance_name), name=self._instance_name + ) + else: self._logger, _ = build_logger( path='./{}/log/{}'.format(self._exp_name, self._instance_name), name=self._instance_name, need_tb=False ) - self._tb_logger = tb_logger - else: - self._logger, self._tb_logger = build_logger( - path='./{}/log/{}'.format(self._exp_name, self._instance_name), name=self._instance_name - ) + self._tb_logger = None + self.reset(policy, env) def reset_env(self, _env: Optional[BaseEnvManager] = None) -> None: @@ -94,8 +106,9 @@ def reset_policy(self, _policy: Optional[namedtuple] = None) -> None: assert hasattr(self, '_env'), "please set env first" if _policy is not None: self._policy = _policy - self._default_n_sample = _policy.get_attribute('cfg').collect.get('n_sample', None) - self._traj_len_inf = _policy.get_attribute('cfg').collect.get('traj_len_inf', False) + self._policy_cfg = self._policy.get_attribute('cfg') + self._default_n_sample = _policy.get_attribute('n_sample') + self._traj_len_inf = self._policy_cfg.traj_len_inf self._unroll_len = _policy.get_attribute('unroll_len') self._on_policy = _policy.get_attribute('on_policy') if self._default_n_sample is not None and not self._traj_len_inf: @@ -131,6 +144,9 @@ def reset(self, _policy: Optional[namedtuple] = None, _env: Optional[BaseEnvMana if _policy is not None: self.reset_policy(_policy) + if self._policy_cfg.type == 'dreamer_command': + self._states = None + self._resets = np.array([False for i in range(self._env_num)]) self._obs_pool = CachePool('obs', self._env_num, deepcopy=self._deepcopy_obs) self._policy_output_pool = CachePool('policy_output', self._env_num) # _traj_buffer is {env_id: TrajBuffer}, is used to store traj_len pieces of transitions @@ -169,10 +185,20 @@ def envstep(self) -> int: Overview: Print the total envstep count. Return: - - envstep (:obj:`int`): the total envstep count + - envstep (:obj:`int`): The total envstep count. """ return self._total_envstep_count + @envstep.setter + def envstep(self, value: int) -> None: + """ + Overview: + Set the total envstep count. + Arguments: + - value (:obj:`int`): The total envstep count. + """ + self._total_envstep_count = value + def close(self) -> None: """ Overview: @@ -183,8 +209,9 @@ def close(self) -> None: return self._end_flag = True self._env.close() - self._tb_logger.flush() - self._tb_logger.close() + if self._tb_logger: + self._tb_logger.flush() + self._tb_logger.close() def __del__(self) -> None: """ @@ -199,6 +226,8 @@ def collect( n_sample: Optional[int] = None, train_iter: int = 0, drop_extra: bool = True, + random_collect: bool = False, + record_random_collect: bool = True, policy_kwargs: Optional[dict] = None, level_seeds: Optional[List] = None, ) -> List[Any]: @@ -209,6 +238,7 @@ def collect( - n_sample (:obj:`int`): The number of collecting data sample. - train_iter (:obj:`int`): The number of training iteration when calling collect method. - drop_extra (:obj:`bool`): Whether to drop extra return_data more than `n_sample`. + - record_random_collect (:obj:`bool`) :Whether to output logs of random collect. - policy_kwargs (:obj:`dict`): The keyword args for policy forward. - level_seeds (:obj:`dict`): Used in PLR, represents the seed of the environment that \ generate the data @@ -228,6 +258,8 @@ def collect( if policy_kwargs is None: policy_kwargs = {} collected_sample = 0 + collected_step = 0 + collected_episode = 0 return_data = [] while collected_sample < n_sample: @@ -238,7 +270,12 @@ def collect( self._obs_pool.update(obs) if self._transform_obs: obs = to_tensor(obs, dtype=torch.float32) - policy_output = self._policy.forward(obs, **policy_kwargs) + if self._policy_cfg.type == 'dreamer_command' and not random_collect: + policy_output = self._policy.forward(obs, **policy_kwargs, reset=self._resets, state=self._states) + #self._states = {env_id: output['state'] for env_id, output in policy_output.items()} + self._states = [output['state'] for output in policy_output.values()] + else: + policy_output = self._policy.forward(obs, **policy_kwargs) self._policy_output_pool.update(policy_output) # Interact with env. actions = {env_id: output['action'] for env_id, output in policy_output.items()} @@ -259,9 +296,9 @@ def collect( self._reset_stat(env_id) self._logger.info('Env{} returns a abnormal step, its info is {}'.format(env_id, timestep.info)) continue - if 'type' in self._policy.get_attribute('cfg') and \ - self._policy.get_attribute('cfg').type == 'ngu_command': - # for NGU policy + if self._policy_cfg.type == 'dreamer_command' and not random_collect: + self._resets[env_id] = timestep.done + if self._policy_cfg.type == 'ngu_command': # for NGU policy transition = self._policy.process_transition( self._obs_pool[env_id], self._policy_output_pool[env_id], timestep, env_id ) @@ -275,7 +312,7 @@ def collect( transition['collect_iter'] = train_iter self._traj_buffer[env_id].append(transition) self._env_info[env_id]['step'] += 1 - self._total_envstep_count += 1 + collected_step += 1 # prepare data if timestep.done or len(self._traj_buffer[env_id]) == self._traj_len: # If policy is r2d2: @@ -289,10 +326,10 @@ def collect( # sequence sample of length (please refer to r2d2.py). # Episode is done or traj_buffer(maxlen=traj_len) is full. - transitions = to_tensor_transitions(self._traj_buffer[env_id]) + # indicate whether to shallow copy next obs, i.e., overlap of s_t and s_t+1 + transitions = to_tensor_transitions(self._traj_buffer[env_id], not self._deepcopy_obs) train_sample = self._policy.get_train_sample(transitions) return_data.extend(train_sample) - self._total_train_sample_count += len(train_sample) self._env_info[env_id]['train_sample'] += len(train_sample) collected_sample += len(train_sample) self._traj_buffer[env_id].clear() @@ -301,8 +338,8 @@ def collect( # If env is done, record episode info and reset if timestep.done: - self._total_episode_count += 1 - reward = timestep.info['final_eval_reward'] + collected_episode += 1 + reward = timestep.info['eval_episode_return'] info = { 'reward': reward, 'time': self._env_info[env_id]['time'], @@ -313,8 +350,23 @@ def collect( # Env reset is done by env_manager automatically self._policy.reset([env_id]) self._reset_stat(env_id) + + collected_duration = sum([d['time'] for d in self._episode_info]) + # reduce data when enables DDP + if self._world_size > 1: + collected_sample = allreduce_data(collected_sample, 'sum') + collected_step = allreduce_data(collected_step, 'sum') + collected_episode = allreduce_data(collected_episode, 'sum') + collected_duration = allreduce_data(collected_duration, 'sum') + self._total_envstep_count += collected_step + self._total_episode_count += collected_episode + self._total_duration += collected_duration + self._total_train_sample_count += collected_sample # log - self._output_log(train_iter) + if record_random_collect: # default is true, but when random collect, record_random_collect is False + self._output_log(train_iter) + else: + self._episode_info.clear() # on-policy reset if self._on_policy: for env_id in range(self._env_num): @@ -328,19 +380,20 @@ def collect( def _output_log(self, train_iter: int) -> None: """ Overview: - Print the output log information. You can refer to Docs/Best Practice/How to understand\ - training generated folders/Serial mode/log/collector for more details. + Print the output log information. You can refer to the docs of `Best Practice` to understand \ + the training generated logs and tensorboards. Arguments: - train_iter (:obj:`int`): the number of training iteration. """ + if self._rank != 0: + return if (train_iter - self._last_train_iter) >= self._collect_print_freq and len(self._episode_info) > 0: self._last_train_iter = train_iter episode_count = len(self._episode_info) envstep_count = sum([d['step'] for d in self._episode_info]) train_sample_count = sum([d['train_sample'] for d in self._episode_info]) duration = sum([d['time'] for d in self._episode_info]) - episode_reward = [d['reward'] for d in self._episode_info] - self._total_duration += duration + episode_return = [d['reward'] for d in self._episode_info] info = { 'episode_count': episode_count, 'envstep_count': envstep_count, @@ -350,16 +403,14 @@ def _output_log(self, train_iter: int) -> None: 'avg_envstep_per_sec': envstep_count / duration, 'avg_train_sample_per_sec': train_sample_count / duration, 'avg_episode_per_sec': episode_count / duration, - 'collect_time': duration, - 'reward_mean': np.mean(episode_reward), - 'reward_std': np.std(episode_reward), - 'reward_max': np.max(episode_reward), - 'reward_min': np.min(episode_reward), + 'reward_mean': np.mean(episode_return), + 'reward_std': np.std(episode_return), + 'reward_max': np.max(episode_return), + 'reward_min': np.min(episode_return), 'total_envstep_count': self._total_envstep_count, 'total_train_sample_count': self._total_train_sample_count, 'total_episode_count': self._total_episode_count, - 'total_duration': self._total_duration, - # 'each_reward': episode_reward, + # 'each_reward': episode_return, } self._episode_info.clear() self._logger.info("collect end:\n{}".format('\n'.join(['{}: {}'.format(k, v) for k, v in info.items()]))) diff --git a/ding/worker/collector/tests/speed_test/fake_env.py b/ding/worker/collector/tests/speed_test/fake_env.py index b9714bc118..731e990e34 100644 --- a/ding/worker/collector/tests/speed_test/fake_env.py +++ b/ding/worker/collector/tests/speed_test/fake_env.py @@ -40,7 +40,7 @@ def reset(self) -> np.ndarray: self._episode_step = int(random_change(self._episode_step_base)) env_sleep(random_change(self._reset_time)) self._step_count = 0 - self._final_eval_reward = 0. + self._eval_episode_return = 0. obs = np.random.randn(self._obs_dim).astype(np.float32) return obs @@ -59,9 +59,9 @@ def step(self, action: np.ndarray) -> BaseEnvTimestep: rew = np.random.randint(2) done = True if self._step_count == self._episode_step else False info = {} - self._final_eval_reward += rew + self._eval_episode_return += rew if done: - info['final_eval_reward'] = self._final_eval_reward + info['eval_episode_return'] = self._eval_episode_return rew = to_ndarray([rew]) # to shape (1,) return BaseEnvTimestep(obs, rew, done, info) diff --git a/ding/worker/collector/tests/speed_test/test_collector_profile.py b/ding/worker/collector/tests/speed_test/test_collector_profile.py index c7019316f9..9f1a268c40 100644 --- a/ding/worker/collector/tests/speed_test/test_collector_profile.py +++ b/ding/worker/collector/tests/speed_test/test_collector_profile.py @@ -169,7 +169,8 @@ def compare_test(cfg: EasyDict, seed: int, test_name: str) -> None: print(template.format(test_name, np.mean(fps), np.std(fps))) -@pytest.mark.benchmark +# TODO(nyz) fix CI bug when py==3.8.15 +@pytest.mark.tmp def test_collector_profile(): # ignore them for clear log collector_log = logging.getLogger('collector_logger') diff --git a/ding/worker/collector/tests/test_base_serial_collector.py b/ding/worker/collector/tests/test_base_serial_collector.py new file mode 100644 index 0000000000..475a6a4b17 --- /dev/null +++ b/ding/worker/collector/tests/test_base_serial_collector.py @@ -0,0 +1,42 @@ +import pytest +import numpy as np +import torch +from ding.worker.collector.base_serial_collector import to_tensor_transitions + + +def get_transition(): + return { + 'obs': np.random.random((2, 3)), + 'action': np.random.randint(0, 6, size=(1, )), + 'reward': np.random.random((1, )), + 'done': False, + 'next_obs': np.random.random((2, 3)), + } + + +@pytest.mark.unittest +def test_to_tensor_transitions(): + # test case when shallow copy is True + transition_list = [get_transition() for _ in range(4)] + tensor_list = to_tensor_transitions(transition_list, shallow_copy_next_obs=True) + for i in range(len(tensor_list)): + tensor = tensor_list[i] + assert isinstance(tensor['obs'], torch.Tensor) + assert isinstance(tensor['action'], torch.Tensor), type(tensor['action']) + assert isinstance(tensor['reward'], torch.Tensor) + assert isinstance(tensor['done'], bool) + assert 'next_obs' in tensor + if i < len(tensor_list) - 1: + assert id(tensor['next_obs']) == id(tensor_list[i + 1]['obs']) + # test case when shallow copy is False + transition_list = [get_transition() for _ in range(4)] + tensor_list = to_tensor_transitions(transition_list, shallow_copy_next_obs=False) + for i in range(len(tensor_list)): + tensor = tensor_list[i] + assert isinstance(tensor['obs'], torch.Tensor) + assert isinstance(tensor['action'], torch.Tensor) + assert isinstance(tensor['reward'], torch.Tensor) + assert isinstance(tensor['done'], bool) + assert 'next_obs' in tensor + if i < len(tensor_list) - 1: + assert id(tensor['next_obs']) != id(tensor_list[i + 1]['obs']) diff --git a/ding/worker/collector/tests/test_marine_parallel_collector.py b/ding/worker/collector/tests/test_marine_parallel_collector.py index be69fbaa25..a4253799e7 100644 --- a/ding/worker/collector/tests/test_marine_parallel_collector.py +++ b/ding/worker/collector/tests/test_marine_parallel_collector.py @@ -48,7 +48,7 @@ def step(self, action: Union[torch.Tensor, np.ndarray, list]) -> BaseEnvTimestep done = False if self._step_times < 20 else True info = {} if done: - info['final_eval_reward'] = np.array([21.]) if self._is_evaluator else np.array([5., -5.]) + info['eval_episode_return'] = np.array([21.]) if self._is_evaluator else np.array([5., -5.]) self._step_times += 1 return BaseEnvTimestep(obs, rew, done, info) diff --git a/ding/worker/collector/tests/test_sample_serial_collector.py b/ding/worker/collector/tests/test_sample_serial_collector.py new file mode 100644 index 0000000000..dbc0994f10 --- /dev/null +++ b/ding/worker/collector/tests/test_sample_serial_collector.py @@ -0,0 +1,39 @@ +import pytest +from ding.worker import SampleSerialCollector +from ding.envs import BaseEnvManager, SyncSubprocessEnvManager, AsyncSubprocessEnvManager +from ding.policy import DQNPolicy +from ding.model import DQN +from dizoo.classic_control.cartpole.envs import CartPoleEnv + + +@pytest.mark.unittest +@pytest.mark.parametrize('env_manager_type', [BaseEnvManager, SyncSubprocessEnvManager]) +def test_collect(env_manager_type): + env = env_manager_type([lambda: CartPoleEnv({}) for _ in range(8)], env_manager_type.default_config()) + env.seed(0) + model = DQN(obs_shape=4, action_shape=1) + policy = DQNPolicy(DQNPolicy.default_config(), model=model).collect_mode + collector = SampleSerialCollector(SampleSerialCollector.default_config(), env, policy) + + collected_sample = collector.collect( + n_sample=1000, + train_iter=collector._collect_print_freq, + record_random_collect=False, + policy_kwargs={'eps': 0.5} + ) + assert len(collected_sample) == 1000 + + +@pytest.mark.unittest +@pytest.mark.parametrize('env_manager_type', [BaseEnvManager, SyncSubprocessEnvManager]) +def test_random_collect(env_manager_type): + env = env_manager_type([lambda: CartPoleEnv({}) for _ in range(8)], env_manager_type.default_config()) + env.seed(0) + model = DQN(obs_shape=4, action_shape=1) + policy = DQNPolicy(DQNPolicy.default_config(), model=model).collect_mode + collector = SampleSerialCollector(SampleSerialCollector.default_config(), env, policy) + + collected_sample = collector.collect( + n_sample=1000, train_iter=collector._collect_print_freq, record_random_collect=True, policy_kwargs={'eps': 0.5} + ) + assert len(collected_sample) == 1000 diff --git a/ding/worker/collector/zergling_parallel_collector.py b/ding/worker/collector/zergling_parallel_collector.py index f15cc368cc..9c1da9c41e 100644 --- a/ding/worker/collector/zergling_parallel_collector.py +++ b/ding/worker/collector/zergling_parallel_collector.py @@ -78,9 +78,10 @@ def policy(self) -> Policy: @policy.setter def policy(self, _policy: Policy) -> None: self._policy = _policy - self._n_episode = _policy.get_attribute('cfg').collect.get('n_episode', None) - self._n_sample = _policy.get_attribute('cfg').collect.get('n_sample', None) - assert any( + self._policy_cfg = self._policy.get_attribute('cfg') + self._n_sample = _policy.get_attribute('n_sample') + self._n_episode = _policy.get_attribute('n_episode') + assert not all( [t is None for t in [self._n_sample, self._n_episode]] ), "n_episode/n_sample in policy cfg can't be not None at the same time" # TODO(nyz) the same definition of traj_len in serial and parallel @@ -182,7 +183,7 @@ def _process_timestep(self, timestep: Dict[int, namedtuple]) -> None: self._obs_pool.reset(env_id) self._policy_output_pool.reset(env_id) self._policy.reset([env_id]) - reward = t.info['final_eval_reward'] + reward = t.info['eval_episode_return'] if isinstance(reward, torch.Tensor): reward = reward.item() self._episode_result[env_id].append(reward) diff --git a/ding/worker/coordinator/one_vs_one_parallel_commander.py b/ding/worker/coordinator/one_vs_one_parallel_commander.py index 45135ea050..82b85420cb 100644 --- a/ding/worker/coordinator/one_vs_one_parallel_commander.py +++ b/ding/worker/coordinator/one_vs_one_parallel_commander.py @@ -237,7 +237,7 @@ def finish_collector_task(self, task_id: str, finished_task: dict) -> bool: print('===', eval_win, difficulty_inc) if eval_stop_value is not None and finished_task['reward_mean'] >= eval_stop_value and is_hardest: self._logger.info( - "[DI-engine parallel pipeline] Current eval_reward: {} is greater than the stop_value: {}". + "[DI-engine parallel pipeline] Current episode_return: {} is greater than the stop_value: {}". format(finished_task['reward_mean'], eval_stop_value) + ", so the total training program is over." ) self._end_flag = True diff --git a/ding/worker/coordinator/solo_parallel_commander.py b/ding/worker/coordinator/solo_parallel_commander.py index 9506db5b5e..ab374ebbc9 100644 --- a/ding/worker/coordinator/solo_parallel_commander.py +++ b/ding/worker/coordinator/solo_parallel_commander.py @@ -174,7 +174,7 @@ def finish_collector_task(self, task_id: str, finished_task: dict) -> bool: eval_stop_value = self._cfg.env.stop_value if eval_stop_value is not None and finished_task['reward_mean'] >= eval_stop_value: self._logger.info( - "[DI-engine parallel pipeline] current eval_reward: {} is greater than the stop_value: {}". + "[DI-engine parallel pipeline] current episode_return: {} is greater than the stop_value: {}". format(finished_task['reward_mean'], eval_stop_value) + ", so the total training program is over." ) self._end_flag = True diff --git a/ding/worker/learner/base_learner.py b/ding/worker/learner/base_learner.py index b7588e5f18..a070e7a58d 100644 --- a/ding/worker/learner/base_learner.py +++ b/ding/worker/learner/base_learner.py @@ -35,6 +35,8 @@ def default_config(cls: type) -> EasyDict: train_iterations=int(1e9), dataloader=dict(num_workers=0, ), log_policy=True, + is_multitask_pipeline=False, + only_monitor_rank0=True, # --- Hooks --- hook=dict( load_ckpt_before_run='', @@ -59,7 +61,9 @@ def __init__( Overview: Initialization method, build common learner components according to cfg, such as hook, wrapper and so on. Arguments: - - cfg (:obj:`EasyDict`): Learner config, you can refer cls.config for details. + - cfg (:obj:`EasyDict`): Learner config, you can refer cls.config for details. It should include \ + `is_multitask_pipeline` to indicate if the pipeline is multitask, default is False, \ + and `only_monitor_rank0` to control whether only rank 0 needs monitor and tb_logger, default is True. - policy (:obj:`namedtuple`): A collection of policy function of learn mode. And policy can also be \ initialized when runtime. - tb_logger (:obj:`SummaryWriter`): Tensorboard summary writer. @@ -78,6 +82,12 @@ def __init__( self._instance_name = instance_name self._ckpt_name = None self._timer = EasyTimer() + self._is_multitask_pipeline = self._cfg.is_multitask_pipeline + self.only_monitor_rank0 = self._cfg.only_monitor_rank0 + + # Adjust only_monitor_rank0 based on is_multitask_pipeline + if self._is_multitask_pipeline: + self.only_monitor_rank0 = False # These 2 attributes are only used in parallel mode. self._end_flag = False @@ -92,8 +102,10 @@ def __init__( self._cfg.hook.log_reduce_after_iter = True # Logger (Monitor will be initialized in policy setter) - # Only rank == 0 learner needs monitor and tb_logger, others only need text_logger to display terminal output. - if self._rank == 0: + # In the multitask pipeline, each rank needs its own tb_logger. + # Otherwise, only rank == 0 learner needs monitor and tb_logger, + # others only need text_logger to display terminal output. + if self._rank == 0 or not self.only_monitor_rank0: if tb_logger is not None: self._logger, _ = build_logger( './{}/log/{}'.format(self._exp_name, self._instance_name), self._instance_name, need_tb=False @@ -108,6 +120,7 @@ def __init__( './{}/log/{}'.format(self._exp_name, self._instance_name), self._instance_name, need_tb=False ) self._tb_logger = None + self._log_buffer = { 'scalar': build_log_buffer(), 'scalars': build_log_buffer(), @@ -122,6 +135,8 @@ def __init__( self._hooks = {'before_run': [], 'before_iter': [], 'after_iter': [], 'after_run': []} # Last iteration. Used to record current iter. self._last_iter = CountVar(init_val=0) + # Collector envstep. Used to record current envstep. + self._collector_envstep = 0 # Setup time wrapper and hook. self._setup_wrapper() @@ -177,6 +192,26 @@ def register_hook(self, hook: LearnerHook) -> None: """ add_learner_hook(self._hooks, hook) + @property + def collector_envstep(self) -> int: + """ + Overview: + Get current collector envstep. + Returns: + - collector_envstep (:obj:`int`): Current collector envstep. + """ + return self._collector_envstep + + @collector_envstep.setter + def collector_envstep(self, value: int) -> None: + """ + Overview: + Set current collector envstep. + Arguments: + - value (:obj:`int`): Current collector envstep. + """ + self._collector_envstep = value + def train(self, data: dict, envstep: int = -1, policy_kwargs: Optional[dict] = None) -> None: """ Overview: @@ -351,7 +386,7 @@ def save_checkpoint(self, ckpt_name: str = None) -> None: - ``auto_checkpoint`` (``torch_utils/checkpoint_helper.py``), which is designed for \ saving checkpoint whenever an exception raises. - ``serial_pipeline`` (``entry/serial_entry.py``). Used to save checkpoint when reaching \ - new highest evaluation reward. + new highest episode return. """ if ckpt_name is not None: self.ckpt_name = ckpt_name @@ -432,7 +467,7 @@ def policy(self, _policy: 'Policy') -> None: # noqa Policy variable monitor is set alongside with policy, because variables are determined by specific policy. """ self._policy = _policy - if self._rank == 0: + if self._rank == 0 or not self.only_monitor_rank0: self._monitor = get_simple_monitor_type(self._policy.monitor_vars())(TickTime(), expire=10) if self._cfg.log_policy: self.info(self._policy.info()) diff --git a/ding/worker/learner/comm/base_comm_learner.py b/ding/worker/learner/comm/base_comm_learner.py index a7234e0e7a..4a9562888f 100644 --- a/ding/worker/learner/comm/base_comm_learner.py +++ b/ding/worker/learner/comm/base_comm_learner.py @@ -116,6 +116,7 @@ def _create_learner(self, task_info: dict) -> 'BaseLearner': # noqa setattr(learner, item, getattr(self, item)) # Set policy in created learner. policy_cfg = task_info['policy'] + policy_cfg = EasyDict(policy_cfg) learner.policy = create_policy(policy_cfg, enable_field=['learn']).learn_mode learner.setup_dataloader() return learner diff --git a/ding/worker/learner/learner_hook.py b/ding/worker/learner/learner_hook.py index 250a8f1950..ffbc42fdf9 100644 --- a/ding/worker/learner/learner_hook.py +++ b/ding/worker/learner/learner_hook.py @@ -6,7 +6,7 @@ from easydict import EasyDict import ding -from ding.utils import allreduce, read_file, save_file, get_rank +from ding.utils import allreduce, read_file, save_file class Hook(ABC): @@ -117,6 +117,9 @@ def __call__(self, engine: 'BaseLearner') -> None: # noqa if 'last_iter' in state_dict: last_iter = state_dict.pop('last_iter') engine.last_iter.update(last_iter) + if 'last_step' in state_dict: + last_step = state_dict.pop('last_step') + engine._collector_envstep = last_step engine.policy.load_state_dict(state_dict) engine.info('{} load ckpt in {}'.format(engine.instance_name, path)) @@ -166,6 +169,7 @@ def __call__(self, engine: 'BaseLearner') -> None: # noqa path = os.path.join(dirname, ckpt_name) state_dict = engine.policy.state_dict() state_dict.update({'last_iter': engine.last_iter.val}) + state_dict.update({'last_step': engine.collector_envstep}) save_file(path, state_dict) engine.info('{} save ckpt in {}'.format(engine.instance_name, path)) @@ -183,29 +187,33 @@ class LogShowHook(LearnerHook): def __init__(self, *args, ext_args: EasyDict = EasyDict(), **kwargs) -> None: """ Overview: - init LogShowHook + Init LogShowHook. Arguments: - - ext_args (:obj:`EasyDict`): extended_args, use ext_args.freq to set freq + - ext_args (:obj:`EasyDict`): Extended arguments, use ext_args.freq to set frequency and \ + ext_args.only_monitor_rank0 to control if only rank 0 should monitor, default is True. """ super().__init__(*args, **kwargs) if ext_args == {}: self._freq = 1 else: self._freq = ext_args.freq + self._only_monitor_rank0 = None def __call__(self, engine: 'BaseLearner') -> None: # noqa """ Overview: Show log, update record and tb_logger if rank is 0 and at interval iterations, - clear the log buffer for all learners regardless of rank + clear the log buffer for all learners regardless of rank. Arguments: - - engine (:obj:`BaseLearner`): the BaseLearner + - engine (:obj:`BaseLearner`): The BaseLearner. """ - # Only show log for rank 0 learner - if engine.rank != 0: + self._only_monitor_rank0 = engine.only_monitor_rank0 + # Only show log for rank 0 learner if _only_monitor_rank0 is True + if engine.rank != 0 and self._only_monitor_rank0: for k in engine.log_buffer: engine.log_buffer[k].clear() return + # For 'scalar' type variables: log_buffer -> tick_monitor -> monitor_time.step for k, v in engine.log_buffer['scalar'].items(): setattr(engine.monitor, k, v) @@ -239,7 +247,7 @@ def __call__(self, engine: 'BaseLearner') -> None: # noqa class LogReduceHook(LearnerHook): """ Overview: - Hook to reduce the distributed(multi-gpu) logs + Hook to reduce the distributed (multi-gpu) logs. Interfaces: __init__, __call__ Property: @@ -249,32 +257,47 @@ class LogReduceHook(LearnerHook): def __init__(self, *args, ext_args: EasyDict = EasyDict(), **kwargs) -> None: """ Overview: - init LogReduceHook + Initialize LogReduceHook. Arguments: - - ext_args (:obj:`EasyDict`): extended_args, use ext_args.freq to set log_reduce_freq + - ext_args (:obj:`EasyDict`): Extended arguments, use ext_args.freq to set log_reduce_freq. """ super().__init__(*args, **kwargs) def __call__(self, engine: 'BaseLearner') -> None: # noqa """ Overview: - reduce the logs from distributed(multi-gpu) learners + Reduce the logs from distributed (multi-gpu) learners. Arguments: - - engine (:obj:`BaseLearner`): the BaseLearner + - engine (:obj:`BaseLearner`): The BaseLearner. """ def aggregate(data): r""" Overview: - aggregate the information from all ranks(usually use sync allreduce) + Aggregate the information from all ranks (usually using sync allreduce). Arguments: - data (:obj:`dict`): Data that needs to be reduced. \ - Could be dict, torch.Tensor, numbers.Integral or numbers.Real. + Could be dict, torch.Tensor, numbers.Integral, or numbers.Real. Returns: - - new_data (:obj:`dict`): data after reduce + - new_data (:obj:`dict`): Data after reduction. """ + + def should_reduce(key): + # Check if the key starts with the "noreduce_" prefix. + # The "noreduce_" prefix is used in the unizero_multitask ddp pipeline + # to indicate data that should not be reduced. + return not key.startswith("noreduce_") + + cuda_device = torch.cuda.current_device() + if isinstance(data, dict): - new_data = {k: aggregate(v) for k, v in data.items()} + new_data = {} + for k, v in data.items(): + if should_reduce(k): + new_data[k] = aggregate(v) # Perform allreduce on data that needs reduction. + else: + new_data[k] = v # Retain data that does not need reduction. + elif isinstance(data, list) or isinstance(data, tuple): new_data = [aggregate(t) for t in data] elif isinstance(data, torch.Tensor): @@ -282,7 +305,7 @@ def aggregate(data): if ding.enable_linklink: allreduce(new_data) else: - new_data = new_data.to(get_rank()) + new_data = new_data.to(cuda_device) allreduce(new_data) new_data = new_data.cpu() elif isinstance(data, numbers.Integral) or isinstance(data, numbers.Real): @@ -290,12 +313,12 @@ def aggregate(data): if ding.enable_linklink: allreduce(new_data) else: - new_data = new_data.to(get_rank()) + new_data = new_data.to(cuda_device) allreduce(new_data) new_data = new_data.cpu() new_data = new_data.item() else: - raise TypeError("invalid type in reduce: {}".format(type(data))) + raise TypeError("Invalid type in reduce: {}".format(type(data))) return new_data engine.log_buffer = aggregate(engine.log_buffer) diff --git a/ding/worker/replay_buffer/__init__.py b/ding/worker/replay_buffer/__init__.py index 8b7009450f..c4f1bf3e87 100644 --- a/ding/worker/replay_buffer/__init__.py +++ b/ding/worker/replay_buffer/__init__.py @@ -1,4 +1,4 @@ from .base_buffer import IBuffer, create_buffer, get_buffer_cls -from .naive_buffer import NaiveReplayBuffer +from .naive_buffer import NaiveReplayBuffer, SequenceReplayBuffer from .advanced_buffer import AdvancedReplayBuffer from .episode_buffer import EpisodeReplayBuffer diff --git a/ding/worker/replay_buffer/advanced_buffer.py b/ding/worker/replay_buffer/advanced_buffer.py index 7196b16b09..b9f700768d 100644 --- a/ding/worker/replay_buffer/advanced_buffer.py +++ b/ding/worker/replay_buffer/advanced_buffer.py @@ -1,11 +1,13 @@ +import os import copy import time -from typing import Union, NoReturn, Any, Optional, List, Dict, Tuple +from typing import Union, Any, Optional, List, Dict, Tuple import numpy as np +import hickle from ding.worker.replay_buffer import IBuffer from ding.utils import SumSegmentTree, MinSegmentTree, BUFFER_REGISTRY -from ding.utils import LockContext, LockContextType, build_logger +from ding.utils import LockContext, LockContextType, build_logger, get_rank from ding.utils.autolog import TickTime from .utils import UsedDataRemover, generate_id, SampledDataAttrMonitor, PeriodicThruputMonitor, ThruputController @@ -104,6 +106,7 @@ def __init__( self._instance_name = instance_name self._end_flag = False self._cfg = cfg + self._rank = get_rank() self._replay_buffer_size = self._cfg.replay_buffer_size self._deepcopy = self._cfg.deepcopy # ``_data`` is a circular queue to store data (full data or meta data) @@ -117,7 +120,7 @@ def __init__( # Is used to generate a unique id for each data: If a new data is inserted, its unique id will be this. self._next_unique_id = 0 # Lock to guarantee thread safe - self._lock = LockContext(type_=LockContextType.THREAD_LOCK) + self._lock = LockContext(lock_type=LockContextType.THREAD_LOCK) # Point to the head of the circular queue. The true data is the stalest(oldest) data in this queue. # Because buffer would remove data due to staleness or use count, and at the beginning when queue is not # filled with data head would always be 0, so ``head`` may be not equal to ``tail``; @@ -161,16 +164,22 @@ def __init__( # Monitor & Logger monitor_cfg = self._cfg.monitor - if tb_logger is not None: + if self._rank == 0: + if tb_logger is not None: + self._logger, _ = build_logger( + './{}/log/{}'.format(self._exp_name, self._instance_name), self._instance_name, need_tb=False + ) + self._tb_logger = tb_logger + else: + self._logger, self._tb_logger = build_logger( + './{}/log/{}'.format(self._exp_name, self._instance_name), + self._instance_name, + ) + else: self._logger, _ = build_logger( './{}/log/{}'.format(self._exp_name, self._instance_name), self._instance_name, need_tb=False ) - self._tb_logger = tb_logger - else: - self._logger, self._tb_logger = build_logger( - './{}/log/{}'.format(self._exp_name, self._instance_name), - self._instance_name, - ) + self._tb_logger = None self._start_time = time.time() # Sampled data attributes. self._cur_learner_iter = -1 @@ -181,9 +190,10 @@ def __init__( ) self._sampled_data_attr_print_freq = monitor_cfg.sampled_data_attr.print_freq # Periodic thruput. - self._periodic_thruput_monitor = PeriodicThruputMonitor( - self._instance_name, monitor_cfg.periodic_thruput, self._logger, self._tb_logger - ) + if self._rank == 0: + self._periodic_thruput_monitor = PeriodicThruputMonitor( + self._instance_name, monitor_cfg.periodic_thruput, self._logger, self._tb_logger + ) # Used data remover self._enable_track_used_data = self._cfg.enable_track_used_data @@ -208,9 +218,10 @@ def close(self) -> None: return self._end_flag = True self.clear() - self._periodic_thruput_monitor.close() - self._tb_logger.flush() - self._tb_logger.close() + if self._rank == 0: + self._periodic_thruput_monitor.close() + self._tb_logger.flush() + self._tb_logger.close() if self._enable_track_used_data: self._used_data_remover.close() @@ -284,6 +295,15 @@ def push(self, data: Union[List[Any], Any], cur_collector_envstep: int) -> None: else: self._append(data, cur_collector_envstep) + def save_data(self, file_name: str): + if not os.path.exists(os.path.dirname(file_name)): + if os.path.dirname(file_name) != "": + os.makedirs(os.path.dirname(file_name)) + hickle.dump(py_obj=self._data, file_obj=file_name) + + def load_data(self, file_name: str): + self.push(hickle.load(file_name), 0) + def _sample_check(self, size: int, cur_learner_iter: int) -> Tuple[bool, str]: r""" Overview: @@ -363,7 +383,8 @@ def _append(self, ori_data: Any, cur_collector_envstep: int = -1) -> None: self._set_weight(data) self._data[self._tail] = data self._valid_count += 1 - self._periodic_thruput_monitor.valid_count = self._valid_count + if self._rank == 0: + self._periodic_thruput_monitor.valid_count = self._valid_count self._tail = (self._tail + 1) % self._replay_buffer_size self._next_unique_id += 1 self._monitor_update_of_push(1, cur_collector_envstep) @@ -424,7 +445,8 @@ def _extend(self, ori_data: List[Any], cur_collector_envstep: int = -1) -> None: data_start = 0 valid_data_start += L self._valid_count += len(valid_data) - self._periodic_thruput_monitor.valid_count = self._valid_count + if self._rank == 0: + self._periodic_thruput_monitor.valid_count = self._valid_count # Update ``tail`` and ``next_unique_id`` after the whole list is pushed into buffer. self._tail = (self._tail + length) % self._replay_buffer_size self._next_unique_id += length @@ -557,8 +579,9 @@ def _remove(self, idx: int, use_too_many_times: bool = False) -> None: if self._enable_track_used_data: self._used_data_remover.add_used_data(self._data[idx]) self._valid_count -= 1 - self._periodic_thruput_monitor.valid_count = self._valid_count - self._periodic_thruput_monitor.remove_data_count += 1 + if self._rank == 0: + self._periodic_thruput_monitor.valid_count = self._valid_count + self._periodic_thruput_monitor.remove_data_count += 1 self._data[idx] = None self._sum_tree[idx] = self._sum_tree.neutral_element self._min_tree[idx] = self._min_tree.neutral_element @@ -613,7 +636,8 @@ def _monitor_update_of_push(self, add_count: int, cur_collector_envstep: int = - - add_count (:obj:`int`): How many datas are added into buffer. - cur_collector_envstep (:obj:`int`): Collector envstep, passed in by collector. """ - self._periodic_thruput_monitor.push_data_count += add_count + if self._rank == 0: + self._periodic_thruput_monitor.push_data_count += add_count if self._use_thruput_controller: self._thruput_controller.history_push_count += add_count self._cur_collector_envstep = cur_collector_envstep @@ -628,7 +652,8 @@ def _monitor_update_of_sample(self, sample_data: list, cur_learner_iter: int) -> e.g. use, priority, staleness, etc. - cur_learner_iter (:obj:`int`): Learner iteration, passed in by learner. """ - self._periodic_thruput_monitor.sample_data_count += len(sample_data) + if self._rank == 0: + self._periodic_thruput_monitor.sample_data_count += len(sample_data) if self._use_thruput_controller: self._thruput_controller.history_sample_count += len(sample_data) self._cur_learner_iter = cur_learner_iter @@ -657,7 +682,7 @@ def _monitor_update_of_sample(self, sample_data: list, cur_learner_iter: int) -> 'staleness_max': self._sampled_data_attr_monitor.max['staleness'](), 'beta': self._beta, } - if self._sampled_data_attr_print_count % self._sampled_data_attr_print_freq == 0: + if self._sampled_data_attr_print_count % self._sampled_data_attr_print_freq == 0 and self._rank == 0: self._logger.info("=== Sample data {} Times ===".format(self._sampled_data_attr_print_count)) self._logger.info(self._logger.get_tabulate_vars_hor(out_dict)) for k, v in out_dict.items(): @@ -710,7 +735,7 @@ def beta(self) -> float: return self._beta @beta.setter - def beta(self, beta: float) -> NoReturn: + def beta(self, beta: float) -> None: self._beta = beta def state_dict(self) -> dict: diff --git a/ding/worker/replay_buffer/base_buffer.py b/ding/worker/replay_buffer/base_buffer.py index 53ada7d484..7231c34067 100644 --- a/ding/worker/replay_buffer/base_buffer.py +++ b/ding/worker/replay_buffer/base_buffer.py @@ -79,6 +79,26 @@ def count(self) -> int: """ raise NotImplementedError + @abstractmethod + def save_data(self, file_name: str): + """ + Overview: + Save buffer data into a file. + Arguments: + - file_name (:obj:`str`): file name of buffer data + """ + raise NotImplementedError + + @abstractmethod + def load_data(self, file_name: str): + """ + Overview: + Load buffer data from a file. + Arguments: + - file_name (:obj:`str`): file name of buffer data + """ + raise NotImplementedError + @abstractmethod def state_dict(self) -> Dict[str, Any]: """ diff --git a/ding/worker/replay_buffer/naive_buffer.py b/ding/worker/replay_buffer/naive_buffer.py index 4a4211bbfd..87b934e6bf 100644 --- a/ding/worker/replay_buffer/naive_buffer.py +++ b/ding/worker/replay_buffer/naive_buffer.py @@ -1,6 +1,9 @@ +import os import copy from typing import Union, Any, Optional, List import numpy as np +import math +import hickle from easydict import EasyDict from ding.worker.replay_buffer import IBuffer @@ -61,7 +64,7 @@ def __init__( # Point to the tail position where next data can be inserted, i.e. latest inserted data's next position. self._tail = 0 # Lock to guarantee thread safe - self._lock = LockContext(type_=LockContextType.THREAD_LOCK) + self._lock = LockContext(lock_type=LockContextType.THREAD_LOCK) self._end_flag = False self._enable_track_used_data = self._cfg.enable_track_used_data if self._enable_track_used_data: @@ -146,6 +149,15 @@ def sample(self, self._periodic_thruput_monitor.sample_data_count += len(sample_data) return sample_data + def save_data(self, file_name: str): + if not os.path.exists(os.path.dirname(file_name)): + if os.path.dirname(file_name) != "": + os.makedirs(os.path.dirname(file_name)) + hickle.dump(py_obj=self._data, file_obj=file_name) + + def load_data(self, file_name: str): + self.push(hickle.load(file_name), 0) + def _append(self, ori_data: Any, cur_collector_envstep: int = -1) -> None: r""" Overview: @@ -454,3 +466,100 @@ def _get_indices(self, size: int, sample_range: slice = None, replace: bool = Fa def update(self, envstep): self._current_buffer_size = self._set_buffer_size(envstep) + + +@BUFFER_REGISTRY.register('sequence') +class SequenceReplayBuffer(NaiveReplayBuffer): + r""" + Overview: + Interface: + start, close, push, update, sample, clear, count, state_dict, load_state_dict, default_config + Property: + replay_buffer_size, push_count + """ + + def sample( + self, + batch: int, + sequence: int, + cur_learner_iter: int, + sample_range: slice = None, + replace: bool = False + ) -> Optional[list]: + """ + Overview: + Sample data with length ``size``. + Arguments: + - size (:obj:`int`): The number of the data that will be sampled. + - sequence (:obj:`int`): The length of the sequence of a data that will be sampled. + - cur_learner_iter (:obj:`int`): Learner's current iteration. \ + Not used in naive buffer, but preserved for compatibility. + - sample_range (:obj:`slice`): Buffer slice for sampling, such as `slice(-10, None)`, which \ + means only sample among the last 10 data + - replace (:obj:`bool`): Whether sample with replacement + Returns: + - sample_data (:obj:`list`): A list of data with length ``size``. + """ + if batch == 0: + return [] + can_sample = self._sample_check(batch * sequence, replace) + if not can_sample: + return None + with self._lock: + indices = self._get_indices(batch, sequence, sample_range, replace) + sample_data = self._sample_with_indices(indices, sequence, cur_learner_iter) + self._periodic_thruput_monitor.sample_data_count += len(sample_data) + return sample_data + + def _get_indices(self, size: int, sequence: int, sample_range: slice = None, replace: bool = False) -> list: + r""" + Overview: + Get the sample index list. + Arguments: + - size (:obj:`int`): The number of the data that will be sampled + - sample_range (:obj:`slice`): Buffer slice for sampling, such as `slice(-10, None)`, which \ + means only sample among the last 10 data + Returns: + - index_list (:obj:`list`): A list including all the sample indices, whose length should equal to ``size``. + """ + assert self._valid_count <= self._replay_buffer_size + if self._valid_count == self._replay_buffer_size: + tail = self._replay_buffer_size + else: + tail = self._tail + episodes = math.ceil(self._valid_count / 500) + batch = 0 + indices = [] + if sample_range is None: + while batch < size: + episode = np.random.choice(episodes) + length = tail - episode * 500 if tail - episode * 500 < 500 else 500 + available = length - sequence + if available < 1: + continue + list(range(episode * 500, episode * 500 + available)) + indices.append(np.random.randint(episode * 500, episode * 500 + available + 1)) + batch += 1 + else: + raise NotImplementedError("sample_range is not implemented in this version") + return indices + + def _sample_with_indices(self, indices: List[int], sequence: int, cur_learner_iter: int) -> list: + r""" + Overview: + Sample data with ``indices``. + Arguments: + - indices (:obj:`List[int]`): A list including all the sample indices. + - cur_learner_iter (:obj:`int`): Not used in this method, but preserved for compatibility. + Returns: + - data (:obj:`list`) Sampled data. + """ + data = [] + for idx in indices: + assert self._data[idx] is not None, idx + if self._deepcopy: + copy_data = copy.deepcopy(self._data[idx:idx + sequence]) + else: + copy_data = self._data[idx:idx + sequence] + data.append(copy_data) + return data diff --git a/ding/worker/replay_buffer/tests/test_advanced_buffer.py b/ding/worker/replay_buffer/tests/test_advanced_buffer.py index b42dd0e14b..e6bb71a6b3 100644 --- a/ding/worker/replay_buffer/tests/test_advanced_buffer.py +++ b/ding/worker/replay_buffer/tests/test_advanced_buffer.py @@ -6,6 +6,7 @@ import os import pickle import time +import tempfile from ding.worker.replay_buffer import AdvancedReplayBuffer from ding.utils import deep_merge_dicts @@ -73,6 +74,28 @@ def test_push(self): assert advanced_buffer._next_unique_id == start_data_id + extend_num * i assert advanced_buffer._valid_count == min(start_data_id + extend_num * i, replay_buffer_size) + def test_save_and_load_data(self): + buffer_cfg = deep_merge_dicts(AdvancedReplayBuffer.default_config(), EasyDict(dict(replay_buffer_size=64))) + advanced_buffer = AdvancedReplayBuffer(buffer_cfg, tb_logger=None, instance_name='test') + start_pointer = advanced_buffer._tail + start_vaildlen = advanced_buffer.count() + start_data_id = advanced_buffer._next_unique_id + valid_count = 0 + for _ in range(100): + if advanced_buffer._data[advanced_buffer._tail] is None: + valid_count += 1 + advanced_buffer.push(generate_data(), 0) + assert (advanced_buffer.replay_buffer_size == 64) + assert (advanced_buffer.count() == 64 == start_vaildlen + valid_count) + with tempfile.TemporaryDirectory() as tmpdirname: + test_file = os.path.join(tmpdirname, "data.hkl") + advanced_buffer.save_data(test_file) + advanced_buffer_new = AdvancedReplayBuffer(buffer_cfg, instance_name='test_new') + advanced_buffer_new.load_data(test_file) + assert (advanced_buffer_new.replay_buffer_size == 64) + assert (advanced_buffer_new.count() == 64 == start_vaildlen + valid_count) + assert (advanced_buffer_new.push_count == 64) + def test_update(self): buffer_cfg = deep_merge_dicts(AdvancedReplayBuffer.default_config(), EasyDict(dict(replay_buffer_size=64))) advanced_buffer = AdvancedReplayBuffer(buffer_cfg, tb_logger=None, instance_name='test') diff --git a/ding/worker/replay_buffer/tests/test_naive_buffer.py b/ding/worker/replay_buffer/tests/test_naive_buffer.py index a4a9de345a..f0122669f9 100644 --- a/ding/worker/replay_buffer/tests/test_naive_buffer.py +++ b/ding/worker/replay_buffer/tests/test_naive_buffer.py @@ -2,6 +2,7 @@ from easydict import EasyDict import os import time +import tempfile from ding.worker.replay_buffer import NaiveReplayBuffer from ding.utils import deep_merge_dicts @@ -37,6 +38,29 @@ def test_push(self): naive_buffer.push(data, 0) assert naive_buffer._tail == (start_pointer + extend_num * i) % replay_buffer_size + def test_save_and_load_data(self): + buffer_cfg = deep_merge_dicts(NaiveReplayBuffer.default_config(), EasyDict(dict(replay_buffer_size=64))) + naive_buffer = NaiveReplayBuffer(buffer_cfg, instance_name='test') + start_pointer = naive_buffer._tail + start_vaildlen = naive_buffer.count() + valid_count = 0 + for _ in range(100): + if naive_buffer._data[naive_buffer._tail] is None: + valid_count += 1 + naive_buffer.push(generate_data(), 0) + assert (naive_buffer.replay_buffer_size == 64) + assert (naive_buffer.count() == 64 == start_vaildlen + valid_count) + assert (naive_buffer.push_count == start_vaildlen + 100) + assert (naive_buffer._tail == (start_pointer + 100) % naive_buffer.replay_buffer_size) + with tempfile.TemporaryDirectory() as tmpdirname: + test_file = os.path.join(tmpdirname, "data.hkl") + naive_buffer.save_data(test_file) + naive_buffer_new = NaiveReplayBuffer(buffer_cfg, instance_name='test_new') + naive_buffer_new.load_data(test_file) + assert (naive_buffer_new.replay_buffer_size == 64) + assert (naive_buffer_new.count() == 64 == start_vaildlen + valid_count) + assert (naive_buffer_new.push_count == 64) + def test_sample(self): buffer_cfg = deep_merge_dicts(NaiveReplayBuffer.default_config(), EasyDict(dict(replay_buffer_size=64))) naive_buffer = NaiveReplayBuffer(buffer_cfg, instance_name='test') diff --git a/ding/worker/replay_buffer/utils.py b/ding/worker/replay_buffer/utils.py index 07e2896085..7264db0115 100644 --- a/ding/worker/replay_buffer/utils.py +++ b/ding/worker/replay_buffer/utils.py @@ -226,8 +226,8 @@ def __init__(self, cfg) -> None: window_seconds = cfg.window_seconds self._decay_factor = 0.01 ** (1 / window_seconds) - self._push_lock = LockContext(type_=LockContextType.THREAD_LOCK) - self._sample_lock = LockContext(type_=LockContextType.THREAD_LOCK) + self._push_lock = LockContext(lock_type=LockContextType.THREAD_LOCK) + self._sample_lock = LockContext(lock_type=LockContextType.THREAD_LOCK) self._history_push_count = 0 self._history_sample_count = 0 diff --git a/ding/world_model/ddppo.py b/ding/world_model/ddppo.py index ceb91180f9..075c2aa63b 100644 --- a/ding/world_model/ddppo.py +++ b/ding/world_model/ddppo.py @@ -1,12 +1,12 @@ +from functools import partial +from ditk import logging import itertools +import copy import numpy as np import multiprocessing -import copy import torch -from torch import nn +import torch.nn as nn -from scipy.spatial import KDTree -from functools import partial from ding.utils import WORLD_MODEL_REGISTRY from ding.utils.data import default_collate from ding.torch_utils import unsqueeze_repeat @@ -27,6 +27,12 @@ def get_neighbor_index(data, k, serial=False): ret: [B, k] """ + try: + from scipy.spatial import KDTree + except ImportError: + import sys + logging.warning("Please install scipy first, such as `pip3 install scipy`.") + sys.exit(1) data = data.cpu().numpy() tree = KDTree(data) diff --git a/ding/world_model/dreamer.py b/ding/world_model/dreamer.py new file mode 100644 index 0000000000..63e65067de --- /dev/null +++ b/ding/world_model/dreamer.py @@ -0,0 +1,302 @@ +import numpy as np +import copy +import torch +from torch import nn + +from ding.utils import WORLD_MODEL_REGISTRY, lists_to_dicts +from ding.utils.data import default_collate +from ding.model import ConvEncoder, FCEncoder +from ding.world_model.base_world_model import WorldModel +from ding.world_model.model.networks import RSSM, ConvDecoder +from ding.torch_utils import to_device, one_hot +from ding.torch_utils.network.dreamer import DenseHead + + +@WORLD_MODEL_REGISTRY.register('dreamer') +class DREAMERWorldModel(WorldModel, nn.Module): + config = dict( + pretrain=100, + train_freq=2, + model=dict( + state_size=None, + action_size=None, + model_lr=1e-4, + reward_size=1, + hidden_size=200, + batch_size=256, + max_epochs_since_update=5, + dyn_stoch=32, + dyn_deter=512, + dyn_hidden=512, + dyn_input_layers=1, + dyn_output_layers=1, + dyn_rec_depth=1, + dyn_shared=False, + dyn_discrete=32, + act='SiLU', + norm='LayerNorm', + grad_heads=['image', 'reward', 'discount'], + units=512, + image_dec_layers=2, + reward_layers=2, + discount_layers=2, + value_layers=2, + actor_layers=2, + cnn_depth=32, + encoder_kernels=[4, 4, 4, 4], + decoder_kernels=[4, 4, 4, 4], + reward_head='twohot_symlog', + kl_lscale=0.1, + kl_rscale=0.5, + kl_free=1.0, + kl_forward=False, + pred_discount=True, + dyn_mean_act='none', + dyn_std_act='sigmoid2', + dyn_temp_post=True, + dyn_min_std=0.1, + dyn_cell='gru_layer_norm', + unimix_ratio=0.01, + device='cuda' if torch.cuda.is_available() else 'cpu', + obs_type='RGB', + action_type='continuous', + encoder_hidden_size_list=[64, 128, 128], + ), + ) + + def __init__(self, cfg, env, tb_logger): + WorldModel.__init__(self, cfg, env, tb_logger) + nn.Module.__init__(self) + + self.pretrain_flag = True + self._cfg = cfg.model + #self._cfg.act = getattr(torch.nn, self._cfg.act), + #self._cfg.norm = getattr(torch.nn, self._cfg.norm), + self._cfg.act = nn.modules.activation.SiLU # nn.SiLU + self._cfg.norm = nn.modules.normalization.LayerNorm # nn.LayerNorm + self.state_size = self._cfg.state_size + self.obs_type = self._cfg.obs_type + self.action_size = self._cfg.action_size + self.action_type = self._cfg.action_type + self.reward_size = self._cfg.reward_size + self.hidden_size = self._cfg.hidden_size + self.batch_size = self._cfg.batch_size + if self.obs_type == 'vector': + self.encoder = FCEncoder(self.state_size, self._cfg.encoder_hidden_size_list, activation=torch.nn.SiLU()) + self.embed_size = self._cfg.encoder_hidden_size_list[-1] + elif self.obs_type == 'RGB': + self.encoder = ConvEncoder( + self.state_size, + hidden_size_list=[32, 64, 128, 256, 4096], # to last layer 128? + activation=torch.nn.SiLU(), + kernel_size=self._cfg.encoder_kernels, + layer_norm=True + ) + self.embed_size = ( + (self.state_size[1] // 2 ** (len(self._cfg.encoder_kernels))) ** 2 * self._cfg.cnn_depth * + 2 ** (len(self._cfg.encoder_kernels) - 1) + ) + + self.dynamics = RSSM( + self._cfg.dyn_stoch, + self._cfg.dyn_deter, + self._cfg.dyn_hidden, + self._cfg.action_type, + self._cfg.dyn_input_layers, + self._cfg.dyn_output_layers, + self._cfg.dyn_rec_depth, + self._cfg.dyn_shared, + self._cfg.dyn_discrete, + self._cfg.act, + self._cfg.norm, + self._cfg.dyn_mean_act, + self._cfg.dyn_std_act, + self._cfg.dyn_temp_post, + self._cfg.dyn_min_std, + self._cfg.dyn_cell, + self._cfg.unimix_ratio, + self._cfg.action_size, + self.embed_size, + self._cfg.device, + ) + self.heads = nn.ModuleDict() + if self._cfg.dyn_discrete: + feat_size = self._cfg.dyn_stoch * self._cfg.dyn_discrete + self._cfg.dyn_deter + else: + feat_size = self._cfg.dyn_stoch + self._cfg.dyn_deter + + if isinstance(self.state_size, int): + self.heads['image'] = DenseHead( + feat_size, + (self.state_size, ), + self._cfg.image_dec_layers, + self._cfg.units, + 'SiLU', # self._cfg.act + 'LN', # self._cfg.norm + dist='binary', + outscale=0.0, + device=self._cfg.device, + ) + elif len(self.state_size) == 3: + self.heads["image"] = ConvDecoder( + feat_size, # pytorch version + self._cfg.cnn_depth, + self._cfg.act, + self._cfg.norm, + self.state_size, + self._cfg.decoder_kernels, + ) + self.heads["reward"] = DenseHead( + feat_size, # dyn_stoch * dyn_discrete + dyn_deter + (255, ), + self._cfg.reward_layers, + self._cfg.units, + 'SiLU', # self._cfg.act + 'LN', # self._cfg.norm + dist=self._cfg.reward_head, + outscale=0.0, + device=self._cfg.device, + ) + if self._cfg.pred_discount: + self.heads["discount"] = DenseHead( + feat_size, # pytorch version + [], + self._cfg.discount_layers, + self._cfg.units, + 'SiLU', # self._cfg.act + 'LN', # self._cfg.norm + dist="binary", + device=self._cfg.device, + ) + + if self._cuda: + self.cuda() + # to do + # grad_clip, weight_decay + self.optimizer = torch.optim.Adam(self.parameters(), lr=self._cfg.model_lr) + + def step(self, obs, act): + pass + + def eval(self, env_buffer, envstep, train_iter): + pass + + def should_pretrain(self): + if self.pretrain_flag: + self.pretrain_flag = False + return True + return False + + def train(self, env_buffer, envstep, train_iter, batch_size, batch_length): + self.last_train_step = envstep + data = env_buffer.sample( + batch_size, batch_length, train_iter + ) # [len=B, ele=[len=T, ele={dict_key: Tensor(any_dims)}]] + data = default_collate(data) # -> [len=T, ele={dict_key: Tensor(B, any_dims)}] + data = lists_to_dicts(data, recursive=True) # -> {some_key: T lists}, each list is [B, some_dim] + data = {k: torch.stack(data[k], dim=1) for k in data} # -> {dict_key: Tensor([B, T, any_dims])} + + data['discount'] = data.get('discount', 1.0 - data['done'].float()) + data['weight'] = data.get('weight', None) + if self.obs_type == 'RGB': + data['image'] = data['obs'] - 0.5 + else: + data['image'] = data['obs'] + if self.action_type == 'continuous': + data['action'] *= (1.0 / torch.clip(torch.abs(data['action']), min=1.0)) + else: + data['action'] = one_hot(data['action'], self.action_size) + data = to_device(data, self._cfg.device) + if len(data['reward'].shape) == 2: + data['reward'] = data['reward'].unsqueeze(-1) + if len(data['action'].shape) == 2: + data['action'] = data['action'].unsqueeze(-1) + if len(data['discount'].shape) == 2: + data['discount'] = data['discount'].unsqueeze(-1) + + self.requires_grad_(requires_grad=True) + + image = data['image'].reshape([-1] + list(data['image'].shape[2:])) + embed = self.encoder(image) + embed = embed.reshape(list(data['image'].shape[:2]) + [embed.shape[-1]]) + + post, prior = self.dynamics.observe(embed, data["action"]) + kl_loss, kl_value, loss_lhs, loss_rhs = self.dynamics.kl_loss( + post, prior, self._cfg.kl_forward, self._cfg.kl_free, self._cfg.kl_lscale, self._cfg.kl_rscale + ) + losses = {} + likes = {} + for name, head in self.heads.items(): + grad_head = name in self._cfg.grad_heads + feat = self.dynamics.get_feat(post) + feat = feat if grad_head else feat.detach() + pred = head(feat) + like = pred.log_prob(data[name]) + likes[name] = like + losses[name] = -torch.mean(like) + model_loss = sum(losses.values()) + kl_loss + + # ==================== + # world model update + # ==================== + self.optimizer.zero_grad() + model_loss.backward() + self.optimizer.step() + + self.requires_grad_(requires_grad=False) + # log + if self.tb_logger is not None: + for name, loss in losses.items(): + self.tb_logger.add_scalar(name + '_loss', loss.detach().cpu().numpy().item(), envstep) + self.tb_logger.add_scalar('kl_free', self._cfg.kl_free, envstep) + self.tb_logger.add_scalar('kl_lscale', self._cfg.kl_lscale, envstep) + self.tb_logger.add_scalar('kl_rscale', self._cfg.kl_rscale, envstep) + self.tb_logger.add_scalar('loss_lhs', loss_lhs.detach().cpu().numpy().item(), envstep) + self.tb_logger.add_scalar('loss_rhs', loss_rhs.detach().cpu().numpy().item(), envstep) + self.tb_logger.add_scalar('kl', torch.mean(kl_value).detach().cpu().numpy().item(), envstep) + + prior_ent = torch.mean(self.dynamics.get_dist(prior).entropy()).detach().cpu().numpy() + post_ent = torch.mean(self.dynamics.get_dist(post).entropy()).detach().cpu().numpy() + + self.tb_logger.add_scalar('prior_ent', prior_ent.item(), envstep) + self.tb_logger.add_scalar('post_ent', post_ent.item(), envstep) + + context = dict( + embed=embed, + feat=self.dynamics.get_feat(post), + kl=kl_value, + postent=self.dynamics.get_dist(post).entropy(), + ) + post = {k: v.detach() for k, v in post.items()} + return post, context + + def _save_states(self, ): + self._states = copy.deepcopy(self.state_dict()) + + def _save_state(self, id): + state_dict = self.state_dict() + for k, v in state_dict.items(): + if 'weight' in k or 'bias' in k: + self._states[k].data[id] = copy.deepcopy(v.data[id]) + + def _load_states(self): + self.load_state_dict(self._states) + + def _save_best(self, epoch, holdout_losses): + updated = False + for i in range(len(holdout_losses)): + current = holdout_losses[i] + _, best = self._snapshots[i] + improvement = (best - current) / best + if improvement > 0.01: + self._snapshots[i] = (epoch, current) + self._save_state(i) + # self._save_state(i) + updated = True + # improvement = (best - current) / best + + if updated: + self._epochs_since_update = 0 + else: + self._epochs_since_update += 1 + return self._epochs_since_update > self.max_epochs_since_update diff --git a/ding/world_model/model/__init__.py b/ding/world_model/model/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/ding/world_model/model/networks.py b/ding/world_model/model/networks.py new file mode 100644 index 0000000000..2e47cc6604 --- /dev/null +++ b/ding/world_model/model/networks.py @@ -0,0 +1,403 @@ +import math +import numpy as np +from typing import Optional, Dict, Union, List + +import torch +from torch import nn +import torch.nn.functional as F +from torch import distributions as torchd +from ding.utils import SequenceType +from ding.torch_utils.network.dreamer import weight_init, uniform_weight_init, static_scan, \ + OneHotDist, ContDist, SymlogDist, DreamerLayerNorm + + +class RSSM(nn.Module): + + def __init__( + self, + stoch=30, + deter=200, + hidden=200, + action_type=None, + layers_input=1, + layers_output=1, + rec_depth=1, + shared=False, + discrete=False, + act=nn.ELU, + norm=nn.LayerNorm, + mean_act="none", + std_act="softplus", + temp_post=True, + min_std=0.1, + cell="gru", + unimix_ratio=0.01, + num_actions=None, + embed=None, + device=None, + ): + super(RSSM, self).__init__() + self._stoch = stoch + self._deter = deter + self._hidden = hidden + self._action_type = action_type + self._min_std = min_std + self._layers_input = layers_input + self._layers_output = layers_output + self._rec_depth = rec_depth + self._shared = shared + self._discrete = discrete + self._act = act + self._norm = norm + self._mean_act = mean_act + self._std_act = std_act + self._temp_post = temp_post + self._unimix_ratio = unimix_ratio + self._embed = embed + self._device = device + + inp_layers = [] + if self._discrete: + inp_dim = self._stoch * self._discrete + num_actions + else: + inp_dim = self._stoch + num_actions + if self._shared: + inp_dim += self._embed + for i in range(self._layers_input): + inp_layers.append(nn.Linear(inp_dim, self._hidden, bias=False)) + inp_layers.append(self._norm(self._hidden, eps=1e-03)) + inp_layers.append(self._act()) + if i == 0: + inp_dim = self._hidden + self._inp_layers = nn.Sequential(*inp_layers) + self._inp_layers.apply(weight_init) + + if cell == "gru": + self._cell = GRUCell(self._hidden, self._deter) + self._cell.apply(weight_init) + elif cell == "gru_layer_norm": + self._cell = GRUCell(self._hidden, self._deter, norm=True) + self._cell.apply(weight_init) + else: + raise NotImplementedError(cell) + + img_out_layers = [] + inp_dim = self._deter + for i in range(self._layers_output): + img_out_layers.append(nn.Linear(inp_dim, self._hidden, bias=False)) + img_out_layers.append(self._norm(self._hidden, eps=1e-03)) + img_out_layers.append(self._act()) + if i == 0: + inp_dim = self._hidden + self._img_out_layers = nn.Sequential(*img_out_layers) + self._img_out_layers.apply(weight_init) + + obs_out_layers = [] + if self._temp_post: + inp_dim = self._deter + self._embed + else: + inp_dim = self._embed + for i in range(self._layers_output): + obs_out_layers.append(nn.Linear(inp_dim, self._hidden, bias=False)) + obs_out_layers.append(self._norm(self._hidden, eps=1e-03)) + obs_out_layers.append(self._act()) + if i == 0: + inp_dim = self._hidden + self._obs_out_layers = nn.Sequential(*obs_out_layers) + self._obs_out_layers.apply(weight_init) + + if self._discrete: + self._ims_stat_layer = nn.Linear(self._hidden, self._stoch * self._discrete) + self._ims_stat_layer.apply(weight_init) + self._obs_stat_layer = nn.Linear(self._hidden, self._stoch * self._discrete) + self._obs_stat_layer.apply(weight_init) + else: + self._ims_stat_layer = nn.Linear(self._hidden, 2 * self._stoch) + self._ims_stat_layer.apply(weight_init) + self._obs_stat_layer = nn.Linear(self._hidden, 2 * self._stoch) + self._obs_stat_layer.apply(weight_init) + + def initial(self, batch_size): + deter = torch.zeros(batch_size, self._deter).to(self._device) + if self._discrete: + state = dict( + logit=torch.zeros([batch_size, self._stoch, self._discrete]).to(self._device), + stoch=torch.zeros([batch_size, self._stoch, self._discrete]).to(self._device), + deter=deter, + ) + else: + state = dict( + mean=torch.zeros([batch_size, self._stoch]).to(self._device), + std=torch.zeros([batch_size, self._stoch]).to(self._device), + stoch=torch.zeros([batch_size, self._stoch]).to(self._device), + deter=deter, + ) + return state + + def observe(self, embed, action, state=None): + swap = lambda x: x.permute([1, 0] + list(range(2, len(x.shape)))) # 交换前两维 + if state is None: + state = self.initial(action.shape[0]) # {logit, stoch, deter} + # (batch, time, ch) -> (time, batch, ch) + embed, action = swap(embed), swap(action) + post, prior = static_scan( + lambda prev_state, prev_act, embed: self.obs_step(prev_state[0], prev_act, embed), + (action, embed), + (state, state), + ) + + # (time, batch, stoch, discrete_num) -> (batch, time, stoch, discrete_num) + post = {k: swap(v) for k, v in post.items()} + prior = {k: swap(v) for k, v in prior.items()} + return post, prior + + def imagine(self, action, state=None): + swap = lambda x: x.permute([1, 0] + list(range(2, len(x.shape)))) + if state is None: + state = self.initial(action.shape[0]) + assert isinstance(state, dict), state + action = action + action = swap(action) + prior = static_scan(self.img_step, [action], state) + prior = prior[0] + prior = {k: swap(v) for k, v in prior.items()} + return prior + + def get_feat(self, state): + stoch = state["stoch"] + if self._discrete: + shape = list(stoch.shape[:-2]) + [self._stoch * self._discrete] + stoch = stoch.reshape(shape) + return torch.cat([stoch, state["deter"]], -1) + + def get_dist(self, state, dtype=None): + if self._discrete: + logit = state["logit"] + dist = torchd.independent.Independent(OneHotDist(logit, unimix_ratio=self._unimix_ratio), 1) + else: + mean, std = state["mean"], state["std"] + dist = ContDist(torchd.independent.Independent(torchd.normal.Normal(mean, std), 1)) + return dist + + def obs_step(self, prev_state, prev_action, embed, sample=True): + # if shared is True, prior and post both use same networks(inp_layers, _img_out_layers, _ims_stat_layer) + # otherwise, post use different network(_obs_out_layers) with prior[deter] and embed as inputs + if self._action_type == 'continuous': + prev_action *= (1.0 / torch.clip(torch.abs(prev_action), min=1.0)).detach() + prior = self.img_step(prev_state, prev_action, None, sample) + if self._shared: + post = self.img_step(prev_state, prev_action, embed, sample) + else: + if self._temp_post: + x = torch.cat([prior["deter"], embed], -1) + else: + x = embed + # (batch_size, prior_deter + embed) -> (batch_size, hidden) + x = self._obs_out_layers(x) + # (batch_size, hidden) -> (batch_size, stoch, discrete_num) + stats = self._suff_stats_layer("obs", x) + if sample: + stoch = self.get_dist(stats).sample() + else: + stoch = self.get_dist(stats).mode() + post = {"stoch": stoch, "deter": prior["deter"], **stats} + return post, prior + + # this is used for making future image + def img_step(self, prev_state, prev_action, embed=None, sample=True): + # (batch, stoch, discrete_num) + if self._action_type == 'continuous': + prev_action *= (1.0 / torch.clip(torch.abs(prev_action), min=1.0)).detach() + prev_stoch = prev_state["stoch"] + if self._discrete: + shape = list(prev_stoch.shape[:-2]) + [self._stoch * self._discrete] + # (batch, stoch, discrete_num) -> (batch, stoch * discrete_num) + prev_stoch = prev_stoch.reshape(shape) + if self._shared: + if embed is None: + shape = list(prev_action.shape[:-1]) + [self._embed] + embed = torch.zeros(shape) + # (batch, stoch * discrete_num) -> (batch, stoch * discrete_num + action, embed) + x = torch.cat([prev_stoch, prev_action, embed], -1) + else: + x = torch.cat([prev_stoch, prev_action], -1) + # (batch, stoch * discrete_num + action, embed) -> (batch, hidden) + x = self._inp_layers(x) + for _ in range(self._rec_depth): # rec depth is not correctly implemented + deter = prev_state["deter"] + # (batch, hidden), (batch, deter) -> (batch, deter), (batch, deter) + x, deter = self._cell(x, [deter]) + deter = deter[0] # Keras wraps the state in a list. + # (batch, deter) -> (batch, hidden) + x = self._img_out_layers(x) + # (batch, hidden) -> (batch_size, stoch, discrete_num) + stats = self._suff_stats_layer("ims", x) + if sample: + stoch = self.get_dist(stats).sample() + else: + stoch = self.get_dist(stats).mode() + prior = {"stoch": stoch, "deter": deter, **stats} # {stoch, deter, logit} + return prior + + def _suff_stats_layer(self, name, x): + if self._discrete: + if name == "ims": + x = self._ims_stat_layer(x) + elif name == "obs": + x = self._obs_stat_layer(x) + else: + raise NotImplementedError + logit = x.reshape(list(x.shape[:-1]) + [self._stoch, self._discrete]) + return {"logit": logit} + else: + if name == "ims": + x = self._ims_stat_layer(x) + elif name == "obs": + x = self._obs_stat_layer(x) + else: + raise NotImplementedError + mean, std = torch.split(x, [self._stoch] * 2, -1) + mean = { + "none": lambda: mean, + "tanh5": lambda: 5.0 * torch.tanh(mean / 5.0), + }[self._mean_act]() + std = { + "softplus": lambda: torch.softplus(std), + "abs": lambda: torch.abs(std + 1), + "sigmoid": lambda: torch.sigmoid(std), + "sigmoid2": lambda: 2 * torch.sigmoid(std / 2), + }[self._std_act]() + std = std + self._min_std + return {"mean": mean, "std": std} + + def kl_loss(self, post, prior, forward, free, lscale, rscale): + kld = torchd.kl.kl_divergence + dist = lambda x: self.get_dist(x) + sg = lambda x: {k: v.detach() for k, v in x.items()} + # forward == false -> (post, prior) + lhs, rhs = (prior, post) if forward else (post, prior) + + # forward == false -> Lrep + value_lhs = value = kld( + dist(lhs) if self._discrete else dist(lhs)._dist, + dist(sg(rhs)) if self._discrete else dist(sg(rhs))._dist, + ) + # forward == false -> Ldyn + value_rhs = kld( + dist(sg(lhs)) if self._discrete else dist(sg(lhs))._dist, + dist(rhs) if self._discrete else dist(rhs)._dist, + ) + # free bits + loss_lhs = torch.mean(torch.clip(value_lhs, min=free)) + loss_rhs = torch.mean(torch.clip(value_rhs, min=free)) + loss = lscale * loss_lhs + rscale * loss_rhs + + return loss, value, loss_lhs, loss_rhs + + +class ConvDecoder(nn.Module): + + def __init__( + self, + inp_depth, # config.dyn_stoch * config.dyn_discrete + config.dyn_deter + depth=32, + act=nn.ELU, + norm=nn.LayerNorm, + shape=(3, 64, 64), + kernels=(3, 3, 3, 3), + outscale=1.0, + ): + super(ConvDecoder, self).__init__() + self._inp_depth = inp_depth + self._act = act + self._norm = norm + self._depth = depth + self._shape = shape + self._kernels = kernels + self._embed_size = ((64 // 2 ** (len(kernels))) ** 2 * depth * 2 ** (len(kernels) - 1)) + + self._linear_layer = nn.Linear(inp_depth, self._embed_size) + inp_dim = self._embed_size // 16 # 除以最后的4*4 feature map来得到channel数 + + layers = [] + h, w = 4, 4 + for i, kernel in enumerate(self._kernels): + depth = self._embed_size // 16 // (2 ** (i + 1)) + act = self._act + bias = False + initializer = weight_init + if i == len(self._kernels) - 1: + depth = self._shape[0] + act = False + bias = True + norm = False + initializer = uniform_weight_init(outscale) + + if i != 0: + inp_dim = 2 ** (len(self._kernels) - (i - 1) - 2) * self._depth + pad_h, outpad_h = self.calc_same_pad(k=kernel, s=2, d=1) + pad_w, outpad_w = self.calc_same_pad(k=kernel, s=2, d=1) + layers.append( + nn.ConvTranspose2d( + inp_dim, + depth, + kernel, + 2, + padding=(pad_h, pad_w), + output_padding=(outpad_h, outpad_w), + bias=bias, + ) + ) + if norm: + layers.append(DreamerLayerNorm(depth)) + if act: + layers.append(act()) + [m.apply(initializer) for m in layers[-3:]] + h, w = h * 2, w * 2 + + self.layers = nn.Sequential(*layers) + + def calc_same_pad(self, k, s, d): + val = d * (k - 1) - s + 1 + pad = math.ceil(val / 2) + outpad = pad * 2 - val + return pad, outpad + + def __call__(self, features): + x = self._linear_layer(features) # feature:[batch, time, stoch*discrete + deter] + x = x.reshape([-1, 4, 4, self._embed_size // 16]) + x = x.permute(0, 3, 1, 2) + x = self.layers(x) + mean = x.reshape(list(features.shape[:-1]) + self._shape) + #mean = mean.permute(0, 1, 3, 4, 2) + return SymlogDist(mean) + + +class GRUCell(nn.Module): + + def __init__(self, inp_size, size, norm=False, act=torch.tanh, update_bias=-1): + super(GRUCell, self).__init__() + self._inp_size = inp_size # hidden + self._size = size # deter + self._act = act + self._norm = norm + self._update_bias = update_bias + self._layer = nn.Linear(inp_size + size, 3 * size, bias=False) + if norm: + self._norm = nn.LayerNorm(3 * size, eps=1e-03) + + @property + def state_size(self): + return self._size + + def forward(self, inputs, state): + state = state[0] # Keras wraps the state in a list. + parts = self._layer(torch.cat([inputs, state], -1)) + if self._norm: + parts = self._norm(parts) + reset, cand, update = torch.split(parts, [self._size] * 3, -1) + reset = torch.sigmoid(reset) + cand = self._act(reset * cand) + update = torch.sigmoid(update + self._update_bias) + output = update * cand + (1 - update) * state + return output, [output] diff --git a/ding/world_model/model/tests/test_networks.py b/ding/world_model/model/tests/test_networks.py new file mode 100644 index 0000000000..c23c94cd3d --- /dev/null +++ b/ding/world_model/model/tests/test_networks.py @@ -0,0 +1,3 @@ +import pytest +import torch +from itertools import product diff --git a/ding/world_model/tests/test_dreamerv3.py b/ding/world_model/tests/test_dreamerv3.py new file mode 100644 index 0000000000..f93673d39b --- /dev/null +++ b/ding/world_model/tests/test_dreamerv3.py @@ -0,0 +1,32 @@ +import pytest +import torch + +from itertools import product +from easydict import EasyDict +from ding.world_model.dreamer import DREAMERWorldModel +from ding.utils import deep_merge_dicts + +# arguments +state_size = [[3, 64, 64]] +action_size = [6, 1] +args = list(product(*[state_size, action_size])) + + +@pytest.mark.unittest +class TestDREAMER: + + def get_world_model(self, state_size, action_size): + cfg = DREAMERWorldModel.default_config() + cfg.model.max_epochs_since_update = 0 + cfg = deep_merge_dicts( + cfg, dict(cuda=False, model=dict(state_size=state_size, action_size=action_size, reward_size=1)) + ) + fake_env = EasyDict(termination_fn=lambda obs: torch.zeros_like(obs.sum(-1)).bool()) + return DREAMERWorldModel(cfg, fake_env, None) + + @pytest.mark.parametrize('state_size, action_size', args) + def test_train(self, state_size, action_size): + states = torch.rand(1280, *state_size) + actions = torch.rand(1280, action_size) + + model = self.get_world_model(state_size, action_size) diff --git a/ding/world_model/tests/test_world_model.py b/ding/world_model/tests/test_world_model.py index a8ed671d09..f8dd620c59 100644 --- a/ding/world_model/tests/test_world_model.py +++ b/ding/world_model/tests/test_world_model.py @@ -1,4 +1,5 @@ import pytest +import os import torch from easydict import EasyDict from ding.world_model.base_world_model import DreamWorldModel, DynaWorldModel @@ -10,8 +11,8 @@ class TestDynaWorldModel: @pytest.mark.parametrize('buffer_type', [NaiveReplayBuffer, EpisodeReplayBuffer]) def test_fill_img_buffer(self, buffer_type): - env_buffer = buffer_type(buffer_type.default_config(), None, 'exp_name', 'env_buffer_for_test') - img_buffer = buffer_type(buffer_type.default_config(), None, 'exp_name', 'img_buffer_for_test') + env_buffer = buffer_type(buffer_type.default_config(), None, 'dyna_exp_name', 'env_buffer_for_test') + img_buffer = buffer_type(buffer_type.default_config(), None, 'dyna_exp_name', 'img_buffer_for_test') fake_config = EasyDict( train_freq=250, # w.r.t environment step eval_freq=250, # w.r.t environment step @@ -51,11 +52,11 @@ def step(self, obs, action): return (torch.zeros(B), torch.rand(B, O), obs.sum(-1) > 0) from ding.policy import SACPolicy - from ding.model import QAC + from ding.model import ContinuousQAC policy_config = SACPolicy.default_config() policy_config.model.update(dict(obs_shape=2, action_shape=2)) - model = QAC(**policy_config.model) + model = ContinuousQAC(**policy_config.model) policy = SACPolicy(policy_config, model=model).collect_mode fake_model = FakeModel(fake_config, None, None) @@ -74,6 +75,7 @@ def step(self, obs, action): ) super(FakeModel, fake_model).fill_img_buffer(policy, env_buffer, img_buffer, 0, 0) + os.popen("rm -rf dyna_exp_name") @pytest.mark.unittest diff --git a/dizoo/atari/config/parallel/qbert_dqn_config.py b/dizoo/atari/config/parallel/qbert_dqn_config.py deleted file mode 100644 index fb8ac8e67f..0000000000 --- a/dizoo/atari/config/parallel/qbert_dqn_config.py +++ /dev/null @@ -1,100 +0,0 @@ -from easydict import EasyDict - -qbert_dqn_config = dict( - exp_name='qbert_dqn', - env=dict( - collector_env_num=16, - collector_episode_num=2, - evaluator_env_num=8, - evaluator_episode_num=1, - stop_value=30000, - env_id='Qbert-v4', - #'ALE/Qbert-v5' is available. But special setting is needed after gym make. - frame_stack=4, - manager=dict(shared_memory=True, ), - ), - policy=dict( - cuda=True, - priority=True, - model=dict( - obs_shape=[4, 84, 84], - action_shape=6, - encoder_hidden_size_list=[128, 128, 512], - ), - nstep=3, - discount_factor=0.99, - learn=dict( - batch_size=32, - learning_rate=0.0001, - learner=dict( - learner_num=1, - send_policy_freq=1, - ), - ), - collect=dict( - n_sample=32, - collector=dict( - collector_num=2, - update_policy_second=3, - ), - ), - eval=dict(evaluator=dict(eval_freq=500, )), - other=dict( - eps=dict( - type='exp', - start=1., - end=0.05, - decay=250000, - ), - replay_buffer=dict( - replay_buffer_size=400000, - enable_track_used_data=True, - ), - commander=dict( - collector_task_space=2, - learner_task_space=1, - eval_interval=300, - ), - ), - ), -) -qbert_dqn_config = EasyDict(qbert_dqn_config) -main_config = qbert_dqn_config - -qbert_dqn_create_config = dict( - env=dict( - type='atari', - import_names=['dizoo.atari.envs.atari_env'], - ), - env_manager=dict(type='subprocess'), - policy=dict(type='dqn_command'), - learner=dict(type='base', import_names=['ding.worker.learner.base_learner']), - collector=dict( - type='zergling', - import_names=['ding.worker.collector.zergling_parallel_collector'], - ), - commander=dict( - type='solo', - import_names=['ding.worker.coordinator.solo_parallel_commander'], - ), - comm_learner=dict( - type='flask_fs', - import_names=['ding.worker.learner.comm.flask_fs_learner'], - ), - comm_collector=dict( - type='flask_fs', - import_names=['ding.worker.collector.comm.flask_fs_collector'], - ), -) -qbert_dqn_create_config = EasyDict(qbert_dqn_create_config) -create_config = qbert_dqn_create_config - -qbert_dqn_system_config = dict( - coordinator=dict(), - path_data='./{}/data'.format(main_config.exp_name), - path_policy='./{}/policy'.format(main_config.exp_name), - communication_mode='auto', - learner_gpu_num=1, -) -qbert_dqn_system_config = EasyDict(qbert_dqn_system_config) -system_config = qbert_dqn_system_config diff --git a/dizoo/atari/config/parallel/qbert_dqn_config_k8s.py b/dizoo/atari/config/parallel/qbert_dqn_config_k8s.py deleted file mode 100644 index aa9abc58ca..0000000000 --- a/dizoo/atari/config/parallel/qbert_dqn_config_k8s.py +++ /dev/null @@ -1,117 +0,0 @@ -from easydict import EasyDict - -qbert_dqn_config = dict( - exp_name='qbert_dqn', - env=dict( - collector_env_num=16, - collector_episode_num=2, - evaluator_env_num=8, - evaluator_episode_num=1, - stop_value=30000, - env_id='Qbert-v4', - #'ALE/Qbert-v5' is available. But special setting is needed after gym make. - frame_stack=4, - manager=dict(shared_memory=False, ), - ), - policy=dict( - cuda=False, - priority=True, - model=dict( - obs_shape=[4, 84, 84], - action_shape=6, - encoder_hidden_size_list=[128, 128, 512], - ), - nstep=3, - discount_factor=0.99, - learn=dict( - batch_size=32, - learning_rate=0.0001, - learner=dict( - learner_num=1, - send_policy_freq=1, - ), - ), - collect=dict( - n_sample=16, - collector=dict( - collector_num=2, - update_policy_second=3, - ), - ), - eval=dict(evaluator=dict(eval_freq=500, )), - other=dict( - eps=dict( - type='exp', - start=1., - end=0.05, - decay=250000, - ), - replay_buffer=dict( - replay_buffer_size=400000, - enable_track_used_data=True, - ), - commander=dict( - # increase collector task space when get rs from server - collector_task_space=0, - learner_task_space=1, - eval_interval=30, - ), - ), - ), -) -qbert_dqn_config = EasyDict(qbert_dqn_config) -main_config = qbert_dqn_config - -qbert_dqn_create_config = dict( - env=dict( - type='atari', - import_names=['dizoo.atari.envs.atari_env'], - ), - env_manager=dict(type='subprocess'), - policy=dict(type='dqn_command'), - learner=dict(type='base', import_names=['ding.worker.learner.base_learner']), - collector=dict( - type='zergling', - import_names=['ding.worker.collector.zergling_parallel_collector'], - ), - commander=dict( - type='solo', - import_names=['ding.worker.coordinator.solo_parallel_commander'], - ), - comm_learner=dict( - type='flask_fs', - import_names=['ding.worker.learner.comm.flask_fs_learner'], - ), - comm_collector=dict( - type='flask_fs', - import_names=['ding.worker.collector.comm.flask_fs_collector'], - ), -) -qbert_dqn_create_config = EasyDict(qbert_dqn_create_config) -create_config = qbert_dqn_create_config - -qbert_dqn_system_config = dict( - coordinator=dict( - operator_server=dict( - system_addr='di-server.di-system:8080', - api_version='/v1alpha1', - init_replicas_request=dict( - collectors={ - "replicas": 2, - }, - learners={ - "gpus": "0", - "replicas": 1, - }, - ), - collector_target_num=2, - learner_target_num=1, - ), - ), - path_data='./{}/data'.format(main_config.exp_name), - path_policy='./{}/policy'.format(main_config.exp_name), - communication_mode='auto', - learner_gpu_num=1, -) -qbert_dqn_system_config = EasyDict(qbert_dqn_system_config) -system_config = qbert_dqn_system_config diff --git a/dizoo/atari/config/serial/__init__.py b/dizoo/atari/config/serial/__init__.py index 8ed40685cc..ce0ea3b805 100644 --- a/dizoo/atari/config/serial/__init__.py +++ b/dizoo/atari/config/serial/__init__.py @@ -1,4 +1,5 @@ from dizoo.atari.config.serial.enduro import * from dizoo.atari.config.serial.pong import * from dizoo.atari.config.serial.qbert import * -from dizoo.atari.config.serial.spaceinvaders import * \ No newline at end of file +from dizoo.atari.config.serial.spaceinvaders import * +from dizoo.atari.config.serial.asterix import * diff --git a/dizoo/atari/config/serial/asterix/__init__.py b/dizoo/atari/config/serial/asterix/__init__.py new file mode 100644 index 0000000000..d25b637d56 --- /dev/null +++ b/dizoo/atari/config/serial/asterix/__init__.py @@ -0,0 +1 @@ +from .asterix_mdqn_config import asterix_mdqn_config, asterix_mdqn_create_config diff --git a/dizoo/atari/config/serial/asterix/asterix_mdqn_config.py b/dizoo/atari/config/serial/asterix/asterix_mdqn_config.py new file mode 100644 index 0000000000..675ae32540 --- /dev/null +++ b/dizoo/atari/config/serial/asterix/asterix_mdqn_config.py @@ -0,0 +1,63 @@ +from copy import deepcopy +from easydict import EasyDict + +asterix_mdqn_config = dict( + exp_name='asterix_mdqn_seed0', + env=dict( + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=20000, + env_id='AsterixNoFrameskip-v0', + #'ALE/SpaceInvaders-v5' is available. But special setting is needed after gym make. + frame_stack=4, + ), + policy=dict( + cuda=True, + priority=False, + model=dict( + obs_shape=[4, 84, 84], + action_shape=9, + encoder_hidden_size_list=[128, 128, 512], + ), + nstep=1, + discount_factor=0.99, + entropy_tau=0.03, + m_alpha=0.9, + learn=dict( + update_per_collect=10, + batch_size=32, + learning_rate=0.0001, + target_update_freq=500, + learner=dict(hook=dict(save_ckpt_after_iter=1000000, )) + ), + collect=dict(n_sample=100, ), + eval=dict(evaluator=dict(eval_freq=4000, )), + other=dict( + eps=dict( + type='exp', + start=1., + end=0.05, + decay=1000000, + ), + replay_buffer=dict(replay_buffer_size=400000, ), + ), + ), +) +asterix_mdqn_config = EasyDict(asterix_mdqn_config) +main_config = asterix_mdqn_config +asterix_mdqn_create_config = dict( + env=dict( + type='atari', + import_names=['dizoo.atari.envs.atari_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='mdqn'), +) +asterix_mdqn_create_config = EasyDict(asterix_mdqn_create_config) +create_config = asterix_mdqn_create_config + +if __name__ == '__main__': + # or you can enter ding -m serial -c asterix_mdqn_config.py -s 0 + from ding.entry import serial_pipeline + serial_pipeline((main_config, create_config), seed=0, max_env_step=int(1e7), dynamic_seed=False) diff --git a/dizoo/atari/config/serial/demon_attack/demon_attack_dqn_config.py b/dizoo/atari/config/serial/demon_attack/demon_attack_dqn_config.py new file mode 100644 index 0000000000..9669bd1392 --- /dev/null +++ b/dizoo/atari/config/serial/demon_attack/demon_attack_dqn_config.py @@ -0,0 +1,60 @@ +from easydict import EasyDict + +demon_attack_dqn_config = dict( + exp_name='DemonAttack_dqn_seed0', + env=dict( + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=1e6, + env_id='DemonAttackNoFrameskip-v4', + frame_stack=4, + ), + policy=dict( + cuda=True, + priority=False, + model=dict( + obs_shape=[4, 84, 84], + action_shape=6, + encoder_hidden_size_list=[128, 128, 512], + noise=True, + ), + nstep=3, + discount_factor=0.99, + learn=dict( + update_per_collect=10, + batch_size=32, + learning_rate=0.0001, + target_update_freq=500, + ), + noisy_net=True, + collect=dict(n_sample=96), + eval=dict(evaluator=dict(eval_freq=4000, )), + other=dict( + eps=dict( + type='exp', + start=1., + end=0.05, + decay=250000, + ), + replay_buffer=dict(replay_buffer_size=100000, ), + ), + ), +) +demon_attack_dqn_config = EasyDict(demon_attack_dqn_config) +main_config = demon_attack_dqn_config +demon_attack_dqn_create_config = dict( + env=dict( + type='atari', + import_names=['dizoo.atari.envs.atari_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='dqn'), +) +demon_attack_dqn_create_config = EasyDict(demon_attack_dqn_create_config) +create_config = demon_attack_dqn_create_config + +if __name__ == '__main__': + # or you can enter `ding -m serial -c demon_attack_dqn_config.py -s 0` + from ding.entry import serial_pipeline + serial_pipeline((main_config, create_config), seed=0, max_env_step=int(10e6)) diff --git a/dizoo/atari/config/serial/enduro/enduro_dqn_config.py b/dizoo/atari/config/serial/enduro/enduro_dqn_config.py index baa5d654d6..eb22c20737 100644 --- a/dizoo/atari/config/serial/enduro/enduro_dqn_config.py +++ b/dizoo/atari/config/serial/enduro/enduro_dqn_config.py @@ -8,7 +8,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=int(1e10), - env_id='Enduro-v4', + env_id='EnduroNoFrameskip-v4', #'ALE/Enduro-v5' is available. But special setting is needed after gym make. frame_stack=4, ), diff --git a/dizoo/atari/config/serial/enduro/enduro_impala_config.py b/dizoo/atari/config/serial/enduro/enduro_impala_config.py index 989bb58dbc..2f13770aad 100644 --- a/dizoo/atari/config/serial/enduro/enduro_impala_config.py +++ b/dizoo/atari/config/serial/enduro/enduro_impala_config.py @@ -8,7 +8,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=10000000000, - env_id='Enduro-v4', + env_id='EnduroNoFrameskip-v4', #'ALE/Enduro-v5' is available. But special setting is needed after gym make. frame_stack=4 ), diff --git a/dizoo/atari/config/serial/enduro/enduro_mdqn_config.py b/dizoo/atari/config/serial/enduro/enduro_mdqn_config.py new file mode 100644 index 0000000000..7b08f4ace9 --- /dev/null +++ b/dizoo/atari/config/serial/enduro/enduro_mdqn_config.py @@ -0,0 +1,63 @@ +from copy import deepcopy +from easydict import EasyDict + +enduro_mdqn_config = dict( + exp_name='enduro_mdqn_seed0', + env=dict( + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=int(1e10), + env_id='EnduroNoFrameskip-v4', + #'ALE/Enduro-v5' is available. But special setting is needed after gym make. + frame_stack=4, + ), + policy=dict( + cuda=True, + priority=False, + model=dict( + obs_shape=[4, 84, 84], + action_shape=9, + encoder_hidden_size_list=[128, 128, 512], + ), + nstep=1, + discount_factor=0.99, + entropy_tau=0.03, + m_alpha=0.9, + learn=dict( + update_per_collect=10, + batch_size=32, + learning_rate=0.0001, + target_update_freq=500, + learner=dict(hook=dict(save_ckpt_after_iter=1000000, )) + ), + collect=dict(n_sample=100, ), + eval=dict(evaluator=dict(eval_freq=4000, )), + other=dict( + eps=dict( + type='exp', + start=0.5, + end=0.05, + decay=1000000, + ), + replay_buffer=dict(replay_buffer_size=400000, ), + ), + ), +) +enduro_mdqn_config = EasyDict(enduro_mdqn_config) +main_config = enduro_mdqn_config +enduro_mdqn_create_config = dict( + env=dict( + type='atari', + import_names=['dizoo.atari.envs.atari_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='mdqn'), +) +enduro_mdqn_create_config = EasyDict(enduro_mdqn_create_config) +create_config = enduro_mdqn_create_config + +if __name__ == '__main__': + # or you can enter ding -m serial -c enduro_mdqn_config.py -s 0 + from ding.entry import serial_pipeline + serial_pipeline((main_config, create_config), seed=0, max_env_step=int(1e7), dynamic_seed=False) diff --git a/dizoo/atari/config/serial/enduro/enduro_onppo_config.py b/dizoo/atari/config/serial/enduro/enduro_onppo_config.py index 58708a50ff..8eee796274 100644 --- a/dizoo/atari/config/serial/enduro/enduro_onppo_config.py +++ b/dizoo/atari/config/serial/enduro/enduro_onppo_config.py @@ -3,49 +3,56 @@ enduro_onppo_config = dict( exp_name='enduro_onppo_seed0', env=dict( - collector_env_num=64, + collector_env_num=8, evaluator_env_num=8, n_evaluator_episode=8, stop_value=10000000000, - env_id='Enduro-v4', + env_id='EnduroNoFrameskip-v4', #'ALE/Enduro-v5' is available. But special setting is needed after gym make. frame_stack=4, manager=dict(shared_memory=False, ) ), policy=dict( cuda=True, + recompute_adv=True, + action_space='discrete', model=dict( obs_shape=[4, 84, 84], action_shape=9, - encoder_hidden_size_list=[32, 64, 64, 128], - actor_head_hidden_size=128, - critic_head_hidden_size=128, - critic_head_layer_num=2, + action_space='discrete', + encoder_hidden_size_list=[32, 64, 64, 512], + actor_head_layer_num=0, + critic_head_layer_num=0, + actor_head_hidden_size=512, + critic_head_hidden_size=512, ), learn=dict( - update_per_collect=24, - batch_size=128, - # (bool) Whether to normalize advantage. Default to False. - adv_norm=False, - learning_rate=0.0001, - # (float) loss weight of the value network, the weight of policy network is set to 1 - value_weight=1.0, - # (float) loss weight of the entropy regularization, the weight of policy network is set to 1 - entropy_weight=0.001, # [0.1, 0.01 ,0.0] - clip_ratio=0.1 + lr_scheduler=dict(epoch_num=5200, min_lr_lambda=0), + epoch_per_collect=4, + batch_size=256, + learning_rate=2.5e-4, + value_weight=0.5, + entropy_weight=0.01, + clip_ratio=0.1, + adv_norm=True, + value_norm=True, + # for onppo, when we recompute adv, we need the key done in data to split traj, so we must + # use ignore_done=False here, + # but when we add key traj_flag in data as the backup for key done, we could choose to use ignore_done=True + # for halfcheetah, the length=1000 + ignore_done=False, + grad_clip_type='clip_norm', + grad_clip_value=0.5, ), collect=dict( # (int) collect n_sample data, train model n_iteration times n_sample=1024, + unroll_len=1, # (float) the trade-off factor lambda to balance 1step td and mc gae_lambda=0.95, discount_factor=0.99, ), - eval=dict(evaluator=dict(eval_freq=1000, )), - other=dict(replay_buffer=dict( - replay_buffer_size=10000, - max_use=3, - ), ), + eval=dict(evaluator=dict(eval_freq=5000, )), ), ) main_config = EasyDict(enduro_onppo_config) diff --git a/dizoo/atari/config/serial/enduro/enduro_qrdqn_config.py b/dizoo/atari/config/serial/enduro/enduro_qrdqn_config.py index 8fc0beac15..ce0409fedb 100644 --- a/dizoo/atari/config/serial/enduro/enduro_qrdqn_config.py +++ b/dizoo/atari/config/serial/enduro/enduro_qrdqn_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=10000000000, - env_id='Enduro-v4', + env_id='EnduroNoFrameskip-v4', #'ALE/Enduro-v5' is available. But special setting is needed after gym make. frame_stack=4, ), diff --git a/dizoo/atari/config/serial/enduro/enduro_rainbow_config.py b/dizoo/atari/config/serial/enduro/enduro_rainbow_config.py index 71997e46c5..0735ceeabc 100644 --- a/dizoo/atari/config/serial/enduro/enduro_rainbow_config.py +++ b/dizoo/atari/config/serial/enduro/enduro_rainbow_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=10000000000, - env_id='Enduro-v4', + env_id='EnduroNoFrameskip-v4', #'ALE/Enduro-v5' is available. But special setting is needed after gym make. frame_stack=4 ), diff --git a/dizoo/atari/config/serial/montezuma/montezuma_ngu_config.py b/dizoo/atari/config/serial/montezuma/montezuma_ngu_config.py index b8dfec3526..1a3bc12594 100644 --- a/dizoo/atari/config/serial/montezuma/montezuma_ngu_config.py +++ b/dizoo/atari/config/serial/montezuma/montezuma_ngu_config.py @@ -3,13 +3,15 @@ collector_env_num = 8 evaluator_env_num = 8 nstep = 5 -montezuma_ppo_rnd_config = dict( +max_env_step = int(10e6) + +montezuma_ngu_config = dict( exp_name='montezuma_ngu_seed0', env=dict( collector_env_num=collector_env_num, evaluator_env_num=evaluator_env_num, n_evaluator_episode=8, - env_id='MontezumaRevenge-v4', + env_id='MontezumaRevengeNoFrameskip-v4', #'ALE/MontezumaRevenge-v5' is available. But special setting is needed after gym make. obs_plus_prev_action_reward=True, # use specific env wrapper for ngu policy stop_value=int(1e5), @@ -107,9 +109,9 @@ ), ), ) -montezuma_ppo_rnd_config = EasyDict(montezuma_ppo_rnd_config) -main_config = montezuma_ppo_rnd_config -montezuma_ppo_rnd_create_config = dict( +montezuma_ngu_config = EasyDict(montezuma_ngu_config) +main_config = montezuma_ngu_config +montezuma_ngu_create_config = dict( env=dict( type='atari', import_names=['dizoo.atari.envs.atari_env'], @@ -119,9 +121,9 @@ rnd_reward_model=dict(type='rnd-ngu'), episodic_reward_model=dict(type='episodic'), ) -montezuma_ppo_rnd_create_config = EasyDict(montezuma_ppo_rnd_create_config) -create_config = montezuma_ppo_rnd_create_config +montezuma_ngu_create_config = EasyDict(montezuma_ngu_create_config) +create_config = montezuma_ngu_create_config if __name__ == "__main__": from ding.entry import serial_pipeline_reward_model_ngu - serial_pipeline_reward_model_ngu([main_config, create_config], seed=0) + serial_pipeline_reward_model_ngu([main_config, create_config], seed=0, max_env_step=max_env_step) diff --git a/dizoo/atari/config/serial/phoenix/phoenix_fqf_config.py b/dizoo/atari/config/serial/phoenix/phoenix_fqf_config.py index 467557fe68..736a7a6cff 100644 --- a/dizoo/atari/config/serial/phoenix/phoenix_fqf_config.py +++ b/dizoo/atari/config/serial/phoenix/phoenix_fqf_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=10000000000, - env_id='Phoenix-v4', + env_id='PhoenixNoFrameskip-v4', #'ALE/Phoenix-v5' is available. But special setting is needed after gym make. frame_stack=4, ), @@ -59,4 +59,4 @@ if __name__ == '__main__': # or you can enter `ding -m serial -c phoenix_fqf_config.py -s 0` from ding.entry import serial_pipeline - serial_pipeline((main_config, create_config), seed=0) \ No newline at end of file + serial_pipeline((main_config, create_config), seed=0) diff --git a/dizoo/atari/config/serial/phoenix/phoenix_iqn_config.py b/dizoo/atari/config/serial/phoenix/phoenix_iqn_config.py index 324f676663..4e2a2036af 100644 --- a/dizoo/atari/config/serial/phoenix/phoenix_iqn_config.py +++ b/dizoo/atari/config/serial/phoenix/phoenix_iqn_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=20, - env_id='Phoenix-v4', + env_id='PhoenixNoFrameskip-v4', #'ALE/Phoenix-v5' is available. But special setting is needed after gym make. frame_stack=4, ), diff --git a/dizoo/atari/config/serial/pitfall/pitfall_ngu_config.py b/dizoo/atari/config/serial/pitfall/pitfall_ngu_config.py index 55f275333a..8e22563adc 100644 --- a/dizoo/atari/config/serial/pitfall/pitfall_ngu_config.py +++ b/dizoo/atari/config/serial/pitfall/pitfall_ngu_config.py @@ -3,20 +3,18 @@ collector_env_num = 32 evaluator_env_num = 5 nstep = 5 +max_env_step = int(10e6) -pitfall_ppo_rnd_config = dict( +pitfall_ngu_config = dict( # Note: # 1. at least 1e10 timesteps, i.e., 10000 million, the reward may increase, please be patient. # 2. the larger unroll_lenth and replay buffer size may have better results, but also require more memory. - # exp_name='debug_pitfall_ngu_ul298_er01_n32_rlbs2e4', - # exp_name='debug_pitfall_ngu_ul98_er01_n32_rlbs2e4', - # exp_name='debug_pitfall_ngu_ul40_er01_n32_rlbs2e4', exp_name='pitfall_ngu_seed0', env=dict( collector_env_num=collector_env_num, evaluator_env_num=evaluator_env_num, n_evaluator_episode=5, - env_id='Pitfall-v4', + env_id='PitfallNoFrameskip-v4', #'ALE/Pitfall-v5' is available. But special setting is needed after gym make. obs_plus_prev_action_reward=True, # use specific env wrapper for ngu policy stop_value=int(1e5), @@ -114,9 +112,9 @@ ), ), ) -pitfall_ppo_rnd_config = EasyDict(pitfall_ppo_rnd_config) -main_config = pitfall_ppo_rnd_config -pitfall_ppo_rnd_create_config = dict( +pitfall_ngu_config = EasyDict(pitfall_ngu_config) +main_config = pitfall_ngu_config +pitfall_ngu_create_config = dict( env=dict( type='atari', import_names=['dizoo.atari.envs.atari_env'], @@ -126,9 +124,9 @@ rnd_reward_model=dict(type='rnd-ngu'), episodic_reward_model=dict(type='episodic'), ) -pitfall_ppo_rnd_create_config = EasyDict(pitfall_ppo_rnd_create_config) -create_config = pitfall_ppo_rnd_create_config +pitfall_ngu_create_config = EasyDict(pitfall_ngu_create_config) +create_config = pitfall_ngu_create_config if __name__ == "__main__": from ding.entry import serial_pipeline_reward_model_ngu - serial_pipeline_reward_model_ngu([main_config, create_config], seed=0) + serial_pipeline_reward_model_ngu([main_config, create_config], seed=0, max_env_step=max_env_step) diff --git a/dizoo/atari/config/serial/pong/__init__.py b/dizoo/atari/config/serial/pong/__init__.py index 5ce3db9a5b..93cba0f3ac 100644 --- a/dizoo/atari/config/serial/pong/__init__.py +++ b/dizoo/atari/config/serial/pong/__init__.py @@ -1,3 +1,3 @@ from .pong_dqn_config import pong_dqn_config, pong_dqn_create_config from .pong_dqn_envpool_config import pong_dqn_envpool_config, pong_dqn_envpool_create_config -from .pong_dqfd_config import pong_dqfd_config, pong_dqfd_create_config \ No newline at end of file +from .pong_dqfd_config import pong_dqfd_config, pong_dqfd_create_config diff --git a/dizoo/atari/config/serial/pong/pong_a2c_config.py b/dizoo/atari/config/serial/pong/pong_a2c_config.py index 50fa251591..eb48f4248e 100644 --- a/dizoo/atari/config/serial/pong/pong_a2c_config.py +++ b/dizoo/atari/config/serial/pong/pong_a2c_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=20, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), @@ -51,11 +51,12 @@ ), env_manager=dict(type='subprocess'), policy=dict(type='a2c'), + replay_buffer=dict(type='naive'), ) pong_a2c_create_config = EasyDict(pong_a2c_create_config) create_config = pong_a2c_create_config if __name__ == '__main__': - # or you can enter `ding -m serial -c pong_a2c_config.py -s 0` - from ding.entry import serial_pipeline - serial_pipeline((main_config, create_config), seed=0) + # or you can enter `ding -m serial_onpolicy -c pong_a2c_config.py -s 0` + from ding.entry import serial_pipeline_onpolicy + serial_pipeline_onpolicy((main_config, create_config), seed=0) diff --git a/dizoo/atari/config/serial/pong/pong_acer_config.py b/dizoo/atari/config/serial/pong/pong_acer_config.py index c0a7edc652..706bb83b3a 100644 --- a/dizoo/atari/config/serial/pong/pong_acer_config.py +++ b/dizoo/atari/config/serial/pong/pong_acer_config.py @@ -6,7 +6,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=20, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), diff --git a/dizoo/atari/config/serial/pong/pong_c51_config.py b/dizoo/atari/config/serial/pong/pong_c51_config.py index 944618f52f..4d9f77fa7b 100644 --- a/dizoo/atari/config/serial/pong/pong_c51_config.py +++ b/dizoo/atari/config/serial/pong/pong_c51_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=20, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), diff --git a/dizoo/atari/config/serial/pong/pong_cql_config.py b/dizoo/atari/config/serial/pong/pong_cql_config.py index ee443e6885..7290863716 100644 --- a/dizoo/atari/config/serial/pong/pong_cql_config.py +++ b/dizoo/atari/config/serial/pong/pong_cql_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=20, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), diff --git a/dizoo/atari/config/serial/pong/pong_dqfd_config.py b/dizoo/atari/config/serial/pong/pong_dqfd_config.py index 926784da33..dde918bd26 100644 --- a/dizoo/atari/config/serial/pong/pong_dqfd_config.py +++ b/dizoo/atari/config/serial/pong/pong_dqfd_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=20, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), diff --git a/dizoo/atari/config/serial/pong/pong_dqn_config.py b/dizoo/atari/config/serial/pong/pong_dqn_config.py index f2bc3e9f6c..1e5f4040f1 100644 --- a/dizoo/atari/config/serial/pong/pong_dqn_config.py +++ b/dizoo/atari/config/serial/pong/pong_dqn_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=20, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), diff --git a/dizoo/atari/config/serial/pong/pong_dqn_ddp_config.py b/dizoo/atari/config/serial/pong/pong_dqn_ddp_config.py new file mode 100644 index 0000000000..ce468a4e0c --- /dev/null +++ b/dizoo/atari/config/serial/pong/pong_dqn_ddp_config.py @@ -0,0 +1,67 @@ +from easydict import EasyDict + +pong_dqn_config = dict( + exp_name='data_pong/pong_dqn_ddp_seed0', + env=dict( + collector_env_num=4, + evaluator_env_num=4, + n_evaluator_episode=8, + stop_value=20, + env_id='PongNoFrameskip-v4', + #'ALE/Pong-v5' is available. But special setting is needed after gym make. + frame_stack=4, + ), + policy=dict( + multi_gpu=True, + cuda=True, + priority=False, + model=dict( + obs_shape=[4, 84, 84], + action_shape=6, + encoder_hidden_size_list=[128, 128, 512], + ), + nstep=3, + discount_factor=0.99, + learn=dict( + update_per_collect=10, + batch_size=32, + learning_rate=0.0001, + target_update_freq=500, + ), + collect=dict(n_sample=96, ), + eval=dict(evaluator=dict(eval_freq=4000, )), + other=dict( + eps=dict( + type='exp', + start=1., + end=0.05, + decay=250000, + ), + replay_buffer=dict(replay_buffer_size=100000, ), + ), + ), +) +pong_dqn_config = EasyDict(pong_dqn_config) +main_config = pong_dqn_config +pong_dqn_create_config = dict( + env=dict( + type='atari', + import_names=['dizoo.atari.envs.atari_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='dqn'), +) +pong_dqn_create_config = EasyDict(pong_dqn_create_config) +create_config = pong_dqn_create_config + +if __name__ == '__main__': + """ + Overview: + This script should be executed with GPUs. + Run the following command to launch the script: + python -m torch.distributed.launch --nproc_per_node=2 ./dizoo/atari/config/serial/pong/pong_dqn_ddp_config.py + """ + from ding.utils import DDPContext + from ding.entry import serial_pipeline + with DDPContext(): + serial_pipeline((main_config, create_config), seed=0, max_env_step=int(3e6)) diff --git a/dizoo/atari/config/serial/pong/pong_dqn_envpool_config.py b/dizoo/atari/config/serial/pong/pong_dqn_envpool_config.py index baa866f48e..0b80e41548 100644 --- a/dizoo/atari/config/serial/pong/pong_dqn_envpool_config.py +++ b/dizoo/atari/config/serial/pong/pong_dqn_envpool_config.py @@ -9,7 +9,7 @@ evaluator_batch_size=8, n_evaluator_episode=8, stop_value=20, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), diff --git a/dizoo/atari/config/serial/pong/pong_dqn_multi_gpu_config.py b/dizoo/atari/config/serial/pong/pong_dqn_multi_gpu_config.py index f528a21476..91b9de260a 100644 --- a/dizoo/atari/config/serial/pong/pong_dqn_multi_gpu_config.py +++ b/dizoo/atari/config/serial/pong/pong_dqn_multi_gpu_config.py @@ -7,13 +7,14 @@ evaluator_env_num=4, n_evaluator_episode=8, stop_value=20, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), policy=dict( cuda=True, priority=False, + multi_gpu=True, model=dict( obs_shape=[4, 84, 84], action_shape=6, @@ -22,7 +23,6 @@ nstep=3, discount_factor=0.99, learn=dict( - multi_gpu=True, update_per_collect=10, batch_size=32, learning_rate=0.0001, diff --git a/dizoo/atari/config/serial/pong/pong_dqn_render_config.py b/dizoo/atari/config/serial/pong/pong_dqn_render_config.py index ce5e5d23eb..b12ed17331 100644 --- a/dizoo/atari/config/serial/pong/pong_dqn_render_config.py +++ b/dizoo/atari/config/serial/pong/pong_dqn_render_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=20, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), diff --git a/dizoo/atari/config/serial/pong/pong_dqn_stdim_config.py b/dizoo/atari/config/serial/pong/pong_dqn_stdim_config.py index 78b54f56dd..7d568ab3bb 100644 --- a/dizoo/atari/config/serial/pong/pong_dqn_stdim_config.py +++ b/dizoo/atari/config/serial/pong/pong_dqn_stdim_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=20, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), diff --git a/dizoo/atari/config/serial/pong/pong_dt_config.py b/dizoo/atari/config/serial/pong/pong_dt_config.py new file mode 100644 index 0000000000..60d795ec29 --- /dev/null +++ b/dizoo/atari/config/serial/pong/pong_dt_config.py @@ -0,0 +1,68 @@ +from easydict import EasyDict +from copy import deepcopy + +Pong_dt_config = dict( + exp_name='dt_log/atari/Pong/Pong_dt_seed0', + env=dict( + env_id='PongNoFrameskip-v4', + collector_env_num=1, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=20, + frame_stack=4, + is_train=False, + episode_num=10000, # stop in breakout + ), + dataset=dict( + env_type='atari', + num_steps=500000, + # num_steps=50, + num_buffers=50, + rtg_scale=None, + context_len=30, + data_dir_prefix='/mnt/nfs/luyd/d4rl_atari/Pong', + trajectories_per_buffer=10, + ), + policy=dict( + cuda=True, + multi_gpu=True, + stop_value=20, + evaluator_env_num=8, + rtg_target=20, # max target return to go + max_eval_ep_len=10000, # max lenght of one episode + wt_decay=1e-4, + clip_grad_norm_p=1.0, + weight_decay=0.1, + warmup_steps=10000, + model=dict( + state_dim=(4, 84, 84), + act_dim=6, + n_blocks=6, + h_dim=128, + context_len=30, + n_heads=8, + drop_p=0.1, + continuous=False, + ), + batch_size=128, + learning_rate=6e-4, + eval=dict(evaluator=dict(eval_freq=100, ), ), + collect=dict( + data_type='d4rl_trajectory', + unroll_len=1, + ), + ), +) + +Pong_dt_config = EasyDict(Pong_dt_config) +main_config = Pong_dt_config +Pong_dt_create_config = dict( + env=dict( + type='atari', + import_names=['dizoo.atari.envs.atari_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='dt'), +) +Pong_dt_create_config = EasyDict(Pong_dt_create_config) +create_config = Pong_dt_create_config diff --git a/dizoo/atari/config/serial/pong/pong_fqf_config.py b/dizoo/atari/config/serial/pong/pong_fqf_config.py index a7e8733f21..e1dbc346e9 100644 --- a/dizoo/atari/config/serial/pong/pong_fqf_config.py +++ b/dizoo/atari/config/serial/pong/pong_fqf_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=10000000000, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), @@ -59,4 +59,4 @@ if __name__ == '__main__': # or you can enter `ding -m serial -c pong_fqf_config.py -s 0` from ding.entry import serial_pipeline - serial_pipeline((main_config, create_config), seed=0) \ No newline at end of file + serial_pipeline((main_config, create_config), seed=0) diff --git a/dizoo/atari/config/serial/pong/pong_gail_dqn_config.py b/dizoo/atari/config/serial/pong/pong_gail_dqn_config.py index d014681a65..201ea414ff 100644 --- a/dizoo/atari/config/serial/pong/pong_gail_dqn_config.py +++ b/dizoo/atari/config/serial/pong/pong_gail_dqn_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=20, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), @@ -87,4 +87,4 @@ max_env_step=1000000, seed=0, collect_data=True - ) \ No newline at end of file + ) diff --git a/dizoo/atari/config/serial/pong/pong_impala_config.py b/dizoo/atari/config/serial/pong/pong_impala_config.py index aa7121b1a7..c8a4136bf9 100644 --- a/dizoo/atari/config/serial/pong/pong_impala_config.py +++ b/dizoo/atari/config/serial/pong/pong_impala_config.py @@ -1,13 +1,13 @@ from easydict import EasyDict pong_impala_config = dict( - exp_name='pong_impala_seed0', + exp_name='impala_log/pong_impala_seed0', env=dict( - collector_env_num=8, + collector_env_num=12, evaluator_env_num=8, n_evaluator_episode=8, - stop_value=20, - env_id='Pong-v4', + stop_value=21, + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), @@ -19,27 +19,29 @@ model=dict( obs_shape=[4, 84, 84], action_shape=6, - encoder_hidden_size_list=[128, 128, 512], - critic_head_hidden_size=512, + encoder_hidden_size_list=[64, 128, 256], + critic_head_hidden_size=256, critic_head_layer_num=2, - actor_head_hidden_size=512, + actor_head_hidden_size=256, actor_head_layer_num=2, + # impala_cnn_encoder=True, ), learn=dict( # (int) collect n_sample data, train model update_per_collect times # here we follow impala serial pipeline - update_per_collect=10, + update_per_collect=2, # (int) the number of data for a train iteration batch_size=128, + # optim_type='rmsprop', grad_clip_type='clip_norm', clip_value=0.5, - learning_rate=0.0003, + learning_rate=0.0006, # (float) loss weight of the value network, the weight of policy network is set to 1 value_weight=0.5, # (float) loss weight of the entropy regularization, the weight of policy network is set to 1 entropy_weight=0.01, # (float) discount factor for future reward, defaults int [0, 1] - discount_factor=0.9, + discount_factor=0.99, # (float) additional discounting parameter lambda_=0.95, # (float) clip ratio of importance weights @@ -54,8 +56,8 @@ n_sample=16, collector=dict(collect_print_freq=1000, ), ), - eval=dict(evaluator=dict(eval_freq=5000, )), - other=dict(replay_buffer=dict(replay_buffer_size=10000, ), ), + eval=dict(evaluator=dict(eval_freq=2000, )), + other=dict(replay_buffer=dict(replay_buffer_size=10000, sliced=False), ), ), ) main_config = EasyDict(pong_impala_config) diff --git a/dizoo/atari/config/serial/pong/pong_iqn_config.py b/dizoo/atari/config/serial/pong/pong_iqn_config.py index 416207b3d7..5eed786a19 100644 --- a/dizoo/atari/config/serial/pong/pong_iqn_config.py +++ b/dizoo/atari/config/serial/pong/pong_iqn_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=20, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), diff --git a/dizoo/atari/config/serial/pong/pong_ngu_config.py b/dizoo/atari/config/serial/pong/pong_ngu_config.py index 231f37d907..215913d20d 100644 --- a/dizoo/atari/config/serial/pong/pong_ngu_config.py +++ b/dizoo/atari/config/serial/pong/pong_ngu_config.py @@ -3,13 +3,15 @@ collector_env_num = 8 evaluator_env_num = 8 nstep = 5 -pong_ppo_rnd_config = dict( +max_env_step = int(10e6) + +pong_ngu_config = dict( exp_name='pong_ngu_seed0', env=dict( collector_env_num=collector_env_num, evaluator_env_num=evaluator_env_num, n_evaluator_episode=evaluator_env_num, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. obs_plus_prev_action_reward=True, # use specific env wrapper for ngu policy stop_value=20, @@ -61,7 +63,7 @@ priority_IS_weight=True, discount_factor=0.997, nstep=nstep, - burnin_step=2, + burnin_step=20, # (int) is the total length of [sequence sample] minus # the length of burnin part in [sequence sample], # i.e., = = + @@ -107,9 +109,9 @@ ), ), ) -pong_ppo_rnd_config = EasyDict(pong_ppo_rnd_config) -main_config = pong_ppo_rnd_config -pong_ppo_rnd_create_config = dict( +pong_ngu_config = EasyDict(pong_ngu_config) +main_config = pong_ngu_config +pong_ngu_create_config = dict( env=dict( type='atari', import_names=['dizoo.atari.envs.atari_env'], @@ -119,10 +121,10 @@ rnd_reward_model=dict(type='rnd-ngu'), episodic_reward_model=dict(type='episodic'), ) -pong_ppo_rnd_create_config = EasyDict(pong_ppo_rnd_create_config) -create_config = pong_ppo_rnd_create_config +pong_ngu_create_config = EasyDict(pong_ngu_create_config) +create_config = pong_ngu_create_config if __name__ == "__main__": # or you can enter `ding -m serial_ngu -c pong_ngu_config.py -s 0` from ding.entry import serial_pipeline_ngu - serial_pipeline_ngu([main_config, create_config], seed=0) + serial_pipeline_ngu([main_config, create_config], seed=0, max_env_step=max_env_step) diff --git a/dizoo/atari/config/serial/pong/pong_ppg_config.py b/dizoo/atari/config/serial/pong/pong_ppg_config.py index 7c3f43f04a..13bc7ed448 100644 --- a/dizoo/atari/config/serial/pong/pong_ppg_config.py +++ b/dizoo/atari/config/serial/pong/pong_ppg_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=20, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), diff --git a/dizoo/atari/config/serial/pong/pong_onppo_config.py b/dizoo/atari/config/serial/pong/pong_ppo_config.py similarity index 74% rename from dizoo/atari/config/serial/pong/pong_onppo_config.py rename to dizoo/atari/config/serial/pong/pong_ppo_config.py index 430a07aee7..0d7ae3ed78 100644 --- a/dizoo/atari/config/serial/pong/pong_onppo_config.py +++ b/dizoo/atari/config/serial/pong/pong_ppo_config.py @@ -1,12 +1,13 @@ from easydict import EasyDict -pong_onppo_config = dict( +pong_ppo_config = dict( + exp_name='pong_ppo_seed0', env=dict( collector_env_num=8, evaluator_env_num=8, n_evaluator_episode=8, stop_value=20, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), @@ -39,6 +40,12 @@ ignore_done=False, grad_clip_type='clip_norm', grad_clip_value=0.5, + # KL divergence regularization between current policy and pretrained policy. + # Supported KL divergence estimators: ['k1', 'k2', 'k3']. + # KL divergence loss will be calculated only when pretrained_model_path is provided. + kl_beta=0.01, + kl_type='k1', + pretrained_model_path=None, ), collect=dict( n_sample=3200, @@ -49,9 +56,9 @@ eval=dict(evaluator=dict(eval_freq=5000, )), ), ) -main_config = EasyDict(pong_onppo_config) +main_config = EasyDict(pong_ppo_config) -pong_onppo_create_config = dict( +pong_ppo_create_config = dict( env=dict( type='atari', import_names=['dizoo.atari.envs.atari_env'], @@ -59,9 +66,9 @@ env_manager=dict(type='subprocess'), policy=dict(type='ppo'), ) -create_config = EasyDict(pong_onppo_create_config) +create_config = EasyDict(pong_ppo_create_config) if __name__ == "__main__": - # or you can enter `ding -m serial_onpolicy -c pong_onppo_config.py -s 0` + # or you can enter `ding -m serial_onpolicy -c pong_ppo_config.py -s 0` from ding.entry import serial_pipeline_onpolicy serial_pipeline_onpolicy((main_config, create_config), seed=0) diff --git a/dizoo/atari/config/serial/pong/pong_ppo_ddp_config.py b/dizoo/atari/config/serial/pong/pong_ppo_ddp_config.py new file mode 100644 index 0000000000..be4200de0f --- /dev/null +++ b/dizoo/atari/config/serial/pong/pong_ppo_ddp_config.py @@ -0,0 +1,76 @@ +from easydict import EasyDict + +pong_ppo_config = dict( + exp_name='data_pong/pong_ppo_ddp_seed0', + env=dict( + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=20, + env_id='PongNoFrameskip-v4', + #'ALE/Pong-v5' is available. But special setting is needed after gym make. + frame_stack=4, + ), + policy=dict( + multi_gpu=True, + cuda=True, + recompute_adv=True, + action_space='discrete', + model=dict( + obs_shape=[4, 84, 84], + action_shape=6, + action_space='discrete', + encoder_hidden_size_list=[64, 64, 128], + actor_head_hidden_size=128, + critic_head_hidden_size=128, + ), + learn=dict( + epoch_per_collect=10, + update_per_collect=1, + batch_size=320, + learning_rate=3e-4, + value_weight=0.5, + entropy_weight=0.001, + clip_ratio=0.2, + adv_norm=True, + value_norm=True, + # for ppo, when we recompute adv, we need the key done in data to split traj, so we must + # use ignore_done=False here, + # but when we add key traj_flag in data as the backup for key done, we could choose to use ignore_done=True + # for halfcheetah, the length=1000 + ignore_done=False, + grad_clip_type='clip_norm', + grad_clip_value=0.5, + ), + collect=dict( + n_sample=3200, + unroll_len=1, + discount_factor=0.99, + gae_lambda=0.95, + ), + eval=dict(evaluator=dict(eval_freq=1000, )), + ), +) +main_config = EasyDict(pong_ppo_config) + +pong_ppo_create_config = dict( + env=dict( + type='atari', + import_names=['dizoo.atari.envs.atari_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='ppo'), +) +create_config = EasyDict(pong_ppo_create_config) + +if __name__ == "__main__": + """ + Overview: + This script should be executed with GPUs. + Run the following command to launch the script: + python -m torch.distributed.launch --nproc_per_node=2 ./dizoo/atari/config/serial/pong/pong_ppo_ddp_config.py + """ + from ding.utils import DDPContext + from ding.entry import serial_pipeline_onpolicy + with DDPContext(): + serial_pipeline_onpolicy((main_config, create_config), seed=0, max_env_step=int(3e6)) diff --git a/dizoo/atari/config/serial/pong/pong_qrdqn_config.py b/dizoo/atari/config/serial/pong/pong_qrdqn_config.py index a779883844..df0ad1ab9d 100644 --- a/dizoo/atari/config/serial/pong/pong_qrdqn_config.py +++ b/dizoo/atari/config/serial/pong/pong_qrdqn_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=20, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), diff --git a/dizoo/atari/config/serial/pong/pong_qrdqn_generation_data_config.py b/dizoo/atari/config/serial/pong/pong_qrdqn_generation_data_config.py index 678b5eef91..32a0346a37 100644 --- a/dizoo/atari/config/serial/pong/pong_qrdqn_generation_data_config.py +++ b/dizoo/atari/config/serial/pong/pong_qrdqn_generation_data_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=20, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), diff --git a/dizoo/atari/config/serial/pong/pong_r2d2_config.py b/dizoo/atari/config/serial/pong/pong_r2d2_config.py index fe3954999c..6160332317 100644 --- a/dizoo/atari/config/serial/pong/pong_r2d2_config.py +++ b/dizoo/atari/config/serial/pong/pong_r2d2_config.py @@ -9,7 +9,7 @@ evaluator_env_num=evaluator_env_num, n_evaluator_episode=evaluator_env_num, stop_value=20, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), diff --git a/dizoo/atari/config/serial/pong/pong_r2d2_gtrxl_config.py b/dizoo/atari/config/serial/pong/pong_r2d2_gtrxl_config.py index 98d23d42e3..860cce4e4c 100644 --- a/dizoo/atari/config/serial/pong/pong_r2d2_gtrxl_config.py +++ b/dizoo/atari/config/serial/pong/pong_r2d2_gtrxl_config.py @@ -9,7 +9,7 @@ evaluator_env_num=evaluator_env_num, n_evaluator_episode=5, stop_value=20, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), diff --git a/dizoo/atari/config/serial/pong/pong_r2d2_residual_config.py b/dizoo/atari/config/serial/pong/pong_r2d2_residual_config.py index bd81825211..db3411102f 100644 --- a/dizoo/atari/config/serial/pong/pong_r2d2_residual_config.py +++ b/dizoo/atari/config/serial/pong/pong_r2d2_residual_config.py @@ -9,7 +9,7 @@ evaluator_env_num=evaluator_env_num, n_evaluator_episode=evaluator_env_num, stop_value=20, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), diff --git a/dizoo/atari/config/serial/pong/pong_r2d3_offppoexpert_config.py b/dizoo/atari/config/serial/pong/pong_r2d3_offppoexpert_config.py index 944079ff36..a61221da87 100644 --- a/dizoo/atari/config/serial/pong/pong_r2d3_offppoexpert_config.py +++ b/dizoo/atari/config/serial/pong/pong_r2d3_offppoexpert_config.py @@ -13,7 +13,7 @@ evaluator_env_num=evaluator_env_num, n_evaluator_episode=evaluator_env_num, stop_value=20, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), @@ -101,7 +101,7 @@ evaluator_env_num=evaluator_env_num, n_evaluator_episode=evaluator_env_num, stop_value=20, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), diff --git a/dizoo/atari/config/serial/pong/pong_r2d3_r2d2expert_config.py b/dizoo/atari/config/serial/pong/pong_r2d3_r2d2expert_config.py index 2108da92df..64d0bd06fc 100644 --- a/dizoo/atari/config/serial/pong/pong_r2d3_r2d2expert_config.py +++ b/dizoo/atari/config/serial/pong/pong_r2d3_r2d2expert_config.py @@ -13,7 +13,7 @@ evaluator_env_num=evaluator_env_num, n_evaluator_episode=evaluator_env_num, stop_value=20, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), @@ -105,7 +105,7 @@ evaluator_env_num=evaluator_env_num, n_evaluator_episode=evaluator_env_num, stop_value=20, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), diff --git a/dizoo/atari/config/serial/pong/pong_rainbow_config.py b/dizoo/atari/config/serial/pong/pong_rainbow_config.py index 2b877c6633..f7b403d57a 100644 --- a/dizoo/atari/config/serial/pong/pong_rainbow_config.py +++ b/dizoo/atari/config/serial/pong/pong_rainbow_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=20, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), diff --git a/dizoo/atari/config/serial/pong/pong_sqil_config.py b/dizoo/atari/config/serial/pong/pong_sqil_config.py index 4d95805b39..3d562795d5 100644 --- a/dizoo/atari/config/serial/pong/pong_sqil_config.py +++ b/dizoo/atari/config/serial/pong/pong_sqil_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=20, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), diff --git a/dizoo/atari/config/serial/pong/pong_sql_config.py b/dizoo/atari/config/serial/pong/pong_sql_config.py index 8bd3bd8973..1bc2ab2ee0 100644 --- a/dizoo/atari/config/serial/pong/pong_sql_config.py +++ b/dizoo/atari/config/serial/pong/pong_sql_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=20, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), diff --git a/dizoo/atari/config/serial/pong/pong_trex_offppo_config.py b/dizoo/atari/config/serial/pong/pong_trex_offppo_config.py index 743fa0571a..3351931380 100644 --- a/dizoo/atari/config/serial/pong/pong_trex_offppo_config.py +++ b/dizoo/atari/config/serial/pong/pong_trex_offppo_config.py @@ -7,7 +7,7 @@ evaluator_env_num=4, n_evaluator_episode=8, stop_value=20, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), diff --git a/dizoo/atari/config/serial/pong/pong_trex_sql_config.py b/dizoo/atari/config/serial/pong/pong_trex_sql_config.py index a1d89af0d7..e1d4991294 100644 --- a/dizoo/atari/config/serial/pong/pong_trex_sql_config.py +++ b/dizoo/atari/config/serial/pong/pong_trex_sql_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=20, - env_id='Pong-v4', + env_id='PongNoFrameskip-v4', #'ALE/Pong-v5' is available. But special setting is needed after gym make. frame_stack=4, ), diff --git a/dizoo/atari/config/serial/qbert/__init__.py b/dizoo/atari/config/serial/qbert/__init__.py index 5032c3a751..ee170aba99 100644 --- a/dizoo/atari/config/serial/qbert/__init__.py +++ b/dizoo/atari/config/serial/qbert/__init__.py @@ -1,2 +1,2 @@ from .qbert_dqn_config import qbert_dqn_config, qbert_dqn_create_config -from .qbert_dqfd_config import qbert_dqfd_config, qbert_dqfd_create_config \ No newline at end of file +from .qbert_dqfd_config import qbert_dqfd_config, qbert_dqfd_create_config diff --git a/dizoo/atari/config/serial/qbert/qbert_a2c_config.py b/dizoo/atari/config/serial/qbert/qbert_a2c_config.py index ba65882f27..f4b0cc7088 100644 --- a/dizoo/atari/config/serial/qbert/qbert_a2c_config.py +++ b/dizoo/atari/config/serial/qbert/qbert_a2c_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=1000000, - env_id='Qbert-v4', + env_id='QbertNoFrameskip-v4', #'ALE/Qbert-v5' is available. But special setting is needed after gym make. frame_stack=4 ), @@ -52,10 +52,11 @@ ), env_manager=dict(type='subprocess'), policy=dict(type='a2c'), + replay_buffer=dict(type='naive'), ) create_config = EasyDict(qbert_a2c_create_config) if __name__ == '__main__': - # or you can enter ding -m serial -c qbert_a2c_config.py -s 0 - from ding.entry import serial_pipeline - serial_pipeline((main_config, create_config), seed=0) + # or you can enter ding -m serial_onpolicy -c qbert_a2c_config.py -s 0 + from ding.entry import serial_pipeline_onpolicy + serial_pipeline_onpolicy((main_config, create_config), seed=0) diff --git a/dizoo/atari/config/serial/qbert/qbert_acer_config.py b/dizoo/atari/config/serial/qbert/qbert_acer_config.py index 505f9c9ae6..46bff881dd 100644 --- a/dizoo/atari/config/serial/qbert/qbert_acer_config.py +++ b/dizoo/atari/config/serial/qbert/qbert_acer_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=int(1e6), - env_id='Qbert-v4', + env_id='QbertNoFrameskip-v4', #'ALE/Qbert-v5' is available. But special setting is needed after gym make. frame_stack=4, manager=dict(shared_memory=False, ) diff --git a/dizoo/atari/config/serial/qbert/qbert_c51_config.py b/dizoo/atari/config/serial/qbert/qbert_c51_config.py index c33ea0dad6..ba0effe2c3 100644 --- a/dizoo/atari/config/serial/qbert/qbert_c51_config.py +++ b/dizoo/atari/config/serial/qbert/qbert_c51_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=30000, - env_id='Qbert-v4', + env_id='QbertNoFrameskip-v4', #'ALE/Qbert-v5' is available. But special setting is needed after gym make. frame_stack=4 ), diff --git a/dizoo/atari/config/serial/qbert/qbert_cql_config.py b/dizoo/atari/config/serial/qbert/qbert_cql_config.py index bd5556be86..de6aef6e89 100644 --- a/dizoo/atari/config/serial/qbert/qbert_cql_config.py +++ b/dizoo/atari/config/serial/qbert/qbert_cql_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=30000, - env_id='Qbert-v4', + env_id='QbertNoFrameskip-v4', #'ALE/Qbert-v5' is available. But special setting is needed after gym make. frame_stack=4 ), diff --git a/dizoo/atari/config/serial/qbert/qbert_dqfd_config.py b/dizoo/atari/config/serial/qbert/qbert_dqfd_config.py index 58787a6974..3907aa3f09 100644 --- a/dizoo/atari/config/serial/qbert/qbert_dqfd_config.py +++ b/dizoo/atari/config/serial/qbert/qbert_dqfd_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=30000, - env_id='Qbert-v4', + env_id='QbertNoFrameskip-v4', #'ALE/Qbert-v5' is available. But special setting is needed after gym make. frame_stack=4 ), diff --git a/dizoo/atari/config/serial/qbert/qbert_dqn_config.py b/dizoo/atari/config/serial/qbert/qbert_dqn_config.py index a5ea56a80c..4b1cd1906f 100644 --- a/dizoo/atari/config/serial/qbert/qbert_dqn_config.py +++ b/dizoo/atari/config/serial/qbert/qbert_dqn_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=30000, - env_id='Qbert-v4', + env_id='QbertNoFrameskip-v4', #'ALE/Qbert-v5' is available. But special setting is needed after gym make. frame_stack=4 ), diff --git a/dizoo/atari/config/serial/qbert/qbert_fqf_config.py b/dizoo/atari/config/serial/qbert/qbert_fqf_config.py index 673c521b58..af1c7f5035 100644 --- a/dizoo/atari/config/serial/qbert/qbert_fqf_config.py +++ b/dizoo/atari/config/serial/qbert/qbert_fqf_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=10000000000, - env_id='Qbert-v4', + env_id='QbertNoFrameskip-v4', #'ALE/Qbert-v5' is available. But special setting is needed after gym make. frame_stack=4, ), @@ -61,4 +61,4 @@ if __name__ == '__main__': # or you can enter `ding -m serial -c qbert_fqf_config.py -s 0` from ding.entry import serial_pipeline - serial_pipeline((main_config, create_config), seed=0) \ No newline at end of file + serial_pipeline((main_config, create_config), seed=0) diff --git a/dizoo/atari/config/serial/qbert/qbert_impala_config.py b/dizoo/atari/config/serial/qbert/qbert_impala_config.py index 6f427717e0..74fbdf4c0c 100644 --- a/dizoo/atari/config/serial/qbert/qbert_impala_config.py +++ b/dizoo/atari/config/serial/qbert/qbert_impala_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=10000000000, - env_id='Qbert-v4', + env_id='QbertNoFrameskip-v4', #'ALE/Qbert-v5' is available. But special setting is needed after gym make. frame_stack=4 ), diff --git a/dizoo/atari/config/serial/qbert/qbert_iqn_config.py b/dizoo/atari/config/serial/qbert/qbert_iqn_config.py index 32c2e91062..00f2f32183 100644 --- a/dizoo/atari/config/serial/qbert/qbert_iqn_config.py +++ b/dizoo/atari/config/serial/qbert/qbert_iqn_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=30000, - env_id='Qbert-v4', + env_id='QbertNoFrameskip-v4', #'ALE/Qbert-v5' is available. But special setting is needed after gym make. frame_stack=4 ), diff --git a/dizoo/atari/config/serial/qbert/qbert_ngu_config.py b/dizoo/atari/config/serial/qbert/qbert_ngu_config.py new file mode 100644 index 0000000000..00bc18f33f --- /dev/null +++ b/dizoo/atari/config/serial/qbert/qbert_ngu_config.py @@ -0,0 +1,129 @@ +from easydict import EasyDict + +collector_env_num = 8 +evaluator_env_num = 8 +nstep = 5 +max_env_step = int(10e6) + +qbert_ngu_config = dict( + exp_name='qbert_ngu_seed0', + env=dict( + collector_env_num=collector_env_num, + evaluator_env_num=evaluator_env_num, + n_evaluator_episode=evaluator_env_num, + env_id='QbertNoFrameskip-v4', + obs_plus_prev_action_reward=True, # use specific env wrapper for ngu policy + stop_value=int(1e6), + frame_stack=4, + ), + rnd_reward_model=dict( + intrinsic_reward_type='add', + learning_rate=1e-4, + obs_shape=[4, 84, 84], + action_shape=6, + batch_size=320, + update_per_collect=10, + only_use_last_five_frames_for_icm_rnd=False, + clear_buffer_per_iters=10, + nstep=nstep, + hidden_size_list=[128, 128, 64], + type='rnd-ngu', + ), + episodic_reward_model=dict( + # means if using rescale trick to the last non-zero reward + # when combing extrinsic and intrinsic reward. + # the rescale trick only used in: + # 1. sparse reward env minigrid, in which the last non-zero reward is a strong positive signal + # 2. the last reward of each episode directly reflects the agent's completion of the task, e.g. lunarlander + # Note that the ngu intrinsic reward is a positive value (max value is 5), in these envs, + # the last non-zero reward should not be overwhelmed by intrinsic rewards, so we need rescale the + # original last nonzero extrinsic reward. + # please refer to ngu_reward_model for details. + last_nonzero_reward_rescale=False, + # means the rescale value for the last non-zero reward, only used when last_nonzero_reward_rescale is True + # please refer to ngu_reward_model for details. + last_nonzero_reward_weight=1, + intrinsic_reward_type='add', + learning_rate=1e-4, + obs_shape=[4, 84, 84], + action_shape=6, + batch_size=320, + update_per_collect=10, + only_use_last_five_frames_for_icm_rnd=False, + clear_buffer_per_iters=10, + nstep=nstep, + hidden_size_list=[128, 128, 64], + type='episodic', + ), + policy=dict( + cuda=True, + on_policy=False, + priority=True, + priority_IS_weight=True, + discount_factor=0.997, + nstep=nstep, + burnin_step=20, + # (int) is the total length of [sequence sample] minus + # the length of burnin part in [sequence sample], + # i.e., = = + + learn_unroll_len=40, # set this key according to the episode length + model=dict( + obs_shape=[4, 84, 84], + action_shape=6, + encoder_hidden_size_list=[128, 128, 512], + collector_env_num=collector_env_num, + ), + learn=dict( + update_per_collect=8, + batch_size=64, + learning_rate=0.0005, + target_update_theta=0.001, + ), + collect=dict( + # NOTE: It is important that set key traj_len_inf=True here, + # to make sure self._traj_len=INF in serial_sample_collector.py. + # In sequence-based policy, for each collect_env, + # we want to collect data of length self._traj_len=INF + # unless the episode enters the 'done' state. + # In each collect phase, we collect a total of sequence samples. + n_sample=32, + traj_len_inf=True, + env_num=collector_env_num, + ), + eval=dict(env_num=evaluator_env_num, ), + other=dict( + eps=dict( + type='exp', + start=0.95, + end=0.05, + decay=1e5, + ), + replay_buffer=dict( + replay_buffer_size=int(2e4), + # (Float type) How much prioritization is used: 0 means no prioritization while 1 means full prioritization + alpha=0.6, + # (Float type) How much correction is used: 0 means no correction while 1 means full correction + beta=0.4, + ) + ), + ), +) +qbert_ngu_config = EasyDict(qbert_ngu_config) +main_config = qbert_ngu_config +qbert_ngu_create_config = dict( + env=dict( + type='atari', + import_names=['dizoo.atari.envs.atari_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='ngu'), + rnd_reward_model=dict(type='rnd-ngu'), + episodic_reward_model=dict(type='episodic'), +) +qbert_ngu_create_config = EasyDict(qbert_ngu_create_config) +create_config = qbert_ngu_create_config + +if __name__ == "__main__": + # or you can enter `ding -m serial_ngu -c qbert_ngu_config.py -s 0` + from ding.entry import serial_pipeline_ngu + serial_pipeline_ngu([main_config, create_config], seed=0, max_env_step=max_env_step) diff --git a/dizoo/atari/config/serial/qbert/qbert_offppo_config.py b/dizoo/atari/config/serial/qbert/qbert_offppo_config.py index b2fcd71805..d0de2b6b43 100644 --- a/dizoo/atari/config/serial/qbert/qbert_offppo_config.py +++ b/dizoo/atari/config/serial/qbert/qbert_offppo_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=10000000000, - env_id='Qbert-v4', + env_id='QbertNoFrameskip-v4', #'ALE/Qbert-v5' is available. But special setting is needed after gym make. frame_stack=4 ), diff --git a/dizoo/atari/config/serial/qbert/qbert_onppo_config.py b/dizoo/atari/config/serial/qbert/qbert_onppo_config.py index 3c08edbd36..3976c6988b 100644 --- a/dizoo/atari/config/serial/qbert/qbert_onppo_config.py +++ b/dizoo/atari/config/serial/qbert/qbert_onppo_config.py @@ -1,13 +1,13 @@ from easydict import EasyDict qbert_onppo_config = dict( - exp_name='enduro_onppo_seed0', + exp_name='qbert_onppo_seed0', env=dict( - collector_env_num=16, + collector_env_num=8, evaluator_env_num=8, n_evaluator_episode=8, stop_value=int(1e10), - env_id='Qbert-v4', + env_id='QbertNoFrameskip-v4', #'ALE/Qbert-v5' is available. But special setting is needed after gym make. frame_stack=4 ), @@ -19,18 +19,20 @@ obs_shape=[4, 84, 84], action_shape=6, action_space='discrete', - encoder_hidden_size_list=[64, 64, 128], - actor_head_hidden_size=128, - critic_head_hidden_size=128, + encoder_hidden_size_list=[32, 64, 64, 512], + actor_head_layer_num=0, + critic_head_layer_num=0, + actor_head_hidden_size=512, + critic_head_hidden_size=512, ), learn=dict( - epoch_per_collect=10, - update_per_collect=1, - batch_size=320, - learning_rate=3e-4, + lr_scheduler=dict(epoch_num=5200, min_lr_lambda=0), + epoch_per_collect=4, + batch_size=256, + learning_rate=2.5e-4, value_weight=0.5, - entropy_weight=0.001, - clip_ratio=0.2, + entropy_weight=0.01, + clip_ratio=0.1, adv_norm=True, value_norm=True, # for onppo, when we recompute adv, we need the key done in data to split traj, so we must @@ -42,7 +44,7 @@ grad_clip_value=0.5, ), collect=dict( - n_sample=3200, + n_sample=1024, unroll_len=1, discount_factor=0.99, gae_lambda=0.95, diff --git a/dizoo/atari/config/serial/qbert/qbert_ppg_config.py b/dizoo/atari/config/serial/qbert/qbert_ppg_config.py index c9baa51866..51a953315a 100644 --- a/dizoo/atari/config/serial/qbert/qbert_ppg_config.py +++ b/dizoo/atari/config/serial/qbert/qbert_ppg_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=1000000, - env_id='Qbert-v4', + env_id='QbertNoFrameskip-v4', #'ALE/Qbert-v5' is available. But special setting is needed after gym make. frame_stack=4 ), diff --git a/dizoo/atari/config/serial/qbert/qbert_qrdqn_config.py b/dizoo/atari/config/serial/qbert/qbert_qrdqn_config.py index 75e6d5c32b..f5485e7f01 100644 --- a/dizoo/atari/config/serial/qbert/qbert_qrdqn_config.py +++ b/dizoo/atari/config/serial/qbert/qbert_qrdqn_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=30000, - env_id='Qbert-v4', + env_id='QbertNoFrameskip-v4', #'ALE/Qbert-v5' is available. But special setting is needed after gym make. frame_stack=4 ), diff --git a/dizoo/atari/config/serial/qbert/qbert_qrdqn_generation_data_config.py b/dizoo/atari/config/serial/qbert/qbert_qrdqn_generation_data_config.py index f68e1530fd..38eaaa7ca6 100644 --- a/dizoo/atari/config/serial/qbert/qbert_qrdqn_generation_data_config.py +++ b/dizoo/atari/config/serial/qbert/qbert_qrdqn_generation_data_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=30000, - env_id='Qbert-v4', + env_id='QbertNoFrameskip-v4', #'ALE/Qbert-v5' is available. But special setting is needed after gym make. frame_stack=4 ), diff --git a/dizoo/atari/config/serial/qbert/qbert_r2d2_config.py b/dizoo/atari/config/serial/qbert/qbert_r2d2_config.py index cbafd8952b..d1c597eb2b 100644 --- a/dizoo/atari/config/serial/qbert/qbert_r2d2_config.py +++ b/dizoo/atari/config/serial/qbert/qbert_r2d2_config.py @@ -9,7 +9,7 @@ evaluator_env_num=evaluator_env_num, n_evaluator_episode=evaluator_env_num, stop_value=int(1e6), - env_id='Qbert-v4', + env_id='QbertNoFrameskip-v4', #'ALE/Qbert-v5' is available. But special setting is needed after gym make. frame_stack=4 ), diff --git a/dizoo/atari/config/serial/qbert/qbert_r2d2_gtrxl_config.py b/dizoo/atari/config/serial/qbert/qbert_r2d2_gtrxl_config.py index d59cd57c28..afc22c2eb0 100644 --- a/dizoo/atari/config/serial/qbert/qbert_r2d2_gtrxl_config.py +++ b/dizoo/atari/config/serial/qbert/qbert_r2d2_gtrxl_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=10000000000, - env_id='Qbert-v4', + env_id='QbertNoFrameskip-v4', #'ALE/Qbert-v5' is available. But special setting is needed after gym make. frame_stack=4 ), diff --git a/dizoo/atari/config/serial/qbert/qbert_rainbow_config.py b/dizoo/atari/config/serial/qbert/qbert_rainbow_config.py index 2f02175231..26f761fd93 100644 --- a/dizoo/atari/config/serial/qbert/qbert_rainbow_config.py +++ b/dizoo/atari/config/serial/qbert/qbert_rainbow_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=30000, - env_id='Qbert-v4', + env_id='QbertNoFrameskip-v4', #'ALE/Qbert-v5' is available. But special setting is needed after gym make. frame_stack=4 ), diff --git a/dizoo/atari/config/serial/qbert/qbert_sqil_config.py b/dizoo/atari/config/serial/qbert/qbert_sqil_config.py index dd7511846f..b030513c7f 100644 --- a/dizoo/atari/config/serial/qbert/qbert_sqil_config.py +++ b/dizoo/atari/config/serial/qbert/qbert_sqil_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=30000, - env_id='Qbert-v4', + env_id='QbertNoFrameskip-v4', #'ALE/Qbert-v5' is available. But special setting is needed after gym make. frame_stack=4 ), diff --git a/dizoo/atari/config/serial/qbert/qbert_sql_config.py b/dizoo/atari/config/serial/qbert/qbert_sql_config.py index 6e89fbc14d..7019318d45 100644 --- a/dizoo/atari/config/serial/qbert/qbert_sql_config.py +++ b/dizoo/atari/config/serial/qbert/qbert_sql_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=30000, - env_id='Qbert-v4', + env_id='QbertNoFrameskip-v4', #'ALE/Qbert-v5' is available. But special setting is needed after gym make. frame_stack=4 ), diff --git a/dizoo/atari/config/serial/qbert/qbert_trex_dqn_config.py b/dizoo/atari/config/serial/qbert/qbert_trex_dqn_config.py index 2a038e37e6..95d1d9716d 100644 --- a/dizoo/atari/config/serial/qbert/qbert_trex_dqn_config.py +++ b/dizoo/atari/config/serial/qbert/qbert_trex_dqn_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=30000, - env_id='Qbert-v4', + env_id='QbertNoFrameskip-v4', #'ALE/Qbert-v5' is available. But special setting is needed after gym make. frame_stack=4 ), diff --git a/dizoo/atari/config/serial/qbert/qbert_trex_offppo_config.py b/dizoo/atari/config/serial/qbert/qbert_trex_offppo_config.py index 8b76e2f728..3621edc462 100644 --- a/dizoo/atari/config/serial/qbert/qbert_trex_offppo_config.py +++ b/dizoo/atari/config/serial/qbert/qbert_trex_offppo_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=10000000000, - env_id='Qbert-v4', + env_id='QbertNoFrameskip-v4', #'ALE/Qbert-v5' is available. But special setting is needed after gym make. frame_stack=4 ), diff --git a/dizoo/atari/config/serial/spaceinvaders/__init__.py b/dizoo/atari/config/serial/spaceinvaders/__init__.py index f4ff222aa6..8be7e58bfe 100644 --- a/dizoo/atari/config/serial/spaceinvaders/__init__.py +++ b/dizoo/atari/config/serial/spaceinvaders/__init__.py @@ -1,2 +1,2 @@ from .spaceinvaders_dqn_config import spaceinvaders_dqn_config, spaceinvaders_dqn_create_config -from .spaceinvaders_dqfd_config import spaceinvaders_dqfd_config, spaceinvaders_dqfd_create_config \ No newline at end of file +from .spaceinvaders_dqfd_config import spaceinvaders_dqfd_config, spaceinvaders_dqfd_create_config diff --git a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_a2c_config.py b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_a2c_config.py index 0b72fd80bd..fb7d4cbb63 100644 --- a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_a2c_config.py +++ b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_a2c_config.py @@ -8,7 +8,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=10000000000, - env_id='SpaceInvaders-v4', + env_id='SpaceInvadersNoFrameskip-v4', #'ALE/SpaceInvaders-v5' is available. But special setting is needed after gym make. frame_stack=4, ), @@ -54,11 +54,12 @@ ), env_manager=dict(type='subprocess'), policy=dict(type='a2c'), + replay_buffer=dict(type='naive'), ) spaceinvaders_a2c_create_config = EasyDict(spaceinvaders_a2c_create_config) create_config = spaceinvaders_a2c_create_config if __name__ == '__main__': # or you can enter ding -m serial_onpolicy -c spaceinvaders_a2c_config.py -s 0 - from ding.entry import serial_pipeline - serial_pipeline((main_config, create_config), seed=0) + from ding.entry import serial_pipeline_onpolicy + serial_pipeline_onpolicy((main_config, create_config), seed=0) diff --git a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_acer_config.py b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_acer_config.py index 6b3f8d1e99..a7afdc3690 100644 --- a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_acer_config.py +++ b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_acer_config.py @@ -8,7 +8,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=int(1e6), - env_id='SpaceInvaders-v4', + env_id='SpaceInvadersNoFrameskip-v4', #'ALE/SpaceInvaders-v5' is available. But special setting is needed after gym make. frame_stack=4, manager=dict(shared_memory=False, ) diff --git a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_c51_config.py b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_c51_config.py index 6e61c0f77f..a12b4846c2 100644 --- a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_c51_config.py +++ b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_c51_config.py @@ -8,7 +8,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=10000000000, - env_id='SpaceInvaders-v4', + env_id='SpaceInvadersNoFrameskip-v4', #'ALE/SpaceInvaders-v5' is available. But special setting is needed after gym make. frame_stack=4, manager=dict(shared_memory=False, ) diff --git a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_dqfd_config.py b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_dqfd_config.py index 849c9a5750..b9d5fde694 100644 --- a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_dqfd_config.py +++ b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_dqfd_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=10000000000, - env_id='SpaceInvaders-v4', + env_id='SpaceInvadersNoFrameskip-v4', #'ALE/SpaceInvaders-v5' is available. But special setting is needed after gym make. frame_stack=4, ), diff --git a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_dqn_config.py b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_dqn_config.py index 5bd7e3c6b9..0d0b810f49 100644 --- a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_dqn_config.py +++ b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_dqn_config.py @@ -7,18 +7,14 @@ collector_env_num=8, evaluator_env_num=8, n_evaluator_episode=8, - stop_value=10000000000, - env_id='SpaceInvaders-v4', + env_id='SpaceInvadersNoFrameskip-v4', #'ALE/SpaceInvaders-v5' is available. But special setting is needed after gym make. frame_stack=4, - manager=dict(shared_memory=False, ), - # The path to save the game replay - replay_path='./spaceinvaders_dqn_seed0/video', ), policy=dict( cuda=True, priority=False, - load_path="./spaceinvaders_dqn_seed0/ckpt/ckpt_best.pth.tar", + random_collect_size=5000, model=dict( obs_shape=[4, 84, 84], action_shape=6, @@ -61,4 +57,4 @@ if __name__ == '__main__': # or you can enter ding -m serial -c spaceinvaders_dqn_config.py -s 0 from ding.entry import serial_pipeline - serial_pipeline((main_config, create_config), seed=0) + serial_pipeline((main_config, create_config), seed=0, max_env_step=int(1e7)) diff --git a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_dqn_config_multi_gpu_ddp.py b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_dqn_config_multi_gpu_ddp.py index 63eabaeb88..093dd7d2ba 100644 --- a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_dqn_config_multi_gpu_ddp.py +++ b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_dqn_config_multi_gpu_ddp.py @@ -7,15 +7,14 @@ collector_env_num=8, evaluator_env_num=8, n_evaluator_episode=8, - stop_value=10000000000, - env_id='SpaceInvaders-v4', - #'ALE/SpaceInvaders-v5' is available. But special setting is needed after gym make. + env_id='SpaceInvadersNoFrameskip-v4', frame_stack=4, - manager=dict(shared_memory=False, ) ), policy=dict( cuda=True, + multi_gpu=True, priority=False, + random_collect_size=5000, model=dict( obs_shape=[4, 84, 84], action_shape=6, @@ -28,7 +27,6 @@ batch_size=32, learning_rate=0.0001, target_update_freq=500, - multi_gpu=True, ), collect=dict(n_sample=100, ), eval=dict(evaluator=dict(eval_freq=4000, )), @@ -58,6 +56,7 @@ if __name__ == '__main__': from ding.entry import serial_pipeline - from ding.utils import DistContext - with DistContext(): - serial_pipeline((main_config, create_config), seed=0) + from ding.utils import DDPContext, to_ddp_config + with DDPContext(): + main_config = to_ddp_config(main_config) + serial_pipeline((main_config, create_config), seed=0, max_env_step=int(1e7)) diff --git a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_dqn_config_multi_gpu_dp.py b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_dqn_config_multi_gpu_dp.py index 5985689adc..5dcbdc8465 100644 --- a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_dqn_config_multi_gpu_dp.py +++ b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_dqn_config_multi_gpu_dp.py @@ -7,11 +7,9 @@ collector_env_num=8, evaluator_env_num=8, n_evaluator_episode=8, - stop_value=10000000000, - env_id='SpaceInvaders-v4', + env_id='SpaceInvadersNoFrameskip-v4', #'ALE/SpaceInvaders-v5' is available. But special setting is needed after gym make. frame_stack=4, - manager=dict(shared_memory=False, ) ), policy=dict( cuda=True, @@ -60,4 +58,4 @@ from ding.model.template.q_learning import DQN from ding.torch_utils import DataParallel model = DataParallel(DQN(obs_shape=[4, 84, 84], action_shape=6)) - serial_pipeline((main_config, create_config), seed=0, model=model) + serial_pipeline((main_config, create_config), seed=0, model=model, max_env_step=int(1e7)) diff --git a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_fqf_config.py b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_fqf_config.py index 5f19ae8ccf..ce3cf7c52b 100644 --- a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_fqf_config.py +++ b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_fqf_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=10000000000, - env_id='SpaceInvaders-v4', + env_id='SpaceInvadersNoFrameskip-v4', #'ALE/SpaceInvaders-v5' is available. But special setting is needed after gym make. frame_stack=4, ), @@ -60,4 +60,4 @@ if __name__ == '__main__': # or you can enter `ding -m serial -c spaceinvaders_fqf_config.py -s 0` from ding.entry import serial_pipeline - serial_pipeline((main_config, create_config), seed=0) \ No newline at end of file + serial_pipeline((main_config, create_config), seed=0) diff --git a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_impala_config.py b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_impala_config.py index 83ae59d261..19be3fc111 100644 --- a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_impala_config.py +++ b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_impala_config.py @@ -1,17 +1,16 @@ -from copy import deepcopy from easydict import EasyDict spaceinvaders_impala_config = dict( - exp_name='spaceinvaders_impala_seed0', + exp_name='impala_log/spaceinvaders_impala_seed0', env=dict( collector_env_num=8, evaluator_env_num=8, n_evaluator_episode=8, stop_value=10000000000, - env_id='SpaceInvaders-v4', + env_id='SpaceInvadersNoFrameskip-v4', #'ALE/SpaceInvaders-v5' is available. But special setting is needed after gym make. frame_stack=4, - manager=dict(shared_memory=False, ) + # manager=dict(shared_memory=False, ) ), policy=dict( cuda=True, @@ -21,21 +20,21 @@ model=dict( obs_shape=[4, 84, 84], action_shape=6, - encoder_hidden_size_list=[128, 128, 256, 512], - critic_head_hidden_size=512, + encoder_hidden_size_list=[128, 128, 256, 256], + critic_head_hidden_size=256, critic_head_layer_num=3, - actor_head_hidden_size=512, + actor_head_hidden_size=256, actor_head_layer_num=3, ), learn=dict( # (int) collect n_sample data, train model update_per_collect times # here we follow impala serial pipeline - update_per_collect=3, # update_per_collect show be in [1, 10] + update_per_collect=2, # update_per_collect show be in [1, 10] # (int) the number of data for a train iteration batch_size=128, grad_clip_type='clip_norm', clip_value=5, - learning_rate=0.0003, + learning_rate=0.0006, # (float) loss weight of the value network, the weight of policy network is set to 1 value_weight=0.5, # (float) loss weight of the entropy regularization, the weight of policy network is set to 1 @@ -56,8 +55,8 @@ n_sample=16, collector=dict(collect_print_freq=1000, ), ), - eval=dict(evaluator=dict(eval_freq=5000, )), - other=dict(replay_buffer=dict(replay_buffer_size=10000, ), ), + eval=dict(evaluator=dict(eval_freq=500, )), + other=dict(replay_buffer=dict(replay_buffer_size=100000, sliced=True), ), ), ) spaceinvaders_impala_config = EasyDict(spaceinvaders_impala_config) diff --git a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_iqn_config.py b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_iqn_config.py index 6eb486114a..1fe4b45903 100644 --- a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_iqn_config.py +++ b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_iqn_config.py @@ -8,7 +8,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=10000000000, - env_id='SpaceInvaders-v4', + env_id='SpaceInvadersNoFrameskip-v4', #'ALE/SpaceInvaders-v5' is available. But special setting is needed after gym make. frame_stack=4, manager=dict(shared_memory=False, ) diff --git a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_mdqn_config.py b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_mdqn_config.py new file mode 100644 index 0000000000..9abf5159d3 --- /dev/null +++ b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_mdqn_config.py @@ -0,0 +1,63 @@ +from copy import deepcopy +from easydict import EasyDict + +spaceinvaders_mdqn_config = dict( + exp_name='spaceinvaders_mdqn_seed0', + env=dict( + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=10000, + env_id='SpaceInvadersNoFrameskip-v0', + #'ALE/SpaceInvaders-v5' is available. But special setting is needed after gym make. + frame_stack=4, + ), + policy=dict( + cuda=True, + priority=False, + model=dict( + obs_shape=[4, 84, 84], + action_shape=6, + encoder_hidden_size_list=[128, 128, 512], + ), + nstep=1, + discount_factor=0.99, + entropy_tau=0.03, + m_alpha=0.9, + learn=dict( + update_per_collect=10, + batch_size=32, + learning_rate=0.0001, + target_update_freq=500, + learner=dict(hook=dict(save_ckpt_after_iter=1000000, )) + ), + collect=dict(n_sample=100, ), + eval=dict(evaluator=dict(eval_freq=4000, )), + other=dict( + eps=dict( + type='exp', + start=1., + end=0.05, + decay=1000000, + ), + replay_buffer=dict(replay_buffer_size=400000, ), + ), + ), +) +spaceinvaders_mdqn_config = EasyDict(spaceinvaders_mdqn_config) +main_config = spaceinvaders_mdqn_config +spaceinvaders_mdqn_create_config = dict( + env=dict( + type='atari', + import_names=['dizoo.atari.envs.atari_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='mdqn'), +) +spaceinvaders_mdqn_create_config = EasyDict(spaceinvaders_mdqn_create_config) +create_config = spaceinvaders_mdqn_create_config + +if __name__ == '__main__': + # or you can enter ding -m serial -c spaceinvaders_mdqn_config.py -s 0 + from ding.entry import serial_pipeline + serial_pipeline((main_config, create_config), seed=0, max_env_step=int(3e7), dynamic_seed=False) diff --git a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_ngu_config.py b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_ngu_config.py new file mode 100644 index 0000000000..f01caa4ac1 --- /dev/null +++ b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_ngu_config.py @@ -0,0 +1,129 @@ +from easydict import EasyDict + +collector_env_num = 8 +evaluator_env_num = 8 +nstep = 5 +max_env_step = int(10e6) + +spaceinvaders_ngu_config = dict( + exp_name='spaceinvaders_ngu_seed0', + env=dict( + collector_env_num=collector_env_num, + evaluator_env_num=evaluator_env_num, + n_evaluator_episode=evaluator_env_num, + env_id='SpaceInvadersNoFrameskip-v4', + obs_plus_prev_action_reward=True, # use specific env wrapper for ngu policy + stop_value=int(1e6), + frame_stack=4, + ), + rnd_reward_model=dict( + intrinsic_reward_type='add', + learning_rate=1e-4, + obs_shape=[4, 84, 84], + action_shape=6, + batch_size=320, + update_per_collect=10, + only_use_last_five_frames_for_icm_rnd=False, + clear_buffer_per_iters=10, + nstep=nstep, + hidden_size_list=[128, 128, 64], + type='rnd-ngu', + ), + episodic_reward_model=dict( + # means if using rescale trick to the last non-zero reward + # when combing extrinsic and intrinsic reward. + # the rescale trick only used in: + # 1. sparse reward env minigrid, in which the last non-zero reward is a strong positive signal + # 2. the last reward of each episode directly reflects the agent's completion of the task, e.g. lunarlander + # Note that the ngu intrinsic reward is a positive value (max value is 5), in these envs, + # the last non-zero reward should not be overwhelmed by intrinsic rewards, so we need rescale the + # original last nonzero extrinsic reward. + # please refer to ngu_reward_model for details. + last_nonzero_reward_rescale=False, + # means the rescale value for the last non-zero reward, only used when last_nonzero_reward_rescale is True + # please refer to ngu_reward_model for details. + last_nonzero_reward_weight=1, + intrinsic_reward_type='add', + learning_rate=1e-4, + obs_shape=[4, 84, 84], + action_shape=6, + batch_size=320, + update_per_collect=10, + only_use_last_five_frames_for_icm_rnd=False, + clear_buffer_per_iters=10, + nstep=nstep, + hidden_size_list=[128, 128, 64], + type='episodic', + ), + policy=dict( + cuda=True, + on_policy=False, + priority=True, + priority_IS_weight=True, + discount_factor=0.997, + nstep=nstep, + burnin_step=20, + # (int) is the total length of [sequence sample] minus + # the length of burnin part in [sequence sample], + # i.e., = = + + learn_unroll_len=40, # set this key according to the episode length + model=dict( + obs_shape=[4, 84, 84], + action_shape=6, + encoder_hidden_size_list=[128, 128, 512], + collector_env_num=collector_env_num, + ), + learn=dict( + update_per_collect=8, + batch_size=64, + learning_rate=0.0005, + target_update_theta=0.001, + ), + collect=dict( + # NOTE: It is important that set key traj_len_inf=True here, + # to make sure self._traj_len=INF in serial_sample_collector.py. + # In sequence-based policy, for each collect_env, + # we want to collect data of length self._traj_len=INF + # unless the episode enters the 'done' state. + # In each collect phase, we collect a total of sequence samples. + n_sample=32, + traj_len_inf=True, + env_num=collector_env_num, + ), + eval=dict(env_num=evaluator_env_num, ), + other=dict( + eps=dict( + type='exp', + start=0.95, + end=0.05, + decay=1e5, + ), + replay_buffer=dict( + replay_buffer_size=int(2e4), + # (Float type) How much prioritization is used: 0 means no prioritization while 1 means full prioritization + alpha=0.6, + # (Float type) How much correction is used: 0 means no correction while 1 means full correction + beta=0.4, + ) + ), + ), +) +spaceinvaders_ngu_config = EasyDict(spaceinvaders_ngu_config) +main_config = spaceinvaders_ngu_config +spaceinvaders_ngu_create_config = dict( + env=dict( + type='atari', + import_names=['dizoo.atari.envs.atari_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='ngu'), + rnd_reward_model=dict(type='rnd-ngu'), + episodic_reward_model=dict(type='episodic'), +) +spaceinvaders_ngu_create_config = EasyDict(spaceinvaders_ngu_create_config) +create_config = spaceinvaders_ngu_create_config + +if __name__ == "__main__": + # or you can enter `ding -m serial_ngu -c spaceinvaders_ngu_config.py -s 0` + from ding.entry import serial_pipeline_ngu + serial_pipeline_ngu([main_config, create_config], seed=0, max_env_step=max_env_step) diff --git a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_offppo_config.py b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_offppo_config.py index aaa7bcd148..2b284d9f5b 100644 --- a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_offppo_config.py +++ b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_offppo_config.py @@ -8,7 +8,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=10000000000, - env_id='SpaceInvaders-v4', + env_id='SpaceInvadersNoFrameskip-v4', #'ALE/SpaceInvaders-v5' is available. But special setting is needed after gym make. frame_stack=4, manager=dict(shared_memory=False, ) diff --git a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_onppo_config.py b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_onppo_config.py index 067a110d65..fb5969282c 100644 --- a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_onppo_config.py +++ b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_onppo_config.py @@ -4,11 +4,11 @@ spaceinvaders_ppo_config = dict( exp_name='spaceinvaders_onppo_seed0', env=dict( - collector_env_num=16, + collector_env_num=8, evaluator_env_num=8, n_evaluator_episode=8, stop_value=int(1e10), - env_id='SpaceInvaders-v4', + env_id='SpaceInvadersNoFrameskip-v4', #'ALE/SpaceInvaders-v5' is available. But special setting is needed after gym make. frame_stack=4, manager=dict(shared_memory=False, ) @@ -21,18 +21,20 @@ obs_shape=[4, 84, 84], action_shape=6, action_space='discrete', - encoder_hidden_size_list=[64, 64, 128], - actor_head_hidden_size=128, - critic_head_hidden_size=128, + encoder_hidden_size_list=[32, 64, 64, 512], + actor_head_layer_num=0, + critic_head_layer_num=0, + actor_head_hidden_size=512, + critic_head_hidden_size=512, ), learn=dict( - epoch_per_collect=10, - update_per_collect=1, - batch_size=320, - learning_rate=3e-4, + lr_scheduler=dict(epoch_num=5200, min_lr_lambda=0), + epoch_per_collect=4, + batch_size=256, + learning_rate=2.5e-4, value_weight=0.5, - entropy_weight=0.001, - clip_ratio=0.2, + entropy_weight=0.01, + clip_ratio=0.1, adv_norm=True, value_norm=True, # for onppo, when we recompute adv, we need the key done in data to split traj, so we must @@ -42,9 +44,15 @@ ignore_done=False, grad_clip_type='clip_norm', grad_clip_value=0.5, + # KL divergence regularization between current policy and pretrained policy. + # Supported KL divergence estimators: ['k1', 'k2', 'k3']. + # KL divergence loss will be calculated only when pretrained_model_path is provided. + kl_beta=0.05, + kl_type='k1', + pretrained_model_path=None, ), collect=dict( - n_sample=3200, + n_sample=1024, unroll_len=1, discount_factor=0.99, gae_lambda=0.95, diff --git a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_ppg_config.py b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_ppg_config.py index b3b7088be9..af111ccf36 100644 --- a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_ppg_config.py +++ b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_ppg_config.py @@ -8,7 +8,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=10000000000, - env_id='SpaceInvaders-v4', + env_id='SpaceInvadersNoFrameskip-v4', #'ALE/SpaceInvaders-v5' is available. But special setting is needed after gym make. frame_stack=4, manager=dict(shared_memory=False, ) diff --git a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_qrdqn_config.py b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_qrdqn_config.py index ac602beab0..c67c33ea43 100644 --- a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_qrdqn_config.py +++ b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_qrdqn_config.py @@ -8,7 +8,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=10000000000, - env_id='SpaceInvaders-v4', + env_id='SpaceInvadersNoFrameskip-v4', #'ALE/SpaceInvaders-v5' is available. But special setting is needed after gym make. frame_stack=4, manager=dict(shared_memory=False, ) @@ -59,4 +59,4 @@ if __name__ == '__main__': # or you can enter ding -m serial -c spaceinvaders_qrdqn_config.py -s 0 from ding.entry import serial_pipeline - serial_pipeline((main_config, create_config), seed=0) + serial_pipeline((main_config, create_config), seed=0, max_env_step=int(1e7)) diff --git a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_r2d2_config.py b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_r2d2_config.py index a0d8ae3f59..5d9145d741 100644 --- a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_r2d2_config.py +++ b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_r2d2_config.py @@ -9,7 +9,7 @@ evaluator_env_num=evaluator_env_num, n_evaluator_episode=8, stop_value=int(1e6), - env_id='SpaceInvaders-v4', + env_id='SpaceInvadersNoFrameskip-v4', #'ALE/SpaceInvaders-v5' is available. But special setting is needed after gym make. frame_stack=4, manager=dict(shared_memory=False, ) diff --git a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_r2d2_gtrxl_config.py b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_r2d2_gtrxl_config.py index ef93a99019..b176597821 100644 --- a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_r2d2_gtrxl_config.py +++ b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_r2d2_gtrxl_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=10000000000, - env_id='SpaceInvaders-v4', + env_id='SpaceInvadersNoFrameskip-v4', #'ALE/SpaceInvaders-v5' is available. But special setting is needed after gym make. frame_stack=4, manager=dict(shared_memory=False) diff --git a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_r2d2_residual_config.py b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_r2d2_residual_config.py index c1315748c1..9341b9f092 100644 --- a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_r2d2_residual_config.py +++ b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_r2d2_residual_config.py @@ -9,7 +9,7 @@ evaluator_env_num=evaluator_env_num, n_evaluator_episode=8, stop_value=10000000000, - env_id='SpaceInvaders-v4', + env_id='SpaceInvadersNoFrameskip-v4', #'ALE/SpaceInvaders-v5' is available. But special setting is needed after gym make. frame_stack=4, manager=dict(shared_memory=False, ) diff --git a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_rainbow_config.py b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_rainbow_config.py index 1b6d651e94..00be159e71 100644 --- a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_rainbow_config.py +++ b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_rainbow_config.py @@ -8,7 +8,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=10000000000, - env_id='SpaceInvaders-v4', + env_id='SpaceInvadersNoFrameskip-v4', #'ALE/SpaceInvaders-v5' is available. But special setting is needed after gym make. frame_stack=4, manager=dict(shared_memory=False, ) diff --git a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_sqil_config.py b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_sqil_config.py index 2cc676122f..84a83a6357 100644 --- a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_sqil_config.py +++ b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_sqil_config.py @@ -7,7 +7,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=10000000000, - env_id='SpaceInvaders-v4', + env_id='SpaceInvadersNoFrameskip-v4', #'ALE/SpaceInvaders-v5' is available. But special setting is needed after gym make. frame_stack=4, ), diff --git a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_sql_config.py b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_sql_config.py index 3c95e5dc77..888eeb1bfe 100644 --- a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_sql_config.py +++ b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_sql_config.py @@ -8,7 +8,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=10000000000, - env_id='SpaceInvaders-v4', + env_id='SpaceInvadersNoFrameskip-v4', #'ALE/SpaceInvaders-v5' is available. But special setting is needed after gym make. frame_stack=4, manager=dict(shared_memory=False, reset_inplace=True) diff --git a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_trex_dqn_config.py b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_trex_dqn_config.py index 0d40d1e26c..1a4491d68a 100644 --- a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_trex_dqn_config.py +++ b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_trex_dqn_config.py @@ -8,7 +8,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=10000000000, - env_id='SpaceInvaders-v4', + env_id='SpaceInvadersNoFrameskip-v4', #'ALE/SpaceInvaders-v5' is available. But special setting is needed after gym make. frame_stack=4, manager=dict(shared_memory=False, ) diff --git a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_trex_offppo_config.py b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_trex_offppo_config.py index 4b75d7289a..7934a67a7d 100644 --- a/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_trex_offppo_config.py +++ b/dizoo/atari/config/serial/spaceinvaders/spaceinvaders_trex_offppo_config.py @@ -8,7 +8,7 @@ evaluator_env_num=8, n_evaluator_episode=8, stop_value=10000000000, - env_id='SpaceInvaders-v4', + env_id='SpaceInvadersNoFrameskip-v4', #'ALE/SpaceInvaders-v5' is available. But special setting is needed after gym make. frame_stack=4, manager=dict(shared_memory=False, ) diff --git a/dizoo/atari/entry/atari_dt_main.py b/dizoo/atari/entry/atari_dt_main.py new file mode 100644 index 0000000000..b89bbaec7e --- /dev/null +++ b/dizoo/atari/entry/atari_dt_main.py @@ -0,0 +1,56 @@ +import torch.nn as nn +import torch.distributed as dist +from ditk import logging +from ding.model import DecisionTransformer +from ding.policy import DTPolicy +from ding.envs import SubprocessEnvManagerV2 +from ding.envs import AllinObsWrapper +from ding.data import create_dataset +from ding.config import compile_config +from ding.framework import task, ding_init +from ding.framework.context import OfflineRLContext +from ding.framework.middleware import interaction_evaluator, trainer, CkptSaver, offline_logger, termination_checker, \ + OfflineMemoryDataFetcher +from ding.utils import set_pkg_seed, DDPContext, to_ddp_config +from dizoo.atari.envs import AtariEnv +from dizoo.atari.config.serial.pong.pong_dt_config import main_config, create_config + + +def main(): + # If you don't have offline data, you need to prepare if first and set the data_path in config + # For demostration, we also can train a RL policy (e.g. SAC) and collect some data + logging.getLogger().setLevel(logging.INFO) + with DDPContext(): + cmain_config = to_ddp_config(main_config) + cfg = compile_config(cmain_config, create_cfg=create_config, auto=True) + ding_init(cfg) + with task.start(async_mode=False, ctx=OfflineRLContext()): + evaluator_env = SubprocessEnvManagerV2( + env_fn=[lambda: AllinObsWrapper(AtariEnv(cfg.env)) for _ in range(cfg.env.evaluator_env_num)], + cfg=cfg.env.manager + ) + + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + dataset = create_dataset(cfg) + cfg.policy.model.max_timestep = dataset.get_max_timestep() + state_encoder = nn.Sequential( + nn.Conv2d(4, 32, 8, stride=4, padding=0), nn.ReLU(), nn.Conv2d(32, 64, 4, stride=2, padding=0), + nn.ReLU(), nn.Conv2d(64, 64, 3, stride=1, padding=0), nn.ReLU(), nn.Flatten(), + nn.Linear(3136, cfg.policy.model.h_dim), nn.Tanh() + ) + + model = DecisionTransformer(**cfg.policy.model, state_encoder=state_encoder) + # model.parallelize() + policy = DTPolicy(cfg.policy, model=model) + + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) + task.use(OfflineMemoryDataFetcher(cfg, dataset)) + task.use(trainer(cfg, policy.learn_mode)) + task.use(termination_checker(max_train_iter=3e4)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) + task.use(offline_logger()) + task.run() + + +if __name__ == "__main__": + main() diff --git a/dizoo/atari/entry/atari_impala_main.py b/dizoo/atari/entry/atari_impala_main.py new file mode 100644 index 0000000000..69dce9f8a4 --- /dev/null +++ b/dizoo/atari/entry/atari_impala_main.py @@ -0,0 +1,51 @@ +import gym +from ditk import logging +from ding.model import VAC +from ding.policy import IMPALAPolicy +from ding.envs import SubprocessEnvManagerV2 +from ding.data import DequeBuffer +from ding.config import compile_config +from ding.framework import task, ding_init +from ding.framework.context import OnlineRLContext +from ding.framework.middleware import OffPolicyLearner, StepCollector, interaction_evaluator, data_pusher, \ + CkptSaver, online_logger, termination_checker +from ding.utils import set_pkg_seed +from dizoo.atari.config.serial.pong.pong_impala_config import main_config, create_config +from dizoo.atari.envs import AtariEnv + + +def main(): + logging.getLogger().setLevel(logging.INFO) + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + ding_init(cfg) + with task.start(async_mode=False, ctx=OnlineRLContext()): + collector_env_cfg = AtariEnv.create_collector_env_cfg(cfg.env) + evaluator_env_cfg = AtariEnv.create_evaluator_env_cfg(cfg.env) + collector_env = SubprocessEnvManagerV2( + env_fn=[lambda: AtariEnv(cfg=c) for c in collector_env_cfg], cfg=cfg.env.manager + ) + evaluator_env = SubprocessEnvManagerV2( + env_fn=[lambda: AtariEnv(cfg=c) for c in evaluator_env_cfg], cfg=cfg.env.manager + ) + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + + model = VAC(**cfg.policy.model) + buffer_ = DequeBuffer( + size=cfg.policy.other.replay_buffer.replay_buffer_size, sliced=cfg.policy.other.replay_buffer.sliced + ) + policy = IMPALAPolicy(cfg.policy, model=model) + + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) + task.use( + StepCollector(cfg, policy.collect_mode, collector_env, random_collect_size=cfg.policy.random_collect_size) + ) + task.use(data_pusher(cfg, buffer_, group_by_env=True)) + task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) + task.use(online_logger(train_show_freq=300)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=10000)) + task.use(termination_checker(max_env_step=1e7)) + task.run() + + +if __name__ == "__main__": + main() diff --git a/dizoo/atari/entry/pong_cql_main.py b/dizoo/atari/entry/pong_cql_main.py index a4b4383ec1..fca48a1ff6 100644 --- a/dizoo/atari/entry/pong_cql_main.py +++ b/dizoo/atari/entry/pong_cql_main.py @@ -1,7 +1,6 @@ import torch from copy import deepcopy -from dizoo.atari.config.serial.pong.pong_qrdqn_generation_data_config import main_config, create_config from ding.entry import serial_pipeline_offline, collect_demo_data, eval, serial_pipeline @@ -15,22 +14,21 @@ def train_cql(args): def eval_ckpt(args): + from dizoo.atari.config.serial.pong.pong_qrdqn_generation_data_config import main_config, create_config main_config.exp_name = 'pong' - main_config.policy.learn.learner.load_path = './pong/ckpt/ckpt_best.pth.tar' - main_config.policy.learn.learner.hook.load_ckpt_before_run = './pong/ckpt/ckpt_best.pth.tar' config = deepcopy([main_config, create_config]) - eval(config, seed=args.seed, load_path=main_config.policy.learn.learner.hook.load_ckpt_before_run) + eval(config, seed=args.seed, load_path='./pong/ckpt/ckpt_best.pth.tar') def generate(args): + from dizoo.atari.config.serial.pong.pong_qrdqn_generation_data_config import main_config, create_config main_config.exp_name = 'pong' - main_config.policy.learn.learner.load_path = './pong/ckpt/ckpt_best.pth.tar' main_config.policy.collect.save_path = './pong/expert.pkl' config = deepcopy([main_config, create_config]) - state_dict = torch.load(main_config.policy.learn.learner.load_path, map_location='cpu') + state_dict = torch.load('./pong/ckpt/ckpt_best.pth.tar', map_location='cpu') collect_demo_data( config, - collect_count=main_config.policy.other.replay_buffer.replay_buffer_size, + collect_count=int(1e5), seed=args.seed, expert_data_path=main_config.policy.collect.save_path, state_dict=state_dict diff --git a/dizoo/atari/entry/spaceinavders_dqn_eval.py b/dizoo/atari/entry/spaceinvaders_dqn_eval.py similarity index 95% rename from dizoo/atari/entry/spaceinavders_dqn_eval.py rename to dizoo/atari/entry/spaceinvaders_dqn_eval.py index d8bfde290d..35e15a578c 100644 --- a/dizoo/atari/entry/spaceinavders_dqn_eval.py +++ b/dizoo/atari/entry/spaceinvaders_dqn_eval.py @@ -15,8 +15,9 @@ from ding.rl_utils import get_epsilon_greedy_fn from dizoo.atari.config.serial.spaceinvaders.spaceinvaders_dqn_config import main_config, create_config + def main(rl_cfg, seed=0): - main_cfg, create_cfg =rl_cfg + main_cfg, create_cfg = rl_cfg cfg = compile_config( main_cfg, BaseEnvManager, @@ -56,4 +57,4 @@ def main(rl_cfg, seed=0): if __name__ == "__main__": - main(rl_cfg=(main_config, create_config),seed=0) + main(rl_cfg=(main_config, create_config), seed=0) diff --git a/dizoo/atari/envs/atari_env.py b/dizoo/atari/envs/atari_env.py index e672f1cf13..c8c3298225 100644 --- a/dizoo/atari/envs/atari_env.py +++ b/dizoo/atari/envs/atari_env.py @@ -43,7 +43,7 @@ def reset(self) -> np.ndarray: self._env.seed(self._seed) obs = self._env.reset() obs = to_ndarray(obs) - self._final_eval_reward = 0. + self._eval_episode_return = 0. return obs def close(self) -> None: @@ -61,11 +61,11 @@ def step(self, action: np.ndarray) -> BaseEnvTimestep: action = action.item() obs, rew, done, info = self._env.step(action) # self._env.render() - self._final_eval_reward += rew + self._eval_episode_return += rew obs = to_ndarray(obs) rew = to_ndarray([rew]).astype(np.float32) # wrapped to be transferred to a Tensor with shape (1,) if done: - info['final_eval_reward'] = self._final_eval_reward + info['eval_episode_return'] = self._eval_episode_return return BaseEnvTimestep(obs, rew, done, info) def enable_save_replay(self, replay_path: Optional[str] = None) -> None: @@ -133,7 +133,7 @@ def reset(self) -> np.ndarray: self._env.seed(self._seed + np_seed) obs = self._env.reset() obs = to_ndarray(obs) - self._final_eval_reward = 0. + self._eval_episode_return = 0. return obs def _make_env(self): diff --git a/dizoo/atari/envs/atari_muzero_env.py b/dizoo/atari/envs/atari_muzero_env.py deleted file mode 100644 index 159a3bb5b9..0000000000 --- a/dizoo/atari/envs/atari_muzero_env.py +++ /dev/null @@ -1,139 +0,0 @@ -""" -Adapt Atari to BaseGameEnv interface -""" - -import sys -from typing import Any, List, Union, Sequence -import copy -import numpy as np -import gym -from ding.envs import BaseEnv, BaseEnvTimestep -from ding.utils import ENV_REGISTRY -from ding.torch_utils import to_ndarray -from dizoo.atari.envs.atari_wrappers import wrap_muzero - - -@ENV_REGISTRY.register('atari-muzero') -class AtariMuZeroEnv(BaseEnv): - - def __init__(self, cfg=None): - self.cfg = cfg - self._init_flag = False - - def _make_env(self): - return wrap_muzero(self.cfg) - - def reset(self): - if not self._init_flag: - self._env = self._make_env() - self._observation_space = self._env.env.observation_space - self._action_space = self._env.env.action_space - self._reward_space = gym.spaces.Box( - low=self._env.env.reward_range[0], high=self._env.env.reward_range[1], shape=(1, ), dtype=np.float32 - ) - - self._init_flag = True - if hasattr(self, '_seed') and hasattr(self, '_dynamic_seed') and self._dynamic_seed: - np_seed = 100 * np.random.randint(1, 1000) - self._env.env.seed(self._seed + np_seed) - elif hasattr(self, '_seed'): - self._env.env.seed(self._seed) - - obs = self._env.reset() - self.obs = to_ndarray(obs) - self._final_eval_reward = 0. - self.has_reset = True - obs = self.observe() - return obs - - def observe(self): - """ - Overview: - add action_mask to obs to adapt with MCTS alg.. - """ - observation = self.obs - action_mask = np.ones(self._action_space.n, 'int8') - return {'observation': observation, 'action_mask': action_mask} - - def step(self, action): - obs, reward, done, info = self._env.step(action) - # self._env.render() - self.obs = to_ndarray(obs) - self.reward = np.array(reward).astype(np.float32) - self._final_eval_reward += self.reward - observation = self.observe() - if done: - info['final_eval_reward'] = self._final_eval_reward - - return BaseEnvTimestep(observation, self.reward, done, info) - - @property - def legal_actions(self): - return np.arange(self._action_space.n) - - def random_action(self): - action_list = self.legal_actions - return np.random.choice(action_list) - - def render(self, mode='human'): - self._env.render() - - def human_to_action(self): - """ - Overview: - For multiplayer games, ask the user for a legal action - and return the corresponding action number. - Returns: - An integer from the action space. - """ - while True: - try: - print(f"Current available actions for the player are:{self.legal_actions}") - choice = int(input(f"Enter the index of next action: ")) - if choice in self.legal_actions: - break - else: - print("Wrong input, try again") - except KeyboardInterrupt: - print("exit") - sys.exit(0) - except Exception as e: - print("Wrong input, try again") - return choice - - def close(self) -> None: - if self._init_flag: - self._env.close() - self._init_flag = False - - def seed(self, seed: int, dynamic_seed: bool = True) -> None: - self._seed = seed - self._dynamic_seed = dynamic_seed - np.random.seed(self._seed) - - @property - def observation_space(self) -> gym.spaces.Space: - return self._observation_space - - @property - def action_space(self) -> gym.spaces.Space: - return self._action_space - - @property - def reward_space(self) -> gym.spaces.Space: - return self._reward_space - - def __repr__(self) -> str: - return "DI-engine Atari MuZero Env({})".format(self.cfg.env_name) - - @staticmethod - def create_collector_envcfg(cfg: dict) -> List[dict]: - collector_env_num = cfg.pop('collector_env_num') - cfg = copy.deepcopy(cfg) - return [cfg for _ in range(collector_env_num)] - - @staticmethod - def create_evaluator_envcfg(cfg: dict) -> List[dict]: - evaluator_env_num = cfg.pop('evaluator_env_num') - cfg = copy.deepcopy(cfg) - return [cfg for _ in range(evaluator_env_num)] diff --git a/dizoo/atari/envs/atari_wrappers.py b/dizoo/atari/envs/atari_wrappers.py index 78de31a4aa..1eb288d1dc 100644 --- a/dizoo/atari/envs/atari_wrappers.py +++ b/dizoo/atari/envs/atari_wrappers.py @@ -7,9 +7,7 @@ ScaledFloatFrameWrapper, \ ClipRewardWrapper, FrameStackWrapper import numpy as np -from ding.rl_utils.efficientzero.game import Game from ding.utils.compression_helper import jpeg_data_compressor -from gym.wrappers import RecordVideo import cv2 @@ -75,42 +73,6 @@ def wrap_deepmind_mr(env_id, episode_life=True, clip_rewards=True, frame_stack=4 return env -""" -The following code is adapted from https://github.com/YeWR/EfficientZero -""" - - -def wrap_muzero(config, warp_frame=True, save_video=False, save_path=None, video_callable=None, uid=None): - """ - Overview: - Configure environment for MuZero-style Atari. The observation is - channel-first: (c, h, w) instead of (h, w, c). - Arguments: - - config (:obj:`Dict`): Dict containing configuration. - :param str env_id: the atari environment id. - :param bool config.episode_life: wrap the episode life wrapper. - :param bool warp_frame: wrap the grayscale + resize observation wrapper. - :return: the wrapped atari environment. - """ - env = gym.make(config.env_name) - assert 'NoFrameskip' in env.spec.id - env = NoopResetWrapper(env, noop_max=30) - env = MaxAndSkipWrapper(env, skip=4) - if config.episode_life: - env = EpisodicLifeWrapper(env) - env = TimeLimit(env, max_episode_steps=config.max_episode_steps) - if warp_frame: - env = WarpFrame(env, width=config.obs_shape[1], height=config.obs_shape[2], grayscale=config.gray_scale) - if save_video: - #env = Monitor(env, directory=save_path, force=True, video_callable=video_callable, uid=uid) - env = RecordVideo(env, video_folder=save_path, episode_trigger=lambda episode_id: True, name_prefix='rl-video-{}'.format(uid)) - env = JpegWrapper(env, cvt_string=config.cvt_string) - if config.game_wrapper: - env = GameWrapper(env) - - return env - - class TimeLimit(gym.Wrapper): def __init__(self, env, max_episode_steps=None): diff --git a/dizoo/atari/envs/test_atari_env.py b/dizoo/atari/envs/test_atari_env.py index 517e5238c9..63c6018bcc 100644 --- a/dizoo/atari/envs/test_atari_env.py +++ b/dizoo/atari/envs/test_atari_env.py @@ -31,10 +31,10 @@ def test_pong(self): assert timestep.obs.shape == (cfg.frame_stack, 84, 84) assert timestep.reward.shape == (1, ) if timestep.done: - assert 'final_eval_reward' in timestep.info, timestep.info + assert 'eval_episode_return' in timestep.info, timestep.info break print(pong_env.observation_space, pong_env.action_space, pong_env.reward_space) - print('final_eval_reward: {}'.format(timestep.info['final_eval_reward'])) + print('eval_episode_return: {}'.format(timestep.info['eval_episode_return'])) pong_env.close() def test_montezuma_revenge(self): @@ -56,8 +56,8 @@ def test_montezuma_revenge(self): assert timestep.obs.shape == (cfg.frame_stack, 84, 84) assert timestep.reward.shape == (1, ) if timestep.done: - assert 'final_eval_reward' in timestep.info, timestep.info + assert 'eval_episode_return' in timestep.info, timestep.info break print(mr_env.observation_space, mr_env.action_space, mr_env.reward_space) - print('final_eval_reward: {}'.format(timestep.info['final_eval_reward'])) + print('eval_episode_return: {}'.format(timestep.info['eval_episode_return'])) mr_env.close() diff --git a/dizoo/atari/envs/test_atari_muzero_env.py b/dizoo/atari/envs/test_atari_muzero_env.py deleted file mode 100644 index 8f38c0c5b6..0000000000 --- a/dizoo/atari/envs/test_atari_muzero_env.py +++ /dev/null @@ -1,33 +0,0 @@ -import pytest -from dizoo.atari.envs.atari_muzero_env import AtariMuZeroEnv -from easydict import EasyDict - -cfg = EasyDict( - env_name='PongNoFrameskip-v4', - frame_skip=4, - frame_stack=4, - episode_life=True, - obs_shape=(12, 96, 96), - gray_scale=False, - discount=0.997, - cvt_string=True, - max_episode_steps=1.08e5, - game_wrapper=True, -) - - -@pytest.mark.envtest -class TestAtariMuZeroEnv: - - def test_naive(self): - env = AtariMuZeroEnv(cfg) - env.reset() - # env.render() - while True: - action = env.random_action() - # action = env.human_to_action() - obs, reward, done, info = env.step(action) - # env.render() - if done: - print(info) - break diff --git a/dizoo/atari/example/atari_dqn.py b/dizoo/atari/example/atari_dqn.py index 4f4be49ecb..660f8576d9 100644 --- a/dizoo/atari/example/atari_dqn.py +++ b/dizoo/atari/example/atari_dqn.py @@ -41,8 +41,8 @@ def main(): task.use(nstep_reward_enhancer(cfg)) task.use(data_pusher(cfg, buffer_)) task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) - task.use(CkptSaver(cfg, policy, train_freq=1000)) - task.use(termination_checker(max_train_iter=int(1e7))) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=1000)) + task.use(termination_checker(max_env_step=int(1e7))) task.run() diff --git a/dizoo/atari/example/atari_dqn_ddp.py b/dizoo/atari/example/atari_dqn_ddp.py new file mode 100644 index 0000000000..3871b3ec11 --- /dev/null +++ b/dizoo/atari/example/atari_dqn_ddp.py @@ -0,0 +1,65 @@ +from copy import deepcopy +from ditk import logging +from ding.model import DQN +from ding.policy import DQNPolicy +from ding.envs import DingEnvWrapper, SubprocessEnvManagerV2 +from ding.data import DequeBuffer +from ding.config import compile_config +from ding.utils import DistContext, get_rank +from ding.framework import task, ding_init +from ding.framework.context import OnlineRLContext +from ding.framework.middleware import OffPolicyLearner, StepCollector, interaction_evaluator, data_pusher, \ + eps_greedy_handler, CkptSaver, nstep_reward_enhancer, online_logger, ddp_termination_checker +from ding.utils import set_pkg_seed +from dizoo.atari.envs.atari_env import AtariEnv +from dizoo.atari.config.serial.pong.pong_dqn_config import main_config, create_config + + +def main(): + logging.getLogger().setLevel(logging.INFO) + main_config.exp_name = 'pong_dqn_seed0_ddp' + main_config.policy.multi_gpu = True + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + ding_init(cfg) + with DistContext(): + rank = get_rank() + with task.start(async_mode=False, ctx=OnlineRLContext()): + collector_cfg = deepcopy(cfg.env) + collector_cfg.is_train = True + evaluator_cfg = deepcopy(cfg.env) + evaluator_cfg.is_train = False + collector_env = SubprocessEnvManagerV2( + env_fn=[lambda: AtariEnv(collector_cfg) for _ in range(cfg.env.collector_env_num)], cfg=cfg.env.manager + ) + evaluator_env = SubprocessEnvManagerV2( + env_fn=[lambda: AtariEnv(evaluator_cfg) for _ in range(cfg.env.evaluator_env_num)], cfg=cfg.env.manager + ) + + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + + model = DQN(**cfg.policy.model) + buffer_ = DequeBuffer(size=cfg.policy.other.replay_buffer.replay_buffer_size) + policy = DQNPolicy(cfg.policy, model=model) + + if rank == 0: + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) + task.use(eps_greedy_handler(cfg)) + task.use(StepCollector(cfg, policy.collect_mode, collector_env)) + task.use(nstep_reward_enhancer(cfg)) + task.use(data_pusher(cfg, buffer_)) + task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) + if rank == 0: + task.use(CkptSaver(policy, cfg.exp_name, train_freq=1000)) + task.use(online_logger(record_train_iter=True)) + task.use(ddp_termination_checker(max_env_step=int(1e7), rank=rank)) + task.run() + + +if __name__ == "__main__": + """ + Overview: + This script should be executed with GPUs. + Run the following command to launch the script: + python -m torch.distributed.launch --nproc_per_node=2 ./dizoo/atari/example/atari_dqn_ddp.py + """ + main() diff --git a/dizoo/atari/example/atari_dqn_dist.py b/dizoo/atari/example/atari_dqn_dist.py new file mode 100644 index 0000000000..8b9f01699a --- /dev/null +++ b/dizoo/atari/example/atari_dqn_dist.py @@ -0,0 +1,85 @@ +from copy import deepcopy +from ditk import logging +from ding.model import DQN +from ding.policy import DQNPolicy +from ding.envs import DingEnvWrapper, SubprocessEnvManagerV2 +from ding.data import DequeBuffer +from ding.config import compile_config +from ding.framework import task, ding_init +from ding.framework.context import OnlineRLContext +from ding.framework.middleware import OffPolicyLearner, StepCollector, interaction_evaluator, data_pusher, \ + eps_greedy_handler, CkptSaver, context_exchanger, model_exchanger, termination_checker, nstep_reward_enhancer, \ + online_logger +from ding.utils import set_pkg_seed +from dizoo.atari.envs.atari_env import AtariEnv +from dizoo.atari.config.serial.pong.pong_dqn_config import main_config, create_config + + +def main(): + logging.getLogger().setLevel(logging.INFO) + main_config.exp_name = 'pong_dqn_seed0_ditask_dist' + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + ding_init(cfg) + with task.start(async_mode=False, ctx=OnlineRLContext()): + assert task.router.is_active, "Please execute this script with ditask! See note in the header." + + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + + model = DQN(**cfg.policy.model) + policy = DQNPolicy(cfg.policy, model=model) + + if 'learner' in task.router.labels: + logging.info("Learner running on node {}".format(task.router.node_id)) + buffer_ = DequeBuffer(size=cfg.policy.other.replay_buffer.replay_buffer_size) + task.use( + context_exchanger( + send_keys=["train_iter"], + recv_keys=["trajectories", "episodes", "env_step", "env_episode"], + skip_n_iter=0 + ) + ) + task.use(model_exchanger(model, is_learner=True)) + task.use(nstep_reward_enhancer(cfg)) + task.use(data_pusher(cfg, buffer_)) + task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=1000)) + + elif 'evaluator' in task.router.labels: + logging.info("Evaluator running on node {}".format(task.router.node_id)) + evaluator_cfg = deepcopy(cfg.env) + evaluator_cfg.is_train = False + evaluator_env = SubprocessEnvManagerV2( + env_fn=[lambda: AtariEnv(evaluator_cfg) for _ in range(cfg.env.evaluator_env_num)], cfg=cfg.env.manager + ) + task.use(context_exchanger(recv_keys=["train_iter", "env_step"], skip_n_iter=1)) + task.use(model_exchanger(model, is_learner=False)) + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) + task.use(CkptSaver(policy, cfg.exp_name, save_finish=False)) + task.use(online_logger(record_train_iter=True)) + + elif 'collector' in task.router.labels: + logging.info("Collector running on node {}".format(task.router.node_id)) + collector_cfg = deepcopy(cfg.env) + collector_cfg.is_train = True + collector_env = SubprocessEnvManagerV2( + env_fn=[lambda: AtariEnv(collector_cfg) for _ in range(cfg.env.collector_env_num)], cfg=cfg.env.manager + ) + task.use( + context_exchanger( + send_keys=["trajectories", "episodes", "env_step", "env_episode"], + recv_keys=["train_iter"], + skip_n_iter=1 + ) + ) + task.use(model_exchanger(model, is_learner=False)) + task.use(eps_greedy_handler(cfg)) + task.use(StepCollector(cfg, policy.collect_mode, collector_env)) + task.use(termination_checker(max_env_step=int(1e7))) + else: + raise KeyError("invalid router labels: {}".format(task.router.labels)) + + task.run() + + +if __name__ == "__main__": + main() diff --git a/dizoo/atari/example/atari_dqn_dist_ddp.py b/dizoo/atari/example/atari_dqn_dist_ddp.py new file mode 100644 index 0000000000..5dbfc4e65c --- /dev/null +++ b/dizoo/atari/example/atari_dqn_dist_ddp.py @@ -0,0 +1,106 @@ +from copy import deepcopy +from ditk import logging +from ding.model import DQN +from ding.policy import DQNPolicy +from ding.envs import DingEnvWrapper, SubprocessEnvManagerV2 +from ding.data import DequeBuffer +from ding.config import compile_config +from ding.framework import task, ding_init +from ding.framework.context import OnlineRLContext +from ding.framework.middleware import OffPolicyLearner, StepCollector, interaction_evaluator, data_pusher, \ + eps_greedy_handler, CkptSaver, context_exchanger, model_exchanger, termination_checker, nstep_reward_enhancer, \ + online_logger +from ding.utils import set_pkg_seed +from dizoo.atari.envs.atari_env import AtariEnv +from dizoo.atari.config.serial.pong.pong_dqn_config import main_config, create_config + +logging.getLogger().setLevel(logging.INFO) +main_config.exp_name = 'pong_dqn_seed0_ditask_dist_ddp' + + +def learner(): + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + ding_init(cfg) + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + + model = DQN(**cfg.policy.model) + policy = DQNPolicy(cfg.policy, model=model, enable_field=['learn']) + buffer_ = DequeBuffer(size=cfg.policy.other.replay_buffer.replay_buffer_size) + + with task.start(async_mode=False, ctx=OnlineRLContext()): + assert task.router.is_active, "Please execute this script with ditask! See note in the header." + logging.info("Learner running on node {}".format(task.router.node_id)) + + from ding.utils import DistContext, get_rank + with DistContext(): + rank = get_rank() + task.use( + context_exchanger( + send_keys=["train_iter"], + recv_keys=["trajectories", "episodes", "env_step", "env_episode"], + skip_n_iter=0 + ) + ) + task.use(model_exchanger(model, is_learner=True)) + task.use(nstep_reward_enhancer(cfg)) + task.use(data_pusher(cfg, buffer_)) + task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) + if rank == 0: + task.use(CkptSaver(policy, cfg.exp_name, train_freq=1000)) + task.run() + + +def collector(): + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + ding_init(cfg) + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + + model = DQN(**cfg.policy.model) + policy = DQNPolicy(cfg.policy, model=model, enable_field=['collect']) + collector_cfg = deepcopy(cfg.env) + collector_cfg.is_train = True + collector_env = SubprocessEnvManagerV2( + env_fn=[lambda: AtariEnv(collector_cfg) for _ in range(cfg.env.collector_env_num)], cfg=cfg.env.manager + ) + + with task.start(async_mode=False, ctx=OnlineRLContext()): + assert task.router.is_active, "Please execute this script with ditask! See note in the header." + logging.info("Collector running on node {}".format(task.router.node_id)) + + task.use( + context_exchanger( + send_keys=["trajectories", "episodes", "env_step", "env_episode"], + recv_keys=["train_iter"], + skip_n_iter=1 + ) + ) + task.use(model_exchanger(model, is_learner=False)) + task.use(eps_greedy_handler(cfg)) + task.use(StepCollector(cfg, policy.collect_mode, collector_env)) + task.use(termination_checker(max_env_step=int(1e7))) + task.run() + + +def evaluator(): + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + ding_init(cfg) + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + + model = DQN(**cfg.policy.model) + policy = DQNPolicy(cfg.policy, model=model, enable_field=['eval']) + evaluator_cfg = deepcopy(cfg.env) + evaluator_cfg.is_train = False + evaluator_env = SubprocessEnvManagerV2( + env_fn=[lambda: AtariEnv(evaluator_cfg) for _ in range(cfg.env.evaluator_env_num)], cfg=cfg.env.manager + ) + + with task.start(async_mode=False, ctx=OnlineRLContext()): + assert task.router.is_active, "Please execute this script with ditask! See note in the header." + logging.info("Evaluator running on node {}".format(task.router.node_id)) + + task.use(context_exchanger(recv_keys=["train_iter", "env_step"], skip_n_iter=1)) + task.use(model_exchanger(model, is_learner=False)) + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) + task.use(CkptSaver(policy, cfg.exp_name, save_finish=False)) + task.use(online_logger(record_train_iter=True)) + task.run() diff --git a/dizoo/atari/example/atari_dqn_dist_rdma.py b/dizoo/atari/example/atari_dqn_dist_rdma.py new file mode 100644 index 0000000000..b364b37966 --- /dev/null +++ b/dizoo/atari/example/atari_dqn_dist_rdma.py @@ -0,0 +1,72 @@ +from copy import deepcopy +from ditk import logging +from ding.model import DQN +from ding.policy import DQNPolicy +from ding.envs import DingEnvWrapper, SubprocessEnvManagerV2 +from ding.data import DequeBuffer +from ding.config import compile_config +from ding.framework import task, ding_init +from ding.framework.context import OnlineRLContext +from ding.framework.middleware import OffPolicyLearner, StepCollector, interaction_evaluator, data_pusher, \ + eps_greedy_handler, CkptSaver, context_exchanger, model_exchanger, termination_checker, nstep_reward_enhancer, \ + online_logger +from ding.utils import set_pkg_seed +from dizoo.atari.envs.atari_env import AtariEnv +from dizoo.atari.config.serial.pong.pong_dqn_config import main_config, create_config + + +def main(): + logging.getLogger().setLevel(logging.INFO) + main_config.exp_name = 'pong_dqn_seed0_dist_rdma' + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + ding_init(cfg) + with task.start(async_mode=False, ctx=OnlineRLContext()): + assert task.router.is_active, "Please execute this script with ditask! See note in the header." + + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + + model = DQN(**cfg.policy.model) + policy = DQNPolicy(cfg.policy, model=model) + + if 'learner' in task.router.labels: + logging.info("Learner running on node {}".format(task.router.node_id)) + buffer_ = DequeBuffer(size=cfg.policy.other.replay_buffer.replay_buffer_size) + task.use( + context_exchanger( + send_keys=["train_iter"], + recv_keys=["trajectories", "episodes", "env_step", "env_episode"], + skip_n_iter=0 + ) + ) + task.use(model_exchanger(model, is_learner=True)) + task.use(nstep_reward_enhancer(cfg)) + task.use(data_pusher(cfg, buffer_)) + task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=1000)) + + elif 'collector' in task.router.labels: + logging.info("Collector running on node {}".format(task.router.node_id)) + collector_cfg = deepcopy(cfg.env) + collector_cfg.is_train = True + collector_env = SubprocessEnvManagerV2( + env_fn=[lambda: AtariEnv(collector_cfg) for _ in range(cfg.env.collector_env_num)], cfg=cfg.env.manager + ) + task.use( + context_exchanger( + send_keys=["trajectories", "episodes", "env_step", "env_episode"], + recv_keys=["train_iter"], + skip_n_iter=1 + ) + ) + task.use(model_exchanger(model, is_learner=False)) + task.use(eps_greedy_handler(cfg)) + task.use(StepCollector(cfg, policy.collect_mode, collector_env)) + task.use(termination_checker(max_env_step=int(1e7))) + else: + raise KeyError("invalid router labels: {}".format(task.router.labels)) + + task.run() + + +if __name__ == "__main__": + main() diff --git a/dizoo/atari/example/atari_dqn_dp.py b/dizoo/atari/example/atari_dqn_dp.py new file mode 100644 index 0000000000..540fb4ab07 --- /dev/null +++ b/dizoo/atari/example/atari_dqn_dp.py @@ -0,0 +1,53 @@ +from copy import deepcopy +from ditk import logging +from ding.model import DQN +from ding.policy import DQNPolicy +from ding.envs import DingEnvWrapper, SubprocessEnvManagerV2 +from ding.data import DequeBuffer +from ding.config import compile_config +from ding.torch_utils import DataParallel +from ding.framework import task +from ding.framework.context import OnlineRLContext +from ding.framework.middleware import OffPolicyLearner, StepCollector, interaction_evaluator, data_pusher, \ + eps_greedy_handler, CkptSaver, nstep_reward_enhancer, termination_checker +from ding.utils import set_pkg_seed +from dizoo.atari.envs.atari_env import AtariEnv +from dizoo.atari.config.serial.pong.pong_dqn_config import main_config, create_config + + +def main(): + logging.getLogger().setLevel(logging.INFO) + main_config.exp_name = 'pong_dqn_seed0_dp' + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + with task.start(async_mode=False, ctx=OnlineRLContext()): + collector_cfg = deepcopy(cfg.env) + collector_cfg.is_train = True + evaluator_cfg = deepcopy(cfg.env) + evaluator_cfg.is_train = False + collector_env = SubprocessEnvManagerV2( + env_fn=[lambda: AtariEnv(collector_cfg) for _ in range(cfg.env.collector_env_num)], cfg=cfg.env.manager + ) + evaluator_env = SubprocessEnvManagerV2( + env_fn=[lambda: AtariEnv(evaluator_cfg) for _ in range(cfg.env.evaluator_env_num)], cfg=cfg.env.manager + ) + + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + + model = DQN(**cfg.policy.model) + model = DataParallel(model) + buffer_ = DequeBuffer(size=cfg.policy.other.replay_buffer.replay_buffer_size) + policy = DQNPolicy(cfg.policy, model=model) + + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) + task.use(eps_greedy_handler(cfg)) + task.use(StepCollector(cfg, policy.collect_mode, collector_env)) + task.use(nstep_reward_enhancer(cfg)) + task.use(data_pusher(cfg, buffer_)) + task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=1000)) + task.use(termination_checker(max_env_step=int(1e7))) + task.run() + + +if __name__ == "__main__": + main() diff --git a/dizoo/atari/example/atari_ppo.py b/dizoo/atari/example/atari_ppo.py new file mode 100644 index 0000000000..cc26fb0a59 --- /dev/null +++ b/dizoo/atari/example/atari_ppo.py @@ -0,0 +1,47 @@ +from copy import deepcopy +from ditk import logging +from ding.model import VAC +from ding.policy import PPOPolicy +from ding.envs import DingEnvWrapper, SubprocessEnvManagerV2 +from ding.data import DequeBuffer +from ding.config import compile_config +from ding.framework import task +from ding.framework.context import OnlineRLContext +from ding.framework.middleware import multistep_trainer, StepCollector, interaction_evaluator, CkptSaver, \ + gae_estimator, termination_checker +from ding.utils import set_pkg_seed +from dizoo.atari.envs.atari_env import AtariEnv +from dizoo.atari.config.serial.pong.pong_ppo_config import main_config, create_config + + +def main(): + logging.getLogger().setLevel(logging.INFO) + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + with task.start(async_mode=False, ctx=OnlineRLContext()): + collector_cfg = deepcopy(cfg.env) + collector_cfg.is_train = True + evaluator_cfg = deepcopy(cfg.env) + evaluator_cfg.is_train = False + collector_env = SubprocessEnvManagerV2( + env_fn=[lambda: AtariEnv(collector_cfg) for _ in range(cfg.env.collector_env_num)], cfg=cfg.env.manager + ) + evaluator_env = SubprocessEnvManagerV2( + env_fn=[lambda: AtariEnv(evaluator_cfg) for _ in range(cfg.env.evaluator_env_num)], cfg=cfg.env.manager + ) + + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + + model = VAC(**cfg.policy.model) + policy = PPOPolicy(cfg.policy, model=model) + + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) + task.use(StepCollector(cfg, policy.collect_mode, collector_env)) + task.use(gae_estimator(cfg, policy.collect_mode)) + task.use(multistep_trainer(cfg, policy.learn_mode)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=1000)) + task.use(termination_checker(max_env_step=int(1e7))) + task.run() + + +if __name__ == "__main__": + main() diff --git a/dizoo/atari/example/atari_ppo_ddp.py b/dizoo/atari/example/atari_ppo_ddp.py new file mode 100644 index 0000000000..b8aeedefa0 --- /dev/null +++ b/dizoo/atari/example/atari_ppo_ddp.py @@ -0,0 +1,63 @@ +from copy import deepcopy +from ditk import logging +from ding.model import VAC +from ding.policy import PPOPolicy +from ding.envs import DingEnvWrapper, SubprocessEnvManagerV2 +from ding.data import DequeBuffer +from ding.config import compile_config +from ding.framework import task, ding_init +from ding.framework.context import OnlineRLContext +from ding.framework.middleware import multistep_trainer, StepCollector, interaction_evaluator, CkptSaver, \ + gae_estimator, ddp_termination_checker, online_logger +from ding.utils import set_pkg_seed, DDPContext, get_rank, get_world_size +from dizoo.atari.envs.atari_env import AtariEnv +from dizoo.atari.config.serial.pong.pong_ppo_config import main_config, create_config + + +def main(): + logging.getLogger().setLevel(logging.INFO) + with DDPContext(): + rank, world_size = get_rank(), get_world_size() + main_config.example = 'pong_ppo_seed0_ddp_avgsplit' + main_config.policy.multi_gpu = True + main_config.policy.learn.batch_size = main_config.policy.learn.batch_size // world_size + main_config.policy.collect.n_sample = main_config.policy.collect.n_sample // world_size + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + ding_init(cfg) + with task.start(async_mode=False, ctx=OnlineRLContext()): + collector_cfg = deepcopy(cfg.env) + collector_cfg.is_train = True + evaluator_cfg = deepcopy(cfg.env) + evaluator_cfg.is_train = False + collector_env = SubprocessEnvManagerV2( + env_fn=[lambda: AtariEnv(collector_cfg) for _ in range(cfg.env.collector_env_num)], cfg=cfg.env.manager + ) + evaluator_env = SubprocessEnvManagerV2( + env_fn=[lambda: AtariEnv(evaluator_cfg) for _ in range(cfg.env.evaluator_env_num)], cfg=cfg.env.manager + ) + + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + + model = VAC(**cfg.policy.model) + policy = PPOPolicy(cfg.policy, model=model) + + if rank == 0: + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) + task.use(StepCollector(cfg, policy.collect_mode, collector_env)) + task.use(gae_estimator(cfg, policy.collect_mode)) + task.use(multistep_trainer(policy.learn_mode)) + if rank == 0: + task.use(CkptSaver(policy, cfg.exp_name, train_freq=1000)) + task.use(online_logger(record_train_iter=True)) + task.use(ddp_termination_checker(max_env_step=int(1e7), rank=rank)) + task.run() + + +if __name__ == "__main__": + """ + Overview: + This script should be executed with GPUs. + Run the following command to launch the script: + python -m torch.distributed.launch --nproc_per_node=2 ./dizoo/atari/example/atari_ppo_ddp.py + """ + main() diff --git a/dizoo/beergame/__init__.py b/dizoo/beergame/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/dizoo/beergame/beergame.png b/dizoo/beergame/beergame.png new file mode 100644 index 0000000000..ec76d6e5ca Binary files /dev/null and b/dizoo/beergame/beergame.png differ diff --git a/dizoo/beergame/config/beergame_onppo_config.py b/dizoo/beergame/config/beergame_onppo_config.py new file mode 100644 index 0000000000..f6bec16a87 --- /dev/null +++ b/dizoo/beergame/config/beergame_onppo_config.py @@ -0,0 +1,70 @@ +from easydict import EasyDict + +beergame_ppo_config = dict( + exp_name='beergame_ppo_seed0', + env=dict( + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=200, + role=0, # 0-3 : retailer, warehouse, distributor, manufacturer + agent_type='bs', + # type of co-player, 'bs'- base stock, 'Strm'- use Sterman formula to model typical human behavior + demandDistribution=0 + # distribution of demand, default=0, '0=uniform, 1=normal distribution, 2=the sequence of 4,4,4,4,8,..., 3= basket data, 4= forecast data' + ), + policy=dict( + cuda=True, + recompute_adv=True, + action_space='discrete', + model=dict( + obs_shape=50, # statedim * multPerdInpt= 5 * 10 + action_shape=5, # the quantity relative to the arriving order + action_space='discrete', + encoder_hidden_size_list=[64, 64, 128], + actor_head_hidden_size=128, + critic_head_hidden_size=128, + ), + learn=dict( + epoch_per_collect=10, + batch_size=320, + learning_rate=3e-4, + entropy_weight=0.001, + adv_norm=True, + value_norm=True, + # for onppo, when we recompute adv, we need the key done in data to split traj, so we must + # use ignore_done=False here, + # but when we add key traj_flag in data as the backup for key done, we could choose to use ignore_done=True + # for halfcheetah, the length=1000 + ignore_done=True, + ), + collect=dict( + n_episode=8, + discount_factor=0.99, + gae_lambda=0.95, + collector=dict( + get_train_sample=True, + reward_shaping=True, # whether use total return to reshape reward + ), + ), + eval=dict(evaluator=dict(eval_freq=500, )), + ), +) +beergame_ppo_config = EasyDict(beergame_ppo_config) +main_config = beergame_ppo_config +beergame_ppo_create_config = dict( + env=dict( + type='beergame', + import_names=['dizoo.beergame.envs.beergame_env'], + ), + env_manager=dict(type='base'), + policy=dict(type='ppo'), + collector=dict(type='episode', ), +) +beergame_ppo_create_config = EasyDict(beergame_ppo_create_config) +create_config = beergame_ppo_create_config + +if __name__ == "__main__": + # or you can enter `ding -m serial_onpolicy -c beergame_onppo_config.py -s 0` + from ding.entry import serial_pipeline_onpolicy + serial_pipeline_onpolicy([main_config, create_config], seed=0) diff --git a/dizoo/beergame/entry/beergame_eval.py b/dizoo/beergame/entry/beergame_eval.py new file mode 100644 index 0000000000..c37a076f67 --- /dev/null +++ b/dizoo/beergame/entry/beergame_eval.py @@ -0,0 +1,42 @@ +import os +import torch +from tensorboardX import SummaryWriter + +from ding.config import compile_config +from ding.worker import InteractionSerialEvaluator +from ding.envs import BaseEnvManager +from ding.policy import PPOPolicy +from ding.model import VAC +from ding.utils import set_pkg_seed +from dizoo.beergame.config.beergame_onppo_config import beergame_ppo_config, beergame_ppo_create_config +from ding.envs import get_vec_env_setting +from functools import partial + + +def main(cfg, seed=0): + env_fn = None + cfg, create_cfg = beergame_ppo_config, beergame_ppo_create_config + cfg = compile_config(cfg, seed=seed, env=env_fn, auto=True, create_cfg=create_cfg, save_cfg=True) + collector_env_num, evaluator_env_num = cfg.env.collector_env_num, cfg.env.evaluator_env_num + + env_fn, collector_env_cfg, evaluator_env_cfg = get_vec_env_setting(cfg.env) + cfg.env.manager.auto_reset = False + evaluator_env = BaseEnvManager(env_fn=[partial(env_fn, cfg=c) for c in evaluator_env_cfg], cfg=cfg.env.manager) + evaluator_env.seed(seed, dynamic_seed=False) + set_pkg_seed(seed, use_cuda=cfg.policy.cuda) + model = VAC(**cfg.policy.model) + tb_logger = SummaryWriter(os.path.join('./{}/log/'.format(cfg.exp_name), 'serial')) + policy = PPOPolicy(cfg.policy, model=model) + # set the path to save figure + cfg.policy.eval.evaluator.figure_path = './' + evaluator = InteractionSerialEvaluator( + cfg.policy.eval.evaluator, evaluator_env, policy.eval_mode, tb_logger, exp_name=cfg.exp_name + ) + # load model + model.load_state_dict(torch.load('model path', map_location='cpu')["model"]) + evaluator.eval(None, -1, -1) + + +if __name__ == "__main__": + beergame_ppo_config.exp_name = 'beergame_evaluate' + main(beergame_ppo_config) diff --git a/dizoo/beergame/envs/BGAgent.py b/dizoo/beergame/envs/BGAgent.py new file mode 100644 index 0000000000..d06866afc2 --- /dev/null +++ b/dizoo/beergame/envs/BGAgent.py @@ -0,0 +1,152 @@ +# Code Reference: https://github.com/OptMLGroup/DeepBeerInventory-RL. +import argparse +import numpy as np + + +# Here we want to define the agent class for the BeerGame +class Agent(object): + # initializes the agents with initial values for IL, OO and saves self.agentNum for recognizing the agents. + def __init__( + self, agentNum: int, IL: int, AO: int, AS: int, c_h: float, c_p: float, eta: int, compuType: str, + config: argparse.Namespace + ) -> None: + self.agentNum = agentNum + self.IL = IL # Inventory level of each agent - changes during the game + self.OO = 0 # Open order of each agent - changes during the game + self.ASInitial = AS # the initial arriving shipment. + self.ILInitial = IL # IL at which we start each game with this number + self.AOInitial = AO # OO at which we start each game with this number + self.config = config # an instance of config is stored inside the class + self.curState = [] # this function gets the current state of the game + self.nextState = [] + self.curReward = 0 # the reward observed at the current step + self.cumReward = 0 # cumulative reward; reset at the beginning of each episode + self.totRew = 0 # it is reward of all players obtained for the current player. + self.c_h = c_h # holding cost + self.c_p = c_p # backorder cost + self.eta = eta # the total cost regulazer + self.AS = np.zeros((1, 1)) # arriced shipment + self.AO = np.zeros((1, 1)) # arrived order + self.action = 0 # the action at time t + self.compType = compuType + # self.compTypeTrain = compuType # rnd -> random / srdqn-> srdqn / Strm-> formula-Rong2008 / bs -> optimal policy if exists + # self.compTypeTest = compuType # rnd -> random / srdqn-> srdqn / Strm-> formula-Rong2008 / bs -> optimal policy if exists + self.alpha_b = self.config.alpha_b[self.agentNum] # parameters for the formula + self.betta_b = self.config.betta_b[self.agentNum] # parameters for the formula + if self.config.demandDistribution == 0: + self.a_b = np.mean((self.config.demandUp, self.config.demandLow)) # parameters for the formula + self.b_b = np.mean((self.config.demandUp, self.config.demandLow)) * ( + np.mean((self.config.leadRecItemLow[self.agentNum], self.config.leadRecItemUp[self.agentNum])) + + np.mean((self.config.leadRecOrderLow[self.agentNum], self.config.leadRecOrderUp[self.agentNum])) + ) # parameters for the formula + elif self.config.demandDistribution == 1 or self.config.demandDistribution == 3 or self.config.demandDistribution == 4: + self.a_b = self.config.demandMu # parameters for the formula + self.b_b = self.config.demandMu * ( + np.mean((self.config.leadRecItemLow[self.agentNum], self.config.leadRecItemUp[self.agentNum])) + + np.mean((self.config.leadRecOrderLow[self.agentNum], self.config.leadRecOrderUp[self.agentNum])) + ) # parameters for the formula + elif self.config.demandDistribution == 2: + self.a_b = 8 # parameters for the formula + self.b_b = (3 / 4.) * 8 * ( + np.mean((self.config.leadRecItemLow[self.agentNum], self.config.leadRecItemUp[self.agentNum])) + + np.mean((self.config.leadRecOrderLow[self.agentNum], self.config.leadRecOrderUp[self.agentNum])) + ) # parameters for the formula + elif self.config.demandDistribution == 3: + self.a_b = 10 # parameters for the formula + self.b_b = 7 * ( + np.mean((self.config.leadRecItemLow[self.agentNum], self.config.leadRecItemUp[self.agentNum])) + + np.mean((self.config.leadRecOrderLow[self.agentNum], self.config.leadRecOrderUp[self.agentNum])) + ) # parameters for the formula + else: + raise Exception('The demand distribution is not defined or it is not a valid type.!') + + self.hist = [] # this is used for plotting - keeps the history for only one game + self.hist2 = [] # this is used for animation usage + self.srdqnBaseStock = [] # this holds the base stock levels that srdqn has came up with. added on Nov 8, 2017 + self.T = 0 + self.bsBaseStock = 0 + self.init_bsBaseStock = 0 + self.nextObservation = [] + + if self.compType == 'srdqn': + # sets the initial input of the network + self.currentState = np.stack( + [self.curState for _ in range(self.config.multPerdInpt)], axis=0 + ) # multPerdInpt observations stacked. each row is an observation + + # reset player information + def resetPlayer(self, T: int): + self.IL = self.ILInitial + self.OO = 0 + self.AS = np.squeeze( + np.zeros((1, T + max(self.config.leadRecItemUp) + max(self.config.leadRecOrderUp) + 10)) + ) # arriced shipment + self.AO = np.squeeze( + np.zeros((1, T + max(self.config.leadRecItemUp) + max(self.config.leadRecOrderUp) + 10)) + ) # arrived order + if self.agentNum != 0: + for i in range(self.config.leadRecOrderUp_aux[self.agentNum - 1]): + self.AO[i] = self.AOInitial[self.agentNum - 1] + for i in range(self.config.leadRecItemUp[self.agentNum]): + self.AS[i] = self.ASInitial + self.curReward = 0 # the reward observed at the current step + self.cumReward = 0 # cumulative reward; reset at the begining of each episode + self.action = [] + self.hist = [] + self.hist2 = [] + self.srdqnBaseStock = [] # this holds the base stock levels that srdqn has came up with. added on Nov 8, 2017 + self.T = T + self.curObservation = self.getCurState(1) # this function gets the current state of the game + self.nextObservation = [] + if self.compType == 'srdqn': + self.currentState = np.stack([self.curObservation for _ in range(self.config.multPerdInpt)], axis=0) + + # updates the IL and OO at time t, after recieving "rec" number of items + def recieveItems(self, time: int) -> None: + self.IL = self.IL + self.AS[time] # inverntory level update + self.OO = self.OO - self.AS[time] # invertory in transient update + + # find action Value associated with the action list + def actionValue(self, curTime: int) -> int: + if self.config.fixedAction: + a = self.config.actionList[np.argmax(self.action)] + else: + # "d + x" rule + if self.compType == 'srdqn': + a = max(0, self.config.actionList[np.argmax(self.action)] * self.config.action_step + self.AO[curTime]) + elif self.compType == 'rnd': + a = max(0, self.config.actionList[np.argmax(self.action)] + self.AO[curTime]) + else: + a = max(0, self.config.actionListOpt[np.argmax(self.action)]) + + return a + + # getReward returns the reward at the current state + def getReward(self) -> None: + # cost (holding + backorder) for one time unit + self.curReward = (self.c_p * max(0, -self.IL) + self.c_h * max(0, self.IL)) / 200. # self.config.Ttest # + self.curReward = -self.curReward + # make reward negative, because it is the cost + + # sum total reward of each agent + self.cumReward = self.config.gamma * self.cumReward + self.curReward + + # This function returns a np.array of the current state of the agent + def getCurState(self, t: int) -> np.ndarray: + if self.config.ifUseASAO: + if self.config.if_use_AS_t_plus_1: + curState = np.array( + [-1 * (self.IL < 0) * self.IL, 1 * (self.IL > 0) * self.IL, self.OO, self.AS[t], self.AO[t]] + ) + else: + curState = np.array( + [-1 * (self.IL < 0) * self.IL, 1 * (self.IL > 0) * self.IL, self.OO, self.AS[t - 1], self.AO[t]] + ) + else: + curState = np.array([-1 * (self.IL < 0) * self.IL, 1 * (self.IL > 0) * self.IL, self.OO]) + + if self.config.ifUseActionInD: + a = self.config.actionList[np.argmax(self.action)] + curState = np.concatenate((curState, np.array([a]))) + + return curState diff --git a/dizoo/beergame/envs/__init__.py b/dizoo/beergame/envs/__init__.py new file mode 100644 index 0000000000..d4ffbfd452 --- /dev/null +++ b/dizoo/beergame/envs/__init__.py @@ -0,0 +1,2 @@ +from .clBeergame import clBeerGame +from .beergame_core import BeerGame diff --git a/dizoo/beergame/envs/beergame_core.py b/dizoo/beergame/envs/beergame_core.py new file mode 100644 index 0000000000..2f0ac61910 --- /dev/null +++ b/dizoo/beergame/envs/beergame_core.py @@ -0,0 +1,112 @@ +from __future__ import print_function +from dizoo.beergame.envs import clBeerGame +from torch import Tensor +import numpy as np +import random +from .utils import get_config, update_config +import gym +import os +from typing import Optional + + +class BeerGame(): + + def __init__(self, role: int, agent_type: str, demandDistribution: int) -> None: + self._cfg, unparsed = get_config() + self._role = role + # prepare loggers and directories + # prepare_dirs_and_logger(self._cfg) + self._cfg = update_config(self._cfg) + + # set agent type + if agent_type == 'bs': + self._cfg.agentTypes = ["bs", "bs", "bs", "bs"] + elif agent_type == 'Strm': + self._cfg.agentTypes = ["Strm", "Strm", "Strm", "Strm"] + self._cfg.agentTypes[role] = "srdqn" + + self._cfg.demandDistribution = demandDistribution + + # load demands:0=uniform, 1=normal distribution, 2=the sequence of 4,4,4,4,8,..., 3= basket data, 4= forecast data + if self._cfg.observation_data: + adsr = 'data/demandTr-obs-' + elif self._cfg.demandDistribution == 3: + if self._cfg.scaled: + adsr = 'data/basket_data/scaled' + else: + adsr = 'data/basket_data' + direc = os.path.realpath(adsr + '/demandTr-' + str(self._cfg.data_id) + '.npy') + self._demandTr = np.load(direc) + print("loaded training set=", direc) + elif self._cfg.demandDistribution == 4: + if self._cfg.scaled: + adsr = 'data/forecast_data/scaled' + else: + adsr = 'data/forecast_data' + direc = os.path.realpath(adsr + '/demandTr-' + str(self._cfg.data_id) + '.npy') + self._demandTr = np.load(direc) + print("loaded training set=", direc) + else: + if self._cfg.demandDistribution == 0: # uniform + self._demandTr = np.random.randint(0, self._cfg.demandUp, size=[self._cfg.demandSize, self._cfg.TUp]) + elif self._cfg.demandDistribution == 1: # normal distribution + self._demandTr = np.round( + np.random.normal( + self._cfg.demandMu, self._cfg.demandSigma, size=[self._cfg.demandSize, self._cfg.TUp] + ) + ).astype(int) + elif self._cfg.demandDistribution == 2: # the sequence of 4,4,4,4,8,... + self._demandTr = np.concatenate( + (4 * np.ones((self._cfg.demandSize, 4)), 8 * np.ones((self._cfg.demandSize, 98))), axis=1 + ).astype(int) + + # initilize an instance of Beergame + self._env = clBeerGame(self._cfg) + self.observation_space = gym.spaces.Box( + low=float("-inf"), + high=float("inf"), + shape=(self._cfg.stateDim * self._cfg.multPerdInpt, ), + dtype=np.float32 + ) # state_space = state_dim * m (considering the reward delay) + self.action_space = gym.spaces.Discrete(self._cfg.actionListLen) # length of action list + self.reward_space = gym.spaces.Box(low=float("-inf"), high=float("inf"), shape=(1, ), dtype=np.float32) + + # get the length of the demand. + self._demand_len = np.shape(self._demandTr)[0] + + def reset(self): + self._env.resetGame(demand=self._demandTr[random.randint(0, self._demand_len - 1)]) + obs = [i for item in self._env.players[self._role].currentState for i in item] + return obs + + def seed(self, seed: int) -> None: + self._seed = seed + np.random.seed(self._seed) + + def close(self) -> None: + pass + + def step(self, action: np.ndarray): + self._env.handelAction(action) + self._env.next() + newstate = np.append( + self._env.players[self._role].currentState[1:, :], [self._env.players[self._role].nextObservation], axis=0 + ) + self._env.players[self._role].currentState = newstate + obs = [i for item in newstate for i in item] + rew = self._env.players[self._role].curReward + done = (self._env.curTime == self._env.T) + info = {} + return obs, rew, done, info + + def reward_shaping(self, reward: Tensor) -> Tensor: + self._totRew, self._cumReward = self._env.distTotReward(self._role) + reward += (self._cfg.distCoeff / 3) * ((self._totRew - self._cumReward) / (self._env.T)) + return reward + + def enable_save_figure(self, figure_path: Optional[str] = None) -> None: + self._cfg.ifSaveFigure = True + if figure_path is None: + figure_path = './' + self._cfg.figure_dir = figure_path + self._env.doTestMid(self._demandTr[random.randint(0, self._demand_len - 1)]) diff --git a/dizoo/beergame/envs/beergame_env.py b/dizoo/beergame/envs/beergame_env.py new file mode 100644 index 0000000000..f48cdd93cf --- /dev/null +++ b/dizoo/beergame/envs/beergame_env.py @@ -0,0 +1,84 @@ +import numpy as np +from dizoo.beergame.envs.beergame_core import BeerGame +from typing import Union, List, Optional + +from ding.envs import BaseEnv, BaseEnvTimestep +from ding.utils import ENV_REGISTRY +from ding.torch_utils import to_ndarray +import copy + + +@ENV_REGISTRY.register('beergame') +class BeerGameEnv(BaseEnv): + + def __init__(self, cfg: dict) -> None: + self._cfg = cfg + self._init_flag = False + + def reset(self) -> np.ndarray: + if not self._init_flag: + self._env = BeerGame(self._cfg.role, self._cfg.agent_type, self._cfg.demandDistribution) + self._observation_space = self._env.observation_space + self._action_space = self._env.action_space + self._reward_space = self._env.reward_space + self._init_flag = True + if hasattr(self, '_seed') and hasattr(self, '_dynamic_seed') and self._dynamic_seed: + np_seed = 100 * np.random.randint(1, 1000) + self._env.seed(self._seed + np_seed) + elif hasattr(self, '_seed'): + self._env.seed(self._seed) + self._eval_episode_return = 0 + obs = self._env.reset() + obs = to_ndarray(obs).astype(np.float32) + return obs + + def close(self) -> None: + if self._init_flag: + self._env.close() + self._init_flag = False + + def seed(self, seed: int, dynamic_seed: bool = True) -> None: + self._seed = seed + self._dynamic_seed = dynamic_seed + np.random.seed(self._seed) + + def step(self, action: Union[int, np.ndarray]) -> BaseEnvTimestep: + if isinstance(action, np.ndarray) and action.shape == (1, ): + action = action.squeeze() # 0-dim array + obs, rew, done, info = self._env.step(action) + self._eval_episode_return += rew + if done: + info['eval_episode_return'] = self._eval_episode_return + obs = to_ndarray(obs).astype(np.float32) + rew = to_ndarray([rew]).astype(np.float32) # wrapped to be transfered to a array with shape (1,) + return BaseEnvTimestep(obs, rew, done, info) + + def reward_shaping(self, transitions: List[dict]) -> List[dict]: + new_transitions = copy.deepcopy(transitions) + for trans in new_transitions: + trans['reward'] = self._env.reward_shaping(trans['reward']) + return new_transitions + + def random_action(self) -> np.ndarray: + random_action = self.action_space.sample() + if isinstance(random_action, int): + random_action = to_ndarray([random_action], dtype=np.int64) + return random_action + + def enable_save_figure(self, figure_path: Optional[str] = None) -> None: + self._env.enable_save_figure(figure_path) + + @property + def observation_space(self) -> int: + return self._observation_space + + @property + def action_space(self) -> int: + return self._action_space + + @property + def reward_space(self) -> int: + return self._reward_space + + def __repr__(self) -> str: + return "DI-engine Beergame Env" diff --git a/dizoo/beergame/envs/clBeergame.py b/dizoo/beergame/envs/clBeergame.py new file mode 100644 index 0000000000..9a237fbb68 --- /dev/null +++ b/dizoo/beergame/envs/clBeergame.py @@ -0,0 +1,439 @@ +# Code Reference: https://github.com/OptMLGroup/DeepBeerInventory-RL. +import numpy as np +from random import randint +from .BGAgent import Agent +from matplotlib import rc +rc('text', usetex=True) +from .plotting import plotting, savePlot +import matplotlib.pyplot as plt +import os +import time +from time import gmtime, strftime + + +class clBeerGame(object): + + def __init__(self, config): + self.config = config + self.curGame = 0 # The number associated with the current game (counter of the game) + self.curTime = 0 + self.totIterPlayed = 0 # total iterations of the game, played so far in this and previous games + self.players = self.createAgent() # create the agents + self.T = 0 + self.demand = [] + self.ifOptimalSolExist = self.config.ifOptimalSolExist + self.getOptimalSol() + self.totRew = 0 # it is reward of all players obtained for the current player. + self.resultTest = [] + self.runnerMidlResults = [] # stores the results to use in runner comparisons + self.runnerFinlResults = [] # stores the results to use in runner comparisons + self.middleTestResult = [ + ] # stores the whole middle results of bs, Strm, and random to avoid doing same tests multiple of times. + self.runNumber = 0 # the runNumber which is used when use runner + self.strNum = 0 # the runNumber which is used when use runner + + # createAgent : Create agent objects (agentNum,IL,OO,c_h,c_p,type,config) + def createAgent(self): + agentTypes = self.config.agentTypes + return [ + Agent( + i, self.config.ILInit[i], self.config.AOInit, self.config.ASInit[i], self.config.c_h[i], + self.config.c_p[i], self.config.eta[i], agentTypes[i], self.config + ) for i in range(self.config.NoAgent) + ] + + # planHorizon : Find a random planning horizon + def planHorizon(self): + # TLow: minimum number for the planning horizon # TUp: maximum number for the planning horizon + # output: The planning horizon which is chosen randomly. + return randint(self.config.TLow, self.config.TUp) + + # this function resets the game for start of the new game + def resetGame(self, demand: np.ndarray): + self.demand = demand + self.curTime = 0 + self.curGame += 1 + self.totIterPlayed += self.T + self.T = self.planHorizon() + # reset the required information of player for each episode + for k in range(0, self.config.NoAgent): + self.players[k].resetPlayer(self.T) + + # update OO when there are initial IL,AO,AS + self.update_OO() + + # correction on cost at time T according to the cost of the other players + def getTotRew(self): + totRew = 0 + for i in range(self.config.NoAgent): + # sum all rewards for the agents and make correction + totRew += self.players[i].cumReward + + for i in range(self.config.NoAgent): + self.players[i].curReward += self.players[i].eta * (totRew - self.players[i].cumReward) # /(self.T) + + # make correction to the rewards in the experience replay for all iterations of current game + def distTotReward(self, role: int): + totRew = 0 + optRew = 0.1 # why? + for i in range(self.config.NoAgent): + # sum all rewards for the agents and make correction + totRew += self.players[i].cumReward + totRew += optRew + + return totRew, self.players[role].cumReward + + def getAction(self, k: int, action: np.ndarray, playType="train"): + if playType == "train": + if self.players[k].compType == "srdqn": + self.players[k].action = np.zeros(self.config.actionListLen) + self.players[k].action[action] = 1 + elif self.players[k].compType == "Strm": + self.players[k].action = np.zeros(self.config.actionListLenOpt) + self.players[k].action[np.argmin(np.abs(np.array(self.config.actionListOpt)\ + - max(0, round(self.players[k].AO[self.curTime] + \ + self.players[k].alpha_b*(self.players[k].IL - self.players[k].a_b) + \ + self.players[k].betta_b*(self.players[k].OO - self.players[k].b_b)))))] = 1 + elif self.players[k].compType == "rnd": + self.players[k].action = np.zeros(self.config.actionListLen) + a = np.random.randint(self.config.actionListLen) + self.players[k].action[a] = 1 + elif self.players[k].compType == "bs": + self.players[k].action = np.zeros(self.config.actionListLenOpt) + if self.config.demandDistribution == 2: + if self.curTime and self.config.use_initial_BS <= 4: + self.players[k].action[np.argmin(np.abs(np.array(self.config.actionListOpt) - \ + max(0, (self.players[k].int_bslBaseStock - (self.players[k].IL + self.players[k].OO - self.players[k].AO[self.curTime])))))] = 1 + else: + self.players[k].action[np.argmin(np.abs(np.array(self.config.actionListOpt) - \ + max(0, (self.players[k].bsBaseStock - (self.players[k].IL + self.players[k].OO - self.players[k].AO[self.curTime])))))] = 1 + else: + self.players[k].action[np.argmin(np.abs(np.array(self.config.actionListOpt) - \ + max(0, (self.players[k].bsBaseStock - (self.players[k].IL + self.players[k].OO - self.players[k].AO[self.curTime])))))] = 1 + elif playType == "test": + if self.players[k].compTypeTest == "srdqn": + self.players[k].action = np.zeros(self.config.actionListLen) + self.players[k].action = self.players[k].brain.getDNNAction(self.playType) + elif self.players[k].compTypeTest == "Strm": + self.players[k].action = np.zeros(self.config.actionListLenOpt) + + self.players[k].action[np.argmin(np.abs(np.array(self.config.actionListOpt)-\ + max(0,round(self.players[k].AO[self.curTime] +\ + self.players[k].alpha_b*(self.players[k].IL - self.players[k].a_b) +\ + self.players[k].betta_b*(self.players[k].OO - self.players[k].b_b)))))] = 1 + elif self.players[k].compTypeTest == "rnd": + self.players[k].action = np.zeros(self.config.actionListLen) + a = np.random.randint(self.config.actionListLen) + self.players[k].action[a] = 1 + elif self.players[k].compTypeTest == "bs": + self.players[k].action = np.zeros(self.config.actionListLenOpt) + + if self.config.demandDistribution == 2: + if self.curTime and self.config.use_initial_BS <= 4: + self.players[k].action [np.argmin(np.abs(np.array(self.config.actionListOpt)-\ + max(0,(self.players[k].int_bslBaseStock - (self.players[k].IL + self.players[k].OO - self.players[k].AO[self.curTime]))) ))] = 1 + else: + self.players[k].action [np.argmin(np.abs(np.array(self.config.actionListOpt)-\ + max(0,(self.players[k].bsBaseStock - (self.players[k].IL + self.players[k].OO - self.players[k].AO[self.curTime]))) ))] = 1 + else: + self.players[k].action [np.argmin(np.abs(np.array(self.config.actionListOpt)-\ + max(0,(self.players[k].bsBaseStock - (self.players[k].IL + self.players[k].OO - self.players[k].AO[self.curTime]))) ))] = 1 + else: + # not a valid player is defined. + raise Exception('The player type is not defined or it is not a valid type.!') + + def next(self): + # get a random leadtime + leadTimeIn = randint( + self.config.leadRecItemLow[self.config.NoAgent - 1], self.config.leadRecItemUp[self.config.NoAgent - 1] + ) + # handle the most upstream recieved shipment + self.players[self.config.NoAgent - 1].AS[self.curTime + + leadTimeIn] += self.players[self.config.NoAgent - + 1].actionValue(self.curTime) + + for k in range(self.config.NoAgent - 1, -1, -1): # [3,2,1,0] + + # get current IL and Backorder + current_IL = max(0, self.players[k].IL) + current_backorder = max(0, -self.players[k].IL) + + # TODO: We have get the AS and AO from the UI and update our AS and AO, so that code update the corresponding variables + + # increase IL and decrease OO based on the action, for the next period + self.players[k].recieveItems(self.curTime) + + # observe the reward + possible_shipment = min( + current_IL + self.players[k].AS[self.curTime], current_backorder + self.players[k].AO[self.curTime] + ) + + # plan arrivals of the items to the downstream agent + if self.players[k].agentNum > 0: + leadTimeIn = randint(self.config.leadRecItemLow[k - 1], self.config.leadRecItemUp[k - 1]) + self.players[k - 1].AS[self.curTime + leadTimeIn] += possible_shipment + + # update IL + self.players[k].IL -= self.players[k].AO[self.curTime] + # observe the reward + self.players[k].getReward() + self.players[k].hist[-1][-2] = self.players[k].curReward + self.players[k].hist2[-1][-2] = self.players[k].curReward + + # update next observation + self.players[k].nextObservation = self.players[k].getCurState(self.curTime + 1) + + if self.config.ifUseTotalReward: + # correction on cost at time T + if self.curTime == self.T: + self.getTotRew() + + self.curTime += 1 + + def handelAction(self, action: np.ndarray, playType="train"): + # get random lead time + leadTime = randint(self.config.leadRecOrderLow[0], self.config.leadRecOrderUp[0]) + # set AO + self.players[0].AO[self.curTime] += self.demand[self.curTime] + for k in range(0, self.config.NoAgent): + self.getAction(k, action, playType) + + self.players[k].srdqnBaseStock += [self.players[k].actionValue( \ + self.curTime) + self.players[k].IL + self.players[k].OO] + + # update hist for the plots + self.players[k].hist += [[self.curTime, self.players[k].IL, self.players[k].OO,\ + self.players[k].actionValue(self.curTime), self.players[k].curReward, self.players[k].srdqnBaseStock[-1]]] + + if self.players[k].compType == "srdqn": + self.players[k].hist2 += [[self.curTime, self.players[k].IL, self.players[k].OO, self.players[k].AO[self.curTime], self.players[k].AS[self.curTime], \ + self.players[k].actionValue(self.curTime), self.players[k].curReward, \ + self.config.actionList[np.argmax(self.players[k].action)]]] + + else: + self.players[k].hist2 += [[self.curTime, self.players[k].IL, self.players[k].OO, self.players[k].AO[self.curTime], self.players[k].AS[self.curTime], \ + self.players[k].actionValue(self.curTime), self.players[k].curReward, 0]] + + # updates OO and AO at time t+1 + self.players[k].OO += self.players[k].actionValue(self.curTime) # open order level update + leadTime = randint(self.config.leadRecOrderLow[k], self.config.leadRecOrderUp[k]) + if self.players[k].agentNum < self.config.NoAgent - 1: + self.players[k + 1].AO[self.curTime + leadTime] += self.players[k].actionValue( + self.curTime + ) # open order level update + + # check the Shang and Song (2003) condition, and if it works, obtains the base stock policy values for each agent + def getOptimalSol(self): + # if self.config.NoAgent !=1: + if self.config.NoAgent != 1 and 1 == 2: + # check the Shang and Song (2003) condition. + for k in range(self.config.NoAgent - 1): + if not (self.players[k].c_h == self.players[k + 1].c_h and self.players[k + 1].c_p == 0): + self.ifOptimalSolExist = False + + # if the Shang and Song (2003) condition satisfied, it runs the algorithm + if self.ifOptimalSolExist == True: + calculations = np.zeros((7, self.config.NoAgent)) + for k in range(self.config.NoAgent): + # DL_high + calculations[0][k] = ((self.config.leadRecItemLow + self.config.leadRecItemUp + 2) / 2 \ + + (self.config.leadRecOrderLow + self.config.leadRecOrderUp + 2) / 2) * \ + (self.config.demandUp - self.config.demandLow - 1) + if k > 0: + calculations[0][k] += calculations[0][k - 1] + # probability_high + nominator_ch = 0 + low_denominator_ch = 0 + for j in range(k, self.config.NoAgent): + if j < self.config.NoAgent - 1: + nominator_ch += self.players[j + 1].c_h + low_denominator_ch += self.players[j].c_h + if k == 0: + high_denominator_ch = low_denominator_ch + calculations[2][k] = (self.players[0].c_p + + nominator_ch) / (self.players[0].c_p + low_denominator_ch + 0.0) + # probability_low + calculations[3][k] = (self.players[0].c_p + + nominator_ch) / (self.players[0].c_p + high_denominator_ch + 0.0) + # S_high + calculations[4] = np.round(np.multiply(calculations[0], calculations[2])) + # S_low + calculations[5] = np.round(np.multiply(calculations[0], calculations[3])) + # S_avg + calculations[6] = np.round(np.mean(calculations[4:6], axis=0)) + # S', set the base stock values into each agent. + for k in range(self.config.NoAgent): + if k == 0: + self.players[k].bsBaseStock = calculations[6][k] + + else: + self.players[k].bsBaseStock = calculations[6][k] - calculations[6][k - 1] + if self.players[k].bsBaseStock < 0: + self.players[k].bsBaseStock = 0 + elif self.config.NoAgent == 1: + if self.config.demandDistribution == 0: + self.players[0].bsBaseStock = np.ceil( + self.config.c_h[0] / (self.config.c_h[0] + self.config.c_p[0] + 0.0) + ) * ((self.config.demandUp - self.config.demandLow - 1) / 2) * self.config.leadRecItemUp + elif 1 == 1: + f = self.config.f + f_init = self.config.f_init + for k in range(self.config.NoAgent): + self.players[k].bsBaseStock = f[k] + self.players[k].int_bslBaseStock = f_init[k] + + def update_OO(self): + for k in range(0, self.config.NoAgent): + if k < self.config.NoAgent - 1: + self.players[k].OO = sum(self.players[k + 1].AO) + sum(self.players[k].AS) + else: + self.players[k].OO = sum(self.players[k].AS) + + def doTestMid(self, demandTs): + self.resultTest = [] + m = strftime("%Y-%m-%d-%H-%M-%S", gmtime()) + self.doTest(m, demandTs) + print("---------------------------------------------------------------------------------------") + resultSummary = np.array(self.resultTest).mean(axis=0).tolist() + + result_srdqn = ', '.join(map("{:.2f}".format, resultSummary[0])) + result_rand = ', '.join(map("{:.2f}".format, resultSummary[1])) + result_strm = ', '.join(map("{:.2f}".format, resultSummary[2])) + if self.ifOptimalSolExist: + result_bs = ', '.join(map("{:.2f}".format, resultSummary[3])) + print( + 'SUMMARY; {0:s}; ITER= {1:d}; OURPOLICY= [{2:s}]; SUM = {3:2.4f}; Rand= [{4:s}]; SUM = {5:2.4f}; STRM= [{6:s}]; SUM = {7:2.4f}; BS= [{8:s}]; SUM = {9:2.4f}' + .format( + strftime("%Y-%m-%d %H:%M:%S", gmtime()), self.curGame, result_srdqn, sum(resultSummary[0]), + result_rand, sum(resultSummary[1]), result_strm, sum(resultSummary[2]), result_bs, + sum(resultSummary[3]) + ) + ) + + else: + print( + 'SUMMARY; {0:s}; ITER= {1:d}; OURPOLICY= [{2:s}]; SUM = {3:2.4f}; Rand= [{4:s}]; SUM = {5:2.4f}; STRM= [{6:s}]; SUM = {7:2.4f}' + .format( + strftime("%Y-%m-%d %H:%M:%S", gmtime()), self.curGame, result_srdqn, sum(resultSummary[0]), + result_rand, sum(resultSummary[1]), result_strm, sum(resultSummary[2]) + ) + ) + + print("=======================================================================================") + + def doTest(self, m, demand): + import matplotlib.pyplot as plt + if self.config.ifSaveFigure: + plt.figure(self.curGame, figsize=(12, 8), dpi=80, facecolor='w', edgecolor='k') + + # self.demand = demand + # use dnn to get output. + Rsltdnn, plt = self.tester(self.config.agentTypes, plt, 'b', 'OurPolicy', m) + baseStockdata = self.players[0].srdqnBaseStock + # # use random to get output. + RsltRnd, plt = self.tester(["rnd", "rnd", "rnd", "rnd"], plt, 'y', 'RAND', m) + + # use formual to get output. + RsltStrm, plt = self.tester(["Strm", "Strm", "Strm", "Strm"], plt, 'g', 'Strm', m) + + # use optimal strategy to get output, if it works. + if self.ifOptimalSolExist: + if self.config.agentTypes == ["srdqn", "Strm", "Strm", "Strm"]: + Rsltbs, plt = self.tester(["bs", "Strm", "Strm", "Strm"], plt, 'r', 'Strm-BS', m) + elif self.config.agentTypes == ["Strm", "srdqn", "Strm", "Strm"]: + Rsltbs, plt = self.tester(["Strm", "bs", "Strm", "Strm"], plt, 'r', 'Strm-BS', m) + elif self.config.agentTypes == ["Strm", "Strm", "srdqn", "Strm"]: + Rsltbs, plt = self.tester(["Strm", "Strm", "bs", "Strm"], plt, 'r', 'Strm-BS', m) + elif self.config.agentTypes == ["Strm", "Strm", "Strm", "srdqn"]: + Rsltbs, plt = self.tester(["Strm", "Strm", "Strm", "bs"], plt, 'r', 'Strm-BS', m) + elif self.config.agentTypes == ["srdqn", "rnd", "rnd", "rnd"]: + Rsltbs, plt = self.tester(["bs", "rnd", "rnd", "rnd"], plt, 'r', 'RND-BS', m) + elif self.config.agentTypes == ["rnd", "srdqn", "rnd", "rnd"]: + Rsltbs, plt = self.tester(["rnd", "bs", "rnd", "rnd"], plt, 'r', 'RND-BS', m) + elif self.config.agentTypes == ["rnd", "rnd", "srdqn", "rnd"]: + Rsltbs, plt = self.tester(["rnd", "rnd", "bs", "rnd"], plt, 'r', 'RND-BS', m) + elif self.config.agentTypes == ["rnd", "rnd", "rnd", "srdqn"]: + Rsltbs, plt = self.tester(["rnd", "rnd", "rnd", "bs"], plt, 'r', 'RND-BS', m) + else: + Rsltbs, plt = self.tester(["bs", "bs", "bs", "bs"], plt, 'r', 'BS', m) + # hold the results of the optimal solution + self.middleTestResult += [[RsltRnd, RsltStrm, Rsltbs]] + else: + self.middleTestResult += [[RsltRnd, RsltStrm]] + + else: + # return the obtained results into their lists + RsltRnd = self.middleTestResult[m][0] + RsltStrm = self.middleTestResult[m][1] + if self.ifOptimalSolExist: + Rsltbs = self.middleTestResult[m][2] + + # save the figure + if self.config.ifSaveFigure: + savePlot(self.players, self.curGame, Rsltdnn, RsltStrm, Rsltbs, RsltRnd, self.config, m) + plt.close() + + result_srdqn = ', '.join(map("{:.2f}".format, Rsltdnn)) + result_rand = ', '.join(map("{:.2f}".format, RsltRnd)) + result_strm = ', '.join(map("{:.2f}".format, RsltStrm)) + if self.ifOptimalSolExist: + result_bs = ', '.join(map("{:.2f}".format, Rsltbs)) + print( + 'output; {0:s}; Iter= {1:s}; SRDQN= [{2:s}]; sum = {3:2.4f}; Rand= [{4:s}]; sum = {5:2.4f}; Strm= [{6:s}]; sum = {7:2.4f}; BS= [{8:s}]; sum = {9:2.4f}' + .format( + strftime("%Y-%m-%d %H:%M:%S", gmtime()), str(str(self.curGame) + "-" + str(m)), result_srdqn, + sum(Rsltdnn), result_rand, sum(RsltRnd), result_strm, sum(RsltStrm), result_bs, sum(Rsltbs) + ) + ) + self.resultTest += [[Rsltdnn, RsltRnd, RsltStrm, Rsltbs]] + + else: + print( + 'output; {0:s}; Iter= {1:s}; SRDQN= [{2:s}]; sum = {3:2.4f}; Rand= [{4:s}]; sum = {5:2.4f}; Strm= [{6:s}]; sum = {7:2.4f}' + .format( + strftime("%Y-%m-%d %H:%M:%S", gmtime()), str(str(self.curGame) + "-" + str(m)), result_srdqn, + sum(Rsltdnn), result_rand, sum(RsltRnd), result_strm, sum(RsltStrm) + ) + ) + + self.resultTest += [[Rsltdnn, RsltRnd, RsltStrm]] + + return sum(Rsltdnn) + + def tester(self, testType, plt, colori, labeli, m): + + # set computation type for test + for k in range(0, self.config.NoAgent): + # self.players[k].compTypeTest = testType[k] + self.players[k].compType = testType[k] + # run the episode to get the results. + if labeli != 'OurPolicy': + result = self.playGame(self.demand) + else: + result = [-1 * self.players[i].cumReward for i in range(0, self.config.NoAgent)] + # add the results into the figure + if self.config.ifSaveFigure: + plt = plotting(plt, [np.array(self.players[i].hist) for i in range(0, self.config.NoAgent)], colori, labeli) + if self.config.ifsaveHistInterval and ((self.curGame == 0) or (self.curGame == 1) or (self.curGame == 2) or (self.curGame == 3) or ((self.curGame - 1) % self.config.saveHistInterval == 0)\ + or ((self.curGame) % self.config.saveHistInterval == 0) or ((self.curGame) % self.config.saveHistInterval == 1) \ + or ((self.curGame) % self.config.saveHistInterval == 2)) : + for k in range(0, self.config.NoAgent): + name = labeli + "-" + str(self.curGame) + "-" + "player" + "-" + str(k) + "-" + str(m) + np.save(os.path.join(self.config.model_dir, name), np.array(self.players[k].hist2)) + + # save the figure of base stocks + # if self.config.ifSaveFigure and (self.curGame in range(self.config.saveFigInt[0],self.config.saveFigInt[1])): + # for k in range(self.config.NoAgent): + # if self.players[k].compTypeTest == 'dnn': + # plotBaseStock(self.players[k].srdqnBaseStock, 'b', 'base stock of agent '+ str(self.players[k].agentNum), self.curGame, self.config, m) + + return result, plt + + def playGame(self, demand): + self.resetGame(demand) + + # run the game + while self.curTime < self.T: + self.handelAction(np.array(0)) # action won't be used. + self.next() + return [-1 * self.players[i].cumReward for i in range(0, self.config.NoAgent)] diff --git a/dizoo/beergame/envs/plotting.py b/dizoo/beergame/envs/plotting.py new file mode 100644 index 0000000000..57776c9641 --- /dev/null +++ b/dizoo/beergame/envs/plotting.py @@ -0,0 +1,72 @@ +# Code Reference: https://github.com/OptMLGroup/DeepBeerInventory-RL. +import os +import numpy as np +import matplotlib +matplotlib.use('Agg') +import matplotlib.pyplot as plt +from pylab import * + + +# plotting +def plotting(plt, data, colori, pltLabel): + # plt.hold(True) + + for i in range(np.shape(data)[0]): + plt.subplot(4, 5, 5 * i + 1) + plt.plot(np.transpose(data[i])[0, :], np.transpose(data[i])[1, :], colori, label=pltLabel) + plt.xlabel('Time') + plt.ylabel('IL') + plt.grid(True) + + plt.subplot(4, 5, 5 * i + 2) + plt.plot(np.transpose(data[i])[0, :], np.transpose(data[i])[2, :], colori, label=pltLabel) + plt.xlabel('Time') + plt.ylabel('OO') + plt.grid(True) + + plt.subplot(4, 5, 5 * i + 3) + plt.plot(np.transpose(data[i])[0, :], np.transpose(data[i])[3, :], colori, label=pltLabel) + plt.xlabel('Time') + plt.ylabel('a') + plt.grid(True) + + plt.subplot(4, 5, 5 * i + 4) + plt.plot(np.transpose(data[i])[0, :], np.transpose(data[i])[5, :], colori, label=pltLabel) + plt.xlabel('Time') + plt.ylabel('OUTL') + plt.grid(True) + + plt.subplot(4, 5, 5 * i + 5) + plt.plot(np.transpose(data[i])[0, :], -1 * np.transpose(data[i])[4, :], colori, label=pltLabel) + plt.xlabel('Time') + plt.ylabel('r') + plt.grid(True) + + return plt + + +def savePlot(players, curGame, Rsltdnn, RsltFrmu, RsltOptm, RsltRnd, config, m): + #add title to plot + if config.if_titled_figure: + plt.suptitle( + "sum OurPolicy=" + str(round(sum(Rsltdnn), 2)) + "; sum Strm=" + str(round(sum(RsltFrmu), 2)) + + "; sum BS=" + str(round(sum(RsltOptm), 2)) + "; sum Rnd=" + str(round(sum(RsltRnd), 2)) + "\n" + + "Ag OurPolicy=" + str([round(Rsltdnn[i], 2) for i in range(config.NoAgent)]) + "; Ag Strm=" + + str([round(RsltFrmu[i], 2) for i in range(config.NoAgent)]) + "; Ag BS=" + + str([round(RsltOptm[i], 2) for i in range(config.NoAgent)]) + "; Ag Rnd=" + + str([round(RsltRnd[i], 2) for i in range(config.NoAgent)]), + fontsize=12 + ) + + #insert legend to the figure + legend = plt.legend(bbox_to_anchor=(-1.4, -.165, 1., -.102), shadow=True, ncol=4) + + # configures spaces between subplots + plt.subplots_adjust(left=None, bottom=None, right=None, top=None, wspace=.5, hspace=.5) + # save the figure + path = os.path.join(config.figure_dir, 'saved_figures/') + if not os.path.exists(path): + os.mkdir(path) + plt.savefig(path + str(curGame) + '-' + str(m) + '.png', format='png') + print("figure" + str(curGame) + ".png saved in folder \"saved_figures\"") + plt.close(curGame) diff --git a/dizoo/beergame/envs/utils.py b/dizoo/beergame/envs/utils.py new file mode 100644 index 0000000000..2a6cf6f83d --- /dev/null +++ b/dizoo/beergame/envs/utils.py @@ -0,0 +1,355 @@ +import argparse +import os +import numpy as np + + +def str2bool(v): + return v.lower() in ('true', '1') + + +arg_lists = [] +parser = argparse.ArgumentParser() + + +def add_argument_group(name): + arg = parser.add_argument_group(name) + arg_lists.append(arg) + return arg + + +# crm +game_arg = add_argument_group('BeerGame') +game_arg.add_argument('--task', type=str, default='bg') +game_arg.add_argument( + '--fixedAction', + type=str2bool, + default='False', + help='if you want to have actions in [0,actionMax] set it to True. with False it will set it [actionLow, actionUp]' +) +game_arg.add_argument( + '--observation_data', + type=str2bool, + default=False, + help='if it is True, then it uses the data that is generated by based on few real world observation' +) +game_arg.add_argument('--data_id', type=int, default=22, help='the default item id for the basket dataset') +game_arg.add_argument('--TLow', type=int, default=100, help='duration of one GAME (lower bound)') +game_arg.add_argument('--TUp', type=int, default=100, help='duration of one GAME (upper bound)') +game_arg.add_argument( + '--demandDistribution', + type=int, + default=0, + help='0=uniform, 1=normal distribution, 2=the sequence of 4,4,4,4,8,..., 3= basket data, 4= forecast data' +) +game_arg.add_argument( + '--scaled', type=str2bool, default=False, help='if true it uses the (if) existing scaled parameters' +) +game_arg.add_argument('--demandSize', type=int, default=6100, help='the size of demand dataset') +game_arg.add_argument('--demandLow', type=int, default=0, help='the lower bound of random demand') +game_arg.add_argument('--demandUp', type=int, default=3, help='the upper bound of random demand') +game_arg.add_argument('--demandMu', type=float, default=10, help='the mu of the normal distribution for demand ') +game_arg.add_argument('--demandSigma', type=float, default=2, help='the sigma of the normal distribution for demand ') +game_arg.add_argument('--actionMax', type=int, default=2, help='it works when fixedAction is True') +game_arg.add_argument( + '--actionUp', type=int, default=2, help='bounds on my decision (upper bound), it works when fixedAction is True' +) +game_arg.add_argument( + '--actionLow', type=int, default=-2, help='bounds on my decision (lower bound), it works when fixedAction is True' +) +game_arg.add_argument( + '--action_step', type=int, default=1, help='The obtained action value by dnn is multiplied by this value' +) +game_arg.add_argument('--actionList', type=list, default=[], help='The list of the available actions') +game_arg.add_argument('--actionListLen', type=int, default=0, help='the length of the action list') +game_arg.add_argument( + '--actionListOpt', type=int, default=0, help='the action list which is used in optimal and sterman' +) +game_arg.add_argument('--actionListLenOpt', type=int, default=0, help='the length of the actionlistopt') +game_arg.add_argument('--agentTypes', type=list, default=['dnn', 'dnn', 'dnn', 'dnn'], help='the player types') +game_arg.add_argument( + '--agent_type1', type=str, default='dnn', help='the player types for agent 1, it can be dnn, Strm, bs, rnd' +) +game_arg.add_argument( + '--agent_type2', type=str, default='dnn', help='the player types for agent 2, it can be dnn, Strm, bs, rnd' +) +game_arg.add_argument( + '--agent_type3', type=str, default='dnn', help='the player types for agent 3, it can be dnn, Strm, bs, rnd' +) +game_arg.add_argument( + '--agent_type4', type=str, default='dnn', help='the player types for agent 4, it can be dnn, Strm, bs, rnd' +) +game_arg.add_argument('--NoAgent', type=int, default=4, help='number of agents, currently it should be in {1,2,3,4}') +game_arg.add_argument('--cp1', type=float, default=2.0, help='shortage cost of player 1') +game_arg.add_argument('--cp2', type=float, default=0.0, help='shortage cost of player 2') +game_arg.add_argument('--cp3', type=float, default=0.0, help='shortage cost of player 3') +game_arg.add_argument('--cp4', type=float, default=0.0, help='shortage cost of player 4') +game_arg.add_argument('--ch1', type=float, default=2.0, help='holding cost of player 1') +game_arg.add_argument('--ch2', type=float, default=2.0, help='holding cost of player 2') +game_arg.add_argument('--ch3', type=float, default=2.0, help='holding cost of player 3') +game_arg.add_argument('--ch4', type=float, default=2.0, help='holding cost of player 4') +game_arg.add_argument('--alpha_b1', type=float, default=-0.5, help='alpha of Sterman formula parameter for player 1') +game_arg.add_argument('--alpha_b2', type=float, default=-0.5, help='alpha of Sterman formula parameter for player 2') +game_arg.add_argument('--alpha_b3', type=float, default=-0.5, help='alpha of Sterman formula parameter for player 3') +game_arg.add_argument('--alpha_b4', type=float, default=-0.5, help='alpha of Sterman formula parameter for player 4') +game_arg.add_argument('--betta_b1', type=float, default=-0.2, help='beta of Sterman formula parameter for player 1') +game_arg.add_argument('--betta_b2', type=float, default=-0.2, help='beta of Sterman formula parameter for player 2') +game_arg.add_argument('--betta_b3', type=float, default=-0.2, help='beta of Sterman formula parameter for player 3') +game_arg.add_argument('--betta_b4', type=float, default=-0.2, help='beta of Sterman formula parameter for player 4') +game_arg.add_argument('--eta', type=list, default=[0, 4, 4, 4], help='the total cost regulazer') +game_arg.add_argument('--distCoeff', type=int, default=20, help='the total cost regulazer') +game_arg.add_argument( + '--ifUseTotalReward', + type=str2bool, + default='False', + help='if you want to have the total rewards in the experience replay, set it to true.' +) +game_arg.add_argument( + '--ifUsedistTotReward', + type=str2bool, + default='True', + help='If use correction to the rewards in the experience replay for all iterations of current game' +) +game_arg.add_argument( + '--ifUseASAO', + type=str2bool, + default='True', + help='if use AS and AO, i.e., received shipment and received orders in the input of DNN' +) +game_arg.add_argument('--ifUseActionInD', type=str2bool, default='False', help='if use action in the input of DNN') +game_arg.add_argument( + '--stateDim', type=int, default=5, help='Number of elements in the state desciptor - Depends on ifUseASAO' +) +game_arg.add_argument('--iftl', type=str2bool, default=False, help='if apply transfer learning') +game_arg.add_argument( + '--ifTransferFromSmallerActionSpace', + type=str2bool, + default=False, + help='if want to transfer knowledge from a network with different action space size.' +) +game_arg.add_argument( + '--baseActionSize', + type=int, + default=5, + help='if ifTransferFromSmallerActionSpace is true, this determines the size of action space of saved network' +) +game_arg.add_argument( + '--tlBaseBrain', + type=int, + default=3, + help='the gameConfig of the base network for re-training with transfer-learning' +) +game_arg.add_argument('--baseDemandDistribution', type=int, default=0, help='same as the demandDistribution') +game_arg.add_argument( + '--MultiAgent', type=str2bool, default=False, help='if run multi-agent RL model, not fully operational' +) +game_arg.add_argument( + '--MultiAgentRun', + type=list, + default=[True, True, True, True], + help='In the multi-RL setting, it determines which agent should get training.' +) +game_arg.add_argument( + '--if_use_AS_t_plus_1', type=str2bool, default='False', help='if use AS[t+1], not AS[t] in the input of DNN' +) +game_arg.add_argument( + '--ifSinglePathExist', + type=str2bool, + default=False, + help='If true it uses the predefined path in pre_model_dir and does not merge it with demandDistribution.' +) +game_arg.add_argument('--gamma', type=float, default=.99, help='discount factor for reward') +game_arg.add_argument( + '--multPerdInpt', type=int, default=10, help='Number of history records which we feed into network' +) + +# parameters of the leadtimes +leadtimes_arg = add_argument_group('leadtimes') +leadtimes_arg.add_argument( + '--leadRecItemLow', type=list, default=[2, 2, 2, 4], help='the min lead time for receiving items' +) +leadtimes_arg.add_argument( + '--leadRecItemUp', type=list, default=[2, 2, 2, 4], help='the max lead time for receiving items' +) +leadtimes_arg.add_argument( + '--leadRecOrderLow', type=int, default=[2, 2, 2, 0], help='the min lead time for receiving orders' +) +leadtimes_arg.add_argument( + '--leadRecOrderUp', type=int, default=[2, 2, 2, 0], help='the max lead time for receiving orders' +) +leadtimes_arg.add_argument('--ILInit', type=list, default=[0, 0, 0, 0], help='') +leadtimes_arg.add_argument('--AOInit', type=list, default=[0, 0, 0, 0], help='') +leadtimes_arg.add_argument('--ASInit', type=list, default=[0, 0, 0, 0], help='the initial shipment of each agent') +leadtimes_arg.add_argument('--leadRecItem1', type=int, default=2, help='the min lead time for receiving items') +leadtimes_arg.add_argument('--leadRecItem2', type=int, default=2, help='the min lead time for receiving items') +leadtimes_arg.add_argument('--leadRecItem3', type=int, default=2, help='the min lead time for receiving items') +leadtimes_arg.add_argument('--leadRecItem4', type=int, default=2, help='the min lead time for receiving items') +leadtimes_arg.add_argument('--leadRecOrder1', type=int, default=2, help='the min lead time for receiving order') +leadtimes_arg.add_argument('--leadRecOrder2', type=int, default=2, help='the min lead time for receiving order') +leadtimes_arg.add_argument('--leadRecOrder3', type=int, default=2, help='the min lead time for receiving order') +leadtimes_arg.add_argument('--leadRecOrder4', type=int, default=2, help='the min lead time for receiving order') +leadtimes_arg.add_argument('--ILInit1', type=int, default=0, help='the initial inventory level of the agent') +leadtimes_arg.add_argument('--ILInit2', type=int, default=0, help='the initial inventory level of the agent') +leadtimes_arg.add_argument('--ILInit3', type=int, default=0, help='the initial inventory level of the agent') +leadtimes_arg.add_argument('--ILInit4', type=int, default=0, help='the initial inventory level of the agent') +leadtimes_arg.add_argument('--AOInit1', type=int, default=0, help='the initial arriving order of the agent') +leadtimes_arg.add_argument('--AOInit2', type=int, default=0, help='the initial arriving order of the agent') +leadtimes_arg.add_argument('--AOInit3', type=int, default=0, help='the initial arriving order of the agent') +leadtimes_arg.add_argument('--AOInit4', type=int, default=0, help='the initial arriving order of the agent') +leadtimes_arg.add_argument('--ASInit1', type=int, default=0, help='the initial arriving shipment of the agent') +leadtimes_arg.add_argument('--ASInit2', type=int, default=0, help='the initial arriving shipment of the agent') +leadtimes_arg.add_argument('--ASInit3', type=int, default=0, help='the initial arriving shipment of the agent') +leadtimes_arg.add_argument('--ASInit4', type=int, default=0, help='the initial arriving shipment of the agent') + +# test +test_arg = add_argument_group('testing') +test_arg.add_argument( + '--testRepeatMid', + type=int, + default=50, + help='it is number of episodes which is going to be used for testing in the middle of training' +) +test_arg.add_argument('--testInterval', type=int, default=100, help='every xx games compute "test error"') +test_arg.add_argument( + '--ifSaveFigure', type=str2bool, default=True, help='if is it True, save the figures in each testing.' +) +test_arg.add_argument( + '--if_titled_figure', + type=str2bool, + default='True', + help='if is it True, save the figures with details in the title.' +) +test_arg.add_argument( + '--ifsaveHistInterval', type=str2bool, default=False, help='if every xx games save details of the episode' +) +test_arg.add_argument('--saveHistInterval', type=int, default=50000, help='every xx games save details of the play') +test_arg.add_argument('--Ttest', type=int, default=100, help='it defines the number of periods in the test cases') +test_arg.add_argument( + '--ifOptimalSolExist', + type=str2bool, + default=True, + help='if the instance has optimal base stock policy, set it to True, otherwise it should be False.' +) +test_arg.add_argument('--f1', type=float, default=8, help='base stock policy decision of player 1') +test_arg.add_argument('--f2', type=float, default=8, help='base stock policy decision of player 2') +test_arg.add_argument('--f3', type=float, default=0, help='base stock policy decision of player 3') +test_arg.add_argument('--f4', type=float, default=0, help='base stock policy decision of player 4') +test_arg.add_argument( + '--f_init', + type=list, + default=[32, 32, 32, 24], + help='base stock policy decision for 4 time-steps on the C(4,8) demand distribution' +) +test_arg.add_argument('--use_initial_BS', type=str2bool, default=False, help='If use f_init set it to True') + +# reporting +reporting_arg = add_argument_group('reporting') +reporting_arg.add_argument('--Rsltdnn', type=list, default=[], help='the result of dnn play tests will be saved here') +reporting_arg.add_argument( + '--RsltRnd', type=list, default=[], help='the result of random play tests will be saved here' +) +reporting_arg.add_argument( + '--RsltStrm', type=list, default=[], help='the result of heuristic fomula play tests will be saved here' +) +reporting_arg.add_argument( + '--Rsltbs', type=list, default=[], help='the result of optimal play tests will be saved here' +) +reporting_arg.add_argument( + '--ifSaveHist', + type=str2bool, + default='False', + help= + 'if it is true, saves history, prediction, and the randBatch in each period, WARNING: just make it True in small runs, it saves huge amount of files.' +) + + +# buildActionList: actions for the beer game problem +def buildActionList(config): + aDiv = 1 # difference in the action list + if config.fixedAction: + actions = list( + range(0, config.actionMax + 1, aDiv) + ) # If you put the second argument =11, creates an actionlist from 0..xx + else: + actions = list(range(config.actionLow, config.actionUp + 1, aDiv)) + return actions + + +# specify the dimension of the state of the game +def getStateDim(config): + if config.ifUseASAO: + stateDim = 5 + else: + stateDim = 3 + + if config.ifUseActionInD: + stateDim += 1 + + return stateDim + + +def set_optimal(config): + if config.demandDistribution == 0: + if config.cp1 == 2 and config.ch1 == 2 and config.ch2 == 2 and config.ch3 == 2 and config.ch4 == 2: + config.f1 = 8. + config.f2 = 8. + config.f3 = 0. + config.f4 = 0. + + +def get_config(): + config, unparsed = parser.parse_known_args() + config = update_config(config) + + return config, unparsed + + +def fill_leadtime_initial_values(config): + config.leadRecItemLow = [config.leadRecItem1, config.leadRecItem2, config.leadRecItem3, config.leadRecItem4] + config.leadRecItemUp = [config.leadRecItem1, config.leadRecItem2, config.leadRecItem3, config.leadRecItem4] + config.leadRecOrderLow = [config.leadRecOrder1, config.leadRecOrder2, config.leadRecOrder3, config.leadRecOrder4] + config.leadRecOrderUp = [config.leadRecOrder1, config.leadRecOrder2, config.leadRecOrder3, config.leadRecOrder4] + config.ILInit = [config.ILInit1, config.ILInit2, config.ILInit3, config.ILInit4] + config.AOInit = [config.AOInit1, config.AOInit2, config.AOInit3, config.AOInit4] + config.ASInit = [config.ASInit1, config.ASInit2, config.ASInit3, config.ASInit4] + + +def get_auxuliary_leadtime_initial_values(config): + config.leadRecOrderUp_aux = [config.leadRecOrder1, config.leadRecOrder2, config.leadRecOrder3, config.leadRecOrder4] + config.leadRecItemUp_aux = [config.leadRecItem1, config.leadRecItem2, config.leadRecItem3, config.leadRecItem4] + + +def fix_lead_time_manufacturer(config): + if config.leadRecOrder4 > 0: + config.leadRecItem4 += config.leadRecOrder4 + config.leadRecOrder4 = 0 + + +def set_sterman_parameters(config): + config.alpha_b = [config.alpha_b1, config.alpha_b2, config.alpha_b3, config.alpha_b4] + config.betta_b = [config.betta_b1, config.betta_b2, config.betta_b3, config.betta_b4] + + +def update_config(config): + config.actionList = buildActionList(config) # The list of the available actions + config.actionListLen = len(config.actionList) # the length of the action list + + set_optimal(config) + config.f = [config.f1, config.f2, config.f3, config.f4] # [6.4, 2.88, 2.08, 0.8] + + config.actionListLen = len(config.actionList) + if config.demandDistribution == 0: + config.actionListOpt = list(range(0, int(max(config.actionUp * 30 + 1, 3 * sum(config.f))), 1)) + else: + config.actionListOpt = list(range(0, int(max(config.actionUp * 30 + 1, 7 * sum(config.f))), 1)) + config.actionListLenOpt = len(config.actionListOpt) + + config.c_h = [config.ch1, config.ch2, config.ch3, config.ch4] + config.c_p = [config.cp1, config.cp2, config.cp3, config.cp4] + + config.stateDim = getStateDim(config) # Number of elements in the state description - Depends on ifUseASAO + get_auxuliary_leadtime_initial_values(config) + fix_lead_time_manufacturer(config) + fill_leadtime_initial_values(config) + set_sterman_parameters(config) + + return config diff --git a/dizoo/bitflip/envs/bitflip_env.py b/dizoo/bitflip/envs/bitflip_env.py index 9255e4cb96..3c74b174be 100644 --- a/dizoo/bitflip/envs/bitflip_env.py +++ b/dizoo/bitflip/envs/bitflip_env.py @@ -19,14 +19,14 @@ def __init__(self, cfg: dict) -> None: self._goal = np.zeros(self._n_bits) self._curr_step = 0 self._maxsize = self._n_bits - self._final_eval_reward = 0 + self._eval_episode_return = 0 self._observation_space = gym.spaces.Box(low=0, high=1, shape=(2 * self._n_bits, ), dtype=np.float32) self._action_space = gym.spaces.Discrete(self._n_bits) self._reward_space = gym.spaces.Box(low=0.0, high=1.0, shape=(1, ), dtype=np.float32) def reset(self) -> np.ndarray: self._curr_step = 0 - self._final_eval_reward = 0 + self._eval_episode_return = 0 if hasattr(self, '_seed') and hasattr(self, '_dynamic_seed') and self._dynamic_seed: random_seed = 100 * random.randint(1, 1000) np.random.seed(self._seed + random_seed) @@ -60,12 +60,12 @@ def step(self, action: np.ndarray) -> BaseEnvTimestep: else: rew = np.array([0]).astype(np.float32) done = False - self._final_eval_reward += float(rew) + self._eval_episode_return += float(rew) if self._curr_step >= self._maxsize - 1: done = True info = {} if done: - info['final_eval_reward'] = self._final_eval_reward + info['eval_episode_return'] = self._eval_episode_return self._curr_step += 1 obs = np.concatenate([self._state, self._goal], axis=0) return BaseEnvTimestep(obs, rew, done, info) diff --git a/dizoo/board_games/base_game_env.py b/dizoo/board_games/base_game_env.py deleted file mode 100644 index 2141543efd..0000000000 --- a/dizoo/board_games/base_game_env.py +++ /dev/null @@ -1,56 +0,0 @@ -from abc import abstractmethod -from ding.envs import BaseEnv - - -class BaseGameEnv(BaseEnv): - """ - Overview: - Base game class for MCTS based method - """ - - @abstractmethod - def current_player(self): - """ - Overview: - Return the current player. - Returns: - - The current player, it should be an element of the players list in the config. - """ - raise NotImplementedError - - @abstractmethod - def to_play(self): - """ - Overview: - In the board game environment, we need to know which player should move on the current state. - and return the corresponding player number. - Returns: - - An integer, indicate which player is to play. - """ - pass - - @abstractmethod - def legal_actions(self): - """ - Overview: - Should return the legal actions at each turn, if it is not available, it can return - the whole action space. At each turn, the game have to be able to handle one of returned actions. - For complex game where calculating legal moves is too long, the idea is to define the legal actions - equal to the action space but to return a negative reward if the action is illegal. - Returns: - - An array of integers, subset of the action space. - """ - pass - - def human_to_action(self): - """ - Overview: - For multiplayer games, ask the user for a legal action - and return the corresponding action number. - Returns: - An integer from the action space. - """ - choice = input(f"Enter the action to play for the player {self.to_play()}: ") - while int(choice) not in self.legal_actions(): - choice = input("Illegal action. Enter another action : ") - return int(choice) diff --git a/dizoo/board_games/chess/envs/chess_env_game.py b/dizoo/board_games/chess/envs/chess_env_game.py deleted file mode 100644 index 46254a48a6..0000000000 --- a/dizoo/board_games/chess/envs/chess_env_game.py +++ /dev/null @@ -1,198 +0,0 @@ -""" -Adapt Chess to BaseGameEnv interface from pettingzoo: https://github.com/Farama-Foundation/PettingZoo -""" - -import chess -from gym import spaces -from pettingzoo.utils.agent_selector import agent_selector -from pettingzoo.classic.chess import chess_utils -import numpy as np -import sys -from dizoo.board_games.base_game_env import BaseGameEnv -from ding.envs import BaseEnvTimestep -from ding.utils import ENV_REGISTRY - - -@ENV_REGISTRY.register('Chess') -class ChessEnv(BaseGameEnv): - - def __init__(self, cfg=None): - self.cfg = cfg - self.current_player_index = 0 - self.next_player_index = 1 - - self.board = chess.Board() - - self.agents = [f"player_{i+1}" for i in range(2)] - self.possible_agents = self.agents[:] - - self._agent_selector = agent_selector(self.agents) - - self._action_spaces = {name: spaces.Discrete(8 * 8 * 73) for name in self.agents} - self._observation_spaces = { - name: spaces.Dict( - { - 'observation': spaces.Box(low=0, high=1, shape=(8, 8, 111), dtype=bool), - 'action_mask': spaces.Box(low=0, high=1, shape=(4672, ), dtype=np.int8) - } - ) - for name in self.agents - } - - self.rewards = None - self.dones = None - self.infos = {name: {} for name in self.agents} - - self.agent_selection = None - - self.board_history = np.zeros((8, 8, 104), dtype=bool) - - @property - def current_player(self): - return self.current_player_index - - def to_play(self): - return self.next_player_index - - def reset(self): - self.has_reset = True - self.agents = self.possible_agents[:] - self.board = chess.Board() - - self._agent_selector = agent_selector(self.agents) - self.agent_selection = self._agent_selector.reset() - - self.rewards = {name: 0 for name in self.agents} - self._cumulative_rewards = {name: 0 for name in self.agents} - self.dones = {name: False for name in self.agents} - self.infos = {name: {} for name in self.agents} - - self.board_history = np.zeros((8, 8, 104), dtype=bool) - self.current_player_index = 0 - - for agent, reward in self.rewards.items(): - self._cumulative_rewards[agent] += reward - - agent = self.agent_selection - current_index = self.agents.index(agent) - self.current_player_index = current_index - obs = self.observe(agent) - return obs - - def observe(self, agent): - observation = chess_utils.get_observation(self.board, self.possible_agents.index(agent)) - observation = np.dstack((observation[:, :, :7], self.board_history)) - action_mask = self.legal_actions - - return {'observation': observation, 'action_mask': action_mask} - - def set_game_result(self, result_val): - for i, name in enumerate(self.agents): - self.dones[name] = True - result_coef = 1 if i == 0 else -1 - self.rewards[name] = result_val * result_coef - self.infos[name] = {'legal_moves': []} - - def step(self, action): - - if self.dones[self.agent_selection]: - return self._was_done_step(action) - - current_agent = self.agent_selection - current_index = self.agents.index(current_agent) - self.current_player_index = current_index - - next_board = chess_utils.get_observation(self.board, current_agent) - self.board_history = np.dstack((next_board[:, :, 7:], self.board_history[:, :, :-13])) - chosen_move = chess_utils.action_to_move(self.board, action, current_index) - assert chosen_move in self.board.legal_moves - self.board.push(chosen_move) # NOTE - - next_legal_moves = chess_utils.legal_moves(self.board) - is_stale_or_checkmate = not any(next_legal_moves) - - # claim draw is set to be true to align with normal tournament rules - is_repetition = self.board.is_repetition(3) - is_50_move_rule = self.board.can_claim_fifty_moves() - is_claimable_draw = is_repetition or is_50_move_rule - game_over = is_claimable_draw or is_stale_or_checkmate - - if game_over: - result = self.board.result(claim_draw=True) - result_val = chess_utils.result_to_int(result) - self.set_game_result(result_val) - - # self._accumulate_rewards() - for agent, reward in self.rewards.items(): - self._cumulative_rewards[agent] += reward - - self.agent_selection = self._agent_selector.next() - agent = self.agent_selection - self.next_player_index = self.agents.index(agent) - - observation = self.observe(agent) - - return BaseEnvTimestep(observation, self._cumulative_rewards[agent], self.dones[agent], self.infos[agent]) - - @property - def legal_actions(self): - action_mask = np.zeros(4672, 'uint8') - action_mask[chess_utils.legal_moves(self.board)] = 1 - return action_mask # 4672 dim {0,1} - - def legal_moves(self): - legal_moves = chess_utils.legal_moves(self.board) - return legal_moves - - def random_action(self): - action_list = self.legal_moves() - return np.random.choice(action_list) - - def expert_action(self): - # TODO - pass - - def human_to_action(self): - """ - For multiplayer games, ask the user for a legal action - and return the corresponding action number. - Returns: - An integer from the action space. - """ - while True: - try: - print(f"Current available actions for the player {self.to_play()} are:{self.legal_moves()}") - choice = int(input(f"Enter the index of next move for the player {self.to_play()}: ")) - if choice in self.legal_moves(): - break - except KeyboardInterrupt: - sys.exit(0) - except Exception as e: - print("Wrong input, try again") - return choice - - def render(self, mode='human'): - print(self.board) - - @property - def observation_space(self): - return self._observation_spaces - - @property - def action_space(self): - return self._action_spaces - - @property - def reward_space(self): - return self._reward_space - - def seed(self, seed: int, dynamic_seed: bool = True) -> None: - self._seed = seed - self._dynamic_seed = dynamic_seed - np.random.seed(self._seed) - - def close(self) -> None: - pass - - def __repr__(self) -> str: - return "DI-engine Chess Env" diff --git a/dizoo/board_games/chess/envs/test_chess_env_game.py b/dizoo/board_games/chess/envs/test_chess_env_game.py deleted file mode 100644 index e20e9ad916..0000000000 --- a/dizoo/board_games/chess/envs/test_chess_env_game.py +++ /dev/null @@ -1,40 +0,0 @@ -import pytest -from dizoo.board_games.chess.envs.chess_env_game import ChessEnv - - -@pytest.mark.envtest -class TestChessEnv: - - def test_naive(self): - env = ChessEnv() - env.reset() - print('init board state: ') - env.render() - for i in range(100): - """player 1""" - # action = env.human_to_action() - action = env.random_action() - print('player 1: ', action) - obs, reward, done, info = env.step(action) - assert isinstance(obs, dict) - assert isinstance(done, bool) - assert isinstance(reward, int) - # env.render() - if done: - if done: - if reward > 0: - print('player 1 (human player) win') - else: - print('draw') - break - """player 2""" - action = env.random_action() - print('player 2 (computer player): ', action) - obs, reward, done, info = env.step(action) - # env.render() - if done: - if reward > 0: - print('player 2 (computer player) win') - else: - print('draw') - break diff --git a/dizoo/board_games/go/envs/go_env_game.py b/dizoo/board_games/go/envs/go_env_game.py deleted file mode 100644 index ab9e068bff..0000000000 --- a/dizoo/board_games/go/envs/go_env_game.py +++ /dev/null @@ -1,318 +0,0 @@ -""" -Adapt Go to BaseGameEnv interface from pettingzoo: https://github.com/Farama-Foundation/PettingZoo -""" - -import sys -from dizoo.board_games.base_game_env import BaseGameEnv -from ding.envs import BaseEnvTimestep -from ding.utils import ENV_REGISTRY -import os -import numpy as np -import pygame -from gym import spaces -from pettingzoo.utils.agent_selector import agent_selector -from pettingzoo.classic.go import coords, go - - -def get_image(path): - from os import path as os_path - - import pygame - cwd = os_path.dirname(__file__) - image = pygame.image.load(cwd + '/' + path) - sfc = pygame.Surface(image.get_size(), flags=pygame.SRCALPHA) - sfc.blit(image, (0, 0)) - return sfc - - -@ENV_REGISTRY.register('Go') -class GoEnv(BaseGameEnv): - - def __init__(self, board_size: int = 19, komi: float = 7.5): - # board_size: a int, representing the board size (board has a board_size x board_size shape) - # komi: a float, representing points given to the second player. - self._overwrite_go_global_variables(board_size=board_size) - self._komi = komi - - self.agents = ['black_0', 'white_0'] - self.num_agents = len(self.agents) - - self.possible_agents = self.agents[:] - self.has_reset = False - - self.screen = None - - self._observation_space = self._convert_to_dict( - [ - spaces.Dict( - { - 'observation': spaces.Box(low=0, high=1, shape=(self._N, self._N, 17), dtype=bool), - 'action_mask': spaces.Box(low=0, high=1, shape=((self._N * self._N) + 1, ), dtype=np.int8) - } - ) for _ in range(self.num_agents) - ] - ) - - self._action_space = self._convert_to_dict( - [spaces.Discrete(self._N * self._N + 1) for _ in range(self.num_agents)] - ) - - self._agent_selector = agent_selector(self.agents) - - self.board_history = np.zeros((self._N, self._N, 16), dtype=bool) - - def _overwrite_go_global_variables(self, board_size: int): - self._N = board_size - go.N = self._N - go.ALL_COORDS = [(i, j) for i in range(self._N) for j in range(self._N)] - go.EMPTY_BOARD = np.zeros([self._N, self._N], dtype=np.int8) - go.NEIGHBORS = { - (x, y): list(filter(self._check_bounds, [(x + 1, y), (x - 1, y), (x, y + 1), (x, y - 1)])) - for x, y in go.ALL_COORDS - } - go.DIAGONALS = { - (x, y): list(filter(self._check_bounds, [(x + 1, y + 1), (x + 1, y - 1), (x - 1, y + 1), (x - 1, y - 1)])) - for x, y in go.ALL_COORDS - } - return - - def _check_bounds(self, c): - return 0 <= c[0] < self._N and 0 <= c[1] < self._N - - def _encode_player_plane(self, agent): - if agent == self.possible_agents[0]: - return np.zeros([self._N, self._N], dtype=bool) - else: - return np.ones([self._N, self._N], dtype=bool) - - def _encode_board_planes(self, agent): - agent_factor = go.BLACK if agent == self.possible_agents[0] else go.WHITE - current_agent_plane_idx = np.where(self._go.board == agent_factor) - opponent_agent_plane_idx = np.where(self._go.board == -agent_factor) - current_agent_plane = np.zeros([self._N, self._N], dtype=bool) - opponent_agent_plane = np.zeros([self._N, self._N], dtype=bool) - current_agent_plane[current_agent_plane_idx] = 1 - opponent_agent_plane[opponent_agent_plane_idx] = 1 - return current_agent_plane, opponent_agent_plane - - def _int_to_name(self, ind): - return self.possible_agents[ind] - - def _name_to_int(self, name): - return self.possible_agents.index(name) - - def _convert_to_dict(self, list_of_list): - return dict(zip(self.possible_agents, list_of_list)) - - def _encode_legal_actions(self, actions): - return np.where(actions == 1)[0] - - def _encode_rewards(self, result): - return [1, -1] if result == 1 else [-1, 1] - - @property - def current_player(self): - return self.current_player_index - - @property - def to_play(self): - return self.current_player_index - - def reset(self): - self.has_reset = True - self._go = go.Position(board=None, komi=self._komi) - self.agents = self.possible_agents[:] - self._agent_selector.reinit(self.agents) - self.agent_selection = self._agent_selector.reset() - self._cumulative_rewards = self._convert_to_dict(np.array([0.0, 0.0])) - self.rewards = self._convert_to_dict(np.array([0.0, 0.0])) - self.dones = self._convert_to_dict([False for _ in range(self.num_agents)]) - self.infos = self._convert_to_dict([{} for _ in range(self.num_agents)]) - self.next_legal_moves = self._encode_legal_actions(self._go.all_legal_moves()) - self._last_obs = self.observe(self.agents[0]) - self.board_history = np.zeros((self._N, self._N, 16), dtype=bool) - - self.current_player_index = 0 - - for agent, reward in self.rewards.items(): - self._cumulative_rewards[agent] += reward - - agent = self.agent_selection - current_index = self.agents.index(agent) - self.current_player_index = current_index - obs = self.observe(agent) - return obs - - def observe(self, agent): - player_plane = self._encode_player_plane(agent) - observation = np.dstack((self.board_history, player_plane)) - legal_moves = self.next_legal_moves if agent == self.agent_selection else [] - action_mask = np.zeros((self._N * self._N) + 1, 'int8') - for i in legal_moves: - action_mask[i] = 1 - - return {'observation': observation, 'action_mask': action_mask} - - def set_game_result(self, result_val): - for i, name in enumerate(self.agents): - self.dones[name] = True - result_coef = 1 if i == 0 else -1 - self.rewards[name] = result_val * result_coef - self.infos[name] = {'legal_moves': []} - - def step(self, action): - if self.dones[self.agent_selection]: - return self._was_done_step(action) - self._go = self._go.play_move(coords.from_flat(action)) - self._last_obs = self.observe(self.agent_selection) - current_agent_plane, opponent_agent_plane = self._encode_board_planes(self.agent_selection) - self.board_history = np.dstack((current_agent_plane, opponent_agent_plane, self.board_history[:, :, :-2])) - next_player = self._agent_selector.next() - - current_agent = next_player # 'black_0', 'white_0' - current_index = self.agents.index(current_agent) # 0, 1 - self.current_player_index = current_index - - if self._go.is_game_over(): - self.dones = self._convert_to_dict([True for _ in range(self.num_agents)]) - self.rewards = self._convert_to_dict(self._encode_rewards(self._go.result())) - self.next_legal_moves = [self._N * self._N] - else: - self.next_legal_moves = self._encode_legal_actions(self._go.all_legal_moves()) - self.agent_selection = next_player if next_player else self._agent_selector.next() - - # self._accumulate_rewards() - for agent, reward in self.rewards.items(): - self._cumulative_rewards[agent] += reward - # observation, reward, done, info = env.last() - agent = self.agent_selection - current_index = self.agents.index(agent) - self.current_player_index = current_index - observation = self.observe(agent) - return BaseEnvTimestep(observation, self._cumulative_rewards[agent], self.dones[agent], self.infos[agent]) - - def legal_actions(self): - pass - - def legal_moves(self): - if self._go.is_game_over(): - self.dones = self._convert_to_dict([True for _ in range(self.num_agents)]) - self.rewards = self._convert_to_dict(self._encode_rewards(self._go.result())) - self.next_legal_moves = [self._N * self._N] - else: - self.next_legal_moves = self._encode_legal_actions(self._go.all_legal_moves()) - - return self.next_legal_moves - - def random_action(self): - action_list = self.legal_moves() - return np.random.choice(action_list) - - def expert_action(self): - # TODO - pass - - def human_to_action(self): - """ - For multiplayer games, ask the user for a legal action - and return the corresponding action number. - Returns: - An integer from the action space. - """ - while True: - try: - print(f"Current available actions for the player {self.to_play()} are:{self.legal_moves()}") - choice = int(input(f"Enter the index of next move for the player {self.to_play()}: ")) - if choice in self.legal_moves(): - break - except KeyboardInterrupt: - sys.exit(0) - except Exception as e: - print("Wrong input, try again") - return choice - - def render(self, mode='human'): - screen_width = 1026 - screen_height = 1026 - - if self.screen is None: - if mode == "human": - pygame.init() - self.screen = pygame.display.set_mode((screen_width, screen_height)) - else: - self.screen = pygame.Surface((screen_width, screen_height)) - if mode == "human": - pygame.event.get() - - size = go.N - - # Load and scale all of the necessary images - tile_size = (screen_width) / size - - black_stone = get_image(os.path.join('../img', 'GoBlackPiece.png')) - black_stone = pygame.transform.scale(black_stone, (int(tile_size * (5 / 6)), int(tile_size * (5 / 6)))) - - white_stone = get_image(os.path.join('../img', 'GoWhitePiece.png')) - white_stone = pygame.transform.scale(white_stone, (int(tile_size * (5 / 6)), int(tile_size * (5 / 6)))) - - tile_img = get_image(os.path.join('../img', 'GO_Tile0.png')) - tile_img = pygame.transform.scale(tile_img, ((int(tile_size * (7 / 6))), int(tile_size * (7 / 6)))) - - # blit board tiles - for i in range(1, size - 1): - for j in range(1, size - 1): - self.screen.blit(tile_img, ((i * (tile_size)), int(j) * (tile_size))) - - for i in range(1, 9): - tile_img = get_image(os.path.join('../img', 'GO_Tile' + str(i) + '.png')) - tile_img = pygame.transform.scale(tile_img, ((int(tile_size * (7 / 6))), int(tile_size * (7 / 6)))) - for j in range(1, size - 1): - if i == 1: - self.screen.blit(tile_img, (0, int(j) * (tile_size))) - elif i == 2: - self.screen.blit(tile_img, ((int(j) * (tile_size)), 0)) - elif i == 3: - self.screen.blit(tile_img, ((size - 1) * (tile_size), int(j) * (tile_size))) - elif i == 4: - self.screen.blit(tile_img, ((int(j) * (tile_size)), (size - 1) * (tile_size))) - if i == 5: - self.screen.blit(tile_img, (0, 0)) - elif i == 6: - self.screen.blit(tile_img, ((size - 1) * (tile_size), 0)) - elif i == 7: - self.screen.blit(tile_img, ((size - 1) * (tile_size), (size - 1) * (tile_size))) - elif i == 8: - self.screen.blit(tile_img, (0, (size - 1) * (tile_size))) - - offset = tile_size * (1 / 6) - # Blit the necessary chips and their positions - for i in range(0, size): - for j in range(0, size): - if self._go.board[i][j] == go.BLACK: - self.screen.blit(black_stone, ((i * (tile_size) + offset), int(j) * (tile_size) + offset)) - elif self._go.board[i][j] == go.WHITE: - self.screen.blit(white_stone, ((i * (tile_size) + offset), int(j) * (tile_size) + offset)) - - if mode == "human": - pygame.display.update() - - observation = np.array(pygame.surfarray.pixels3d(self.screen)) - - return np.transpose(observation, axes=(1, 0, 2)) if mode == "rgb_array" else None - - def observation_space(self): - return self.observation_spaces - - def action_space(self): - return self._action_space - - def seed(self, seed: int, dynamic_seed: bool = True) -> None: - self._seed = seed - self._dynamic_seed = dynamic_seed - np.random.seed(self._seed) - - def close(self) -> None: - pass - - def __repr__(self) -> str: - return "DI-engine Go Env" diff --git a/dizoo/board_games/go/envs/test_go_env_game.py b/dizoo/board_games/go/envs/test_go_env_game.py deleted file mode 100644 index 8749bf896e..0000000000 --- a/dizoo/board_games/go/envs/test_go_env_game.py +++ /dev/null @@ -1,39 +0,0 @@ -import pytest -from dizoo.board_games.go.envs.go_env_game import GoEnv - - -@pytest.mark.envtest -class TestGoEnv: - - def test_naive(self): - env = GoEnv(board_size=9, komi=7.5) - print('NOTE：actions are counted by column, such as action 9, which is the second column and the first row') - env.reset() - # env.render() - for i in range(100): - """player 1""" - # action = env.human_to_action() - action = env.random_action() - print('player 1 (black_0): ', action) - obs, reward, done, info = env.step(action) - assert isinstance(obs, dict) - assert isinstance(done, bool) - assert isinstance(reward, float) - # env.render() - if done: - if reward > 0: - print('player 1 (black_0) win') - else: - print('draw') - break - """player 2""" - action = env.random_action() - print('player 2 (white_0): ', action) - obs, reward, done, info = env.step(action) - # env.render() - if done: - if reward > 0: - print('player 2 (white_0) win') - else: - print('draw') - break diff --git a/dizoo/board_games/gomoku/envs/gomoku_env.py b/dizoo/board_games/gomoku/envs/gomoku_env.py deleted file mode 100644 index bf41dc31fb..0000000000 --- a/dizoo/board_games/gomoku/envs/gomoku_env.py +++ /dev/null @@ -1,246 +0,0 @@ -from ditk import logging -import gym -import numpy as np -import sys - -from dizoo.board_games.base_game_env import BaseGameEnv -from ding.envs import BaseEnvTimestep -from ding.utils import ENV_REGISTRY -from dizoo.board_games.gomoku.envs.gomoku_expert import GomokuExpert - - -@ENV_REGISTRY.register('gomoku') -class GomokuEnv(BaseGameEnv): - - def __init__(self, cfg: dict = None, board_size: int = 15): - self.cfg = cfg - self.board_size = board_size - self.players = [1, 2] - self.board_markers = [str(i + 1) for i in range(self.board_size)] - self.total_num_actions = self.board_size * self.board_size - self.expert = GomokuExpert() - - @property - def current_player(self): - return self._current_player - - @property - def to_play(self): - return self.players[0] if self.current_player == self.players[1] else self.players[1] - - @property - def legal_actions(self): - legal_actions = [] - for i in range(self.board_size): - for j in range(self.board_size): - if self.board[i][j] == 0: - legal_actions.append(self.coord_to_action(i, j)) - return legal_actions - - def reset(self, start_player=0): - self._observation_space = gym.spaces.Box( - low=0, high=2, shape=(self.board_size, self.board_size, 3), dtype=np.int32 - ) - self._action_space = gym.spaces.Discrete(self.board_size ** 2) - self._reward_space = gym.spaces.Box(low=0, high=1, shape=(1, ), dtype=np.float32) - - self._current_player = self.players[start_player] - self.board = np.zeros((self.board_size, self.board_size), dtype="int32") - action_mask = np.zeros(self.total_num_actions, 'int8') - action_mask[self.legal_actions] = 1 - obs = {'observation': self.current_state(), 'action_mask': action_mask} - return obs - - def step(self, action): - if action in self.legal_actions: - row, col = self.action_to_coord(action) - self.board[row, col] = self.current_player - else: - logging.warning( - f"You input illegal action: {action}, the legal_actions are {self.legal_actions}. " - f"Now we randomly choice a action from self.legal_actions." - ) - action = np.random.choice(self.legal_actions) - row, col = self.action_to_coord(action) - self.board[row, col] = self.current_player - - # Check whether the game is ended or not and give the winner - done, winner = self.have_winner() - - reward = np.array(float(winner == self.current_player)).astype(np.float32) - info = {'next player to play': self.to_play} - # NOTE: exchange the player - self.current_player = self.to_play - - if done: - info['final_eval_reward'] = reward - # print('gomoku one episode done: ', info) - - action_mask = np.zeros(self.total_num_actions, 'int8') - action_mask[self.legal_actions] = 1 - obs = {'observation': self.current_state(), 'action_mask': action_mask} - return BaseEnvTimestep(obs, reward, done, info) - - def current_state(self): - """ - Overview: - self.board is nd-array, 0 indicates that no stones is placed here, - 1 indicates that player 1's stone is placed here, 2 indicates player 2's stone is placed here - Arguments: - - obs (:obj:`array`): the 0 dim means which positions is occupied by self.current_player, - the 1 dim indicates which positions are occupied by self.to_play, - the 2 dim indicates which player is the to_play player, 1 means player 1, 2 means player 2 - """ - board_curr_player = np.where(self.board == self.current_player, 1, 0) - board_opponent_player = np.where(self.board == self.to_play, 1, 0) - board_to_play = np.full((self.board_size, self.board_size), self.current_player) - return np.array([board_curr_player, board_opponent_player, board_to_play], dtype=np.float32) - - def coord_to_action(self, i, j): - """ - Overview: - convert coordinate i, j to action index a in [0, board_size**2) - """ - return i * self.board_size + j - - def action_to_coord(self, a): - """ - Overview: - convert action index a in [0, board_size**2) to coordinate (i, j) - """ - return a // self.board_size, a % self.board_size - - def have_winner(self): - has_legal_actions = False - directions = ((1, -1), (1, 0), (1, 1), (0, 1)) - for i in range(self.board_size): - for j in range(self.board_size): - # if no stone is on the position, don't need to consider this position - if self.board[i][j] == 0: - has_legal_actions = True - continue - # value-value at a coord, i-row, j-col - player = self.board[i][j] - # check if there exist 5 in a line - for d in directions: - x, y = i, j - count = 0 - for _ in range(5): - if (x not in range(self.board_size)) or (y not in range(self.board_size)): - break - if self.board[x][y] != player: - break - x += d[0] - y += d[1] - count += 1 - # if 5 in a line, store positions of all stones, return value - if count == 5: - return True, player - # if the players don't have legal actions, return done=True - return not has_legal_actions, -1 - - def seed(self, seed: int, dynamic_seed: bool = True) -> None: - self._seed = seed - self._dynamic_seed = dynamic_seed - np.random.seed(self._seed) - - def random_action(self): - action_list = self.legal_actions - return np.random.choice(action_list) - - def expert_action(self): - action_mask = np.zeros(self.total_num_actions, 'int8') - action_mask[self.legal_actions] = 1 - obs = {'observation': self.current_state(), 'action_mask': action_mask} - return self.expert.get_move(obs) - - def human_to_action(self): - """ - Overview: - For multiplayer games, ask the user for a legal action - and return the corresponding action number. - Returns: - An integer from the action space. - """ - while True: - try: - row = int( - input( - f"Enter the row (1, 2, ...,{self.board_size}, from up to bottom) to play for the player {self.current_player}: " - ) - ) - col = int( - input( - f"Enter the column (1, 2, ...,{self.board_size}, from left to right) to play for the player {self.current_player}: " - ) - ) - choice = self.coord_to_action(row - 1, col - 1) - if (choice in self.legal_actions and 1 <= row and 1 <= col and row <= self.board_size - and col <= self.board_size): - break - else: - print("Wrong input, try again") - except KeyboardInterrupt: - print("exit") - sys.exit(0) - except Exception as e: - print("Wrong input, try again") - return choice - - def render(self, mode="human"): - marker = " " - for i in range(self.board_size): - if i <= 8: - marker = marker + self.board_markers[i] + " " - else: - marker = marker + self.board_markers[i] + " " - print(marker) - for row in range(self.board_size): - if row <= 8: - print(str(1 + row) + ' ', end=" ") - else: - print(str(1 + row), end=" ") - for col in range(self.board_size): - ch = self.board[row][col] - if ch == 0: - print(".", end=" ") - elif ch == 1: - print("X", end=" ") - elif ch == 2: - print("O", end=" ") - print() - - def action_to_string(self, action_number): - """ - Overview: - Convert an action number to a string representing the action. - Arguments: - - action_number: an integer from the action space. - Returns: - - String representing the action. - """ - row = action_number // self.board_size + 1 - col = action_number % self.board_size + 1 - return f"Play row {row}, column {col}" - - @property - def observation_space(self) -> gym.spaces.Space: - return self._observation_space - - @property - def action_space(self) -> gym.spaces.Space: - return self._action_space - - @property - def reward_space(self) -> gym.spaces.Space: - return self._reward_space - - def close(self) -> None: - pass - - @current_player.setter - def current_player(self, value): - self._current_player = value - - def __repr__(self) -> str: - return "DI-engine Gomoku Env" diff --git a/dizoo/board_games/gomoku/envs/gomoku_expert.py b/dizoo/board_games/gomoku/envs/gomoku_expert.py deleted file mode 100644 index 4caf9a85c2..0000000000 --- a/dizoo/board_games/gomoku/envs/gomoku_expert.py +++ /dev/null @@ -1,384 +0,0 @@ -# Reference link: -# https://github.com/LouisCaixuran/gomoku/blob/c1b6d508522d9e8c78be827f326bbee54c4dfd8b/gomoku/expert.py - -from collections import defaultdict -import logging -import numpy as np - - -class GomokuExpert(object): - """ - Overview: - The ``GomokuExpert`` used to output rule-based gomoku expert actions. \ - Input is gomoku board obs(:obj:`np,array`) of shape ``(board_w, board_h)`` and returns action (:obj:`Int`) \ - The output action is the sequence number i*board_w+j of the placement position (i, j) - Interfaces: - ``__init__``, ``get_action``. - """ - def __init__(self): - """ - Overview: - Init the ``GomokuExpert``. - """ - # The weight of pieces of the same color - self.grade = 100 - # Indicates that it has already reached 5, reaching the maximum weight - self.max_value = 10012345 - self.init_board_flag = False - - def location_to_action(self, i, j): - """ - Overview: - Convert coordinate to serial number. - Arguments: - - i (:obj:`Int`): X-axis. - - j (:obj:`Int`): Y-axis. - Returns: - - action (:obj:`Int`): The serial number of the entered coordinates on the chessboard. - """ - # location = (i,j), action=j+i*width - return j + i * self.board_width - - def action_to_location(self, action): - """ - Overview: - Convert serial number to coordinate. - Arguments: - - action (:obj:`Int`): The serial number of the entered coordinates on the chessboard. - - Returns: - - [i, j]. - """ - # location = (i,j), action=j+i*width - j = action % self.board_width - i = action // self.board_width - return [i, j] - - def get_loc_player(self, i, j): - """ - Overview: - Returns the state of the pawn at the given coordinates. - Arguments: - - [i, j](:obj:`[Int, Int]`): The coordinate on the chessboard. - - Returns: - - board_status: \ - 0: no pawns, \ - 1: player 1, \ - 2: player 2. - """ - action = self.location_to_action(i, j) - return self.board_status[action] - - def scan_leftright(self, i, j, player): - """ - Overview: - Calculate the estimated value of the pawn from left to right when the player moves at (i,j) - Arguments: - - i (:obj:`Int`): X-axis. - - j (:obj:`Int`): Y-axis. - - player (:obj:`Int`): Current player. - - Returns: - - value: Situation valuation in this direction. - """ - # Count the number of consecutive pieces or empty pieces of the current player - # and get the value in this direction when moving pieces (i, j) - value = 0 - count = 0 - grade = self.grade - # scan left - m, n = i, j - 1 - while n >= 0: - is_continue, value, grade = self.caculate_once_value(m, n, - player, - value, - grade) - if is_continue: - # Continue to move one step to the left - n = n - 1 - else: - break - count += 1 - # Change the direction, the weights are returned to the initial values - grade = self.grade - n = j + 1 - while n < self.board_width: - is_continue, value, grade = self.caculate_once_value(m, n, - player, - value, - grade) - if is_continue: - # Continue to move one step to the right - n = n + 1 - else: - break - count += 1 - # Returns the value if there are four consecutive pawn in this direction, otherwise 0 - return value if count >= 4 else 0 - - def scan_updown(self, i, j, player): - """ - Overview: - Calculate the estimated value of the pawn from up to down when the player moves at (i,j) - Arguments: - - i (:obj:`Int`): X-axis. - - j (:obj:`Int`): Y-axis. - - player (:obj:`Int`): Current player. - - Returns: - - value: Situation valuation in this direction. - """ - value = 0 - count = 0 - # Count the number of consecutive pieces or empty pieces of the current player - # and get the value in this direction when moving pieces (i, j) - grade = self.grade - - m, n = i - 1, j - # scan up - while m >= 0: - is_continue, value, grade = self.caculate_once_value(m, n, - player, - value, - grade) - if is_continue: - # Continue to move one step to the up - m = m - 1 - else: - break - count += 1 - # Change the direction and change the weight back to the initial value - grade = self.grade - m = i + 1 - # scan down - while m < self.board_height: - is_continue, value, grade = self.caculate_once_value(m, n, - player, - value, - grade) - if is_continue: - # Continue to move one step to the down - m = m + 1 - else: - break - count += 1 - # Returns the value if there are four consecutive pawn in this direction, otherwise 0 - return value if count >= 4 else 0 - - def scan_left_updown(self, i, j, player): - """ - Overview: - Calculate the estimated value of the pawn from top left to bottom right when the player moves at (i,j) - Arguments: - - i (:obj:`Int`): X-axis. - - j (:obj:`Int`): Y-axis. - - player (:obj:`Int`): Current player. - - Returns: - - value: Situation valuation in this direction. - """ - # Count the number of consecutive pieces or empty pieces of the current player - # and get the value in this direction when moving pieces (i, j) - value = 0 - count = 0 - - grade = self.grade - m, n = i - 1, j - 1 - # scan left up - while m >= 0 and n >= 0: - is_continue, value, grade = self.caculate_once_value(m, n, - player, - value, - grade) - if is_continue: - # Continue to move one step to the left up - m, n = m - 1, n - 1 - else: - break - count += 1 - - grade = self.grade - # right down - m, n = i + 1, j + 1 - while m < self.board_height and n < self.board_width: - is_continue, value, grade = self.caculate_once_value(m, n, - player, - value, - grade) - if is_continue: - # Continue to move one step down to the right - m, n = m + 1, n + 1 - else: - break - count += 1 - # Returns the value if there are four consecutive pawn in this direction, otherwise 0 - return value if count >= 4 else 0 - - def scan_right_updown(self, i, j, player): - """ - Overview: - Calculate the estimated value of the pawn from top right to bottom left when the player moves at (i,j) - Arguments: - - i (:obj:`Int`): X-axis. - - j (:obj:`Int`): Y-axis. - - player (:obj:`Int`): Current player. - - Returns: - - value: Situation valuation in this direction. - """ - # Count the number of consecutive pieces or empty pieces of the current player - # and get the value in this direction when moving pieces (i, j) - value = 0 - count = 0 - grade = self.grade - # scan left down - m, n = i + 1, j - 1 - while m < self.board_height and n >= 0: - is_continue, value, grade = self.caculate_once_value(m, n, - player, - value, - grade) - if is_continue: - m, n = m + 1, n - 1 - else: - break - count += 1 - grade = self.grade - # scan right up - m, n = i - 1, j + 1 - while m >= 0 and n < self.board_width: - is_continue, value, grade = self.caculate_once_value(m, n, - player, - value, - grade) - if is_continue: - # Continue to move up one step to the right - m, n = m - 1, n + 1 - else: - break - count += 1 - # Returns the value if there are four consecutive pawn in this direction, otherwise 0 - return value if count >= 4 else 0 - - def caculate_once_value(self, m, n, player, value, grade): - """ - Overview: - Calculate the income brought by the pawns adjacent to the position (m,n) \ - when the player places the pawn at the specified position (i,j) \ - in the current situation - Arguments: - - m (:obj:`Int`): x. - - n (:obj:`Int`): y. - - player (:obj:`Int`): current chess player. - - value (:obj:`Int`): The current position (the pawn is at (i,j)) is evaluated. - - grade (:obj:`Int`): The weight of the pawn in the current position - - Returns: - - is_continue: Whether there is a pawn of current_player at the current position - - value: The evaluation value of the move to (i,j) - - grade: The weight of a single one of our chess pieces - """ - loc_player = self.get_loc_player(m, n) - if loc_player == player: - value += grade - elif loc_player == 0: - value += 1 - # When encountering an empty chess, reduce the weight of subsequent black chess - grade = grade / 10 - else: - # opponent's pawn - value -= 5 - # When encountering an opponent's chess, return - return 0, value, grade - # 1 means 'is_continue' - return 1, value, grade - - def evaluate_all_value(self, player): - """ - Overview: - Calculate the estimates of all possible positions, \ - and choose the most favorable position from them. - Arguments: - - player (:obj:`Int`): current chess player. - - Returns: - - action: the most favorable action - - self.action_value[action][4]: Situation valuation under this action - """ - self.action_value = defaultdict(lambda: [0, 0, 0, 0, 0]) - for action in self.available: - i, j = self.action_to_location(action) - - self.action_value[action][0] = self.scan_updown(i, j, player) - self.action_value[action][1] = self.scan_leftright(i, j, player) - self.action_value[action][2] = self.scan_left_updown(i, j, player) - self.action_value[action][3] = self.scan_right_updown(i, j, player) - - # Indicates that one direction can already be rushed to 4 - for k in range(4): - if self.action_value[action][k] >= 390: - self.action_value[action][k] = 2000 - elif self.action_value[action][k] >= 302: - self.action_value[action][k] = 1000 - - # Comprehensive score in all directions - self.action_value[action][4] = ( - self.action_value[action][0] + - self.action_value[action][1] + - self.action_value[action][2] + - self.action_value[action][3]) - - action = max(self.available, key=lambda x: self.action_value[x][4]) - - return action, self.action_value[action][4] - - def get_action(self, obs): - """ - Overview: - Returns a rule-based expert action. - Arguments: - - obs (:obj:`np.array`) - - Returns: - - expert_action - """ - self.board = obs - - if self.init_board_flag is False: - self.board_width = self.board['observation'][0].shape[0] - self.board_height = self.board['observation'][0].shape[1] - # the 2 dim indicates which player is the to_play player, 1 means player 1, 2 means player 2 - self.init_board_flag = True - if self.board['observation'][2][0][0] == 1: - self.m_player_id = 1 - self.s_player_id = 2 - else: - self.m_player_id = 2 - self.s_player_id = 1 - # transform observation,action_mask to self.available, self.board_status - self.available = [] - self.board_status = np.zeros(self.board_width * self.board_height, 'int8') - for i in range(self.board_width): - for j in range(self.board_height): - action = self.location_to_action(i, j) - if self.board['action_mask'][action] == 1: - self.available.append(action) - if self.board['observation'][0][i][j] == 1: - self.board_status[action] = self.m_player_id - elif self.board['observation'][1][i][j] == 1: - self.board_status[action] = self.s_player_id - - m_action, m_value = self.evaluate_all_value(self.m_player_id) - - logging.info("loaction:{loc},value:{value}".format( - loc=self.action_to_location(m_action), value=m_value)) - - s_action, s_value = self.evaluate_all_value(self.s_player_id) - logging.info("O_loaction:{loc},value:{value}".format( - loc=self.action_to_location(s_action), value=s_value)) - # Block this position if it is better for the opponent to play here - if m_value >= s_value: - return m_action - else: - return s_action diff --git a/dizoo/board_games/gomoku/envs/test_gomoku_env.py b/dizoo/board_games/gomoku/envs/test_gomoku_env.py deleted file mode 100644 index cec6ddebcb..0000000000 --- a/dizoo/board_games/gomoku/envs/test_gomoku_env.py +++ /dev/null @@ -1,36 +0,0 @@ -import pytest -from dizoo.board_games.gomoku.envs.gomoku_env import GomokuEnv - - -@pytest.mark.envtest -class TestGomokuEnv: - - def test_naive(self): - env = GomokuEnv() - obs = env.reset() - print('init board state: ') - env.render() - done = False - while True: - action = env.random_action() - # action = env.human_to_action() - print('player 1: ' + env.action_to_string(action)) - obs, reward, done, info = env.step(action) - env.render() - if done: - if reward > 0: - print('player 1 (human player) win') - else: - print('draw') - break - - action = env.random_action() - print('player 2 (computer player): ' + env.action_to_string(action)) - obs, reward, done, info = env.step(action) - env.render() - if done: - if reward > 0: - print('player 2 (computer player) win') - else: - print('draw') - break diff --git a/dizoo/board_games/gomoku/envs/test_gomoku_expert_action.py b/dizoo/board_games/gomoku/envs/test_gomoku_expert_action.py deleted file mode 100644 index a169fc9ef3..0000000000 --- a/dizoo/board_games/gomoku/envs/test_gomoku_expert_action.py +++ /dev/null @@ -1,37 +0,0 @@ -import pytest -from dizoo.board_games.gomoku.envs.gomoku_env import GomokuEnv - - -@pytest.mark.envtest -class TestExpertAction: - def test_naive(): - env = GomokuEnv() - obs = env.reset() - print('init board state: ', obs) - env.render() - done = False - while True: - # action = env.random_action() - action = env.human_to_action() - # action = env.expert_action() - print('original player 1: ',action) - print('player 1: ' + env.action_to_string(action)) - obs, reward, done, info = env.step(action) - env.render() - if done: - if reward > 0: - print('player 1 (human player) win') - else: - print('draw') - break - - action = env.expert_action() - print('player 2 (computer player): ' + env.action_to_string(action)) - obs, reward, done, info = env.step(action) - env.render() - if done: - if reward > 0: - print('player 2 (computer player) win') - else: - print('draw') - break diff --git a/dizoo/board_games/tictactoe/envs/test_tictactoe_env.py b/dizoo/board_games/tictactoe/envs/test_tictactoe_env.py deleted file mode 100644 index 7b6afe3c45..0000000000 --- a/dizoo/board_games/tictactoe/envs/test_tictactoe_env.py +++ /dev/null @@ -1,36 +0,0 @@ -import pytest -from dizoo.board_games.tictactoe.envs.tictactoe_env import TicTacToeEnv - - -@pytest.mark.envtest -class TestTicTacToeEnv: - - def test_naive(self): - env = TicTacToeEnv() - env.reset() - print('init board state: ') - env.render() - while True: - """player 1""" - # action = env.human_to_action() - action = env.random_action() - print('player 1: ' + env.action_to_string(action)) - obs, reward, done, info = env.step(action) - env.render() - if done: - if reward > 0: - print('player 1 (human player) win') - else: - print('draw') - break - """player 2""" - action = env.expert_action() - print('player 2 (computer player): ' + env.action_to_string(action)) - obs, reward, done, info = env.step(action) - env.render() - if done: - if reward > 0: - print('player 2 (computer player) win') - else: - print('draw') - break diff --git a/dizoo/board_games/tictactoe/envs/test_tictactoe_expert_action.py b/dizoo/board_games/tictactoe/envs/test_tictactoe_expert_action.py deleted file mode 100644 index 1794758780..0000000000 --- a/dizoo/board_games/tictactoe/envs/test_tictactoe_expert_action.py +++ /dev/null @@ -1,30 +0,0 @@ -import pytest -from dizoo.board_games.tictactoe.envs.tictactoe_env import TicTacToeEnv -import numpy as np - - -@pytest.mark.envtest -class TestExpertAction: - - def test_expert_action(self): - env = TicTacToeEnv() - env.reset() - print('init board state: ') - env.render() - # TODO(pu): How to fully test all cases - # case 1 - env.board = np.array([[1, 2, 1], [1, 2, 0], [0, 0, 2]]) - env.current_player = 1 - assert 6 == env.expert_action() - # case 2 - env.board = np.array([[1, 2, 1], [2, 2, 0], [1, 0, 0]]) - env.current_player = 1 - assert env.expert_action() in [5, 7] - # case 3 - env.board = np.array([[1, 2, 1], [1, 2, 2], [0, 0, 1]]) - env.current_player = 2 - assert 7 == env.expert_action() - # case 4 - env.board = np.array([[1, 2, 1], [1, 0, 2], [0, 0, 0]]) - env.current_player = 2 - assert 6 == env.expert_action() diff --git a/dizoo/board_games/tictactoe/envs/tictactoe_env.py b/dizoo/board_games/tictactoe/envs/tictactoe_env.py deleted file mode 100644 index 9cabc68114..0000000000 --- a/dizoo/board_games/tictactoe/envs/tictactoe_env.py +++ /dev/null @@ -1,265 +0,0 @@ -""" -Adapt TicTacToe to BaseGameEnv interface from https://github.com/werner-duvaud/muzero-general -""" -from ditk import logging -import sys -import gym -import copy -import numpy as np - -from ding.envs.env.base_env import BaseEnvTimestep -from ding.utils.registry_factory import ENV_REGISTRY -from dizoo.board_games.base_game_env import BaseGameEnv - - -@ENV_REGISTRY.register('tictactoe') -class TicTacToeEnv(BaseGameEnv): - - def __init__(self): - self.board_size = 3 - self.players = [1, 2] - self.total_num_actions = 9 - - @property - def current_player(self): - return self._current_player - - @property - def to_play(self): - return self.players[0] if self.current_player == self.players[1] else self.players[1] - - @property - def current_player_to_compute_expert_action(self): - """ - Overview: to compute expert action easily. - """ - return -1 if self.current_player == 1 else 1 - - def reset(self, start_player=0): - self._observation_space = gym.spaces.Box( - low=0, high=2, shape=(self.board_size, self.board_size, 3), dtype=np.uint8 - ) - self._action_space = gym.spaces.Discrete(self.board_size ** 2) - self._reward_space = gym.spaces.Box(low=0, high=1, shape=(1, ), dtype=np.float32) - - self._current_player = self.players[start_player] - self.board = np.zeros((self.board_size, self.board_size), dtype="int32") - action_mask = np.zeros(self.total_num_actions, 'int8') - action_mask[self.legal_actions] = 1 - obs = {'observation': self.current_state(), 'action_mask': action_mask} - return obs - - def step(self, action): - if action in self.legal_actions: - row, col = self.action_to_coord(action) - self.board[row, col] = self.current_player - else: - logging.warning( - f"You input illegal action: {action}, the legal_actions are {self.legal_actions}. " - f"Now we randomly choice a action from self.legal_actions." - ) - action = np.random.choice(self.legal_actions) - row, col = self.action_to_coord(action) - self.board[row, col] = self.current_player - - # Check whether the game is ended or not and give the winner - have_winner, winner = self.have_winner() - if have_winner: - done, winner = True, winner - elif len(self.legal_actions) == 0: - # the agent don't have legal_actions to move, so episode is done - # winner=-1 indicates draw - done, winner = True, -1 - else: - # episode is not done - done, winner = False, -1 - - reward = np.array(float(winner == self.current_player)).astype(np.float32) - info = {'next player to play': self.to_play} - # NOTE: exchange the player - self.current_player = self.to_play - - if done: - info['final_eval_reward'] = reward - # print('tictactoe one episode done: ', info) - - action_mask = np.zeros(self.total_num_actions, 'int8') - action_mask[self.legal_actions] = 1 - obs = {'observation': self.current_state(), 'action_mask': action_mask} - return BaseEnvTimestep(obs, reward, done, info) - - def current_state(self): - board_curr_player = np.where(self.board == self.current_player, 1, 0) - board_opponent_player = np.where(self.board == self.to_play, 1, 0) - board_to_play = np.full((self.board_size, self.board_size), self.current_player) - return np.array([board_curr_player, board_opponent_player, board_to_play], dtype=np.float32) - - def coord_to_action(self, i, j): - """ - Overview: - convert coordinate i, j to action index a in [0, board_size**2) - """ - return i * self.board_size + j - - def action_to_coord(self, a): - """ - Overview: - convert action index a in [0, board_size**2) to coordinate (i, j) - """ - return a // self.board_size, a % self.board_size - - def have_winner(self): - # Horizontal and vertical checks - for i in range(self.board_size): - if len(set(self.board[i, :])) == 1 and (self.board[i, 0] != 0): - return True, self.board[i, 0] - if len(set(self.board[:, i])) == 1 and (self.board[0, i] != 0): - return True, self.board[0, i] - - # Diagonal checks - if self.board[0, 0] == self.board[1, 1] == self.board[2, 2] != 0: - return True, self.board[0, 0] - if self.board[2, 0] == self.board[1, 1] == self.board[0, 2] != 0: - return True, self.board[2, 0] - - return False, -1 - - @property - def legal_actions(self): - legal_actions = [] - for i in range(self.board_size): - for j in range(self.board_size): - if self.board[i][j] == 0: - legal_actions.append(self.coord_to_action(i, j)) - return legal_actions - - def random_action(self): - action_list = self.legal_actions - return np.random.choice(action_list) - - def expert_action(self): - """ - Overview: - Hard coded expert agent for tictactoe env - Returns: - - action (:obj:`int`): the expert action to take in the current game state. - """ - # To easily calculate expert action, we convert the chessboard notation: - # from player 1: 1, player 2: 2 - # to player 1: -1, player 2: 1 - # TODO: more elegant implementation - board = copy.deepcopy(self.board) - for i in range(board.shape[0]): - for j in range(board.shape[1]): - if board[i][j] == 1: - board[i][j] = -1 - elif board[i][j] == 2: - board[i][j] = 1 - - action = np.random.choice(self.legal_actions) - # Horizontal and vertical checks - for i in range(3): - if abs(sum(board[i, :])) == 2: - ind = np.where(board[i, :] == 0)[0][0] - action = np.ravel_multi_index((np.array([i]), np.array([ind])), (3, 3))[0] - if self.current_player_to_compute_expert_action * sum(board[i, :]) > 0: - return action - - if abs(sum(board[:, i])) == 2: - ind = np.where(board[:, i] == 0)[0][0] - action = np.ravel_multi_index((np.array([ind]), np.array([i])), (3, 3))[0] - if self.current_player_to_compute_expert_action * sum(board[:, i]) > 0: - return action - - # Diagonal checks - diag = board.diagonal() - anti_diag = np.fliplr(board).diagonal() - if abs(sum(diag)) == 2: - ind = np.where(diag == 0)[0][0] - action = np.ravel_multi_index((np.array([ind]), np.array([ind])), (3, 3))[0] - if self.current_player_to_compute_expert_action * sum(diag) > 0: - return action - - if abs(sum(anti_diag)) == 2: - ind = np.where(anti_diag == 0)[0][0] - action = np.ravel_multi_index((np.array([ind]), np.array([2 - ind])), (3, 3))[0] - if self.current_player_to_compute_expert_action * sum(anti_diag) > 0: - return action - - return action - - def human_to_action(self): - """ - Overview: - For multiplayer games, ask the user for a legal action - and return the corresponding action number. - Returns: - An integer from the action space. - """ - while True: - try: - row = int( - input( - f"Enter the row (1, 2, or 3, from up to bottom) to play for the player {self.current_player}: " - ) - ) - col = int( - input( - f"Enter the column (1, 2 or 3, from left to right) to play for the player {self.current_player}: " - ) - ) - choice = self.coord_to_action(row - 1, col - 1) - if (choice in self.legal_actions and 1 <= row and 1 <= col and row <= self.board_size - and col <= self.board_size): - break - else: - print("Wrong input, try again") - except KeyboardInterrupt: - print("exit") - sys.exit(0) - except Exception as e: - print("Wrong input, try again") - return choice - - def render(self, mode="human"): - print(self.board) - - def action_to_string(self, action_number): - """ - Overview: - Convert an action number to a string representing the action. - Arguments: - - action_number: an integer from the action space. - Returns: - - String representing the action. - """ - row = action_number // self.board_size + 1 - col = action_number % self.board_size + 1 - return f"Play row {row}, column {col}" - - @property - def observation_space(self) -> gym.spaces.Space: - return self._observation_space - - @property - def action_space(self) -> gym.spaces.Space: - return self._action_space - - @property - def reward_space(self) -> gym.spaces.Space: - return self._reward_space - - def seed(self, seed: int, dynamic_seed: bool = True) -> None: - self._seed = seed - self._dynamic_seed = dynamic_seed - np.random.seed(self._seed) - - def close(self) -> None: - pass - - @current_player.setter - def current_player(self, value): - self._current_player = value - - def __repr__(self) -> str: - return "DI-engine TicTacToe Env" diff --git a/dizoo/box2d/bipedalwalker/README.md b/dizoo/box2d/bipedalwalker/README.md deleted file mode 100644 index 078b2da8c9..0000000000 --- a/dizoo/box2d/bipedalwalker/README.md +++ /dev/null @@ -1,11 +0,0 @@ -## BipedalWalker Environment - -Bipedalwalker is a continuous action space environment in openAI gym. With 24 dim input and 4 dim action space.Reward is given for moving forward, total 300+ points up to the far end. If the robot falls, it gets -100. Applying motor torque costs a small amount of points, more optimal agent will get better score. State consists of hull angle speed, angular velocity, horizontal speed, vertical speed, position of joints and joints angular speed, legs contact with ground, and 10 lidar rangefinder measurements. - -![original](./original.gif) - -## Train BipedalWalker with DI-engine - -DI-engine can achive 300+ return on average within 2000 episodes by TD3. The tuned example can be found in `dizoo/box2d/bipedalwalker/config/bipedalwalker_td3_config.py`. The training episode reward is as follows. - -![tb](./bipedalwalkertb.png) diff --git a/dizoo/box2d/bipedalwalker/bipedalwalkertb.png b/dizoo/box2d/bipedalwalker/bipedalwalkertb.png deleted file mode 100644 index 821630493d..0000000000 Binary files a/dizoo/box2d/bipedalwalker/bipedalwalkertb.png and /dev/null differ diff --git a/dizoo/box2d/bipedalwalker/config/bipedalwalker_a2c_config.py b/dizoo/box2d/bipedalwalker/config/bipedalwalker_a2c_config.py new file mode 100644 index 0000000000..c82542597f --- /dev/null +++ b/dizoo/box2d/bipedalwalker/config/bipedalwalker_a2c_config.py @@ -0,0 +1,64 @@ +from easydict import EasyDict + +bipedalwalker_a2c_config = dict( + exp_name='bipedalwalker_a2c_seed0', + env=dict( + env_id='BipedalWalker-v3', + collector_env_num=8, + evaluator_env_num=8, + # (bool) Scale output action into legal range. + act_scale=True, + n_evaluator_episode=8, + stop_value=300, + rew_clip=True, + # The path to save the game replay + # replay_path='./bipedalwalker_a2c_seed0/video', + ), + policy=dict( + cuda=True, + # load_path="./bipedalwalker_a2c_seed0/ckpt/ckpt_best.pth.tar", + action_space='continuous', + model=dict( + action_space='continuous', + obs_shape=24, + action_shape=4, + ), + learn=dict( + # (int) the number of data for a train iteration + batch_size=256, + learning_rate=0.0003, + # (float) loss weight of the value network, the weight of policy network is set to 1 + value_weight=0.5, + # (float) loss weight of the entropy regularization, the weight of policy network is set to 1 + entropy_weight=0.001, + # (float) discount factor for future reward, defaults int [0, 1] + discount_factor=0.99, + adv_norm=True, + ), + collect=dict( + # (int) collect n_sample data, train model n_iteration times + n_sample=512, + discount_factor=0.99, + collector=dict(collect_print_freq=100, ), + ), + eval=dict(evaluator=dict(eval_freq=100, )), + ), +) +bipedalwalker_a2c_config = EasyDict(bipedalwalker_a2c_config) +main_config = bipedalwalker_a2c_config +bipedalwalker_a2c_create_config = dict( + env=dict( + type='bipedalwalker', + import_names=['dizoo.box2d.bipedalwalker.envs.bipedalwalker_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='a2c'), + replay_buffer=dict(type='naive'), +) +bipedalwalker_a2c_create_config = EasyDict(bipedalwalker_a2c_create_config) +create_config = bipedalwalker_a2c_create_config + +if __name__ == "__main__": + # or you can enter `ding -m serial_onpolicy -c bipedalwalker_a2c_config.py -s 0` + from ding.entry import serial_pipeline_onpolicy + serial_pipeline_onpolicy([main_config, create_config], seed=0) diff --git a/dizoo/box2d/bipedalwalker/config/bipedalwalker_ddpg_config.py b/dizoo/box2d/bipedalwalker/config/bipedalwalker_ddpg_config.py new file mode 100644 index 0000000000..de70a09c86 --- /dev/null +++ b/dizoo/box2d/bipedalwalker/config/bipedalwalker_ddpg_config.py @@ -0,0 +1,55 @@ +from easydict import EasyDict + +bipedalwalker_ddpg_config = dict( + exp_name='bipedalwalker_ddpg_seed0', + env=dict( + env_id='BipedalWalker-v3', + collector_env_num=8, + evaluator_env_num=5, + # (bool) Scale output action into legal range. + act_scale=True, + n_evaluator_episode=5, + rew_clip=True, + ), + policy=dict( + cuda=True, + random_collect_size=10000, + model=dict( + obs_shape=24, + action_shape=4, + twin_critic=False, + action_space='regression', + actor_head_hidden_size=400, + critic_head_hidden_size=400, + ), + learn=dict( + update_per_collect=64, + batch_size=256, + learning_rate_actor=0.0003, + learning_rate_critic=0.0003, + target_theta=0.005, + discount_factor=0.99, + learner=dict(hook=dict(log_show_after_iter=1000, )) + ), + collect=dict(n_sample=64, ), + other=dict(replay_buffer=dict(replay_buffer_size=300000, ), ), + ), +) +bipedalwalker_ddpg_config = EasyDict(bipedalwalker_ddpg_config) +main_config = bipedalwalker_ddpg_config + +bipedalwalker_ddpg_create_config = dict( + env=dict( + type='bipedalwalker', + import_names=['dizoo.box2d.bipedalwalker.envs.bipedalwalker_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='ddpg'), +) +bipedalwalker_ddpg_create_config = EasyDict(bipedalwalker_ddpg_create_config) +create_config = bipedalwalker_ddpg_create_config + +if __name__ == "__main__": + # or you can enter `ding -m serial -c bipedalwalker_ddpg_config.py -s 0` + from ding.entry import serial_pipeline + serial_pipeline([main_config, create_config], seed=0, max_env_step=int(1e5)) diff --git a/dizoo/box2d/bipedalwalker/config/bipedalwalker_impala_config.py b/dizoo/box2d/bipedalwalker/config/bipedalwalker_impala_config.py new file mode 100644 index 0000000000..ee58dbe818 --- /dev/null +++ b/dizoo/box2d/bipedalwalker/config/bipedalwalker_impala_config.py @@ -0,0 +1,73 @@ +from easydict import EasyDict + +bipedalwalker_impala_config = dict( + exp_name='bipedalwalker_impala_seed0', + env=dict( + env_id='BipedalWalker-v3', + collector_env_num=8, + evaluator_env_num=8, + # (bool) Scale output action into legal range. + act_scale=True, + n_evaluator_episode=8, + stop_value=300, + rew_clip=True, + # The path to save the game replay + # replay_path='./bipedalwalker_impala_seed0/video', + ), + policy=dict( + cuda=True, + # (int) the trajectory length to calculate v-trace target + unroll_len=32, + random_collect_size=256, + # load_path="./bipedalwalker_impala_seed0/ckpt/ckpt_best.pth.tar", + action_space='continuous', + model=dict( + action_space='continuous', + obs_shape=24, + action_shape=4, + ), + learn=dict( + # (int) collect n_sample data, train model update_per_collect times + # here we follow impala serial pipeline + update_per_collect=3, # update_per_collect show be in [1, 10] + # (int) the number of data for a train iteration + batch_size=64, + grad_clip_type='clip_norm', + clip_value=5, + learning_rate=0.0003, + # (float) loss weight of the value network, the weight of policy network is set to 1 + value_weight=0.5, + # (float) loss weight of the entropy regularization, the weight of policy network is set to 1 + entropy_weight=0.01, + # (float) discount factor for future reward, defaults int [0, 1] + discount_factor=0.99, + # (float) additional discounting parameter + lambda_=0.99, + ), + collect=dict( + # (int) collect n_sample data, train model n_iteration times + n_sample=32, + collector=dict(collect_print_freq=1000, ), + ), + eval=dict(evaluator=dict(eval_freq=100, )), + other=dict(replay_buffer=dict(replay_buffer_size=10000, ), ), + ), +) +bipedalwalker_impala_config = EasyDict(bipedalwalker_impala_config) +main_config = bipedalwalker_impala_config +bipedalwalker_impala_create_config = dict( + env=dict( + type='bipedalwalker', + import_names=['dizoo.box2d.bipedalwalker.envs.bipedalwalker_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='impala'), + replay_buffer=dict(type='naive'), +) +bipedalwalker_impala_create_config = EasyDict(bipedalwalker_impala_create_config) +create_config = bipedalwalker_impala_create_config + +if __name__ == "__main__": + # or you can enter `ding -m serial_onpolicy -c bipedalwalker_impala_config.py -s 0` + from ding.entry import serial_pipeline + serial_pipeline([main_config, create_config], seed=0) diff --git a/dizoo/box2d/bipedalwalker/config/bipedalwalker_pg_config.py b/dizoo/box2d/bipedalwalker/config/bipedalwalker_pg_config.py new file mode 100644 index 0000000000..96aa08aee8 --- /dev/null +++ b/dizoo/box2d/bipedalwalker/config/bipedalwalker_pg_config.py @@ -0,0 +1,52 @@ +from easydict import EasyDict + +bipedalwalker_pg_config = dict( + exp_name='bipedalwalker_pg_seed0', + env=dict( + env_id='BipedalWalker-v3', + collector_env_num=8, + evaluator_env_num=8, + act_scale=True, + n_evaluator_episode=8, + stop_value=300, + rew_clip=True, + ), + policy=dict( + cuda=True, + action_space='continuous', + model=dict( + action_space='continuous', + obs_shape=24, + action_shape=4, + ), + learn=dict( + batch_size=64, + learning_rate=0.001, + entropy_weight=0.001, + ), + collect=dict( + n_episode=8, + unroll_len=1, + discount_factor=0.99, + ), + eval=dict(evaluator=dict(eval_freq=200, )) + ), +) +bipedalwalker_pg_config = EasyDict(bipedalwalker_pg_config) +main_config = bipedalwalker_pg_config +bipedalwalker_pg_create_config = dict( + env=dict( + type='bipedalwalker', + import_names=['dizoo.box2d.bipedalwalker.envs.bipedalwalker_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='pg'), + collector=dict(type='episode'), +) +bipedalwalker_pg_create_config = EasyDict(bipedalwalker_pg_create_config) +create_config = bipedalwalker_pg_create_config + +if __name__ == "__main__": + # or you can enter `ding -m serial_onpolicy -c bipedalwalker_pg_config.py -s 0` + from ding.entry import serial_pipeline_onpolicy + serial_pipeline_onpolicy([main_config, create_config], seed=0) diff --git a/dizoo/box2d/bipedalwalker/config/bipedalwalker_ppopg_config.py b/dizoo/box2d/bipedalwalker/config/bipedalwalker_ppopg_config.py index 410d31a595..4e2e7df403 100644 --- a/dizoo/box2d/bipedalwalker/config/bipedalwalker_ppopg_config.py +++ b/dizoo/box2d/bipedalwalker/config/bipedalwalker_ppopg_config.py @@ -53,6 +53,7 @@ class PPOPGContinuousModel(nn.Module): + def __init__(self, obs_shape, action_shape): super(PPOPGContinuousModel, self).__init__() self.encoder = nn.Sequential(nn.Linear(obs_shape, 64), nn.Tanh()) @@ -78,4 +79,6 @@ def forward(self, inputs): new_main_config = deepcopy(main_config) new_main_config.exp_name += "_seed{}".format(seed) model = PPOPGContinuousModel(new_main_config.policy.model.obs_shape, new_main_config.policy.model.action_shape) - serial_pipeline_onpolicy([new_main_config, deepcopy(create_config)], seed=seed, max_env_step=int(5e6), model=model) + serial_pipeline_onpolicy( + [new_main_config, deepcopy(create_config)], seed=seed, max_env_step=int(5e6), model=model + ) diff --git a/dizoo/box2d/bipedalwalker/config/bipedalwalker_sac_config.py b/dizoo/box2d/bipedalwalker/config/bipedalwalker_sac_config.py index 643ee06ce2..f905c4031b 100644 --- a/dizoo/box2d/bipedalwalker/config/bipedalwalker_sac_config.py +++ b/dizoo/box2d/bipedalwalker/config/bipedalwalker_sac_config.py @@ -1,7 +1,7 @@ from easydict import EasyDict bipedalwalker_sac_config = dict( - exp_name='bipedalwalker_sac_seed0', + exp_name='bipedalwalker_sac_config0', env=dict( env_id='BipedalWalker-v3', collector_env_num=8, @@ -9,15 +9,11 @@ # (bool) Scale output action into legal range. act_scale=True, n_evaluator_episode=5, - stop_value=300, rew_clip=True, - # The path to save the game replay - replay_path=None, ), policy=dict( - cuda=False, - priority=False, - random_collect_size=1000, + cuda=True, + random_collect_size=10000, model=dict( obs_shape=24, action_shape=4, @@ -27,22 +23,18 @@ critic_head_hidden_size=128, ), learn=dict( - update_per_collect=1, - batch_size=128, - learning_rate_q=0.001, - learning_rate_policy=0.001, + update_per_collect=64, + batch_size=256, + learning_rate_q=0.0003, + learning_rate_policy=0.0003, learning_rate_alpha=0.0003, - ignore_done=True, target_theta=0.005, discount_factor=0.99, auto_alpha=True, - value_network=False, + learner=dict(hook=dict(log_show_after_iter=1000, )) ), - collect=dict( - n_sample=128, - unroll_len=1, - ), - other=dict(replay_buffer=dict(replay_buffer_size=100000, ), ), + collect=dict(n_sample=64, ), + other=dict(replay_buffer=dict(replay_buffer_size=300000, ), ), ), ) bipedalwalker_sac_config = EasyDict(bipedalwalker_sac_config) @@ -53,10 +45,7 @@ import_names=['dizoo.box2d.bipedalwalker.envs.bipedalwalker_env'], ), env_manager=dict(type='subprocess'), - policy=dict( - type='sac', - import_names=['ding.policy.sac'], - ), + policy=dict(type='sac', ), replay_buffer=dict(type='naive', ), ) bipedalwalker_sac_create_config = EasyDict(bipedalwalker_sac_create_config) @@ -65,4 +54,4 @@ if __name__ == "__main__": # or you can enter `ding -m serial -c bipedalwalker_sac_config.py -s 0` from ding.entry import serial_pipeline - serial_pipeline([main_config, create_config], seed=0) + serial_pipeline([main_config, create_config], seed=0, max_env_step=int(1e5)) diff --git a/dizoo/box2d/bipedalwalker/config/bipedalwalker_td3_config.py b/dizoo/box2d/bipedalwalker/config/bipedalwalker_td3_config.py index 6cb431b143..09cc3d1bf1 100644 --- a/dizoo/box2d/bipedalwalker/config/bipedalwalker_td3_config.py +++ b/dizoo/box2d/bipedalwalker/config/bipedalwalker_td3_config.py @@ -4,35 +4,31 @@ exp_name='bipedalwalker_td3_seed0', env=dict( env_id='BipedalWalker-v3', - collector_env_num=1, + collector_env_num=8, evaluator_env_num=5, # (bool) Scale output action into legal range. act_scale=True, n_evaluator_episode=5, - stop_value=300, rew_clip=True, - # The path to save the game replay - replay_path=None, ), policy=dict( cuda=True, - priority=False, + random_collect_size=10000, model=dict( obs_shape=24, action_shape=4, twin_critic=True, + action_space='regression', actor_head_hidden_size=400, critic_head_hidden_size=400, - action_space='regression', ), learn=dict( - update_per_collect=4, - discount_factor=0.99, - batch_size=128, - learning_rate_actor=0.001, - learning_rate_critic=0.001, + update_per_collect=64, + batch_size=256, + learning_rate_actor=0.0003, + learning_rate_critic=0.0003, target_theta=0.005, - ignore_done=False, + discount_factor=0.99, actor_update_freq=2, noise=True, noise_sigma=0.2, @@ -40,14 +36,10 @@ min=-0.5, max=0.5, ), + learner=dict(hook=dict(log_show_after_iter=1000, )) ), - collect=dict( - n_sample=256, - noise_sigma=0.1, - collector=dict(collect_print_freq=1000, ), - ), - eval=dict(evaluator=dict(eval_freq=100, ), ), - other=dict(replay_buffer=dict(replay_buffer_size=50000, ), ), + collect=dict(n_sample=64, ), + other=dict(replay_buffer=dict(replay_buffer_size=300000, ), ), ), ) bipedalwalker_td3_config = EasyDict(bipedalwalker_td3_config) @@ -67,4 +59,4 @@ if __name__ == "__main__": # or you can enter `ding -m serial -c bipedalwalker_td3_config.py -s 0` from ding.entry import serial_pipeline - serial_pipeline([main_config, create_config], seed=0) + serial_pipeline([main_config, create_config], seed=0, max_env_step=int(1e5)) diff --git a/dizoo/box2d/bipedalwalker/dt_data/bipedalwalker_collect_data.py b/dizoo/box2d/bipedalwalker/dt_data/bipedalwalker_collect_data.py deleted file mode 100644 index 75054bbaa0..0000000000 --- a/dizoo/box2d/bipedalwalker/dt_data/bipedalwalker_collect_data.py +++ /dev/null @@ -1,33 +0,0 @@ -from dizoo.box2d.bipedalwalker.dt_data.collect_sac_data_config import main_config, create_config -from ding.entry import collect_episodic_demo_data, eval -import torch -import copy - - -def eval_ckpt(args): - config = copy.deepcopy([main_config, create_config]) - eval(config, seed=args.seed, load_path=main_config.policy.learn.learner.hook.load_ckpt_before_run) - - -def generate(args): - print('start generate data') - config = copy.deepcopy([main_config, create_config]) - state_dict = torch.load(main_config.policy.learn.learner.load_path, map_location='cpu') - collect_episodic_demo_data( - config, - collect_count=main_config.policy.other.replay_buffer.replay_buffer_size, - seed=args.seed, - expert_data_path=main_config.policy.collect.save_path, - state_dict=state_dict - ) - - -if __name__ == "__main__": - import argparse - - parser = argparse.ArgumentParser() - parser.add_argument('--seed', '-s', type=int, default=0) - args = parser.parse_args() - - eval_ckpt(args) - generate(args) diff --git a/dizoo/box2d/bipedalwalker/dt_data/bipedalwalker_sac_data_config.py b/dizoo/box2d/bipedalwalker/dt_data/bipedalwalker_sac_data_config.py deleted file mode 100644 index 643ee06ce2..0000000000 --- a/dizoo/box2d/bipedalwalker/dt_data/bipedalwalker_sac_data_config.py +++ /dev/null @@ -1,68 +0,0 @@ -from easydict import EasyDict - -bipedalwalker_sac_config = dict( - exp_name='bipedalwalker_sac_seed0', - env=dict( - env_id='BipedalWalker-v3', - collector_env_num=8, - evaluator_env_num=5, - # (bool) Scale output action into legal range. - act_scale=True, - n_evaluator_episode=5, - stop_value=300, - rew_clip=True, - # The path to save the game replay - replay_path=None, - ), - policy=dict( - cuda=False, - priority=False, - random_collect_size=1000, - model=dict( - obs_shape=24, - action_shape=4, - twin_critic=True, - action_space='reparameterization', - actor_head_hidden_size=128, - critic_head_hidden_size=128, - ), - learn=dict( - update_per_collect=1, - batch_size=128, - learning_rate_q=0.001, - learning_rate_policy=0.001, - learning_rate_alpha=0.0003, - ignore_done=True, - target_theta=0.005, - discount_factor=0.99, - auto_alpha=True, - value_network=False, - ), - collect=dict( - n_sample=128, - unroll_len=1, - ), - other=dict(replay_buffer=dict(replay_buffer_size=100000, ), ), - ), -) -bipedalwalker_sac_config = EasyDict(bipedalwalker_sac_config) -main_config = bipedalwalker_sac_config -bipedalwalker_sac_create_config = dict( - env=dict( - type='bipedalwalker', - import_names=['dizoo.box2d.bipedalwalker.envs.bipedalwalker_env'], - ), - env_manager=dict(type='subprocess'), - policy=dict( - type='sac', - import_names=['ding.policy.sac'], - ), - replay_buffer=dict(type='naive', ), -) -bipedalwalker_sac_create_config = EasyDict(bipedalwalker_sac_create_config) -create_config = bipedalwalker_sac_create_config - -if __name__ == "__main__": - # or you can enter `ding -m serial -c bipedalwalker_sac_config.py -s 0` - from ding.entry import serial_pipeline - serial_pipeline([main_config, create_config], seed=0) diff --git a/dizoo/box2d/bipedalwalker/dt_data/collect_sac_data_config.py b/dizoo/box2d/bipedalwalker/dt_data/collect_sac_data_config.py deleted file mode 100644 index 405773848e..0000000000 --- a/dizoo/box2d/bipedalwalker/dt_data/collect_sac_data_config.py +++ /dev/null @@ -1,122 +0,0 @@ -from email import policy -from easydict import EasyDict - -# learn learner -# collect -# other replay buffer - -# td step -bipedalwalker_sac_config = dict( - exp_name='bipedalwalker_sac_dt_dataset_seed0', - env=dict( - env_id='BipedalWalker-v3', - collector_env_num=8, - evaluator_env_num=5, - # act scale - act_scale=True, - n_evaluator_episode=5, - stop_value=300, - rew_clip=True, - # save game play - replay_path=None, - ), - policy=dict( - cuda=True, - priority=False, - random_collect_size=1000, - model=dict( - obs_shape=24, - action_shape=4, - twin_critic=True, - action_space='reparameterization', - actor_head_hidden_size=128, - critic_head_hidden_size=128, - ), - learn=dict( - # NOTE - learner=dict( - train_iterations=int(1e9), - dataloader=dict(num_workers=0, ), - log_policy=True, - hook=dict( - load_ckpt_before_run='/mnt/nfs/wangzilin/bipedalwalker_sac_seed0/ckpt/ckpt_best.pth.tar', - log_show_after_iter=100, - save_ckpt_after_iter=10000, - save_ckpt_after_run=False, - ), - cfg_type='BaseLearnerDict', - load_path='/mnt/nfs/wangzilin/bipedalwalker_sac_seed0/ckpt/ckpt_best.pth.tar', - ), - update_per_collect=1, - batch_size=128, - learning_rate_q=0.001, - learning_rate_policy=0.001, - learning_rate_alpha=0.0003, - ignore_done=True, - target_theta=0.005, - discount_factor=0.99, - auto_alpha=True, - value_network=False, - ), - collect=dict( - n_sample=128, - unroll_len=1, - data_type='naive', - save_path= - '/home/wangzilin/projects/decision_transformer/DI-engine/dizoo/box2d/bipedalwalker/dt_data/data/sac_data_1000eps.pkl', - data_path= - '/home/wangzilin/projects/decision_transformer/DI-engine/dizoo/box2d/bipedalwalker/dt_data/data/sac_data_10eps.pkl', - ), - other=dict( - replay_buffer=dict( - type='advanced', - replay_buffer_size=1000, - max_use=float('inf'), - max_staleness=float('inf'), - alpha=0.6, - beta=0.4, - anneal_step=int(1e5), - enable_track_used_data=False, - deepcopy=False, - thruput_controller=dict( - push_sample_rate_limit=dict( - max=float('inf'), - min=0, - ), - window_seconds=30, - sample_min_limit_ratio=1, - ), - monitor=dict( - sample_data_attr=dict( - average_range=5, - print_freq=200, - ), - periodic_thruput=dict(seconds=60, ), - ), - cfg_type='AdvancedReplayBufferDict', - ), - ), - ), -) -bipedalwalker_sac_config = EasyDict(bipedalwalker_sac_config) -main_config = bipedalwalker_sac_config -bipedalwalker_sac_create_config = dict( - env=dict( - type='bipedalwalker', - import_names=['dizoo.box2d.bipedalwalker.envs.bipedalwalker_env'], - ), - env_manager=dict(type='subprocess'), - policy=dict( - type='sac', - import_names=['ding.policy.sac'], - ), - replay_buffer=dict(type='naive', ), -) - -bipedalwalker_sac_create_config = EasyDict(bipedalwalker_sac_create_config) -create_config = bipedalwalker_sac_create_config - -if __name__ == "__main__": - # or you can enter `ding -m serial -c bipedalwalker_sac_config.py -s 0` - from ding.entry import serial_pipeline - serial_pipeline([main_config, create_config], seed=0) diff --git a/dizoo/box2d/bipedalwalker/entry/bipedalwalker_ppo_eval.py b/dizoo/box2d/bipedalwalker/entry/bipedalwalker_ppo_eval.py index 0261a05587..1423c8c27f 100644 --- a/dizoo/box2d/bipedalwalker/entry/bipedalwalker_ppo_eval.py +++ b/dizoo/box2d/bipedalwalker/entry/bipedalwalker_ppo_eval.py @@ -15,8 +15,9 @@ from ding.rl_utils import get_epsilon_greedy_fn from dizoo.box2d.bipedalwalker.config.bipedalwalker_ppo_config import main_config, create_config + def main(rl_cfg, seed=0): - main_cfg, create_cfg =rl_cfg + main_cfg, create_cfg = rl_cfg cfg = compile_config( main_cfg, BaseEnvManager, @@ -56,4 +57,4 @@ def main(rl_cfg, seed=0): if __name__ == "__main__": - main(rl_cfg=(main_config, create_config),seed=0) + main(rl_cfg=(main_config, create_config), seed=0) diff --git a/dizoo/box2d/bipedalwalker/envs/bipedalwalker_env.py b/dizoo/box2d/bipedalwalker/envs/bipedalwalker_env.py index 4b280fc8c7..ae4d526032 100644 --- a/dizoo/box2d/bipedalwalker/envs/bipedalwalker_env.py +++ b/dizoo/box2d/bipedalwalker/envs/bipedalwalker_env.py @@ -42,7 +42,7 @@ def reset(self) -> np.ndarray: episode_trigger=lambda episode_id: True, name_prefix='rl-video-{}'.format(id(self)) ) - self._final_eval_reward = 0 + self._eval_episode_return = 0 obs = self._env.reset() obs = to_ndarray(obs).astype(np.float32) return obs @@ -68,13 +68,13 @@ def step(self, action: np.ndarray) -> BaseEnvTimestep: action = affine_transform(action, min_val=self.action_space.low, max_val=self.action_space.high) obs, rew, done, info = self._env.step(action) - self._final_eval_reward += rew + self._eval_episode_return += rew if self._rew_clip: rew = max(-10, rew) rew = np.float32(rew) if done: - info['final_eval_reward'] = self._final_eval_reward + info['eval_episode_return'] = self._eval_episode_return obs = to_ndarray(obs).astype(np.float32) rew = to_ndarray([rew]) # wrapped to be transfered to a array with shape (1,) return BaseEnvTimestep(obs, rew, done, info) diff --git a/dizoo/box2d/carracing/__init__.py b/dizoo/box2d/carracing/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/dizoo/box2d/carracing/car_racing.gif b/dizoo/box2d/carracing/car_racing.gif new file mode 100644 index 0000000000..5d3bdd38e9 Binary files /dev/null and b/dizoo/box2d/carracing/car_racing.gif differ diff --git a/dizoo/box2d/carracing/config/__init__.py b/dizoo/box2d/carracing/config/__init__.py new file mode 100644 index 0000000000..1571e58a64 --- /dev/null +++ b/dizoo/box2d/carracing/config/__init__.py @@ -0,0 +1 @@ +from .carracing_dqn_config import carracing_dqn_config, carracing_dqn_create_config diff --git a/dizoo/box2d/carracing/config/carracing_dqn_config.py b/dizoo/box2d/carracing/config/carracing_dqn_config.py new file mode 100644 index 0000000000..1792056a83 --- /dev/null +++ b/dizoo/box2d/carracing/config/carracing_dqn_config.py @@ -0,0 +1,60 @@ +from easydict import EasyDict + +nstep = 3 +carracing_dqn_config = dict( + exp_name='carracing_dqn_seed0', + env=dict( + collector_env_num=8, + evaluator_env_num=8, + env_id='CarRacing-v2', + continuous=False, + n_evaluator_episode=8, + stop_value=900, + # replay_path='./carracing_dqn_seed0/video', + ), + policy=dict( + cuda=True, + # load_path='carracing_dqn_seed0/ckpt/ckpt_best.pth.tar', + model=dict( + obs_shape=[3, 96, 96], + action_shape=5, + encoder_hidden_size_list=[64, 64, 128], + dueling=True, + ), + discount_factor=0.99, + nstep=nstep, + learn=dict( + update_per_collect=10, + batch_size=64, + learning_rate=0.0001, + target_update_freq=100, + ), + collect=dict(n_sample=64, ), + other=dict( + eps=dict( + type='exp', + start=0.95, + end=0.1, + decay=50000, + ), replay_buffer=dict(replay_buffer_size=100000, ) + ), + ), +) +carracing_dqn_config = EasyDict(carracing_dqn_config) +main_config = carracing_dqn_config + +carracing_dqn_create_config = dict( + env=dict( + type='carracing', + import_names=['dizoo.box2d.carracing.envs.carracing_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='dqn'), +) +carracing_dqn_create_config = EasyDict(carracing_dqn_create_config) +create_config = carracing_dqn_create_config + +if __name__ == "__main__": + # or you can enter `ding -m serial -c carracing_dqn_config.py -s 0` + from ding.entry import serial_pipeline + serial_pipeline([main_config, create_config], seed=0) diff --git a/dizoo/box2d/carracing/envs/__init__.py b/dizoo/box2d/carracing/envs/__init__.py new file mode 100644 index 0000000000..a36760ccf7 --- /dev/null +++ b/dizoo/box2d/carracing/envs/__init__.py @@ -0,0 +1 @@ +from .carracing_env import CarRacingEnv diff --git a/dizoo/box2d/carracing/envs/carracing_env.py b/dizoo/box2d/carracing/envs/carracing_env.py new file mode 100644 index 0000000000..60ebaa97d1 --- /dev/null +++ b/dizoo/box2d/carracing/envs/carracing_env.py @@ -0,0 +1,160 @@ +from typing import Optional +import copy +import os + +import gym +import numpy as np +from easydict import EasyDict + +from ding.envs import BaseEnv, BaseEnvTimestep +from ding.envs import ObsPlusPrevActRewWrapper +from ding.envs.common import affine_transform, save_frames_as_gif +from ding.torch_utils import to_ndarray +from ding.utils import ENV_REGISTRY + + +@ENV_REGISTRY.register('carracing') +class CarRacingEnv(BaseEnv): + + config = dict( + replay_path=None, + save_replay_gif=False, + replay_path_gif=None, + action_clip=False, + ) + + @classmethod + def default_config(cls: type) -> EasyDict: + cfg = EasyDict(copy.deepcopy(cls.config)) + cfg.cfg_type = cls.__name__ + 'Dict' + return cfg + + def __init__(self, cfg: dict) -> None: + self._cfg = cfg + self._init_flag = False + # env_id:CarRacing-v2 + self._env_id = cfg.env_id + self._replay_path = None + self._replay_path_gif = cfg.replay_path_gif + self._save_replay_gif = cfg.save_replay_gif + self._save_replay_count = 0 + if cfg.continuous: + self._act_scale = cfg.act_scale # act_scale only works in continuous env + self._action_clip = cfg.action_clip + else: + self._act_scale = False + + def reset(self) -> np.ndarray: + if not self._init_flag: + self._env = gym.make(self._cfg.env_id, continuous=self._cfg.continuous) + if self._replay_path is not None: + self._env = gym.wrappers.RecordVideo( + self._env, + video_folder=self._replay_path, + episode_trigger=lambda episode_id: True, + name_prefix='rl-video-{}'.format(id(self)) + ) + self._observation_space = gym.spaces.Box( + low=np.min(self._env.observation_space.low.astype(np.float32) / 255), + high=np.max(self._env.observation_space.high.astype(np.float32) / 255), + shape=( + self._env.observation_space.shape[2], self._env.observation_space.shape[0], + self._env.observation_space.shape[1] + ), + dtype=np.float32 + ) + self._action_space = self._env.action_space + self._reward_space = gym.spaces.Box( + low=self._env.reward_range[0], high=self._env.reward_range[1], shape=(1, ), dtype=np.float32 + ) + self._init_flag = True + if hasattr(self, '_seed') and hasattr(self, '_dynamic_seed') and self._dynamic_seed: + np_seed = 100 * np.random.randint(1, 1000) + self._env.seed(self._seed + np_seed) + elif hasattr(self, '_seed'): + self._env.seed(self._seed) + self._eval_episode_return = 0 + obs = self._env.reset() + obs = obs.astype(np.float32) / 255 + obs = obs.transpose(2, 0, 1) + obs = to_ndarray(obs) + if self._save_replay_gif: + self._frames = [] + return obs + + def close(self) -> None: + if self._init_flag: + self._env.close() + self._init_flag = False + + def render(self) -> None: + self._env.render() + + def seed(self, seed: int, dynamic_seed: bool = True) -> None: + self._seed = seed + self._dynamic_seed = dynamic_seed + np.random.seed(self._seed) + + def step(self, action: np.ndarray) -> BaseEnvTimestep: + assert isinstance(action, np.ndarray), type(action) + if action.shape == (1, ): + action = action.item() # 0-dim array + if self._act_scale: + action = affine_transform(action, action_clip=self._action_clip, min_val=-1, max_val=1) + if self._save_replay_gif: + self._frames.append(self._env.render(mode='rgb_array')) + obs, rew, done, info = self._env.step(action) + obs = obs.astype(np.float32) / 255 + obs = obs.transpose(2, 0, 1) + self._eval_episode_return += rew + if done: + info['eval_episode_return'] = self._eval_episode_return + if self._save_replay_gif: + if not os.path.exists(self._replay_path_gif): + os.makedirs(self._replay_path_gif) + path = os.path.join( + self._replay_path_gif, '{}_episode_{}.gif'.format(self._env_id, self._save_replay_count) + ) + save_frames_as_gif(self._frames, path) + self._save_replay_count += 1 + + obs = to_ndarray(obs) + rew = to_ndarray([rew]).astype(np.float32) # wrapped to be transferred to a array with shape (1,) + return BaseEnvTimestep(obs, rew, done, info) + + def enable_save_replay(self, replay_path: Optional[str] = None) -> None: + if replay_path is None: + replay_path = './video' + self._replay_path = replay_path + self._save_replay_gif = True + self._save_replay_count = 0 + # this function can lead to the meaningless result + self._env = gym.wrappers.RecordVideo( + self._env, + video_folder=self._replay_path, + episode_trigger=lambda episode_id: True, + name_prefix='rl-video-{}'.format(id(self)) + ) + + def random_action(self) -> np.ndarray: + random_action = self.action_space.sample() + if isinstance(random_action, np.ndarray): + pass + elif isinstance(random_action, int): + random_action = to_ndarray([random_action], dtype=np.int64) + return random_action + + @property + def observation_space(self) -> gym.spaces.Space: + return self._observation_space + + @property + def action_space(self) -> gym.spaces.Space: + return self._action_space + + @property + def reward_space(self) -> gym.spaces.Space: + return self._reward_space + + def __repr__(self) -> str: + return "DI-engine CarRacing Env" diff --git a/dizoo/box2d/carracing/envs/test_carracing_env.py b/dizoo/box2d/carracing/envs/test_carracing_env.py new file mode 100644 index 0000000000..47a5fa4638 --- /dev/null +++ b/dizoo/box2d/carracing/envs/test_carracing_env.py @@ -0,0 +1,28 @@ +import pytest +import numpy as np +from easydict import EasyDict +from carracing_env import CarRacingEnv + + +@pytest.mark.envtest +@pytest.mark.parametrize('cfg', [EasyDict({'env_id': 'CarRacing-v2', 'continuous': False, 'act_scale': False})]) +class TestCarRacing: + + def test_naive(self, cfg): + env = CarRacingEnv(cfg) + env.seed(314) + assert env._seed == 314 + obs = env.reset() + assert obs.shape == (3, 96, 96) + for i in range(10): + random_action = env.random_action() + timestep = env.step(random_action) + print(timestep) + assert isinstance(timestep.obs, np.ndarray) + assert isinstance(timestep.done, bool) + assert timestep.obs.shape == (3, 96, 96) + assert timestep.reward.shape == (1, ) + assert timestep.reward >= env.reward_space.low + assert timestep.reward <= env.reward_space.high + print(env.observation_space, env.action_space, env.reward_space) + env.close() diff --git a/dizoo/box2d/lunarlander/config/lunarlander_c51_config.py b/dizoo/box2d/lunarlander/config/lunarlander_c51_config.py index 8ca965dfb8..8a843f838c 100644 --- a/dizoo/box2d/lunarlander/config/lunarlander_c51_config.py +++ b/dizoo/box2d/lunarlander/config/lunarlander_c51_config.py @@ -10,26 +10,25 @@ stop_value=200, ), policy=dict( - cuda=False, - priority=True, + cuda=True, model=dict( obs_shape=8, action_shape=4, - encoder_hidden_size_list=[128, 128, 64], - v_min=-10, - v_max=10, + encoder_hidden_size_list=[512, 64], + v_min=-30, + v_max=30, n_atom=51, ), - discount_factor=0.97, + discount_factor=0.99, nstep=3, learn=dict( - update_per_collect=3, + update_per_collect=10, batch_size=64, learning_rate=0.001, target_update_freq=100, ), collect=dict( - n_sample=80, + n_sample=64, unroll_len=1, ), other=dict( @@ -37,8 +36,8 @@ type='exp', start=0.95, end=0.1, - decay=10000, - ), replay_buffer=dict(replay_buffer_size=20000, ) + decay=50000, + ), replay_buffer=dict(replay_buffer_size=100000, ) ), ), ) diff --git a/dizoo/box2d/lunarlander/config/lunarlander_cont_sac_config.py b/dizoo/box2d/lunarlander/config/lunarlander_cont_sac_config.py new file mode 100644 index 0000000000..f8a8ab47e7 --- /dev/null +++ b/dizoo/box2d/lunarlander/config/lunarlander_cont_sac_config.py @@ -0,0 +1,53 @@ +from easydict import EasyDict + +lunarlander_sac_config = dict( + exp_name='lunarlander_cont_sac_seed0', + env=dict( + env_id='LunarLanderContinuous-v2', + collector_env_num=4, + evaluator_env_num=8, + # (bool) Scale output action into legal range. + act_scale=True, + n_evaluator_episode=8, + stop_value=200, + ), + policy=dict( + cuda=False, + random_collect_size=10000, + model=dict( + obs_shape=8, + action_shape=2, + twin_critic=True, + action_space='reparameterization', + ), + learn=dict( + update_per_collect=256, + batch_size=128, + learning_rate_q=1e-3, + learning_rate_policy=3e-4, + learning_rate_alpha=3e-4, + auto_alpha=True, + ), + collect=dict(n_sample=256, ), + eval=dict(evaluator=dict(eval_freq=1000, ), ), + other=dict(replay_buffer=dict(replay_buffer_size=int(1e5), ), ), + ), +) +lunarlander_sac_config = EasyDict(lunarlander_sac_config) +main_config = lunarlander_sac_config + +lunarlander_sac_create_config = dict( + env=dict( + type='lunarlander', + import_names=['dizoo.box2d.lunarlander.envs.lunarlander_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='sac'), +) +lunarlander_sac_create_config = EasyDict(lunarlander_sac_create_config) +create_config = lunarlander_sac_create_config + +if __name__ == '__main__': + # or you can enter `ding -m serial -c lunarlander_cont_sac_config.py -s 0` + from ding.entry import serial_pipeline + serial_pipeline([main_config, create_config], seed=0) diff --git a/dizoo/box2d/lunarlander/config/lunarlander_cont_td3_config.py b/dizoo/box2d/lunarlander/config/lunarlander_cont_td3_config.py index e57c296555..d95932f237 100644 --- a/dizoo/box2d/lunarlander/config/lunarlander_cont_td3_config.py +++ b/dizoo/box2d/lunarlander/config/lunarlander_cont_td3_config.py @@ -4,7 +4,7 @@ exp_name='lunarlander_cont_td3_seed0', env=dict( env_id='LunarLanderContinuous-v2', - collector_env_num=8, + collector_env_num=4, evaluator_env_num=8, # (bool) Scale output action into legal range. act_scale=True, @@ -13,8 +13,7 @@ ), policy=dict( cuda=False, - priority=False, - random_collect_size=0, + random_collect_size=10000, model=dict( obs_shape=8, action_shape=2, @@ -25,8 +24,7 @@ update_per_collect=256, batch_size=128, learning_rate_actor=3e-4, - learning_rate_critic=3e-4, - ignore_done=False, + learning_rate_critic=1e-3, actor_update_freq=2, noise=True, noise_sigma=0.1, @@ -38,9 +36,8 @@ collect=dict( n_sample=256, noise_sigma=0.1, - collector=dict(collect_print_freq=1000, ), ), - eval=dict(evaluator=dict(eval_freq=100, ), ), + eval=dict(evaluator=dict(eval_freq=1000, ), ), other=dict(replay_buffer=dict(replay_buffer_size=int(1e5), ), ), ), ) diff --git a/dizoo/box2d/lunarlander/config/lunarlander_dqn_deque_config.py b/dizoo/box2d/lunarlander/config/lunarlander_dqn_deque_config.py index 778c911b4c..ab02cfcb64 100644 --- a/dizoo/box2d/lunarlander/config/lunarlander_dqn_deque_config.py +++ b/dizoo/box2d/lunarlander/config/lunarlander_dqn_deque_config.py @@ -54,7 +54,7 @@ end=0.1, decay=50000, ), - replay_buffer=dict(replay_buffer_size=100000, priority=True, priority_IS_weight=False) + replay_buffer=dict(replay_buffer_size=100000, ) ), ), ) diff --git a/dizoo/box2d/lunarlander/config/lunarlander_decision_transformer.py b/dizoo/box2d/lunarlander/config/lunarlander_dt_config.py similarity index 57% rename from dizoo/box2d/lunarlander/config/lunarlander_decision_transformer.py rename to dizoo/box2d/lunarlander/config/lunarlander_dt_config.py index 01d5fac4ea..0a3ff0a165 100644 --- a/dizoo/box2d/lunarlander/config/lunarlander_decision_transformer.py +++ b/dizoo/box2d/lunarlander/config/lunarlander_dt_config.py @@ -5,29 +5,25 @@ lunarlander_dt_config = dict( exp_name='data_dt/lunarlander_dt_1000eps_rtgt300_meel1000_seed0_debug', env=dict( - env_name='LunarLander-v2', - collector_env_num=8, + env_id='LunarLander-v2', evaluator_env_num=8, n_evaluator_episode=8, stop_value=200, ), policy=dict( stop_value=200, + state_mean=None, + state_std=None, device='cuda', env_name='LunarLander-v2', rtg_target=300, # max target reward_to_go + rtg_scale=150, max_eval_ep_len=1000, # max len of one episode # TODO - num_eval_ep=10, # num of evaluation episodes - batch_size=64, # training batch size wt_decay=1e-4, warmup_steps=10000, - num_updates_per_iter=100, context_len=20, # TODO - n_blocks=3, - embed_dim=128, - n_heads=1, - dropout_p=0.1, - log_dir='/home/puyuan/DI-engine/dizoo/box2d/lunarlander/dt_log_1000eps', + evaluator_env_num=8, + log_dir='DI-engine/dizoo/box2d/lunarlander/dt_log_1000eps', model=dict( state_dim=8, act_dim=4, @@ -41,22 +37,17 @@ discount_factor=0.999, nstep=3, learn=dict( - dataset_path='/home/puyuan/DI-engine/dizoo/box2d/lunarlander/dt_data/data/dqn_data_1000eps.pkl', # TODO - learning_rate=1e-4, + dataset_path='DI-engine/dizoo/box2d/lunarlander/offline_data/dt_data/dqn_data_1000eps.pkl', # TODO + learning_rate=3e-4, + batch_size=64, # training batch size target_update_freq=100, - kappa=1.0, - min_q_weight=4.0, ), - collect=dict(unroll_len=1, ), - eval=dict(evaluator=dict(eval_freq=100, )), - other=dict( - eps=dict( - type='exp', - start=0.95, - end=0.1, - decay=10000, - ), replay_buffer=dict(replay_buffer_size=1000, ) + collect=dict( + data_type='d4rl_trajectory', + data_path='DI-engine/dizoo/box2d/lunarlander/offline_data/dt_data/dqn_data_1000eps.pkl', + unroll_len=1, ), + eval=dict(evaluator=dict(eval_freq=100, )), ), ) lunarlander_dt_config = EasyDict(lunarlander_dt_config) @@ -71,8 +62,3 @@ ) lunarlander_dt_create_config = EasyDict(lunarlander_dt_create_config) create_config = lunarlander_dt_create_config - -if __name__ == "__main__": - from ding.entry import serial_pipeline_dt, collect_demo_data, eval, serial_pipeline - config = deepcopy([main_config, create_config]) - serial_pipeline_dt(config, seed=0, max_train_iter=1000) diff --git a/dizoo/box2d/lunarlander/config/lunarlander_hpt_config.py b/dizoo/box2d/lunarlander/config/lunarlander_hpt_config.py new file mode 100644 index 0000000000..b58166d566 --- /dev/null +++ b/dizoo/box2d/lunarlander/config/lunarlander_hpt_config.py @@ -0,0 +1,78 @@ +from easydict import EasyDict + +nstep = 3 +lunarlander_hpt_config = dict( + exp_name='lunarlander_hpt_seed0', + env=dict( + # Whether to use shared memory. Only effective if "env_manager_type" is 'subprocess' + # Env number respectively for collector and evaluator. + collector_env_num=8, + evaluator_env_num=8, + env_id='LunarLander-v2', + n_evaluator_episode=8, + stop_value=200, + # The path to save the game replay + # replay_path='./lunarlander_hpt_seed0/video', + ), + policy=dict( + # Whether to use cuda for network. + cuda=True, + load_path="./lunarlander_hpt_seed0/ckpt/ckpt_best.pth.tar", + model=dict( + obs_shape=8, + action_shape=4, + ), + # Reward's future discount factor, aka. gamma. + discount_factor=0.99, + # How many steps in td error. + nstep=nstep, + # learn_mode config + learn=dict( + update_per_collect=10, + batch_size=64, + learning_rate=0.0005, + # Frequency of target network update. + target_update_freq=100, + ), + # collect_mode config + collect=dict( + # You can use either "n_sample" or "n_episode" in collector.collect. + # Get "n_sample" samples per collect. + n_sample=64, + # Cut trajectories into pieces with length "unroll_len". + unroll_len=1, + ), + # command_mode config + other=dict( + # Epsilon greedy with decay. + eps=dict( + # Decay type. Support ['exp', 'linear']. + type='exp', + start=0.95, + end=0.1, + decay=50000, + ), + replay_buffer=dict(replay_buffer_size=100000, ) + ), + ), +) +lunarlander_hpt_config = EasyDict(lunarlander_hpt_config) +main_config = lunarlander_hpt_config + +lunarlander_hpt_create_config = dict( + env=dict( + type='lunarlander', + import_names=['dizoo.box2d.lunarlander.envs.lunarlander_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='dqn'), +) +lunarlander_hpt_create_config = EasyDict(lunarlander_hpt_create_config) +create_config = lunarlander_hpt_create_config +""" +This is a configuration file for LunarLander environment with HPT (Hindsight Policy Transformer). +To run this config, please use the lunarlander_hpt_example.py script. + +Example: + python lunarlander_hpt_example.py +""" diff --git a/dizoo/box2d/lunarlander/config/lunarlander_impala_config.py b/dizoo/box2d/lunarlander/config/lunarlander_impala_config.py new file mode 100644 index 0000000000..8725f5bd83 --- /dev/null +++ b/dizoo/box2d/lunarlander/config/lunarlander_impala_config.py @@ -0,0 +1,74 @@ +from easydict import EasyDict + +lunarlander_impala_config = dict( + exp_name='impala_log/lunarlander_impala_seed0', + env=dict( + env_id='LunarLander-v2', + collector_env_num=8, + evaluator_env_num=5, + n_evaluator_episode=5, + stop_value=3000, + ), + policy=dict( + cuda=True, + # (int) the trajectory length to calculate v-trace target + unroll_len=32, + random_collect_size=256, + model=dict( + obs_shape=8, + action_shape=4, + encoder_hidden_size_list=[64, 64], + ), + learn=dict( + # (int) collect n_sample data, train model update_per_collect times + # here we follow ppo serial pipeline + update_per_collect=10, + # (int) the number of data for a train iteration + batch_size=128, + grad_clip_type='clip_norm', + clip_value=5, + learning_rate=0.0003, + # (float) loss weight of the value network, the weight of policy network is set to 1 + value_weight=0.5, + # (float) loss weight of the entropy regularization, the weight of policy network is set to 1 + entropy_weight=0.0001, + # (float) discount factor for future reward, defaults int [0, 1] + discount_factor=0.99, + # (float) additional discounting parameter + lambda_=0.95, + # (float) clip ratio of importance weights + rho_clip_ratio=1.0, + # (float) clip ratio of importance weights + c_clip_ratio=1.0, + # (float) clip ratio of importance sampling + rho_pg_clip_ratio=1.0, + ), + collect=dict( + # (int) collect n_sample data, train model update_per_collect times + n_sample=32, + ), + eval=dict(evaluator=dict(eval_freq=500, )), + other=dict(replay_buffer=dict(replay_buffer_size=1000, sliced=True), ), + ), +) + +lunarlander_impala_config = EasyDict(lunarlander_impala_config) +main_config = lunarlander_impala_config + +lunarlander_impala_create_config = dict( + env=dict( + type='lunarlander', + import_names=['dizoo.box2d.lunarlander.envs.lunarlander_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='impala'), + replay_buffer=dict(type='naive'), +) + +lunarlander_impala_create_config = EasyDict(lunarlander_impala_create_config) +create_config = lunarlander_impala_create_config + +if __name__ == "__main__": + # or you can enter `ding -m serial -c lunarlander_impala_config.py -s 0` + from ding.entry import serial_pipeline + serial_pipeline((main_config, create_config), seed=0) diff --git a/dizoo/box2d/lunarlander/config/lunarlander_pg_config.py b/dizoo/box2d/lunarlander/config/lunarlander_pg_config.py new file mode 100644 index 0000000000..fa66ef4ae9 --- /dev/null +++ b/dizoo/box2d/lunarlander/config/lunarlander_pg_config.py @@ -0,0 +1,45 @@ +from easydict import EasyDict + +lunarlander_pg_config = dict( + exp_name='lunarlander_pg_seed0', + env=dict( + env_id='LunarLander-v2', + collector_env_num=4, + evaluator_env_num=4, + n_evaluator_episode=4, + stop_value=200, + ), + policy=dict( + cuda=True, + model=dict( + obs_shape=8, + action_shape=4, + ), + learn=dict( + batch_size=320, + learning_rate=3e-4, + entropy_weight=0.001, + grad_norm=0.5, + ), + collect=dict(n_episode=8, discount_factor=0.99), + eval=dict(evaluator=dict(eval_freq=1000, ), ), + ), +) +lunarlander_pg_config = EasyDict(lunarlander_pg_config) +main_config = lunarlander_pg_config +lunarlander_pg_create_config = dict( + env=dict( + type='lunarlander', + import_names=['dizoo.box2d.lunarlander.envs.lunarlander_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='pg'), + collector=dict(type='episode'), +) +lunarlander_pg_create_config = EasyDict(lunarlander_pg_create_config) +create_config = lunarlander_pg_create_config + +if __name__ == "__main__": + # or you can enter `ding -m serial_onpolicy -c lunarlander_pg_config.py -s 0` + from ding.entry import serial_pipeline_onpolicy + serial_pipeline_onpolicy((main_config, create_config), seed=0, max_env_step=int(1e7)) diff --git a/dizoo/box2d/lunarlander/config/lunarlander_ppo_config.py b/dizoo/box2d/lunarlander/config/lunarlander_ppo_config.py new file mode 100644 index 0000000000..ad622c444d --- /dev/null +++ b/dizoo/box2d/lunarlander/config/lunarlander_ppo_config.py @@ -0,0 +1,50 @@ +from easydict import EasyDict + +lunarlander_ppo_config = dict( + exp_name='lunarlander_ppo_seed0', + env=dict( + collector_env_num=8, + evaluator_env_num=5, + env_id='LunarLander-v2', + n_evaluator_episode=5, + stop_value=200, + ), + policy=dict( + recompute_adv=True, + cuda=True, + action_space='discrete', + model=dict( + obs_shape=8, + action_shape=4, + action_space='discrete', + ), + learn=dict( + epoch_per_collect=10, + batch_size=64, + learning_rate=3e-4, + entropy_weight=0.01, + adv_norm=True, + value_norm=True, + ), + collect=dict( + n_sample=512, + discount_factor=0.99, + ), + ), +) +lunarlander_ppo_config = EasyDict(lunarlander_ppo_config) +main_config = lunarlander_ppo_config +lunarlander_ppo_create_config = dict( + env=dict( + type='lunarlander', + import_names=['dizoo.box2d.lunarlander.envs.lunarlander_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='ppo'), +) +lunarlander_ppo_create_config = EasyDict(lunarlander_ppo_create_config) +create_config = lunarlander_ppo_create_config + +if __name__ == "__main__": + from ding.entry import serial_pipeline_onpolicy + serial_pipeline_onpolicy([main_config, create_config], seed=0) diff --git a/dizoo/box2d/lunarlander/config/lunarlander_ppo_continuous_config.py b/dizoo/box2d/lunarlander/config/lunarlander_ppo_continuous_config.py new file mode 100644 index 0000000000..fa89c0d878 --- /dev/null +++ b/dizoo/box2d/lunarlander/config/lunarlander_ppo_continuous_config.py @@ -0,0 +1,50 @@ +from easydict import EasyDict + +lunarlander_ppo_config = dict( + exp_name='lunarlander_ppo_continuous_seed0', + env=dict( + collector_env_num=8, + evaluator_env_num=5, + env_id='LunarLander-v2', + n_evaluator_episode=5, + stop_value=200, + ), + policy=dict( + recompute_adv=True, + cuda=True, + action_space='continuous', + model=dict( + obs_shape=8, + action_shape=2, + action_space='continuous', + ), + learn=dict( + epoch_per_collect=5, + batch_size=64, + learning_rate=3e-4, + entropy_weight=0.005, + adv_norm=True, + value_norm=True, + ), + collect=dict( + n_sample=512, + discount_factor=0.99, + ), + ), +) +lunarlander_ppo_config = EasyDict(lunarlander_ppo_config) +main_config = lunarlander_ppo_config +lunarlander_ppo_create_config = dict( + env=dict( + type='lunarlander', + import_names=['dizoo.box2d.lunarlander.envs.lunarlander_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='ppo'), +) +lunarlander_ppo_create_config = EasyDict(lunarlander_ppo_create_config) +create_config = lunarlander_ppo_create_config + +if __name__ == "__main__": + from ding.entry import serial_pipeline_onpolicy + serial_pipeline_onpolicy([main_config, create_config], seed=0) diff --git a/dizoo/box2d/lunarlander/config/lunarlander_sqil_config.py b/dizoo/box2d/lunarlander/config/lunarlander_sqil_config.py index 638c2d2981..b5deffc897 100644 --- a/dizoo/box2d/lunarlander/config/lunarlander_sqil_config.py +++ b/dizoo/box2d/lunarlander/config/lunarlander_sqil_config.py @@ -63,4 +63,4 @@ from dizoo.box2d.lunarlander.config import lunarlander_dqn_config, lunarlander_dqn_create_config expert_main_config = lunarlander_dqn_config expert_create_config = lunarlander_dqn_create_config - serial_pipeline_sqil([main_config, create_config], [expert_main_config, expert_create_config], seed=0) \ No newline at end of file + serial_pipeline_sqil([main_config, create_config], [expert_main_config, expert_create_config], seed=0) diff --git a/dizoo/box2d/lunarlander/entry/lunarlander_dqn_eval.py b/dizoo/box2d/lunarlander/entry/lunarlander_dqn_eval.py index 0bc754376d..a87a80f8e8 100644 --- a/dizoo/box2d/lunarlander/entry/lunarlander_dqn_eval.py +++ b/dizoo/box2d/lunarlander/entry/lunarlander_dqn_eval.py @@ -15,8 +15,9 @@ from ding.rl_utils import get_epsilon_greedy_fn from dizoo.box2d.lunarlander.config.lunarlander_dqn_config import main_config, create_config + def main(rl_cfg, seed=0): - main_cfg, create_cfg =rl_cfg + main_cfg, create_cfg = rl_cfg cfg = compile_config( main_cfg, BaseEnvManager, @@ -56,4 +57,4 @@ def main(rl_cfg, seed=0): if __name__ == "__main__": - main(rl_cfg=(main_config, create_config),seed=0) + main(rl_cfg=(main_config, create_config), seed=0) diff --git a/dizoo/box2d/lunarlander/entry/lunarlander_dqn_example.py b/dizoo/box2d/lunarlander/entry/lunarlander_dqn_example.py new file mode 100644 index 0000000000..ae9b30d1fd --- /dev/null +++ b/dizoo/box2d/lunarlander/entry/lunarlander_dqn_example.py @@ -0,0 +1,75 @@ +import gym +import torch +from ditk import logging +from ding.data.model_loader import FileModelLoader +from ding.data.storage_loader import FileStorageLoader +from ding.model import DQN +from ding.policy import DQNPolicy +from ding.envs import DingEnvWrapper, SubprocessEnvManagerV2 +from ding.data import DequeBuffer +from ding.config import compile_config +from ding.framework import task, ding_init +from ding.framework.context import OnlineRLContext +from ding.framework.middleware import OffPolicyLearner, StepCollector, interaction_evaluator, data_pusher, \ + eps_greedy_handler, CkptSaver, ContextExchanger, ModelExchanger, online_logger, termination_checker, \ + nstep_reward_enhancer +from ding.utils import set_pkg_seed +from dizoo.box2d.lunarlander.config.lunarlander_dqn_config import main_config, create_config + + +def main(): + logging.getLogger().setLevel(logging.INFO) + cfg = compile_config(main_config, create_cfg=create_config, auto=True, save_cfg=task.router.node_id == 0) + ding_init(cfg) + + with task.start(async_mode=False, ctx=OnlineRLContext()): + collector_env = SubprocessEnvManagerV2( + env_fn=[lambda: DingEnvWrapper(gym.make("LunarLander-v2")) for _ in range(cfg.env.collector_env_num)], + cfg=cfg.env.manager + ) + evaluator_env = SubprocessEnvManagerV2( + env_fn=[lambda: DingEnvWrapper(gym.make("LunarLander-v2")) for _ in range(cfg.env.evaluator_env_num)], + cfg=cfg.env.manager + ) + + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + + # Migrating models to the GPU + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = DQN(**cfg.policy.model).to(device) + + buffer_ = DequeBuffer(size=cfg.policy.other.replay_buffer.replay_buffer_size) + + # Pass the model into Policy + policy = DQNPolicy(cfg.policy, model=model) + + # Consider the case with multiple processes + if task.router.is_active: + # You can use labels to distinguish between workers with different roles, + # here we use node_id to distinguish. + if task.router.node_id == 0: + task.add_role(task.role.LEARNER) + elif task.router.node_id == 1: + task.add_role(task.role.EVALUATOR) + else: + task.add_role(task.role.COLLECTOR) + + # Sync their context and model between each worker. + task.use(ContextExchanger(skip_n_iter=1)) + task.use(ModelExchanger(model)) + + # Here is the part of single process pipeline. + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) + task.use(eps_greedy_handler(cfg)) + task.use(StepCollector(cfg, policy.collect_mode, collector_env)) + task.use(nstep_reward_enhancer(cfg)) + task.use(data_pusher(cfg, buffer_)) + task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) + task.use(online_logger(train_show_freq=50)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=1000)) + task.use(termination_checker(max_env_step=int(3e6))) + task.run() + + +if __name__ == "__main__": + main() diff --git a/dizoo/box2d/lunarlander/entry/lunarlander_hpt_example.py b/dizoo/box2d/lunarlander/entry/lunarlander_hpt_example.py new file mode 100644 index 0000000000..f5da04eaf9 --- /dev/null +++ b/dizoo/box2d/lunarlander/entry/lunarlander_hpt_example.py @@ -0,0 +1,77 @@ +import gym +import torch +import torch.nn as nn +from ditk import logging +from ding.data.model_loader import FileModelLoader +from ding.data.storage_loader import FileStorageLoader +from ding.model.common.head import DuelingHead +from ding.model.template.hpt import HPT +from ding.policy import DQNPolicy +from ding.envs import DingEnvWrapper, SubprocessEnvManagerV2 +from ding.data import DequeBuffer +from ding.config import compile_config +from ding.framework import task, ding_init +from ding.framework.context import OnlineRLContext +from ding.framework.middleware import OffPolicyLearner, StepCollector, interaction_evaluator, data_pusher, \ + eps_greedy_handler, CkptSaver, ContextExchanger, ModelExchanger, online_logger, termination_checker, \ + nstep_reward_enhancer +from ding.utils import set_pkg_seed +from dizoo.box2d.lunarlander.config.lunarlander_hpt_config import main_config, create_config + + +def main(): + logging.getLogger().setLevel(logging.INFO) + cfg = compile_config(main_config, create_cfg=create_config, auto=True, save_cfg=task.router.node_id == 0) + ding_init(cfg) + + with task.start(async_mode=False, ctx=OnlineRLContext()): + collector_env = SubprocessEnvManagerV2( + env_fn=[lambda: DingEnvWrapper(gym.make("LunarLander-v2")) for _ in range(cfg.env.collector_env_num)], + cfg=cfg.env.manager + ) + evaluator_env = SubprocessEnvManagerV2( + env_fn=[lambda: DingEnvWrapper(gym.make("LunarLander-v2")) for _ in range(cfg.env.evaluator_env_num)], + cfg=cfg.env.manager + ) + + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + + # Migrating models to the GPU + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + # HPT introduces a Policy Stem module, which processes the input features using Cross-Attention. + model = HPT(cfg.policy.model.obs_shape, cfg.policy.model.action_shape).to(device) + buffer_ = DequeBuffer(size=cfg.policy.other.replay_buffer.replay_buffer_size) + + # Pass the model into Policy + policy = DQNPolicy(cfg.policy, model=model) + + # Consider the case with multiple processes + if task.router.is_active: + # You can use labels to distinguish between workers with different roles, + # here we use node_id to distinguish. + if task.router.node_id == 0: + task.add_role(task.role.LEARNER) + elif task.router.node_id == 1: + task.add_role(task.role.EVALUATOR) + else: + task.add_role(task.role.COLLECTOR) + + # Sync their context and model between each worker. + task.use(ContextExchanger(skip_n_iter=1)) + task.use(ModelExchanger(model)) + + # Here is the part of single process pipeline. + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) + task.use(eps_greedy_handler(cfg)) + task.use(StepCollector(cfg, policy.collect_mode, collector_env)) + task.use(nstep_reward_enhancer(cfg)) + task.use(data_pusher(cfg, buffer_)) + task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) + task.use(online_logger(train_show_freq=50)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=1000)) + task.use(termination_checker(max_env_step=int(3e6))) + task.run() + + +if __name__ == "__main__": + main() diff --git a/dizoo/box2d/lunarlander/envs/lunarlander_env.py b/dizoo/box2d/lunarlander/envs/lunarlander_env.py index d9d65f8eef..235d888115 100644 --- a/dizoo/box2d/lunarlander/envs/lunarlander_env.py +++ b/dizoo/box2d/lunarlander/envs/lunarlander_env.py @@ -1,25 +1,46 @@ -from typing import Any, List, Union, Optional -import time +import copy +import os +from typing import Optional + import gym import numpy as np +from easydict import EasyDict + from ding.envs import BaseEnv, BaseEnvTimestep -from ding.torch_utils import to_ndarray, to_list -from ding.utils import ENV_REGISTRY -from ding.envs.common import affine_transform from ding.envs import ObsPlusPrevActRewWrapper +from ding.envs.common import affine_transform, save_frames_as_gif +from ding.torch_utils import to_ndarray +from ding.utils import ENV_REGISTRY @ENV_REGISTRY.register('lunarlander') class LunarLanderEnv(BaseEnv): + config = dict( + replay_path=None, + save_replay_gif=False, + replay_path_gif=None, + action_clip=False, + ) + + @classmethod + def default_config(cls: type) -> EasyDict: + cfg = EasyDict(copy.deepcopy(cls.config)) + cfg.cfg_type = cls.__name__ + 'Dict' + return cfg + def __init__(self, cfg: dict) -> None: self._cfg = cfg self._init_flag = False # env_id: LunarLander-v2, LunarLanderContinuous-v2 self._env_id = cfg.env_id self._replay_path = None + self._replay_path_gif = cfg.replay_path_gif + self._save_replay_gif = cfg.save_replay_gif + self._save_replay_count = 0 if 'Continuous' in self._env_id: - self._act_scale = cfg.act_scale # act_scale only works in continous env + self._act_scale = cfg.act_scale # act_scale only works in continuous env + self._action_clip = cfg.action_clip else: self._act_scale = False @@ -46,10 +67,11 @@ def reset(self) -> np.ndarray: self._env.seed(self._seed + np_seed) elif hasattr(self, '_seed'): self._env.seed(self._seed) - self._final_eval_reward = 0 + self._eval_episode_return = 0 obs = self._env.reset() obs = to_ndarray(obs) - + if self._save_replay_gif: + self._frames = [] return obs def close(self) -> None: @@ -70,11 +92,21 @@ def step(self, action: np.ndarray) -> BaseEnvTimestep: if action.shape == (1, ): action = action.item() # 0-dim array if self._act_scale: - action = affine_transform(action, min_val=-1, max_val=1) + action = affine_transform(action, action_clip=self._action_clip, min_val=-1, max_val=1) + if self._save_replay_gif: + self._frames.append(self._env.render(mode='rgb_array')) obs, rew, done, info = self._env.step(action) - self._final_eval_reward += rew + self._eval_episode_return += rew if done: - info['final_eval_reward'] = self._final_eval_reward + info['eval_episode_return'] = self._eval_episode_return + if self._save_replay_gif: + if not os.path.exists(self._replay_path_gif): + os.makedirs(self._replay_path_gif) + path = os.path.join( + self._replay_path_gif, '{}_episode_{}.gif'.format(self._env_id, self._save_replay_count) + ) + save_frames_as_gif(self._frames, path) + self._save_replay_count += 1 obs = to_ndarray(obs) rew = to_ndarray([rew]).astype(np.float32) # wrapped to be transferred to a array with shape (1,) @@ -84,6 +116,8 @@ def enable_save_replay(self, replay_path: Optional[str] = None) -> None: if replay_path is None: replay_path = './video' self._replay_path = replay_path + self._save_replay_gif = True + self._save_replay_count = 0 # this function can lead to the meaningless result self._env = gym.wrappers.RecordVideo( self._env, diff --git a/dizoo/box2d/lunarlander/dt_data/collect_dqn_data_config.py b/dizoo/box2d/lunarlander/offline_data/collect_dqn_data_config.py similarity index 86% rename from dizoo/box2d/lunarlander/dt_data/collect_dqn_data_config.py rename to dizoo/box2d/lunarlander/offline_data/collect_dqn_data_config.py index 0beb38fbd5..e7cc7b383d 100644 --- a/dizoo/box2d/lunarlander/dt_data/collect_dqn_data_config.py +++ b/dizoo/box2d/lunarlander/offline_data/collect_dqn_data_config.py @@ -34,15 +34,16 @@ dataloader=dict(num_workers=0, ), log_policy=True, hook=dict( - # load_ckpt_before_run='./lunarlander/ckpt/ckpt_best.pth.tar', - load_ckpt_before_run='/home/puyuan/DI-engine/dizoo/box2d/lunarlander/dt_data/ckpt/ckpt_best.pth.tar', + load_ckpt_before_run= + './ckpt_best.pth.tar', # TODO: syspath modeified in other place, have to use abs path. May be fix in next version. + # load_ckpt_before_run='DI-engine/dizoo/box2d/lunarlander/dt_data/ckpt/ckpt_best.pth.tar', log_show_after_iter=100, save_ckpt_after_iter=10000, save_ckpt_after_run=False, ), cfg_type='BaseLearnerDict', - # load_path='./cartpole/ckpt/ckpt_best.pth.tar', - load_path='/home/puyuan/DI-engine/dizoo/box2d/lunarlander/dt_data/ckpt/ckpt_best.pth.tar', + load_path='./ckpt_best.pth.tar', # TODO: same like last path. + # load_path='DI-engine/dizoo/box2d/lunarlander/dt_data/ckpt/ckpt_best.pth.tar', ), update_per_collect=10, batch_size=64, @@ -61,9 +62,9 @@ # save # data_type='hdf5', data_type='naive', - save_path='/home/puyuan/DI-engine/dizoo/box2d/lunarlander/dt_data/data/dqn_data_1000eps.pkl', # TODO(pu) + save_path='./dt_data/dqn_data_1000eps.pkl', # TODO(pu) # load - data_path='/home/puyuan/DI-engine/dizoo/box2d/lunarlander/dt_data/data/dqn_data_10eps.pkl', # TODO(pu) + data_path='./dt_data/dqn_data_10eps.pkl', # TODO(pu) ), # command_mode config other=dict( diff --git a/dizoo/box2d/lunarlander/dt_data/lunarlander_collect_data.py b/dizoo/box2d/lunarlander/offline_data/lunarlander_collect_data.py similarity index 91% rename from dizoo/box2d/lunarlander/dt_data/lunarlander_collect_data.py rename to dizoo/box2d/lunarlander/offline_data/lunarlander_collect_data.py index fda635b41f..58256b758e 100644 --- a/dizoo/box2d/lunarlander/dt_data/lunarlander_collect_data.py +++ b/dizoo/box2d/lunarlander/offline_data/lunarlander_collect_data.py @@ -1,4 +1,4 @@ -from dizoo.box2d.lunarlander.dt_data.collect_dqn_data_config import main_config, create_config +from dizoo.box2d.lunarlander.offline_data.collect_dqn_data_config import main_config, create_config from ding.entry import collect_episodic_demo_data, eval import torch import copy diff --git a/dizoo/box2d/lunarlander/dt_data/lunarlander_show_data.py b/dizoo/box2d/lunarlander/offline_data/lunarlander_show_data.py similarity index 93% rename from dizoo/box2d/lunarlander/dt_data/lunarlander_show_data.py rename to dizoo/box2d/lunarlander/offline_data/lunarlander_show_data.py index 94bbf89d97..c3a63b071d 100644 --- a/dizoo/box2d/lunarlander/dt_data/lunarlander_show_data.py +++ b/dizoo/box2d/lunarlander/offline_data/lunarlander_show_data.py @@ -1,4 +1,4 @@ -from dizoo.classic_control.cartpole.dt_data.collect_dqn_data_config import main_config, create_config +from dizoo.classic_control.cartpole.offline_data.collect_dqn_data_config import main_config, create_config from ding.entry import serial_pipeline_offline import os diff --git a/dizoo/bsuite/envs/bsuite_env.py b/dizoo/bsuite/envs/bsuite_env.py index 45d5406e68..915411f57d 100644 --- a/dizoo/bsuite/envs/bsuite_env.py +++ b/dizoo/bsuite/envs/bsuite_env.py @@ -36,7 +36,7 @@ def reset(self) -> np.ndarray: self._env.seed(self._seed + np_seed) elif hasattr(self, '_seed'): self._env.seed(self._seed) - self._final_eval_reward = 0 + self._eval_episode_return = 0 obs = self._env.reset() if obs.shape[0] == 1: obs = obs[0] @@ -58,9 +58,9 @@ def step(self, action: np.ndarray) -> BaseEnvTimestep: if action.shape[0] == 1: action = action[0] obs, rew, done, info = self._env.step(action) - self._final_eval_reward += rew + self._eval_episode_return += rew if done: - info['final_eval_reward'] = self._final_eval_reward + info['eval_episode_return'] = self._eval_episode_return if obs.shape[0] == 1: obs = obs[0] obs = to_ndarray(obs) diff --git a/dizoo/bsuite/envs/test_bsuite_env.py b/dizoo/bsuite/envs/test_bsuite_env.py index 09f8e9339b..93a330bfb8 100644 --- a/dizoo/bsuite/envs/test_bsuite_env.py +++ b/dizoo/bsuite/envs/test_bsuite_env.py @@ -21,7 +21,7 @@ def test_memory_len(self): assert timestep.obs.shape == (3, ) assert timestep.reward.shape == (1, ) if timestep.done: - assert 'final_eval_reward' in timestep.info, timestep.info + assert 'eval_episode_return' in timestep.info, timestep.info break memory_len_env.close() @@ -38,6 +38,6 @@ def test_cartpole_swingup(self): assert timestep.obs.shape == (8, ) assert timestep.reward.shape == (1, ) if timestep.done: - assert 'final_eval_reward' in timestep.info, timestep.info + assert 'eval_episode_return' in timestep.info, timestep.info break bandit_noise_env.close() diff --git a/dizoo/classic_control/acrobot/__init__.py b/dizoo/classic_control/acrobot/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/dizoo/classic_control/acrobot/acrobot.gif b/dizoo/classic_control/acrobot/acrobot.gif new file mode 100644 index 0000000000..3eef302afe Binary files /dev/null and b/dizoo/classic_control/acrobot/acrobot.gif differ diff --git a/dizoo/classic_control/acrobot/config/__init__.py b/dizoo/classic_control/acrobot/config/__init__.py new file mode 100644 index 0000000000..036dbf6a93 --- /dev/null +++ b/dizoo/classic_control/acrobot/config/__init__.py @@ -0,0 +1 @@ +from .acrobot_dqn_config import acrobot_dqn_config, acrobot_dqn_create_config diff --git a/dizoo/classic_control/acrobot/config/acrobot_dqn_config.py b/dizoo/classic_control/acrobot/config/acrobot_dqn_config.py new file mode 100644 index 0000000000..4957db987f --- /dev/null +++ b/dizoo/classic_control/acrobot/config/acrobot_dqn_config.py @@ -0,0 +1,55 @@ +from easydict import EasyDict + +acrobot_dqn_config = dict( + exp_name='acrobot_dqn_seed0', + env=dict( + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=-60, + env_id='Acrobot-v1', + replay_path='acrobot_dqn_seed0/video', + ), + policy=dict( + cuda=True, + model=dict( + obs_shape=6, + action_shape=3, + encoder_hidden_size_list=[256, 256], + dueling=True, + ), + nstep=3, + discount_factor=0.99, + learn=dict( + update_per_collect=10, + batch_size=128, + learning_rate=0.0001, + target_update_freq=250, + ), + collect=dict(n_sample=96, ), + eval=dict(evaluator=dict(eval_freq=2000, )), + other=dict( + eps=dict( + type='exp', + start=1., + end=0.05, + decay=250000, + ), + replay_buffer=dict(replay_buffer_size=100000, ), + ), + ), +) +acrobot_dqn_config = EasyDict(acrobot_dqn_config) +main_config = acrobot_dqn_config +acrobot_dqn_create_config = dict( + env=dict(type='acrobot', import_names=['dizoo.classic_control.acrobot.envs.acrobot_env']), + env_manager=dict(type='subprocess'), + policy=dict(type='dqn'), + replay_buffer=dict(type='deque', import_names=['ding.data.buffer.deque_buffer_wrapper']), +) +acrobot_dqn_create_config = EasyDict(acrobot_dqn_create_config) +create_config = acrobot_dqn_create_config + +if __name__ == "__main__": + from ding.entry import serial_pipeline + serial_pipeline((main_config, create_config), seed=0) diff --git a/dizoo/classic_control/acrobot/envs/__init__.py b/dizoo/classic_control/acrobot/envs/__init__.py new file mode 100644 index 0000000000..be6537f2c9 --- /dev/null +++ b/dizoo/classic_control/acrobot/envs/__init__.py @@ -0,0 +1 @@ +from .acrobot_env import AcroBotEnv diff --git a/dizoo/classic_control/acrobot/envs/acrobot_env.py b/dizoo/classic_control/acrobot/envs/acrobot_env.py new file mode 100644 index 0000000000..3c26323315 --- /dev/null +++ b/dizoo/classic_control/acrobot/envs/acrobot_env.py @@ -0,0 +1,98 @@ +from typing import Any, List, Union, Optional +import time +import gym +import copy +import numpy as np +from easydict import EasyDict +from ding.envs import BaseEnv, BaseEnvTimestep +from ding.torch_utils import to_ndarray, to_list +from ding.utils import ENV_REGISTRY +from ding.envs import ObsPlusPrevActRewWrapper + + +@ENV_REGISTRY.register('acrobot') +class AcroBotEnv(BaseEnv): + + def __init__(self, cfg: dict = {}) -> None: + self._cfg = cfg + self._init_flag = False + self._replay_path = None + self._observation_space = gym.spaces.Box( + low=np.array([-1.0, -1.0, -1.0, -1.0, -12.57, -28.27]), + high=np.array([1.0, 1.0, 1.0, 1.0, 12.57, 28.27]), + shape=(6, ), + dtype=np.float32 + ) + self._action_space = gym.spaces.Discrete(3) + self._action_space.seed(0) # default seed + self._reward_space = gym.spaces.Box(low=-1.0, high=0.0, shape=(1, ), dtype=np.float32) + + def reset(self) -> np.ndarray: + if not self._init_flag: + self._env = gym.make('Acrobot-v1') + if self._replay_path is not None: + self._env = gym.wrappers.RecordVideo( + self._env, + video_folder=self._replay_path, + episode_trigger=lambda episode_id: True, + name_prefix='rl-video-{}'.format(id(self)) + ) + self._init_flag = True + if hasattr(self, '_seed') and hasattr(self, '_dynamic_seed') and self._dynamic_seed: + np_seed = 100 * np.random.randint(1, 1000) + self._env.seed(self._seed + np_seed) + self._action_space.seed(self._seed + np_seed) + elif hasattr(self, '_seed'): + self._env.seed(self._seed) + self._action_space.seed(self._seed) + self._observation_space = self._env.observation_space + self._eval_episode_return = 0 + obs = self._env.reset() + obs = to_ndarray(obs) + return obs + + def close(self) -> None: + if self._init_flag: + self._env.close() + self._init_flag = False + + def seed(self, seed: int, dynamic_seed: bool = True) -> None: + self._seed = seed + self._dynamic_seed = dynamic_seed + np.random.seed(self._seed) + + def step(self, action: Union[int, np.ndarray]) -> BaseEnvTimestep: + if isinstance(action, np.ndarray) and action.shape == (1, ): + action = action.squeeze() # 0-dim array + obs, rew, done, info = self._env.step(action) + self._eval_episode_return += rew + if done: + info['eval_episode_return'] = self._eval_episode_return + obs = to_ndarray(obs) + rew = to_ndarray([rew]).astype(np.float32) # wrapped to be transfered to a array with shape (1,) + return BaseEnvTimestep(obs, rew, done, info) + + def enable_save_replay(self, replay_path: Optional[str] = None) -> None: + if replay_path is None: + replay_path = './video' + self._replay_path = replay_path + + def random_action(self) -> np.ndarray: + random_action = self.action_space.sample() + random_action = to_ndarray([random_action], dtype=np.int64) + return random_action + + @property + def observation_space(self) -> gym.spaces.Space: + return self._observation_space + + @property + def action_space(self) -> gym.spaces.Space: + return self._action_space + + @property + def reward_space(self) -> gym.spaces.Space: + return self._reward_space + + def __repr__(self) -> str: + return "DI-engine Acrobot Env" diff --git a/dizoo/classic_control/acrobot/envs/test_acrobot_env.py b/dizoo/classic_control/acrobot/envs/test_acrobot_env.py new file mode 100644 index 0000000000..fba0914cfa --- /dev/null +++ b/dizoo/classic_control/acrobot/envs/test_acrobot_env.py @@ -0,0 +1,35 @@ +import pytest +import numpy as np +from dizoo.classic_control.acrobot.envs import AcroBotEnv + + +@pytest.mark.envtest +class TestAcrobotEnv: + + def test_naive(self): + env = AcroBotEnv({}) + env.seed(314, dynamic_seed=False) + assert env._seed == 314 + obs = env.reset() + assert obs.shape == (6, ) + for _ in range(5): + env.reset() + np.random.seed(314) + print('=' * 60) + for i in range(10): + # Both ``env.random_action()``, and utilizing ``np.random`` as well as action space, + # can generate legal random action. + if i < 5: + random_action = np.array([env.action_space.sample()]) + else: + random_action = env.random_action() + timestep = env.step(random_action) + print(timestep) + assert isinstance(timestep.obs, np.ndarray) + assert isinstance(timestep.done, bool) + assert timestep.obs.shape == (6, ) + assert timestep.reward.shape == (1, ) + assert timestep.reward >= env.reward_space.low + assert timestep.reward <= env.reward_space.high + print(env.observation_space, env.action_space, env.reward_space) + env.close() diff --git a/dizoo/classic_control/cartpole/config/__init__.py b/dizoo/classic_control/cartpole/config/__init__.py index 64fd50b9b2..3d6d124274 100644 --- a/dizoo/classic_control/cartpole/config/__init__.py +++ b/dizoo/classic_control/cartpole/config/__init__.py @@ -7,7 +7,7 @@ from .cartpole_gcl_config import cartpole_gcl_ppo_onpolicy_config, cartpole_gcl_ppo_onpolicy_create_config from .cartpole_impala_config import cartpole_impala_config, cartpole_impala_create_config from .cartpole_iqn_config import cartpole_iqn_config, cartpole_iqn_create_config -from .cartpole_offppo_config import cartpole_offppo_config, cartpole_offppo_create_config +from .cartpole_ppo_offpolicy_config import cartpole_ppo_offpolicy_config, cartpole_ppo_offpolicy_create_config from .cartpole_ppg_config import cartpole_ppg_config, cartpole_ppg_create_config from .cartpole_ppo_config import cartpole_ppo_config, cartpole_ppo_create_config from .cartpole_qrdqn_config import cartpole_qrdqn_config, cartpole_qrdqn_create_config @@ -19,4 +19,5 @@ from .cartpole_trex_dqn_config import cartpole_trex_dqn_config, cartpole_trex_dqn_create_config from .cartpole_trex_offppo_config import cartpole_trex_offppo_config, cartpole_trex_offppo_create_config from .cartpole_trex_onppo_config import cartpole_trex_ppo_onpolicy_config, cartpole_trex_ppo_onpolicy_create_config +from .cartpole_mdqn_config import cartpole_mdqn_config, cartpole_mdqn_create_config # from .cartpole_ppo_default_loader import cartpole_ppo_default_loader diff --git a/dizoo/classic_control/cartpole/config/cartpole_bc_config.py b/dizoo/classic_control/cartpole/config/cartpole_bc_config.py new file mode 100644 index 0000000000..b1975718f3 --- /dev/null +++ b/dizoo/classic_control/cartpole/config/cartpole_bc_config.py @@ -0,0 +1,44 @@ +from easydict import EasyDict + +cartpole_bc_config = dict( + exp_name='cartpole_bc_seed0', + env=dict( + collector_env_num=8, + evaluator_env_num=5, + n_evaluator_episode=5, + stop_value=195, + ), + policy=dict( + cuda=False, + continuous=False, + model=dict( + obs_shape=4, + action_shape=2, + encoder_hidden_size_list=[64, 64, 128], + ), + learn=dict( + batch_size=64, + learning_rate=0.01, + learner=dict(hook=dict(save_ckpt_after_iter=1000)), + train_epoch=20, + ), + eval=dict(evaluator=dict(eval_freq=40, )) + ), +) +cartpole_bc_config = EasyDict(cartpole_bc_config) +main_config = cartpole_bc_config +cartpole_bc_create_config = dict( + env=dict( + type='cartpole', + import_names=['dizoo.classic_control.cartpole.envs.cartpole_env'], + ), + env_manager=dict(type='base'), + policy=dict(type='bc'), +) +cartpole_bc_create_config = EasyDict(cartpole_bc_create_config) +create_config = cartpole_bc_create_config + +if __name__ == "__main__": + # Note: Users need to generate expert data, and save the data to ``expert_data_path`` + from ding.entry import serial_pipeline_bc + serial_pipeline_bc([main_config, create_config], seed=0, data_path=expert_data_path) diff --git a/dizoo/classic_control/cartpole/config/cartpole_c51_config.py b/dizoo/classic_control/cartpole/config/cartpole_c51_config.py index dfb4733a1c..a9f8557292 100644 --- a/dizoo/classic_control/cartpole/config/cartpole_c51_config.py +++ b/dizoo/classic_control/cartpole/config/cartpole_c51_config.py @@ -31,6 +31,7 @@ n_sample=80, unroll_len=1, ), + eval=dict(evaluator=dict(eval_freq=40, )), other=dict( eps=dict( type='exp', diff --git a/dizoo/classic_control/cartpole/config/cartpole_cql_config.py b/dizoo/classic_control/cartpole/config/cartpole_cql_config.py index ec804a6339..0b1932e5ad 100644 --- a/dizoo/classic_control/cartpole/config/cartpole_cql_config.py +++ b/dizoo/classic_control/cartpole/config/cartpole_cql_config.py @@ -42,7 +42,7 @@ import_names=['dizoo.classic_control.cartpole.envs.cartpole_env'], ), env_manager=dict(type='base'), - policy=dict(type='cql_discrete'), + policy=dict(type='discrete_cql'), ) cartpole_discrete_cql_create_config = EasyDict(cartpole_discrete_cql_create_config) create_config = cartpole_discrete_cql_create_config diff --git a/dizoo/classic_control/cartpole/config/cartpole_dqn_config.py b/dizoo/classic_control/cartpole/config/cartpole_dqn_config.py index ca34d7d879..1e2cb85433 100644 --- a/dizoo/classic_control/cartpole/config/cartpole_dqn_config.py +++ b/dizoo/classic_control/cartpole/config/cartpole_dqn_config.py @@ -17,10 +17,12 @@ action_shape=2, encoder_hidden_size_list=[128, 128, 64], dueling=True, + # dropout=0.1, ), nstep=1, discount_factor=0.97, learn=dict( + update_per_collect=5, batch_size=64, learning_rate=0.001, ), diff --git a/dizoo/classic_control/cartpole/config/cartpole_dqn_ddp_config.py b/dizoo/classic_control/cartpole/config/cartpole_dqn_ddp_config.py new file mode 100644 index 0000000000..a80662941a --- /dev/null +++ b/dizoo/classic_control/cartpole/config/cartpole_dqn_ddp_config.py @@ -0,0 +1,65 @@ +from easydict import EasyDict + +cartpole_dqn_config = dict( + exp_name='cartpole_dqn_seed0', + env=dict( + collector_env_num=8, + evaluator_env_num=5, + n_evaluator_episode=5, + stop_value=195, + replay_path='cartpole_dqn_seed0/video', + ), + policy=dict( + multi_gpu=True, + cuda=True, + model=dict( + obs_shape=4, + action_shape=2, + encoder_hidden_size_list=[128, 128, 64], + dueling=True, + # dropout=0.1, + ), + nstep=1, + discount_factor=0.97, + learn=dict( + update_per_collect=5, + batch_size=64, + learning_rate=0.001, + ), + collect=dict(n_sample=8), + eval=dict(evaluator=dict(eval_freq=40, )), + other=dict( + eps=dict( + type='exp', + start=0.95, + end=0.1, + decay=10000, + ), + replay_buffer=dict(replay_buffer_size=20000, ), + ), + ), +) +cartpole_dqn_config = EasyDict(cartpole_dqn_config) +main_config = cartpole_dqn_config +cartpole_dqn_create_config = dict( + env=dict( + type='cartpole', + import_names=['dizoo.classic_control.cartpole.envs.cartpole_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='dqn'), +) +cartpole_dqn_create_config = EasyDict(cartpole_dqn_create_config) +create_config = cartpole_dqn_create_config + +if __name__ == "__main__": + """ + Overview: + This script should be executed with GPUs. + Run the following command to launch the script: + python -m torch.distributed.launch --nproc_per_node=2 --master_port=29501 ./dizoo/classic_control/cartpole/config/cartpole_dqn_ddp_config.py + """ + from ding.utils import DDPContext + from ding.entry import serial_pipeline + with DDPContext(): + serial_pipeline((main_config, create_config), seed=0) diff --git a/dizoo/classic_control/cartpole/config/cartpole_drex_dqn_config.py b/dizoo/classic_control/cartpole/config/cartpole_drex_dqn_config.py index 91005cceaf..c898528a39 100644 --- a/dizoo/classic_control/cartpole/config/cartpole_drex_dqn_config.py +++ b/dizoo/classic_control/cartpole/config/cartpole_drex_dqn_config.py @@ -79,11 +79,7 @@ ), env_manager=dict(type='subprocess'), policy=dict(type='dqn'), + collector=dict(type='episode'), ) cartpole_drex_dqn_create_config = EasyDict(cartpole_drex_dqn_create_config) create_config = cartpole_drex_dqn_create_config - -if __name__ == "__main__": - # or you can enter `ding -m serial -c cartpole_drex_dqn_config.py -s 0` - from ding.entry import serial_pipeline - serial_pipeline((main_config, create_config), seed=0) diff --git a/dizoo/classic_control/cartpole/config/cartpole_dt_config.py b/dizoo/classic_control/cartpole/config/cartpole_dt_config.py new file mode 100644 index 0000000000..7eebd77428 --- /dev/null +++ b/dizoo/classic_control/cartpole/config/cartpole_dt_config.py @@ -0,0 +1,61 @@ +from easydict import EasyDict + +cartpole_discrete_dt_config = dict( + exp_name='cartpole_dt_seed0', + env=dict( + collector_env_num=8, + evaluator_env_num=5, + n_evaluator_episode=5, + stop_value=195, + ), + dataset=dict( + data_dir_prefix='./cartpole_qrdqn_generation_data_seed0/expert_demos.hdf5', + rtg_scale=None, + context_len=20, + env_type='classic', + ), + policy=dict( + cuda=False, + rtg_target=10, + evaluator_env_num=5, + clip_grad_norm_p=1.0, + state_mean=1, + state_std=0, + model=dict( + state_dim=4, + act_dim=2, + n_blocks=6, + h_dim=128, + context_len=20, + n_heads=8, + drop_p=0.1, + continuous=False, + ), + max_timestep=1000, + discount_factor=0.97, + nstep=3, + batch_size=64, + learning_rate=0.001, + target_update_freq=100, + kappa=1.0, + min_q_weight=4.0, + collect=dict( + data_type='hdf5', + data_path='./cartpole_qrdqn_generation_data_seed0/expert_demos.hdf5', + ), + eval=dict(evaluator=dict(eval_freq=100, )), + ), +) +cartpole_discrete_dt_config = EasyDict(cartpole_discrete_dt_config) +main_config = cartpole_discrete_dt_config +cartpole_discrete_dt_create_config = dict( + env=dict( + type='cartpole', + import_names=['dizoo.classic_control.cartpole.envs.cartpole_env'], + ), + env_manager=dict(type='base'), + policy=dict(type='dt'), +) +cartpole_discrete_dt_create_config = EasyDict(cartpole_discrete_dt_create_config) +create_config = cartpole_discrete_dt_create_config +# You can run this config with the entry file like `ding/example/dt.py` diff --git a/dizoo/classic_control/cartpole/config/cartpole_fqf_config.py b/dizoo/classic_control/cartpole/config/cartpole_fqf_config.py index ca65670d81..3c1103ba75 100644 --- a/dizoo/classic_control/cartpole/config/cartpole_fqf_config.py +++ b/dizoo/classic_control/cartpole/config/cartpole_fqf_config.py @@ -60,4 +60,4 @@ if __name__ == '__main__': # or you can enter `ding -m serial -c cartpole_fqf_config.py -s 0` from ding.entry import serial_pipeline - serial_pipeline((main_config, create_config), seed=0) \ No newline at end of file + serial_pipeline((main_config, create_config), seed=0) diff --git a/dizoo/classic_control/cartpole/config/cartpole_impala_config.py b/dizoo/classic_control/cartpole/config/cartpole_impala_config.py index 32b7c4d2ca..c78d9392af 100644 --- a/dizoo/classic_control/cartpole/config/cartpole_impala_config.py +++ b/dizoo/classic_control/cartpole/config/cartpole_impala_config.py @@ -15,6 +15,8 @@ action_shape=2, encoder_hidden_size_list=[64, 64], ), + # (int) the trajectory length to calculate v-trace target + unroll_len=8, learn=dict( # (int) collect n_sample data, train model update_per_collect times # here we follow ppo serial pipeline @@ -30,8 +32,6 @@ discount_factor=0.9, # (float) additional discounting parameter lambda_=0.95, - # (int) the trajectory length to calculate v-trace target - unroll_len=32, # (float) clip ratio of importance weights rho_clip_ratio=1.0, # (float) clip ratio of importance weights @@ -42,8 +42,6 @@ collect=dict( # (int) collect n_sample data, train model n_iteration times n_sample=16, - # (int) the trajectory length to calculate v-trace target - unroll_len=32, # (float) discount factor for future reward, defaults int [0, 1] discount_factor=0.9, gae_lambda=0.95, diff --git a/dizoo/classic_control/cartpole/config/cartpole_iqn_config.py b/dizoo/classic_control/cartpole/config/cartpole_iqn_config.py index d869005b65..d6fca73e93 100644 --- a/dizoo/classic_control/cartpole/config/cartpole_iqn_config.py +++ b/dizoo/classic_control/cartpole/config/cartpole_iqn_config.py @@ -29,6 +29,7 @@ n_sample=80, unroll_len=1, ), + eval=dict(evaluator=dict(eval_freq=40, )), other=dict( eps=dict( type='exp', diff --git a/dizoo/classic_control/cartpole/config/cartpole_mdqn_config.py b/dizoo/classic_control/cartpole/config/cartpole_mdqn_config.py new file mode 100644 index 0000000000..b72a6375a1 --- /dev/null +++ b/dizoo/classic_control/cartpole/config/cartpole_mdqn_config.py @@ -0,0 +1,58 @@ +from easydict import EasyDict + +cartpole_mdqn_config = dict( + exp_name='cartpole_mdqn_seed0', + env=dict( + collector_env_num=8, + evaluator_env_num=5, + n_evaluator_episode=5, + stop_value=195, + ), + policy=dict( + cuda=False, + model=dict( + obs_shape=4, + action_shape=2, + encoder_hidden_size_list=[128, 128, 64], + dueling=True, + ), + nstep=1, + discount_factor=0.97, + entropy_tau=0.03, + m_alpha=0.9, + learn=dict( + update_per_collect=5, + batch_size=64, + learning_rate=0.001, + ), + collect=dict(n_sample=8), + eval=dict(evaluator=dict(eval_freq=40, )), + other=dict( + eps=dict( + type='exp', + start=0.95, + end=0.1, + decay=10000, + ), + replay_buffer=dict(replay_buffer_size=20000, ), + ), + ), +) +cartpole_mdqn_config = EasyDict(cartpole_mdqn_config) +main_config = cartpole_mdqn_config +cartpole_mdqn_create_config = dict( + env=dict( + type='cartpole', + import_names=['dizoo.classic_control.cartpole.envs.cartpole_env'], + ), + env_manager=dict(type='base'), + policy=dict(type='mdqn'), + replay_buffer=dict(type='deque', import_names=['ding.data.buffer.deque_buffer_wrapper']), +) +cartpole_mdqn_create_config = EasyDict(cartpole_mdqn_create_config) +create_config = cartpole_mdqn_create_config + +if __name__ == "__main__": + # or you can enter `ding -m serial -c cartpole_mdqn_config.py -s 0` + from ding.entry import serial_pipeline + serial_pipeline((main_config, create_config), seed=0, dynamic_seed=False) diff --git a/dizoo/classic_control/cartpole/config/cartpole_pg_config.py b/dizoo/classic_control/cartpole/config/cartpole_pg_config.py new file mode 100644 index 0000000000..af3ee5ba04 --- /dev/null +++ b/dizoo/classic_control/cartpole/config/cartpole_pg_config.py @@ -0,0 +1,43 @@ +from easydict import EasyDict + +cartpole_pg_config = dict( + exp_name='cartpole_pg_seed0', + env=dict( + collector_env_num=8, + evaluator_env_num=5, + n_evaluator_episode=5, + stop_value=195, + ), + policy=dict( + cuda=False, + model=dict( + obs_shape=4, + action_shape=2, + ), + learn=dict( + batch_size=64, + learning_rate=0.001, + entropy_weight=0.001, + ), + collect=dict(n_episode=80, unroll_len=1, discount_factor=0.9), + eval=dict(evaluator=dict(eval_freq=100, ), ), + ), +) +cartpole_pg_config = EasyDict(cartpole_pg_config) +main_config = cartpole_pg_config +cartpole_pg_create_config = dict( + env=dict( + type='cartpole', + import_names=['dizoo.classic_control.cartpole.envs.cartpole_env'], + ), + env_manager=dict(type='base'), + policy=dict(type='pg'), + collector=dict(type='episode'), +) +cartpole_pg_create_config = EasyDict(cartpole_pg_create_config) +create_config = cartpole_pg_create_config + +if __name__ == "__main__": + # or you can enter `ding -m serial_onpolicy -c cartpole_pg_config.py -s 0` + from ding.entry import serial_pipeline_onpolicy + serial_pipeline_onpolicy((main_config, create_config), seed=0) diff --git a/dizoo/classic_control/cartpole/config/cartpole_ppo_config.py b/dizoo/classic_control/cartpole/config/cartpole_ppo_config.py index e8a5108721..4c1333bbc8 100644 --- a/dizoo/classic_control/cartpole/config/cartpole_ppo_config.py +++ b/dizoo/classic_control/cartpole/config/cartpole_ppo_config.py @@ -26,7 +26,15 @@ value_weight=0.5, entropy_weight=0.01, clip_ratio=0.2, - learner=dict(hook=dict(save_ckpt_after_iter=100)), + # Path to the pretrained checkpoint (ckpt). + # If set to an empty string (''), no pretrained model will be loaded. + # To load a pretrained ckpt, specify the path like this: + # learner=dict(hook=dict(load_ckpt_before_run='/path/to/your/ckpt/iteration_100.pth.tar')), + + # If True, the environment step count (collector.envstep) and training iteration (train_iter) + # will be loaded from the pretrained checkpoint, allowing training to resume seamlessly + # from where the ckpt left off. + resume_training=False, ), collect=dict( n_sample=256, diff --git a/dizoo/classic_control/cartpole/config/cartpole_ppo_ddp_config.py b/dizoo/classic_control/cartpole/config/cartpole_ppo_ddp_config.py new file mode 100644 index 0000000000..101698ea91 --- /dev/null +++ b/dizoo/classic_control/cartpole/config/cartpole_ppo_ddp_config.py @@ -0,0 +1,64 @@ +from easydict import EasyDict + +cartpole_ppo_config = dict( + exp_name='cartpole_ppo_seed0', + env=dict( + collector_env_num=8, + evaluator_env_num=5, + n_evaluator_episode=5, + stop_value=195, + ), + policy=dict( + multi_gpu=True, + cuda=True, + action_space='discrete', + model=dict( + obs_shape=4, + action_shape=2, + action_space='discrete', + encoder_hidden_size_list=[64, 64, 128], + critic_head_hidden_size=128, + actor_head_hidden_size=128, + ), + learn=dict( + epoch_per_collect=2, + batch_size=64, + learning_rate=0.001, + value_weight=0.5, + entropy_weight=0.01, + clip_ratio=0.2, + learner=dict(hook=dict(save_ckpt_after_iter=100)), + ), + collect=dict( + n_sample=256, + unroll_len=1, + discount_factor=0.9, + gae_lambda=0.95, + ), + eval=dict(evaluator=dict(eval_freq=100, ), ), + ), +) +cartpole_ppo_config = EasyDict(cartpole_ppo_config) +main_config = cartpole_ppo_config +cartpole_ppo_create_config = dict( + env=dict( + type='cartpole', + import_names=['dizoo.classic_control.cartpole.envs.cartpole_env'], + ), + env_manager=dict(type='base'), + policy=dict(type='ppo'), +) +cartpole_ppo_create_config = EasyDict(cartpole_ppo_create_config) +create_config = cartpole_ppo_create_config + +if __name__ == "__main__": + """ + Overview: + This script should be executed with GPUs. + Run the following command to launch the script: + python -m torch.distributed.launch --nproc_per_node=2 --master_port=29501 ./dizoo/classic_control/cartpole/config/cartpole_ppo_ddp_config.py + """ + from ding.utils import DDPContext + from ding.entry import serial_pipeline_onpolicy + with DDPContext(): + serial_pipeline_onpolicy((main_config, create_config), seed=0) diff --git a/dizoo/classic_control/cartpole/config/cartpole_offppo_config.py b/dizoo/classic_control/cartpole/config/cartpole_ppo_offpolicy_config.py similarity index 74% rename from dizoo/classic_control/cartpole/config/cartpole_offppo_config.py rename to dizoo/classic_control/cartpole/config/cartpole_ppo_offpolicy_config.py index a720fec55d..f952ecadaa 100644 --- a/dizoo/classic_control/cartpole/config/cartpole_offppo_config.py +++ b/dizoo/classic_control/cartpole/config/cartpole_ppo_offpolicy_config.py @@ -1,7 +1,7 @@ from easydict import EasyDict -cartpole_offppo_config = dict( - exp_name='cartpole_offppo_seed0', +cartpole_ppo_offpolicy_config = dict( + exp_name='cartpole_ppo_offpolicy_seed0', env=dict( collector_env_num=8, evaluator_env_num=5, @@ -37,9 +37,9 @@ other=dict(replay_buffer=dict(replay_buffer_size=5000)) ), ) -cartpole_offppo_config = EasyDict(cartpole_offppo_config) -main_config = cartpole_offppo_config -cartpole_offppo_create_config = dict( +cartpole_ppo_offpolicy_config = EasyDict(cartpole_ppo_offpolicy_config) +main_config = cartpole_ppo_offpolicy_config +cartpole_ppo_offpolicy_create_config = dict( env=dict( type='cartpole', import_names=['dizoo.classic_control.cartpole.envs.cartpole_env'], @@ -47,10 +47,10 @@ env_manager=dict(type='base'), policy=dict(type='ppo_offpolicy'), ) -cartpole_offppo_create_config = EasyDict(cartpole_offppo_create_config) -create_config = cartpole_offppo_create_config +cartpole_ppo_offpolicy_create_config = EasyDict(cartpole_ppo_offpolicy_create_config) +create_config = cartpole_ppo_offpolicy_create_config if __name__ == "__main__": - # or you can enter `ding -m serial -c cartpole_offppo_config.py -s 0` + # or you can enter `ding -m serial -c cartpole_ppo_offpolicy_config.py -s 0` from ding.entry import serial_pipeline serial_pipeline((main_config, create_config), seed=0) diff --git a/dizoo/classic_control/cartpole/config/cartpole_ppopg_config.py b/dizoo/classic_control/cartpole/config/cartpole_ppopg_config.py index efe5ad66c4..623a3b5048 100644 --- a/dizoo/classic_control/cartpole/config/cartpole_ppopg_config.py +++ b/dizoo/classic_control/cartpole/config/cartpole_ppopg_config.py @@ -18,9 +18,9 @@ epoch_per_collect=1, batch_size=64, learning_rate=0.001, - entropy_weight=0.01, + entropy_weight=0.001, ), - collect=dict(n_episode=160, unroll_len=1, discount_factor=0.9, collector=dict(get_train_sample=True)), + collect=dict(n_episode=80, unroll_len=1, discount_factor=0.9, collector=dict(get_train_sample=True)), eval=dict(evaluator=dict(eval_freq=100, ), ), ), ) diff --git a/dizoo/classic_control/cartpole/config/cartpole_r2d2_gtrxl_config.py b/dizoo/classic_control/cartpole/config/cartpole_r2d2_gtrxl_config.py index 5a2cdac020..a165245fc1 100644 --- a/dizoo/classic_control/cartpole/config/cartpole_r2d2_gtrxl_config.py +++ b/dizoo/classic_control/cartpole/config/cartpole_r2d2_gtrxl_config.py @@ -18,11 +18,11 @@ obs_shape=4, action_shape=2, memory_len=5, # length of transformer memory (can be 0) - hidden_size=256, + hidden_size=64, gru_bias=2., att_layer_num=3, dropout=0., - att_head_num=8, + att_head_num=4, ), discount_factor=0.99, nstep=3, @@ -31,7 +31,7 @@ seq_len=8, # transformer input segment # training sequence: unroll_len - burnin_step - nstep learn=dict( - update_per_collect=8, + update_per_collect=16, batch_size=64, learning_rate=0.0005, target_update_freq=500, diff --git a/dizoo/classic_control/cartpole/config/cartpole_rnd_onppo_config.py b/dizoo/classic_control/cartpole/config/cartpole_rnd_onppo_config.py index 9516efd852..ed4e89560a 100644 --- a/dizoo/classic_control/cartpole/config/cartpole_rnd_onppo_config.py +++ b/dizoo/classic_control/cartpole/config/cartpole_rnd_onppo_config.py @@ -10,12 +10,7 @@ ), reward_model=dict( intrinsic_reward_type='add', - intrinsic_reward_weight=None, - # means the relative weight of RND intrinsic_reward. - # If intrinsic_reward_weight=None, we will automatically set it based on - # the absolute value of the difference between max and min extrinsic reward in the sampled mini-batch - # please refer to rnd_reward_model for details. - intrinsic_reward_rescale=0.001, + intrinsic_reward_weight=0.001, # means the rescale value of RND intrinsic_reward only used when intrinsic_reward_weight is None # please refer to rnd_reward_model for details. learning_rate=5e-4, diff --git a/dizoo/classic_control/cartpole/config/cartpole_sac_config.py b/dizoo/classic_control/cartpole/config/cartpole_sac_config.py index 736c7ee930..74c5281577 100644 --- a/dizoo/classic_control/cartpole/config/cartpole_sac_config.py +++ b/dizoo/classic_control/cartpole/config/cartpole_sac_config.py @@ -13,8 +13,7 @@ random_collect_size=0, multi_agent=False, model=dict( - agent_obs_shape=4, - global_obs_shape=4, + obs_shape=4, action_shape=2, twin_critic=True, actor_head_hidden_size=64, @@ -62,7 +61,7 @@ import_names=['dizoo.classic_control.cartpole.envs.cartpole_env'], ), env_manager=dict(type='base'), - policy=dict(type='sac_discrete'), + policy=dict(type='discrete_sac'), ) cartpole_sac_create_config = EasyDict(cartpole_sac_create_config) create_config = cartpole_sac_create_config diff --git a/dizoo/classic_control/cartpole/config/cartpole_sqil_config.py b/dizoo/classic_control/cartpole/config/cartpole_sqil_config.py index ffff606c71..7e4553f729 100644 --- a/dizoo/classic_control/cartpole/config/cartpole_sqil_config.py +++ b/dizoo/classic_control/cartpole/config/cartpole_sqil_config.py @@ -24,7 +24,7 @@ # Users should add their own model path here. Model path should lead to a model. # Absolute path is recommended. # In DI-engine, it is ``exp_name/ckpt/ckpt_best.pth.tar``. - model_path='model_path_placeholder' + model_path='cartpole_dqn_seed0/ckpt/eval.pth.tar' ), # note: this is the times after which you learns to evaluate eval=dict(evaluator=dict(eval_freq=50, )), diff --git a/dizoo/classic_control/cartpole/config/cartpole_trex_dqn_config.py b/dizoo/classic_control/cartpole/config/cartpole_trex_dqn_config.py index 724b9f64fa..306cadd6f2 100644 --- a/dizoo/classic_control/cartpole/config/cartpole_trex_dqn_config.py +++ b/dizoo/classic_control/cartpole/config/cartpole_trex_dqn_config.py @@ -7,7 +7,6 @@ evaluator_env_num=5, n_evaluator_episode=5, stop_value=195, - replay_path='cartpole_dqn/video', ), reward_model=dict( type='trex', @@ -20,12 +19,9 @@ update_per_collect=1, num_trajs=6, num_snippets=6000, - expert_model_path='abs model path', - reward_model_path='abs data path + ./cartpole.params', - data_path='abs data path', + expert_model_path='cartpole_dqn_seed0', # expert model experiment directory path ), policy=dict( - load_path='', cuda=False, model=dict( obs_shape=4, @@ -36,6 +32,7 @@ nstep=1, discount_factor=0.97, learn=dict( + update_per_collect=5, batch_size=64, learning_rate=0.001, ), diff --git a/dizoo/classic_control/cartpole/config/cartpole_trex_offppo_config.py b/dizoo/classic_control/cartpole/config/cartpole_trex_offppo_config.py index c817e783fa..b58535f900 100644 --- a/dizoo/classic_control/cartpole/config/cartpole_trex_offppo_config.py +++ b/dizoo/classic_control/cartpole/config/cartpole_trex_offppo_config.py @@ -13,8 +13,8 @@ min_snippet_length=5, max_snippet_length=100, checkpoint_min=0, - checkpoint_max=1000, - checkpoint_step=1000, + checkpoint_max=100, + checkpoint_step=100, learning_rate=1e-5, update_per_collect=1, expert_model_path='abs model path', diff --git a/dizoo/classic_control/cartpole/dt_data/cartpole_collect_data.py b/dizoo/classic_control/cartpole/dt_data/cartpole_collect_data.py deleted file mode 100644 index efe272dbb0..0000000000 --- a/dizoo/classic_control/cartpole/dt_data/cartpole_collect_data.py +++ /dev/null @@ -1,33 +0,0 @@ -from dizoo.classic_control.cartpole.dt_data.collect_qrdqn_data_config import main_config, create_config -from ding.entry import collect_episodic_demo_data, eval -import torch -import copy - - -def eval_ckpt(args): - config = copy.deepcopy([main_config, create_config]) - # eval(config, seed=args.seed, load_path=main_config.policy.learn.learner.hook.load_ckpt_before_run, replay_path='./replay') - eval(config, seed=args.seed, load_path=main_config.policy.learn.learner.hook.load_ckpt_before_run) - - -def generate(args): - config = copy.deepcopy([main_config, create_config]) - state_dict = torch.load(main_config.policy.learn.learner.load_path, map_location='cpu') - collect_episodic_demo_data( - config, - collect_count=main_config.policy.other.replay_buffer.replay_buffer_size, - seed=args.seed, - expert_data_path=main_config.policy.collect.save_path, - state_dict=state_dict - ) - - -if __name__ == "__main__": - import argparse - - parser = argparse.ArgumentParser() - parser.add_argument('--seed', '-s', type=int, default=0) - args = parser.parse_args() - - eval_ckpt(args) - generate(args) diff --git a/dizoo/classic_control/cartpole/dt_data/cartpole_show_data.py b/dizoo/classic_control/cartpole/dt_data/cartpole_show_data.py deleted file mode 100644 index 9634168fcb..0000000000 --- a/dizoo/classic_control/cartpole/dt_data/cartpole_show_data.py +++ /dev/null @@ -1,50 +0,0 @@ -from dizoo.classic_control.cartpole.dt_data_cartpole.collect_demo_data_config import main_config, create_config - -from ding.entry import serial_pipeline_offline -import os -import torch -from torch.utils.data import DataLoader -from ding.config import read_config, compile_config -from ding.utils.data import create_dataset - - -def train(args): - config = [main_config, create_config] - input_cfg = config - if isinstance(input_cfg, str): - cfg, create_cfg = read_config(input_cfg) - else: - cfg, create_cfg = input_cfg - create_cfg.policy.type = create_cfg.policy.type + '_command' - cfg = compile_config(cfg, seed=args.seed, auto=True, create_cfg=create_cfg) - - # Dataset - dataset = create_dataset(cfg) - print(dataset.__len__()) - - # print(dataset.__getitem__(0)) - print(dataset.__getitem__(0)[0]['action']) - - # episode_action = [] - # for i in range(dataset.__getitem__(0).__len__()): # length of the firse collected episode - # episode_action.append(dataset.__getitem__(0)[i]['action']) - - # stacked action of the first collected episode - episode_action = torch.stack( - [dataset.__getitem__(0)[i]['action'] for i in range(dataset.__getitem__(0).__len__())], axis=0 - ) - - # dataloader = DataLoader(dataset, cfg.policy.learn.batch_size, shuffle=True, collate_fn=lambda x: x) - # for i, train_data in enumerate(dataloader): - # print(i, train_data) - # serial_pipeline_offline(config, seed=args.seed) - - -if __name__ == "__main__": - import argparse - - parser = argparse.ArgumentParser() - parser.add_argument('--seed', '-s', type=int, default=0) - args = parser.parse_args() - - train(args) diff --git a/dizoo/classic_control/cartpole/dt_data/collect_qrdqn_data_config.py b/dizoo/classic_control/cartpole/dt_data/collect_qrdqn_data_config.py deleted file mode 100644 index 124dc6d763..0000000000 --- a/dizoo/classic_control/cartpole/dt_data/collect_qrdqn_data_config.py +++ /dev/null @@ -1,136 +0,0 @@ -from easydict import EasyDict - -main_config = dict( - exp_name='cartpole', - env=dict( - manager=dict( - episode_num=float('inf'), - max_retry=1, - retry_type='reset', - auto_reset=True, - step_timeout=None, - reset_timeout=None, - retry_waiting_time=0.1, - cfg_type='BaseEnvManagerDict', - type='base', - ), - # collector_env_num=8, - # evaluator_env_num=5, - # n_evaluator_episode=5, - collector_env_num=1, - evaluator_env_num=1, - n_evaluator_episode=5, - stop_value=195, - ), - policy=dict( - model=dict( - obs_shape=4, - action_shape=2, - encoder_hidden_size_list=[128, 128, 64], - num_quantiles=64, - ), - learn=dict( - learner=dict( - train_iterations=1000000000, - dataloader=dict(num_workers=0, ), - log_policy=True, - hook=dict( - # load_ckpt_before_run='./cartpole/ckpt/ckpt_best.pth.tar', - load_ckpt_before_run= - '/home/puyuan/DI-engine/dizoo/classic_control/cartpole/dt_data_cartpole/ckpt/ckpt_best.pth.tar', - log_show_after_iter=100, - save_ckpt_after_iter=10000, - save_ckpt_after_run=False, - ), - cfg_type='BaseLearnerDict', - # load_path='./cartpole/ckpt/ckpt_best.pth.tar', - load_path= - '/home/puyuan/DI-engine/dizoo/classic_control/cartpole/dt_data_cartpole/ckpt/ckpt_best.pth.tar', - ), - multi_gpu=False, - update_per_collect=3, - batch_size=64, - learning_rate=0.001, - target_update_freq=100, - ignore_done=False, - kappa=1.0, - ), - collect=dict( - collector=dict( - deepcopy_obs=False, - transform_obs=False, - collect_print_freq=100, - cfg_type='SampleSerialCollectorDict', - type='sample', - ), - unroll_len=1, - n_sample=80, - # save - # data_type='hdf5', - data_type='naive', # TODO - save_path= - '/home/puyuan/DI-engine/dizoo/classic_control/cartpole/dt_data_cartpole/cartpole/qrdqn_data_10eps.pkl', - # load - data_path='/home/puyuan/DI-engine/dizoo/classic_control/cartpole/dt_data_cartpole/cartpole/qrdqn_data.pkl', - ), - eval=dict( - evaluator=dict( - eval_freq=1000, - cfg_type='InteractionSerialEvaluatorDict', - stop_value=195, - n_episode=5, - ), - ), - other=dict( - replay_buffer=dict( - type='advanced', - # replay_buffer_size=100000, - replay_buffer_size=10, #TODO(pu) - max_use=float('inf'), - max_staleness=float('inf'), - alpha=0.6, - beta=0.4, - anneal_step=100000, - enable_track_used_data=False, - deepcopy=False, - thruput_controller=dict( - push_sample_rate_limit=dict( - max=float('inf'), - min=0, - ), - window_seconds=30, - sample_min_limit_ratio=1, - ), - monitor=dict( - sampled_data_attr=dict( - average_range=5, - print_freq=200, - ), - periodic_thruput=dict(seconds=60, ), - ), - cfg_type='AdvancedReplayBufferDict', - ), - ), - cuda=False, - on_policy=False, - priority=True, - discount_factor=0.97, - nstep=3, - cfg_type='QRDQNCommandModePolicyDict', - ), -) -main_config = EasyDict(main_config) -main_config = main_config -create_config = dict( - env=dict( - type='cartpole', - import_names=['dizoo.classic_control.cartpole.envs.cartpole_env'], - ), - env_manager=dict( - cfg_type='BaseEnvManagerDict', - type='base', - ), - policy=dict(type='qrdqn'), -) -create_config = EasyDict(create_config) -create_config = create_config diff --git a/dizoo/classic_control/cartpole/entry/cartpole_c51_deploy.py b/dizoo/classic_control/cartpole/entry/cartpole_c51_deploy.py new file mode 100644 index 0000000000..bf1025dd04 --- /dev/null +++ b/dizoo/classic_control/cartpole/entry/cartpole_c51_deploy.py @@ -0,0 +1,33 @@ +import gym +import torch +from easydict import EasyDict +from ding.config import compile_config +from ding.envs import DingEnvWrapper +from ding.policy import C51Policy, single_env_forward_wrapper +from ding.model import C51DQN +from dizoo.classic_control.cartpole.config.cartpole_c51_config import cartpole_c51_config, cartpole_c51_create_config + + +def main(main_config: EasyDict, create_config: EasyDict, ckpt_path: str): + main_config.exp_name = 'cartpole_c51_deploy' + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + env = DingEnvWrapper(gym.make('CartPole-v0'), EasyDict(env_wrapper='default')) + model = C51DQN(**cfg.policy.model) + state_dict = torch.load(ckpt_path, map_location='cpu') + model.load_state_dict(state_dict['model']) + policy = C51Policy(cfg.policy, model=model).eval_mode + forward_fn = single_env_forward_wrapper(policy.forward) + + obs = env.reset() + returns = 0. + while True: + action = forward_fn(obs) + obs, rew, done, info = env.step(action) + returns += rew + if done: + break + print(f'Deploy is finished, final epsiode return is: {returns}') + + +if __name__ == "__main__": + main(cartpole_c51_config, cartpole_c51_create_config, 'cartpole_c51_seed0/ckpt/ckpt_best.pth.tar') diff --git a/dizoo/classic_control/cartpole/entry/cartpole_cql_main.py b/dizoo/classic_control/cartpole/entry/cartpole_cql_main.py index 5610a879be..311a2c7c11 100644 --- a/dizoo/classic_control/cartpole/entry/cartpole_cql_main.py +++ b/dizoo/classic_control/cartpole/entry/cartpole_cql_main.py @@ -1,7 +1,6 @@ import torch from copy import deepcopy -from dizoo.classic_control.cartpole.config.cartpole_qrdqn_generation_data_config import main_config, create_config from ding.entry import serial_pipeline_offline, collect_demo_data, eval, serial_pipeline @@ -15,23 +14,23 @@ def train_cql(args): def eval_ckpt(args): + from dizoo.classic_control.cartpole.config.cartpole_qrdqn_config import main_config, create_config + main_config, create_config = deepcopy(main_config), deepcopy(create_config) main_config.exp_name = 'cartpole' - main_config.policy.learn.learner.load_path = './cartpole/ckpt/ckpt_best.pth.tar' - main_config.policy.learn.learner.hook.load_ckpt_before_run = './cartpole/ckpt/ckpt_best.pth.tar' config = deepcopy([main_config, create_config]) - eval(config, seed=args.seed, load_path=main_config.policy.learn.learner.hook.load_ckpt_before_run) + eval(config, seed=args.seed, load_path='./cartpole/ckpt/ckpt_best.pth.tar') def generate(args): + from dizoo.classic_control.cartpole.config.cartpole_qrdqn_generation_data_config import main_config, create_config main_config.exp_name = 'cartpole' - main_config.policy.learn.learner.load_path = './cartpole/ckpt/ckpt_best.pth.tar' main_config.policy.collect.save_path = './cartpole/expert.pkl' main_config.policy.collect.data_type = 'hdf5' config = deepcopy([main_config, create_config]) - state_dict = torch.load(main_config.policy.learn.learner.load_path, map_location='cpu') + state_dict = torch.load('./cartpole/ckpt/ckpt_best.pth.tar', map_location='cpu') collect_demo_data( config, - collect_count=main_config.policy.other.replay_buffer.replay_buffer_size, + collect_count=10000, seed=args.seed, expert_data_path=main_config.policy.collect.save_path, state_dict=state_dict @@ -40,6 +39,7 @@ def generate(args): def train_expert(args): from dizoo.classic_control.cartpole.config.cartpole_qrdqn_config import main_config, create_config + main_config, create_config = deepcopy(main_config), deepcopy(create_config) main_config.exp_name = 'cartpole' config = deepcopy([main_config, create_config]) serial_pipeline(config, seed=args.seed) diff --git a/dizoo/classic_control/cartpole/entry/cartpole_dqn_pwil_main.py b/dizoo/classic_control/cartpole/entry/cartpole_dqn_pwil_main.py new file mode 100644 index 0000000000..b88262a2e4 --- /dev/null +++ b/dizoo/classic_control/cartpole/entry/cartpole_dqn_pwil_main.py @@ -0,0 +1,44 @@ +from easydict import EasyDict +from copy import deepcopy + +from dizoo.classic_control.cartpole.config.cartpole_ppo_offpolicy_config import cartpole_ppo_offpolicy_config, cartpole_ppo_offpolicy_create_config # noqa +from dizoo.classic_control.cartpole.config.cartpole_dqn_config import cartpole_dqn_config, cartpole_dqn_create_config +from ding.entry import serial_pipeline, collect_demo_data, serial_pipeline_reward_model_offpolicy + + +def cartpole_dqn_pwil_main(): + reward_model_config = { + 'type': 'pwil', + 's_size': 4, + 'a_size': 2, + 'sample_size': 500, + } + + # train a expert policy (PPO offpolicy) + reward_model_config = EasyDict(reward_model_config) + config = deepcopy(cartpole_ppo_offpolicy_config), deepcopy(cartpole_ppo_offpolicy_create_config) + expert_policy = serial_pipeline(config, seed=0) + + # (optional) collect expert demo data + collect_count = 10000 + expert_data_path = 'expert_data.pkl' + state_dict = expert_policy.collect_mode.state_dict() + config = deepcopy(cartpole_ppo_offpolicy_config), deepcopy(cartpole_ppo_offpolicy_create_config) + collect_demo_data( + config, seed=0, state_dict=state_dict, expert_data_path=expert_data_path, collect_count=collect_count + ) + + # irl + rl training + cp_cartpole_dqn_config = deepcopy(cartpole_dqn_config) + cp_cartpole_dqn_create_config = deepcopy(cartpole_dqn_create_config) + cp_cartpole_dqn_create_config.reward_model = dict(type=reward_model_config.type) + reward_model_config['expert_data_path'] = expert_data_path + cp_cartpole_dqn_config.exp_name = 'cartpole_dqn_pwil' + cp_cartpole_dqn_config.reward_model = reward_model_config + cp_cartpole_dqn_config.policy.collect.n_sample = 128 + + serial_pipeline_reward_model_offpolicy((cp_cartpole_dqn_config, cp_cartpole_dqn_create_config), seed=0) + + +if __name__ == "__main__": + cartpole_dqn_pwil_main() diff --git a/dizoo/classic_control/cartpole/envs/cartpole_env.py b/dizoo/classic_control/cartpole/envs/cartpole_env.py index 2912692c24..7cd36702b0 100644 --- a/dizoo/classic_control/cartpole/envs/cartpole_env.py +++ b/dizoo/classic_control/cartpole/envs/cartpole_env.py @@ -48,7 +48,7 @@ def reset(self) -> np.ndarray: self._env.seed(self._seed) self._action_space.seed(self._seed) self._observation_space = self._env.observation_space - self._final_eval_reward = 0 + self._eval_episode_return = 0 obs = self._env.reset() obs = to_ndarray(obs) return obs @@ -67,9 +67,9 @@ def step(self, action: Union[int, np.ndarray]) -> BaseEnvTimestep: if isinstance(action, np.ndarray) and action.shape == (1, ): action = action.squeeze() # 0-dim array obs, rew, done, info = self._env.step(action) - self._final_eval_reward += rew + self._eval_episode_return += rew if done: - info['final_eval_reward'] = self._final_eval_reward + info['eval_episode_return'] = self._eval_episode_return obs = to_ndarray(obs) rew = to_ndarray([rew]).astype(np.float32) # wrapped to be transfered to a array with shape (1,) return BaseEnvTimestep(obs, rew, done, info) diff --git a/dizoo/classic_control/mountain_car/config/mtcar_rainbow_config.py b/dizoo/classic_control/mountain_car/config/mtcar_rainbow_config.py index c6c4fb4db0..b293d44494 100644 --- a/dizoo/classic_control/mountain_car/config/mtcar_rainbow_config.py +++ b/dizoo/classic_control/mountain_car/config/mtcar_rainbow_config.py @@ -1,58 +1,63 @@ from easydict import EasyDict # DI-Engine uses EasyDict for configuration, by convention -mtcar_rainbow_config = EasyDict(dict( - exp_name='mtcar_rainbow_seed0', - env=dict( - collector_env_num=8, - evaluator_env_num=5, - n_evaluator_episode=5, - stop_value=195, - ), - policy=dict( - cuda=False, - priority=True, - discount_factor=0.97, - nstep=3, - model=dict( - obs_shape=2, - action_shape=3, - encoder_hidden_size_list=[128, 128, 64], +mtcar_rainbow_config = EasyDict( + dict( + exp_name='mtcar_rainbow_seed0', + env=dict( + collector_env_num=8, + evaluator_env_num=5, + n_evaluator_episode=5, + stop_value=195, ), - learn=dict( - update_per_collect=3, - batch_size=64, - learning_rate=0.001, - target_update_freq=100, + policy=dict( + cuda=False, + priority=True, + discount_factor=0.97, + nstep=3, + model=dict( + obs_shape=2, + action_shape=3, + encoder_hidden_size_list=[128, 128, 64], + ), + learn=dict( + update_per_collect=3, + batch_size=64, + learning_rate=0.001, + target_update_freq=100, + ), + collect=dict( + n_sample=80, + unroll_len=1, + ), + other=dict( + eps=dict( + type='exp', + start=0.95, + end=0.1, + decay=10000, + ), + replay_buffer=dict(replay_buffer_size=20000, ) + ), ), - collect=dict( - n_sample=80, - unroll_len=1, - ), - other=dict( - eps=dict( - type='exp', - start=0.95, - end=0.1, - decay=10000, - ), replay_buffer=dict(replay_buffer_size=20000, ) - ), - ), -)) + ) +) main_config = mtcar_rainbow_config -mtcar_rainbow_create_config = EasyDict(dict( - env=dict( - type='mountain_car', - import_names=['dizoo.classic_control.mountain_car.envs.mtcar_env'], - ), - env_manager=dict(type='base'), - policy=dict(type='rainbow'), -)) +mtcar_rainbow_create_config = EasyDict( + dict( + env=dict( + type='mountain_car', + import_names=['dizoo.classic_control.mountain_car.envs.mtcar_env'], + ), + env_manager=dict(type='base'), + policy=dict(type='rainbow'), + ) +) create_config = mtcar_rainbow_create_config if __name__ == "__main__": from ding.entry import serial_pipeline - serial_pipeline((main_config, create_config), seed=0) \ No newline at end of file + serial_pipeline((main_config, create_config), seed=0) diff --git a/dizoo/classic_control/mountain_car/envs/__init__.py b/dizoo/classic_control/mountain_car/envs/__init__.py index 19f7eaf1cc..9e8ca86d5f 100644 --- a/dizoo/classic_control/mountain_car/envs/__init__.py +++ b/dizoo/classic_control/mountain_car/envs/__init__.py @@ -1 +1 @@ -from .mtcar_env import MountainCarEnv \ No newline at end of file +from .mtcar_env import MountainCarEnv diff --git a/dizoo/classic_control/mountain_car/envs/mtcar_env.py b/dizoo/classic_control/mountain_car/envs/mtcar_env.py index 9f8d6a7233..515b424765 100644 --- a/dizoo/classic_control/mountain_car/envs/mtcar_env.py +++ b/dizoo/classic_control/mountain_car/envs/mtcar_env.py @@ -70,7 +70,7 @@ def reset(self) -> np.ndarray: # Init final reward : cumulative sum of the real rewards obtained by a whole episode, # used to evaluate the agent Performance on this environment, not used for training. - self._final_eval_reward = 0. + self._eval_episode_return = 0. return obs def step(self, action: np.ndarray) -> BaseEnvTimestep: @@ -85,11 +85,11 @@ def step(self, action: np.ndarray) -> BaseEnvTimestep: obs, rew, done, info = self._env.step(action) # Cummulate reward - self._final_eval_reward += rew + self._eval_episode_return += rew # Save final cummulative reward when done. if done: - info['final_eval_reward'] = self._final_eval_reward + info['eval_episode_return'] = self._eval_episode_return # Making sure we conform to di-engine conventions obs = to_ndarray(obs) diff --git a/dizoo/classic_control/pendulum/config/__init__.py b/dizoo/classic_control/pendulum/config/__init__.py index e7c2988f06..642386797f 100644 --- a/dizoo/classic_control/pendulum/config/__init__.py +++ b/dizoo/classic_control/pendulum/config/__init__.py @@ -4,4 +4,4 @@ from .pendulum_d4pg_config import pendulum_d4pg_config, pendulum_d4pg_create_config from .pendulum_ppo_config import pendulum_ppo_config, pendulum_ppo_create_config from .pendulum_td3_bc_config import pendulum_td3_bc_config, pendulum_td3_bc_create_config -from .pendulum_td3_data_generation_config import pendulum_td3_generation_config, pendulum_td3_generation_create_config \ No newline at end of file +from .pendulum_td3_data_generation_config import pendulum_td3_generation_config, pendulum_td3_generation_create_config diff --git a/dizoo/classic_control/pendulum/config/pendulum_a2c_config.py b/dizoo/classic_control/pendulum/config/pendulum_a2c_config.py new file mode 100644 index 0000000000..d66e8b1a17 --- /dev/null +++ b/dizoo/classic_control/pendulum/config/pendulum_a2c_config.py @@ -0,0 +1,51 @@ +from easydict import EasyDict + +pendulum_a2c_config = dict( + exp_name='pendulum_a2c_seed0', + env=dict( + collector_env_num=10, + evaluator_env_num=5, + act_scale=True, + n_evaluator_episode=5, + stop_value=-200, + ), + policy=dict( + cuda=False, + action_space='continuous', + model=dict( + action_space='continuous', + obs_shape=3, + action_shape=1, + ), + learn=dict( + epoch_per_collect=10, + batch_size=32, + learning_rate=3e-5, + value_weight=0.5, + entropy_weight=0.0, + ), + collect=dict( + n_sample=200, + unroll_len=1, + discount_factor=0.99, + ), + eval=dict(evaluator=dict(eval_freq=200, )) + ), +) +pendulum_a2c_config = EasyDict(pendulum_a2c_config) +main_config = pendulum_a2c_config +pendulum_a2c_create_config = dict( + env=dict( + type='pendulum', + import_names=['dizoo.classic_control.pendulum.envs.pendulum_env'], + ), + env_manager=dict(type='base'), + policy=dict(type='a2c'), +) +pendulum_a2c_create_config = EasyDict(pendulum_a2c_create_config) +create_config = pendulum_a2c_create_config + +if __name__ == "__main__": + # or you can enter `ding -m serial_onpolicy -c pendulum_a2c_config.py -s 0` + from ding.entry import serial_pipeline_onpolicy + serial_pipeline_onpolicy([main_config, create_config], seed=0) diff --git a/dizoo/classic_control/pendulum/config/pendulum_bdq_config.py b/dizoo/classic_control/pendulum/config/pendulum_bdq_config.py new file mode 100644 index 0000000000..59bffb05b9 --- /dev/null +++ b/dizoo/classic_control/pendulum/config/pendulum_bdq_config.py @@ -0,0 +1,62 @@ +from easydict import EasyDict +import sys +sys.path.insert(0, "/mnt/lustre/chenyun/bdq_implement1/DI-engine") +pendulum_bdq_config = dict( + exp_name='pendulum_bdq_seed0', + env=dict( + collector_env_num=10, + evaluator_env_num=5, + # (bool) Scale output action into legal range. + act_scale=True, + n_evaluator_episode=5, + stop_value=-250, + continuous=False, + # The path to save the game replay + # replay_path='./pendulum_bdq_seed0/video', + ), + policy=dict( + cuda=False, + load_path='pendulum_bdq_seed0/ckpt/ckpt_best.pth.tar', # necessary for eval + model=dict( + obs_shape=3, + num_branches=1, + action_bins_per_branch=11, + encoder_hidden_size_list=[128, 128, 64], + ), + nstep=1, + discount_factor=0.97, + learn=dict( + batch_size=64, + learning_rate=0.001, + ), + collect=dict(n_sample=8), + eval=dict(evaluator=dict(eval_freq=40, )), + other=dict( + eps=dict( + type='exp', + start=0.95, + end=0.1, + decay=10000, + ), + replay_buffer=dict(replay_buffer_size=20000, ), + ), + ), +) +pendulum_bdq_config = EasyDict(pendulum_bdq_config) +main_config = pendulum_bdq_config +pendulum_bdq_create_config = dict( + env=dict( + type='pendulum', + import_names=['dizoo.classic_control.pendulum.envs.pendulum_env'], + ), + env_manager=dict(type='base'), + policy=dict(type='bdq'), + replay_buffer=dict(type='deque', import_names=['ding.data.buffer.deque_buffer_wrapper']), +) +pendulum_bdq_create_config = EasyDict(pendulum_bdq_create_config) +create_config = pendulum_bdq_create_config + +if __name__ == "__main__": + # or you can enter `ding -m serial -c pendulum_bdq_config.py -s 0` + from ding.entry import serial_pipeline + serial_pipeline((main_config, create_config), seed=0) diff --git a/dizoo/classic_control/pendulum/config/pendulum_ibc_config.py b/dizoo/classic_control/pendulum/config/pendulum_ibc_config.py index 247fdad045..7c56f283fe 100644 --- a/dizoo/classic_control/pendulum/config/pendulum_ibc_config.py +++ b/dizoo/classic_control/pendulum/config/pendulum_ibc_config.py @@ -13,16 +13,15 @@ ), policy=dict( cuda=cuda, - model=dict( - obs_shape=3, - action_shape=1, - stochastic_optim=dict(type='mcmc', cuda=cuda,) - ), + model=dict(obs_shape=3, action_shape=1, stochastic_optim=dict( + type='mcmc', + cuda=cuda, + )), learn=dict( multi_gpu=multi_gpu, train_epoch=15, batch_size=256, - optim=dict(learning_rate=1e-5,), + optim=dict(learning_rate=1e-5, ), learner=dict(hook=dict(log_show_after_iter=1000)), ), collect=dict( @@ -30,7 +29,7 @@ data_path='./pendulum_sac_data_generation/expert_demos.hdf5', collector_logit=False, ), - eval=dict(evaluator=dict(eval_freq=-1,)), + eval=dict(evaluator=dict(eval_freq=-1, )), ), ) pendulum_ibc_config = EasyDict(main_config) diff --git a/dizoo/classic_control/pendulum/config/pendulum_pg_config.py b/dizoo/classic_control/pendulum/config/pendulum_pg_config.py new file mode 100644 index 0000000000..b512548398 --- /dev/null +++ b/dizoo/classic_control/pendulum/config/pendulum_pg_config.py @@ -0,0 +1,50 @@ +from easydict import EasyDict + +pendulum_pg_config = dict( + exp_name='pendulum_pg_seed0', + env=dict( + collector_env_num=8, + evaluator_env_num=5, + act_scale=True, + n_evaluator_episode=5, + stop_value=-200, + ), + policy=dict( + cuda=False, + action_space='continuous', + model=dict( + action_space='continuous', + obs_shape=3, + action_shape=1, + ), + learn=dict( + batch_size=400, + learning_rate=0.001, + entropy_weight=0.001, + ), + collect=dict( + n_episode=2, + unroll_len=1, + discount_factor=0.99, + ), + eval=dict(evaluator=dict(eval_freq=200, )) + ), +) +pendulum_pg_config = EasyDict(pendulum_pg_config) +main_config = pendulum_pg_config +pendulum_pg_create_config = dict( + env=dict( + type='pendulum', + import_names=['dizoo.classic_control.pendulum.envs.pendulum_env'], + ), + env_manager=dict(type='base'), + policy=dict(type='pg'), + collector=dict(type='episode'), +) +pendulum_pg_create_config = EasyDict(pendulum_pg_create_config) +create_config = pendulum_pg_create_config + +if __name__ == "__main__": + # or you can enter `ding -m serial_onpolicy -c pendulum_pg_config.py -s 0` + from ding.entry import serial_pipeline_onpolicy + serial_pipeline_onpolicy([main_config, create_config], seed=0) diff --git a/dizoo/classic_control/pendulum/config/pendulum_ppo_config.py b/dizoo/classic_control/pendulum/config/pendulum_ppo_config.py index 2431a5aa2b..151455aec1 100644 --- a/dizoo/classic_control/pendulum/config/pendulum_ppo_config.py +++ b/dizoo/classic_control/pendulum/config/pendulum_ppo_config.py @@ -1,4 +1,5 @@ from easydict import EasyDict +import torch.nn as nn pendulum_ppo_config = dict( exp_name='pendulum_ppo_seed0', @@ -20,25 +21,26 @@ action_space='continuous', actor_head_layer_num=0, critic_head_layer_num=0, - sigma_type='conditioned', + sigma_type='independent', + activation=nn.Tanh(), bound_type='tanh', ), learn=dict( epoch_per_collect=10, batch_size=32, - learning_rate=3e-5, + learning_rate=1e-3, value_weight=0.5, entropy_weight=0.0, clip_ratio=0.2, - adv_norm=False, + adv_norm=True, value_norm=True, ignore_done=True, ), collect=dict( - n_sample=200, + n_sample=5000, unroll_len=1, discount_factor=0.9, - gae_lambda=1., + gae_lambda=.95, ), eval=dict(evaluator=dict(eval_freq=200, )) ), diff --git a/dizoo/classic_control/pendulum/config/pendulum_sqil_sac_config.py b/dizoo/classic_control/pendulum/config/pendulum_sqil_sac_config.py index cecd4a057c..f80b85235f 100644 --- a/dizoo/classic_control/pendulum/config/pendulum_sqil_sac_config.py +++ b/dizoo/classic_control/pendulum/config/pendulum_sqil_sac_config.py @@ -36,12 +36,13 @@ target_theta=0.005, discount_factor=0.99, auto_alpha=True, - value_network=False, ), collect=dict( n_sample=10, - model_path='model_path_placeholder', - unroll_len=1, + # Users should add their own model path here. Model path should lead to a model. + # Absolute path is recommended. + # In DI-engine, it is ``exp_name/ckpt/ckpt_best.pth.tar``. + model_path='pendulum_sac_seed0/ckpt/eval.pth.tar', ), eval=dict(evaluator=dict(eval_freq=100, )), other=dict(replay_buffer=dict(replay_buffer_size=100000, ), ), diff --git a/dizoo/classic_control/pendulum/config/pendulum_td3_bc_config.py b/dizoo/classic_control/pendulum/config/pendulum_td3_bc_config.py index 82a44f034e..8583fc6ada 100644 --- a/dizoo/classic_control/pendulum/config/pendulum_td3_bc_config.py +++ b/dizoo/classic_control/pendulum/config/pendulum_td3_bc_config.py @@ -6,7 +6,7 @@ collector_env_num=8, evaluator_env_num=5, norm_obs=dict( - use_norm=True, + use_norm=True, offline_stats=dict(use_offline_stats=True, ), ), # (bool) Scale output action into legal range. diff --git a/dizoo/classic_control/pendulum/entry/pendulum_ddpg_main.py b/dizoo/classic_control/pendulum/entry/pendulum_ddpg_main.py index 3ad6304f50..0153071ed7 100644 --- a/dizoo/classic_control/pendulum/entry/pendulum_ddpg_main.py +++ b/dizoo/classic_control/pendulum/entry/pendulum_ddpg_main.py @@ -6,7 +6,7 @@ from ding.worker import BaseLearner, SampleSerialCollector, InteractionSerialEvaluator, AdvancedReplayBuffer from ding.envs import BaseEnvManager, DingEnvWrapper from ding.policy import DDPGPolicy -from ding.model import QAC +from ding.model import ContinuousQAC from ding.utils import set_pkg_seed from dizoo.classic_control.pendulum.envs import PendulumEnv from dizoo.classic_control.pendulum.config.pendulum_ddpg_config import pendulum_ddpg_config @@ -50,7 +50,7 @@ def main(cfg, seed=0): set_pkg_seed(seed, use_cuda=cfg.policy.cuda) # Set up RL Policy - model = QAC(**cfg.policy.model) + model = ContinuousQAC(**cfg.policy.model) policy = DDPGPolicy(cfg.policy, model=model) # Set up collection, training and evaluation utilities diff --git a/dizoo/classic_control/pendulum/entry/pendulum_dqn_eval.py b/dizoo/classic_control/pendulum/entry/pendulum_dqn_eval.py index a5a7b9ab32..fb80ad42ad 100644 --- a/dizoo/classic_control/pendulum/entry/pendulum_dqn_eval.py +++ b/dizoo/classic_control/pendulum/entry/pendulum_dqn_eval.py @@ -15,8 +15,9 @@ from ding.rl_utils import get_epsilon_greedy_fn from dizoo.classic_control.pendulum.config.pendulum_dqn_config import main_config, create_config + def main(rl_cfg, seed=0): - main_cfg, create_cfg =rl_cfg + main_cfg, create_cfg = rl_cfg cfg = compile_config( main_cfg, BaseEnvManager, @@ -56,4 +57,4 @@ def main(rl_cfg, seed=0): if __name__ == "__main__": - main(rl_cfg=(main_config, create_config),seed=0) + main(rl_cfg=(main_config, create_config), seed=0) diff --git a/dizoo/classic_control/pendulum/entry/pendulum_td3_main.py b/dizoo/classic_control/pendulum/entry/pendulum_td3_main.py index 354c13da0c..36bfce14f3 100644 --- a/dizoo/classic_control/pendulum/entry/pendulum_td3_main.py +++ b/dizoo/classic_control/pendulum/entry/pendulum_td3_main.py @@ -7,7 +7,7 @@ from ding.worker import BaseLearner, SampleSerialCollector, InteractionSerialEvaluator, AdvancedReplayBuffer from ding.envs import BaseEnvManager, DingEnvWrapper from ding.policy import DDPGPolicy -from ding.model import QAC +from ding.model import ContinuousQAC from ding.utils import set_pkg_seed from dizoo.classic_control.pendulum.envs import PendulumEnv from dizoo.classic_control.pendulum.config.pendulum_td3_config import pendulum_td3_config @@ -40,7 +40,7 @@ def main(cfg, seed=0): set_pkg_seed(seed, use_cuda=cfg.policy.cuda) # Set up RL Policy - model = QAC(**cfg.policy.model) + model = ContinuousQAC(**cfg.policy.model) policy = DDPGPolicy(cfg.policy, model=model) # lr_scheduler demo lr_scheduler = LambdaLR( diff --git a/dizoo/classic_control/pendulum/envs/pendulum_env.py b/dizoo/classic_control/pendulum/envs/pendulum_env.py index b848144383..f3265cbaae 100644 --- a/dizoo/classic_control/pendulum/envs/pendulum_env.py +++ b/dizoo/classic_control/pendulum/envs/pendulum_env.py @@ -54,7 +54,7 @@ def reset(self) -> np.ndarray: self._action_space.seed(self._seed) obs = self._env.reset() obs = to_ndarray(obs).astype(np.float32) - self._final_eval_reward = 0. + self._eval_episode_return = 0. return obs def close(self) -> None: @@ -76,12 +76,12 @@ def step(self, action: np.ndarray) -> BaseEnvTimestep: if self._act_scale: action = affine_transform(action, min_val=self._env.action_space.low, max_val=self._env.action_space.high) obs, rew, done, info = self._env.step(action) - self._final_eval_reward += rew + self._eval_episode_return += rew obs = to_ndarray(obs).astype(np.float32) # wrapped to be transfered to a array with shape (1,) rew = to_ndarray([rew]).astype(np.float32) if done: - info['final_eval_reward'] = self._final_eval_reward + info['eval_episode_return'] = self._eval_episode_return return BaseEnvTimestep(obs, rew, done, info) def enable_save_replay(self, replay_path: Optional[str] = None) -> None: diff --git a/dizoo/cliffwalking/__init__.py b/dizoo/cliffwalking/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/dizoo/cliffwalking/cliff_walking.gif b/dizoo/cliffwalking/cliff_walking.gif new file mode 100644 index 0000000000..29679ad468 Binary files /dev/null and b/dizoo/cliffwalking/cliff_walking.gif differ diff --git a/dizoo/cliffwalking/config/cliffwalking_dqn_config.py b/dizoo/cliffwalking/config/cliffwalking_dqn_config.py new file mode 100644 index 0000000000..c852858ab7 --- /dev/null +++ b/dizoo/cliffwalking/config/cliffwalking_dqn_config.py @@ -0,0 +1,60 @@ +from easydict import EasyDict + +cliffwalking_dqn_config = dict( + exp_name='cliffwalking_dqn_seed0', + env=dict( + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=-13, # the optimal value of cliffwalking env + max_episode_steps=300, + ), + policy=dict( + cuda=True, + load_path="./cliffwalking_dqn_seed0/ckpt/ckpt_best.pth.tar", + model=dict( + obs_shape=48, + action_shape=4, + encoder_hidden_size_list=[512, 64], + ), + discount_factor=0.99, + nstep=1, + learn=dict( + update_per_collect=10, + batch_size=128, + learning_rate=0.001, + target_update_freq=100, + ), + collect=dict( + n_sample=64, + unroll_len=1, + ), + other=dict( + eps=dict( + type='linear', + start=0.95, + end=0.25, + decay=50000, + ), + replay_buffer=dict(replay_buffer_size=100000, ) + ), + ), +) +cliffwalking_dqn_config = EasyDict(cliffwalking_dqn_config) +main_config = cliffwalking_dqn_config + +cliffwalking_dqn_create_config = dict( + env=dict( + type='cliffwalking', + import_names=['dizoo.cliffwalking.envs.cliffwalking_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='dqn'), +) +cliffwalking_dqn_create_config = EasyDict(cliffwalking_dqn_create_config) +create_config = cliffwalking_dqn_create_config + +if __name__ == "__main__": + # or you can enter `ding -m serial -c cliffwalking_dqn_config.py -s 0` + from ding.entry import serial_pipeline + serial_pipeline([main_config, create_config], seed=0) diff --git a/dizoo/cliffwalking/entry/cliffwalking_dqn_deploy.py b/dizoo/cliffwalking/entry/cliffwalking_dqn_deploy.py new file mode 100644 index 0000000000..02fe49a0a7 --- /dev/null +++ b/dizoo/cliffwalking/entry/cliffwalking_dqn_deploy.py @@ -0,0 +1,39 @@ +import gym +import torch +from easydict import EasyDict + +from ding.config import compile_config +from ding.envs import DingEnvWrapper +from ding.model import DQN +from ding.policy import DQNPolicy, single_env_forward_wrapper +from dizoo.cliffwalking.config.cliffwalking_dqn_config import create_config, main_config +from dizoo.cliffwalking.envs.cliffwalking_env import CliffWalkingEnv + + +def main(main_config: EasyDict, create_config: EasyDict, ckpt_path: str): + main_config.exp_name = f'cliffwalking_dqn_seed0_deploy' + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + env = CliffWalkingEnv(cfg.env) + env.enable_save_replay(replay_path=f'./{main_config.exp_name}/video') + model = DQN(**cfg.policy.model) + state_dict = torch.load(ckpt_path, map_location='cpu') + model.load_state_dict(state_dict['model']) + policy = DQNPolicy(cfg.policy, model=model).eval_mode + forward_fn = single_env_forward_wrapper(policy.forward) + obs = env.reset() + returns = 0. + while True: + action = forward_fn(obs) + obs, rew, done, info = env.step(action) + returns += rew + if done: + break + print(f'Deploy is finished, final epsiode return is: {returns}') + + +if __name__ == "__main__": + main( + main_config=main_config, + create_config=create_config, + ckpt_path=f'./cliffwalking_dqn_seed0/ckpt/ckpt_best.pth.tar' + ) diff --git a/dizoo/cliffwalking/entry/cliffwalking_dqn_main.py b/dizoo/cliffwalking/entry/cliffwalking_dqn_main.py new file mode 100644 index 0000000000..8a8c082155 --- /dev/null +++ b/dizoo/cliffwalking/entry/cliffwalking_dqn_main.py @@ -0,0 +1,50 @@ +import gym +from ditk import logging + +from ding.config import compile_config +from ding.data import DequeBuffer +from ding.envs import BaseEnvManagerV2, DingEnvWrapper +from ding.framework import ding_init, task +from ding.framework.context import OnlineRLContext +from ding.framework.middleware import CkptSaver, OffPolicyLearner, StepCollector, data_pusher, eps_greedy_handler, \ + interaction_evaluator, online_logger +from ding.model import DQN +from ding.policy import DQNPolicy +from ding.utils import set_pkg_seed +from dizoo.cliffwalking.config.cliffwalking_dqn_config import create_config, main_config +from dizoo.cliffwalking.envs.cliffwalking_env import CliffWalkingEnv + + +def main(): + filename = '{}/log.txt'.format(main_config.exp_name) + logging.getLogger(with_files=[filename]).setLevel(logging.INFO) + + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + ding_init(cfg) + + collector_env = BaseEnvManagerV2( + env_fn=[lambda: CliffWalkingEnv(cfg.env) for _ in range(cfg.env.collector_env_num)], cfg=cfg.env.manager + ) + evaluator_env = BaseEnvManagerV2( + env_fn=[lambda: CliffWalkingEnv(cfg.env) for _ in range(cfg.env.collector_env_num)], cfg=cfg.env.manager + ) + + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + + model = DQN(**cfg.policy.model) + buffer = DequeBuffer(size=cfg.policy.other.replay_buffer.replay_buffer_size) + policy = DQNPolicy(cfg.policy, model=model) + + with task.start(async_mode=False, ctx=OnlineRLContext()): + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) + task.use(eps_greedy_handler(cfg)) + task.use(StepCollector(cfg, policy.collect_mode, collector_env)) + task.use(data_pusher(cfg, buffer)) + task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer)) + task.use(online_logger(train_show_freq=10)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=1000)) + task.run() + + +if __name__ == '__main__': + main() diff --git a/dizoo/cliffwalking/envs/__init__.py b/dizoo/cliffwalking/envs/__init__.py new file mode 100644 index 0000000000..d90c30675f --- /dev/null +++ b/dizoo/cliffwalking/envs/__init__.py @@ -0,0 +1 @@ +from .cliffwalking_env import CliffWalkingEnv diff --git a/dizoo/cliffwalking/envs/cliffwalking_env.py b/dizoo/cliffwalking/envs/cliffwalking_env.py new file mode 100644 index 0000000000..af5d094aa6 --- /dev/null +++ b/dizoo/cliffwalking/envs/cliffwalking_env.py @@ -0,0 +1,111 @@ +import copy +from typing import List, Union, Optional + +import gym +import numpy as np +from easydict import EasyDict + +from ding.envs.env.base_env import BaseEnv, BaseEnvTimestep +from ding.torch_utils import to_ndarray +from ding.utils import ENV_REGISTRY + + +@ENV_REGISTRY.register('cliffwalking') +class CliffWalkingEnv(BaseEnv): + + def __init__(self, cfg: dict) -> None: + self._cfg = EasyDict( + env_id='CliffWalking', + render_mode='rgb_array', + max_episode_steps=300, # default max trajectory length to truncate possible infinite attempts + ) + self._cfg.update(cfg) + self._init_flag = False + self._replay_path = None + self._observation_space = gym.spaces.Box(low=0, high=1, shape=(48, ), dtype=np.float32) + self._env = gym.make( + "CliffWalking", render_mode=self._cfg.render_mode, max_episode_steps=self._cfg.max_episode_steps + ) + self._action_space = self._env.action_space + self._reward_space = gym.spaces.Box( + low=self._env.reward_range[0], high=self._env.reward_range[1], shape=(1, ), dtype=np.float32 + ) + + def reset(self) -> np.ndarray: + if not self._init_flag: + self._env = gym.make( + "CliffWalking", render_mode=self._cfg.render_mode, max_episode_steps=self._cfg.max_episode_steps + ) + self._init_flag = True + if hasattr(self, '_seed') and hasattr(self, '_dynamic_seed') and self._dynamic_seed: + dy_seed = self._seed + 100 * np.random.randint(1, 1000) + self._env.seed(dy_seed) + elif hasattr(self, '_seed'): + self._env.seed(self._seed) + if self._replay_path is not None: + self._env = gym.wrappers.RecordVideo( + self._env, + video_folder=self._replay_path, + episode_trigger=lambda episode_id: True, + name_prefix='cliffwalking-{}'.format(id(self)) + ) + obs = self._env.reset() + obs_encode = self._encode_obs(obs) + self._eval_episode_return = 0. + return obs_encode + + def close(self) -> None: + try: + self._env.close() + del self._env + except: + pass + + def seed(self, seed: int, dynamic_seed: bool = True) -> None: + self._seed = seed + self._dynamic_seed = dynamic_seed + np.random.seed(seed) + + def step(self, action: Union[int, np.ndarray]) -> BaseEnvTimestep: + if isinstance(action, np.ndarray): + if action.shape == (1, ): + action = action.squeeze() # 0-dim array + action = action.item() + obs, reward, done, info = self._env.step(action) + obs_encode = self._encode_obs(obs) + self._eval_episode_return += reward + reward = to_ndarray([reward], dtype=np.float32) + if done: + info['eval_episode_return'] = self._eval_episode_return + return BaseEnvTimestep(obs_encode, reward, done, info) + + def enable_save_replay(self, replay_path: Optional[str] = None) -> None: + if replay_path is None: + replay_path = './video' + self._replay_path = replay_path + + def random_action(self) -> np.ndarray: + random_action = self.action_space.sample() + if isinstance(random_action, int): + random_action = to_ndarray([random_action], dtype=np.int64) + return random_action + + def _encode_obs(self, obs) -> np.ndarray: + onehot = np.zeros(48, dtype=np.float32) + onehot[int(obs)] = 1 + return onehot + + @property + def observation_space(self) -> gym.spaces.Space: + return self._observation_space + + @property + def action_space(self) -> gym.spaces.Space: + return self._action_space + + @property + def reward_space(self) -> gym.spaces.Space: + return self._reward_space + + def __repr__(self) -> str: + return "DI-engine CliffWalking Env" diff --git a/dizoo/cliffwalking/envs/test_cliffwalking_env.py b/dizoo/cliffwalking/envs/test_cliffwalking_env.py new file mode 100644 index 0000000000..b378d1a1a8 --- /dev/null +++ b/dizoo/cliffwalking/envs/test_cliffwalking_env.py @@ -0,0 +1,35 @@ +import numpy as np +import pytest +from dizoo.cliffwalking.envs import CliffWalkingEnv + + +@pytest.mark.envtest +class TestCliffWalkingEnv: + + def test_naive(self): + env = CliffWalkingEnv({}) + env.seed(314, dynamic_seed=False) + assert env._seed == 314 + obs = env.reset() + assert obs.shape == (48, ) + for _ in range(5): + env.reset() + np.random.seed(314) + print('=' * 60) + for i in range(10): + # Both ``env.random_action()``, and utilizing ``np.random`` as well as action space, + # can generate legal random action. + if i < 5: + random_action = np.array([env.action_space.sample()]) + else: + random_action = env.random_action() + timestep = env.step(random_action) + print(timestep) + assert isinstance(timestep.obs, np.ndarray) + assert isinstance(timestep.done, bool) + assert timestep.obs.shape == (48, ) + assert timestep.reward.shape == (1, ) + assert timestep.reward >= env.reward_space.low + assert timestep.reward <= env.reward_space.high + print(env.observation_space, env.action_space, env.reward_space) + env.close() diff --git a/dizoo/common/policy/md_dqn.py b/dizoo/common/policy/md_dqn.py index a36992cb15..01cc9e13e9 100644 --- a/dizoo/common/policy/md_dqn.py +++ b/dizoo/common/policy/md_dqn.py @@ -87,7 +87,7 @@ def _forward_learn(self, data: dict) -> Dict[str, Any]: # ==================== self._optimizer.zero_grad() loss.backward() - if self._cfg.learn.multi_gpu: + if self._cfg.multi_gpu: self.sync_gradients(self._learn_model) self._optimizer.step() diff --git a/dizoo/competitive_rl/envs/competitive_rl_env.py b/dizoo/competitive_rl/envs/competitive_rl_env.py index f6ae7e5265..4db8964c9d 100644 --- a/dizoo/competitive_rl/envs/competitive_rl_env.py +++ b/dizoo/competitive_rl/envs/competitive_rl_env.py @@ -92,9 +92,9 @@ def reset(self) -> np.ndarray: obs = self.process_obs(obs) # process if self._builtin_wrap: - self._final_eval_reward = np.array([0.]) + self._eval_episode_return = np.array([0.]) else: - self._final_eval_reward = np.array([0., 0.]) + self._eval_episode_return = np.array([0., 0.]) return obs def close(self) -> None: @@ -116,13 +116,13 @@ def step(self, action: Union[np.ndarray, list]) -> BaseEnvTimestep: if not isinstance(rew, tuple): rew = [rew] rew = np.array(rew) - self._final_eval_reward += rew + self._eval_episode_return += rew obs = to_ndarray(obs) obs = self.process_obs(obs) # process if done: - info['final_eval_reward'] = self._final_eval_reward + info['eval_episode_return'] = self._eval_episode_return return BaseEnvTimestep(obs, rew, done, info) diff --git a/dizoo/d4rl/config/__init__.py b/dizoo/d4rl/config/__init__.py index 450f069392..92bca79cc2 100644 --- a/dizoo/d4rl/config/__init__.py +++ b/dizoo/d4rl/config/__init__.py @@ -1,3 +1,3 @@ -from .hopper_cql_config import hopper_cql_config -from .hopper_expert_cql_config import hopper_expert_cql_config -from .hopper_medium_cql_config import hopper_medium_cql_config +# from .hopper_cql_config import hopper_cql_config +# from .hopper_expert_cql_config import hopper_expert_cql_config +# from .hopper_medium_cql_config import hopper_medium_cql_config diff --git a/dizoo/d4rl/config/antmaze_umaze_pd_config.py b/dizoo/d4rl/config/antmaze_umaze_pd_config.py new file mode 100755 index 0000000000..5f111f58fd --- /dev/null +++ b/dizoo/d4rl/config/antmaze_umaze_pd_config.py @@ -0,0 +1,99 @@ +from easydict import EasyDict + +main_config = dict( + exp_name="antmaze_umaze_pd_seed0", + env=dict( + env_id='antmaze-umaze-v0', + collector_env_num=1, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + returns_scale=1.0, + termination_penalty=-100, + max_path_length=1000, + use_padding=True, + include_returns=True, + normed=False, + stop_value=8000, + horizon=256, + obs_dim=29, + action_dim=8, + ), + policy=dict( + cuda=True, + model=dict( + diffuser_model='GaussianDiffusion', + diffuser_model_cfg=dict( + model='DiffusionUNet1d', + model_cfg=dict( + transition_dim=37, + dim=32, + dim_mults=[1, 2, 4, 8], + returns_condition=False, + kernel_size=5, + attention=False, + ), + horizon=256, + obs_dim=29, + action_dim=8, + n_timesteps=20, + predict_epsilon=False, + loss_discount=1, + action_weight=10, + ), + value_model='ValueDiffusion', + value_model_cfg=dict( + model='TemporalValue', + model_cfg=dict( + horizon=256, + transition_dim=37, + dim=32, + dim_mults=[1, 2, 4, 8], + kernel_size=5, + ), + horizon=256, + obs_dim=29, + action_dim=8, + n_timesteps=20, + predict_epsilon=True, + loss_discount=1, + ), + n_guide_steps=2, + scale=0.1, + t_stopgrad=2, + scale_grad_by_std=True, + ), + normalizer='GaussianNormalizer', + learn=dict( + data_path=None, + train_epoch=60000, + gradient_accumulate_every=2, + batch_size=32, + learning_rate=2e-4, + discount_factor=0.99, + plan_batch_size=64, + learner=dict(hook=dict(save_ckpt_after_iter=1000000000, )), + ), + collect=dict(data_type='diffuser_traj', ), + eval=dict( + evaluator=dict(eval_freq=500, ), + test_ret=0.9, + ), + other=dict(replay_buffer=dict(replay_buffer_size=2000000, ), ), + ), +) + +main_config = EasyDict(main_config) +main_config = main_config + +create_config = dict( + env=dict( + type='d4rl', + import_names=['dizoo.d4rl.envs.d4rl_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='pd', ), + replay_buffer=dict(type='naive', ), +) +create_config = EasyDict(create_config) +create_config = create_config diff --git a/dizoo/d4rl/config/halfcheetah_expert_cql_config.py b/dizoo/d4rl/config/halfcheetah_expert_cql_config.py old mode 100644 new mode 100755 index bb4b321d80..58ead98dc5 --- a/dizoo/d4rl/config/halfcheetah_expert_cql_config.py +++ b/dizoo/d4rl/config/halfcheetah_expert_cql_config.py @@ -5,7 +5,7 @@ main_config = dict( exp_name="halfcheetah_expert_cql_seed0", env=dict( - env_id='halfcheetah-expert-v0', + env_id='halfcheetah-expert-v2', collector_env_num=1, evaluator_env_num=8, use_act_scale=True, @@ -15,8 +15,8 @@ policy=dict( cuda=True, model=dict( - obs_shape=11, - action_shape=3, + obs_shape=17, + action_shape=6, ), learn=dict( data_path=None, diff --git a/dizoo/d4rl/config/halfcheetah_expert_dt_config.py b/dizoo/d4rl/config/halfcheetah_expert_dt_config.py index 6c8e219bc0..617d17bc73 100644 --- a/dizoo/d4rl/config/halfcheetah_expert_dt_config.py +++ b/dizoo/d4rl/config/halfcheetah_expert_dt_config.py @@ -2,37 +2,38 @@ from copy import deepcopy halfcheetah_dt_config = dict( - exp_name='halfcheetah_expert_dt_seed0', + exp_name='dt_log/d4rl/halfcheetah/halfcheetah_expert_dt_seed0', env=dict( env_id='HalfCheetah-v3', - norm_obs=dict(use_norm=False, ), - norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, use_act_scale=True, n_evaluator_episode=8, stop_value=6000, ), + dataset=dict( + env_type='mujoco', + rtg_scale=1000, + context_len=30, + data_dir_prefix='d4rl/halfcheetah_expert-v2.pkl', + ), policy=dict( - stop_value=6000, cuda=True, + stop_value=6000, + state_mean=None, + state_std=None, + evaluator_env_num=8, env_name='HalfCheetah-v3', rtg_target=6000, # max target return to go max_eval_ep_len=1000, # max lenght of one episode - num_eval_ep=10, # num of evaluation episode - batch_size=64, wt_decay=1e-4, warmup_steps=10000, - num_updates_per_iter=100, context_len=20, - n_blocks=3, - embed_dim=128, - n_heads=1, - dropout_p=0.1, - log_dir='/home/wangzilin/research/dt/DI-engine/dizoo/d4rl/dt_data/halfcheetah_expert_dt_log', + weight_decay=0.1, + clip_grad_norm_p=0.25, model=dict( - state_dim=17, - act_dim=6, + state_dim=11, + act_dim=3, n_blocks=3, h_dim=128, context_len=20, @@ -40,26 +41,13 @@ drop_p=0.1, continuous=True, ), - discount_factor=0.999, - nstep=3, - learn=dict( - dataset_path='/mnt/lustre/wangzilin/d4rl_data/halfcheetah-expert-v2.pkl', - learning_rate=0.0001, - target_update_freq=100, - kappa=1.0, - min_q_weight=4.0, - ), - collect=dict(unroll_len=1, ), - eval=dict(evaluator=dict(evalu_freq=100, ), ), - other=dict( - eps=dict( - type='exp', - start=0.95, - end=0.1, - decay=10000, - ), - replay_buffer=dict(replay_buffer_size=1000, ), + batch_size=64, + learning_rate=1e-4, + collect=dict( + data_type='d4rl_trajectory', + unroll_len=1, ), + eval=dict(evaluator=dict(eval_freq=100, ), ), ), ) diff --git a/dizoo/d4rl/config/halfcheetah_expert_td3bc_config.py b/dizoo/d4rl/config/halfcheetah_expert_td3bc_config.py old mode 100644 new mode 100755 index e798cf66e3..3cae080e92 --- a/dizoo/d4rl/config/halfcheetah_expert_td3bc_config.py +++ b/dizoo/d4rl/config/halfcheetah_expert_td3bc_config.py @@ -3,11 +3,11 @@ from easydict import EasyDict main_config = dict( - exp_name='walker2d_expert_td3-bc_seed0', + exp_name='halfcheetah_expert_td3-bc_seed0', env=dict( - env_id='walker2d-expert-v0', + env_id='halfcheetah-expert-v2', norm_obs=dict( - use_norm=True, + use_norm=True, offline_stats=dict(use_offline_stats=True, ), ), collector_env_num=1, @@ -17,9 +17,10 @@ stop_value=6000, ), policy=dict( + cuda=True, model=dict( - obs_shape=11, - action_shape=3, + obs_shape=17, + action_shape=6, ), learn=dict( train_epoch=30000, diff --git a/dizoo/d4rl/config/halfcheetah_medium_bcq_config.py b/dizoo/d4rl/config/halfcheetah_medium_bcq_config.py new file mode 100755 index 0000000000..c0199dcb09 --- /dev/null +++ b/dizoo/d4rl/config/halfcheetah_medium_bcq_config.py @@ -0,0 +1,55 @@ +from easydict import EasyDict + +main_config = dict( + exp_name="halfcheetah_medium_bcq_seed0", + env=dict( + env_id='halfcheetah-medium-v2', + collector_env_num=1, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + stop_value=7000, + ), + policy=dict( + cuda=True, + model=dict( + obs_shape=17, + action_shape=6, + actor_head_hidden_size=[400, 300], + critic_head_hidden_size=[400, 300], + phi=0.05, + ), + learn=dict( + data_path=None, + train_epoch=30000, + batch_size=100, + learning_rate_q=3e-3, + learning_rate_policy=3e-3, + learning_rate_alpha=3e-3, + lmbda=0.75, + learner=dict(hook=dict(save_ckpt_after_iter=1000000000, )), + ), + collect=dict(data_type='d4rl', ), + eval=dict(evaluator=dict(eval_freq=500, )), + other=dict(replay_buffer=dict(replay_buffer_size=2000000, ), ), + ), + seed=123, +) + +main_config = EasyDict(main_config) +main_config = main_config + +create_config = dict( + env=dict( + type='d4rl', + import_names=['dizoo.d4rl.envs.d4rl_env'], + ), + env_manager=dict(type='base'), + policy=dict( + type='bcq', + import_names=['ding.policy.bcq'], + ), + replay_buffer=dict(type='naive', ), +) +create_config = EasyDict(create_config) +create_config = create_config diff --git a/dizoo/d4rl/config/halfcheetah_medium_cql_config.py b/dizoo/d4rl/config/halfcheetah_medium_cql_config.py old mode 100644 new mode 100755 index a879599fcf..84a504cadb --- a/dizoo/d4rl/config/halfcheetah_medium_cql_config.py +++ b/dizoo/d4rl/config/halfcheetah_medium_cql_config.py @@ -5,7 +5,7 @@ main_config = dict( exp_name="halfcheetah_medium_cql_seed0", env=dict( - env_id='halfcheetah-medium-v0', + env_id='halfcheetah-medium-v2', collector_env_num=1, evaluator_env_num=8, use_act_scale=True, @@ -15,8 +15,8 @@ policy=dict( cuda=True, model=dict( - obs_shape=11, - action_shape=3, + obs_shape=17, + action_shape=6, ), learn=dict( data_path=None, diff --git a/dizoo/d4rl/config/halfcheetah_medium_dt_config.py b/dizoo/d4rl/config/halfcheetah_medium_dt_config.py index f61ca5e0c2..7521af6dd5 100644 --- a/dizoo/d4rl/config/halfcheetah_medium_dt_config.py +++ b/dizoo/d4rl/config/halfcheetah_medium_dt_config.py @@ -2,37 +2,38 @@ from copy import deepcopy halfcheetah_dt_config = dict( - exp_name='halfcheetah_medium_dt_seed0', + exp_name='dt_log/d4rl/halfcheetah/halfcheetah_medium_dt_seed0', env=dict( env_id='HalfCheetah-v3', - norm_obs=dict(use_norm=False, ), - norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, use_act_scale=True, n_evaluator_episode=8, stop_value=6000, ), + dataset=dict( + env_type='mujoco', + rtg_scale=1000, + context_len=30, + data_dir_prefix='d4rl/halfcheetah_medium-v2.pkl', + ), policy=dict( - stop_value=6000, cuda=True, + stop_value=6000, + state_mean=None, + state_std=None, + evaluator_env_num=8, env_name='HalfCheetah-v3', rtg_target=6000, # max target return to go max_eval_ep_len=1000, # max lenght of one episode - num_eval_ep=10, # num of evaluation episode - batch_size=64, wt_decay=1e-4, warmup_steps=10000, - num_updates_per_iter=100, context_len=20, - n_blocks=3, - embed_dim=128, - n_heads=1, - dropout_p=0.1, - log_dir='/home/wangzilin/research/dt/DI-engine/dizoo/d4rl/dt_data/halfcheetah_medium_dt_log', + weight_decay=0.1, + clip_grad_norm_p=0.25, model=dict( - state_dim=17, - act_dim=6, + state_dim=11, + act_dim=3, n_blocks=3, h_dim=128, context_len=20, @@ -40,26 +41,13 @@ drop_p=0.1, continuous=True, ), - discount_factor=0.999, - nstep=3, - learn=dict( - dataset_path='/mnt/lustre/wangzilin/d4rl_data/halfcheetah-medium-v2.pkl', - learning_rate=0.0001, - target_update_freq=100, - kappa=1.0, - min_q_weight=4.0, - ), - collect=dict(unroll_len=1, ), - eval=dict(evaluator=dict(evalu_freq=100, ), ), - other=dict( - eps=dict( - type='exp', - start=0.95, - end=0.1, - decay=10000, - ), - replay_buffer=dict(replay_buffer_size=1000, ), + batch_size=64, + learning_rate=1e-4, + collect=dict( + data_type='d4rl_trajectory', + unroll_len=1, ), + eval=dict(evaluator=dict(eval_freq=100, ), ), ), ) diff --git a/dizoo/d4rl/config/halfcheetah_medium_edac_config.py b/dizoo/d4rl/config/halfcheetah_medium_edac_config.py new file mode 100755 index 0000000000..66ea8039dc --- /dev/null +++ b/dizoo/d4rl/config/halfcheetah_medium_edac_config.py @@ -0,0 +1,58 @@ +from easydict import EasyDict + +main_config = dict( + exp_name="halfcheetah_medium_edac_seed0", + env=dict( + env_id='halfcheetah-medium-v2', + collector_env_num=1, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + stop_value=7600, + ), + policy=dict( + cuda=True, + model=dict( + obs_shape=17, + action_shape=6, + ensemble_num=10, + actor_head_hidden_size=256, + actor_head_layer_num=3, + critic_head_hidden_size=256, + critic_head_layer_num=3, + ), + learn=dict( + data_path=None, + train_epoch=30000, + batch_size=256, + learning_rate_q=3e-4, + learning_rate_policy=3e-4, + learning_rate_alpha=3e-4, + alpha=1, + auto_alpha=True, + eta=1.0, + with_q_entropy=False, + learner=dict(hook=dict(save_ckpt_after_iter=100000, )), + ), + collect=dict(data_type='d4rl', ), + eval=dict(evaluator=dict(eval_freq=500, )), + ), + seed=0, +) + +main_config = EasyDict(main_config) +main_config = main_config + +create_config = dict( + env=dict( + type='d4rl', + import_names=['dizoo.d4rl.envs.d4rl_env'], + ), + env_manager=dict(type='base'), + policy=dict( + type='edac', + import_names=['ding.policy.edac'], + ), +) +create_config = EasyDict(create_config) +create_config = create_config diff --git a/dizoo/d4rl/config/halfcheetah_medium_expert_bcq_config.py b/dizoo/d4rl/config/halfcheetah_medium_expert_bcq_config.py new file mode 100755 index 0000000000..6c3ac39c18 --- /dev/null +++ b/dizoo/d4rl/config/halfcheetah_medium_expert_bcq_config.py @@ -0,0 +1,55 @@ +from easydict import EasyDict + +main_config = dict( + exp_name="halfcheetah_medium_expert_bcq_seed0", + env=dict( + env_id='halfcheetah-medium-expert-v2', + collector_env_num=1, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + stop_value=12000, + ), + policy=dict( + cuda=True, + model=dict( + obs_shape=17, + action_shape=6, + actor_head_hidden_size=[400, 300], + critic_head_hidden_size=[400, 300], + phi=0.05, + ), + learn=dict( + data_path=None, + train_epoch=30000, + batch_size=100, + learning_rate_q=3e-3, + learning_rate_policy=3e-3, + learning_rate_alpha=3e-3, + lmbda=0.75, + learner=dict(hook=dict(save_ckpt_after_iter=1000000000, )), + ), + collect=dict(data_type='d4rl', ), + eval=dict(evaluator=dict(eval_freq=500, )), + other=dict(replay_buffer=dict(replay_buffer_size=2000000, ), ), + ), + seed=123, +) + +main_config = EasyDict(main_config) +main_config = main_config + +create_config = dict( + env=dict( + type='d4rl', + import_names=['dizoo.d4rl.envs.d4rl_env'], + ), + env_manager=dict(type='base'), + policy=dict( + type='bcq', + import_names=['ding.policy.bcq'], + ), + replay_buffer=dict(type='naive', ), +) +create_config = EasyDict(create_config) +create_config = create_config diff --git a/dizoo/d4rl/config/halfcheetah_medium_expert_cql_config.py b/dizoo/d4rl/config/halfcheetah_medium_expert_cql_config.py old mode 100644 new mode 100755 index 2e74c43ac2..05aa2d1752 --- a/dizoo/d4rl/config/halfcheetah_medium_expert_cql_config.py +++ b/dizoo/d4rl/config/halfcheetah_medium_expert_cql_config.py @@ -5,7 +5,7 @@ main_config = dict( exp_name="halfcheetah_medium_expert_cql_seed0", env=dict( - env_id='halfcheetah-medium-expert-v0', + env_id='halfcheetah-medium-expert-v2', collector_env_num=1, evaluator_env_num=8, use_act_scale=True, @@ -15,8 +15,8 @@ policy=dict( cuda=True, model=dict( - obs_shape=11, - action_shape=3, + obs_shape=17, + action_shape=6, ), learn=dict( data_path=None, diff --git a/dizoo/d4rl/config/halfcheetah_medium_expert_dt_config.py b/dizoo/d4rl/config/halfcheetah_medium_expert_dt_config.py index 9d5d5c6aa2..1f9c636d20 100644 --- a/dizoo/d4rl/config/halfcheetah_medium_expert_dt_config.py +++ b/dizoo/d4rl/config/halfcheetah_medium_expert_dt_config.py @@ -2,37 +2,38 @@ from copy import deepcopy halfcheetah_dt_config = dict( - exp_name='halfcheetah_medium_expert_dt_seed0', + exp_name='dt_log/d4rl/halfcheetah/halfcheetah_medium_expert_dt_seed0', env=dict( env_id='HalfCheetah-v3', - norm_obs=dict(use_norm=False, ), - norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, use_act_scale=True, n_evaluator_episode=8, stop_value=6000, ), + dataset=dict( + env_type='mujoco', + rtg_scale=1000, + context_len=30, + data_dir_prefix='d4rl/halfcheetah_medium_expert-v2.pkl', + ), policy=dict( - stop_value=6000, cuda=True, + stop_value=6000, + state_mean=None, + state_std=None, + evaluator_env_num=8, env_name='HalfCheetah-v3', rtg_target=6000, # max target return to go max_eval_ep_len=1000, # max lenght of one episode - num_eval_ep=10, # num of evaluation episode - batch_size=64, wt_decay=1e-4, warmup_steps=10000, - num_updates_per_iter=100, context_len=20, - n_blocks=3, - embed_dim=128, - n_heads=1, - dropout_p=0.1, - log_dir='/home/wangzilin/research/dt/DI-engine/dizoo/d4rl/dt_data/halfcheetah_medium_expert_dt_log', + weight_decay=0.1, + clip_grad_norm_p=0.25, model=dict( - state_dim=17, - act_dim=6, + state_dim=11, + act_dim=3, n_blocks=3, h_dim=128, context_len=20, @@ -40,26 +41,13 @@ drop_p=0.1, continuous=True, ), - discount_factor=0.999, - nstep=3, - learn=dict( - dataset_path='/mnt/lustre/wangzilin/d4rl_data/halfcheetah-medium-expert-v2.pkl', - learning_rate=0.0001, - target_update_freq=100, - kappa=1.0, - min_q_weight=4.0, - ), - collect=dict(unroll_len=1, ), - eval=dict(evaluator=dict(evalu_freq=100, ), ), - other=dict( - eps=dict( - type='exp', - start=0.95, - end=0.1, - decay=10000, - ), - replay_buffer=dict(replay_buffer_size=1000, ), + batch_size=64, + learning_rate=1e-4, + collect=dict( + data_type='d4rl_trajectory', + unroll_len=1, ), + eval=dict(evaluator=dict(eval_freq=100, ), ), ), ) diff --git a/dizoo/d4rl/config/halfcheetah_medium_expert_edac_config.py b/dizoo/d4rl/config/halfcheetah_medium_expert_edac_config.py new file mode 100755 index 0000000000..17e897f048 --- /dev/null +++ b/dizoo/d4rl/config/halfcheetah_medium_expert_edac_config.py @@ -0,0 +1,58 @@ +from easydict import EasyDict + +main_config = dict( + exp_name="halfcheetah_medium_expert_edac_seed123", + env=dict( + env_id='halfcheetah-medium-expert-v2', + collector_env_num=1, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + stop_value=13000, + ), + policy=dict( + cuda=True, + model=dict( + obs_shape=17, + action_shape=6, + ensemble_num=10, + actor_head_hidden_size=256, + actor_head_layer_num=3, + critic_head_hidden_size=256, + critic_head_layer_num=3, + ), + learn=dict( + data_path=None, + train_epoch=30000, + batch_size=256, + learning_rate_q=3e-4, + learning_rate_policy=3e-4, + learning_rate_alpha=3e-4, + alpha=1, + auto_alpha=True, + eta=5.0, + with_q_entropy=False, + learner=dict(hook=dict(save_ckpt_after_iter=100000, )), + ), + collect=dict(data_type='d4rl', ), + eval=dict(evaluator=dict(eval_freq=500, )), + ), + seed=123, +) + +main_config = EasyDict(main_config) +main_config = main_config + +create_config = dict( + env=dict( + type='d4rl', + import_names=['dizoo.d4rl.envs.d4rl_env'], + ), + env_manager=dict(type='base'), + policy=dict( + type='edac', + import_names=['ding.policy.edac'], + ), +) +create_config = EasyDict(create_config) +create_config = create_config diff --git a/dizoo/d4rl/config/halfcheetah_medium_expert_iql_config.py b/dizoo/d4rl/config/halfcheetah_medium_expert_iql_config.py new file mode 100644 index 0000000000..e3aa855afe --- /dev/null +++ b/dizoo/d4rl/config/halfcheetah_medium_expert_iql_config.py @@ -0,0 +1,53 @@ +# You can conduct Experiments on D4RL with this config file through the following command: +# cd ../entry && python d4rl_iql_main.py +from easydict import EasyDict + +main_config = dict( + exp_name="halfcheetah_medium_expert_iql_seed0", + env=dict( + env_id='halfcheetah-medium-expert-v2', + collector_env_num=1, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + stop_value=6000, + reward_norm="iql_locomotion", + ), + policy=dict( + cuda=True, + model=dict( + obs_shape=17, + action_shape=6, + ), + learn=dict( + data_path=None, + train_epoch=30000, + batch_size=4096, + learning_rate_q=3e-4, + learning_rate_policy=1e-4, + beta=0.05, + tau=0.7, + ), + collect=dict(data_type='d4rl', ), + eval=dict(evaluator=dict(eval_freq=5000, )), + other=dict(replay_buffer=dict(replay_buffer_size=2000000, ), ), + ), +) + +main_config = EasyDict(main_config) +main_config = main_config + +create_config = dict( + env=dict( + type='d4rl', + import_names=['dizoo.d4rl.envs.d4rl_env'], + ), + env_manager=dict(type='base'), + policy=dict( + type='iql', + import_names=['ding.policy.iql'], + ), + replay_buffer=dict(type='naive', ), +) +create_config = EasyDict(create_config) +create_config = create_config diff --git a/dizoo/d4rl/config/halfcheetah_medium_expert_pd_config.py b/dizoo/d4rl/config/halfcheetah_medium_expert_pd_config.py new file mode 100755 index 0000000000..d42692bc63 --- /dev/null +++ b/dizoo/d4rl/config/halfcheetah_medium_expert_pd_config.py @@ -0,0 +1,99 @@ +from easydict import EasyDict + +main_config = dict( + exp_name="halfcheetah_medium_expert_pd_seed0", + env=dict( + env_id='halfcheetah-medium-expert-v2', + collector_env_num=1, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + returns_scale=1.0, + termination_penalty=-100, + max_path_length=1000, + use_padding=True, + include_returns=True, + normed=False, + stop_value=12000, + horizon=4, + obs_dim=17, + action_dim=6, + ), + policy=dict( + cuda=True, + model=dict( + diffuser_model='GaussianDiffusion', + diffuser_model_cfg=dict( + model='DiffusionUNet1d', + model_cfg=dict( + transition_dim=23, + dim=32, + dim_mults=[1, 4, 8], + returns_condition=False, + kernel_size=5, + attention=True, + ), + horizon=4, + obs_dim=17, + action_dim=6, + n_timesteps=20, + predict_epsilon=False, + loss_discount=1, + action_weight=10, + ), + value_model='ValueDiffusion', + value_model_cfg=dict( + model='TemporalValue', + model_cfg=dict( + horizon=4, + transition_dim=23, + dim=32, + dim_mults=[1, 4, 8], + kernel_size=5, + ), + horizon=4, + obs_dim=17, + action_dim=6, + n_timesteps=20, + predict_epsilon=True, + loss_discount=1, + ), + n_guide_steps=2, + scale=0.001, + t_stopgrad=4, + scale_grad_by_std=True, + ), + normalizer='GaussianNormalizer', + learn=dict( + data_path=None, + train_epoch=60000, + gradient_accumulate_every=2, + batch_size=32, + learning_rate=2e-4, + discount_factor=0.99, + plan_batch_size=64, + learner=dict(hook=dict(save_ckpt_after_iter=1000000000, )), + ), + collect=dict(data_type='diffuser_traj', ), + eval=dict( + evaluator=dict(eval_freq=500, ), + test_ret=0.9, + ), + other=dict(replay_buffer=dict(replay_buffer_size=2000000, ), ), + ), +) + +main_config = EasyDict(main_config) +main_config = main_config + +create_config = dict( + env=dict( + type='d4rl', + import_names=['dizoo.d4rl.envs.d4rl_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='pd', ), + replay_buffer=dict(type='naive', ), +) +create_config = EasyDict(create_config) +create_config = create_config diff --git a/dizoo/d4rl/config/halfcheetah_medium_expert_qgpo_config.py b/dizoo/d4rl/config/halfcheetah_medium_expert_qgpo_config.py new file mode 100644 index 0000000000..c8b75891ca --- /dev/null +++ b/dizoo/d4rl/config/halfcheetah_medium_expert_qgpo_config.py @@ -0,0 +1,47 @@ +from easydict import EasyDict + +main_config = dict( + exp_name='halfcheetah_medium_expert_v2_QGPO_seed0', + env=dict( + env_id="halfcheetah-medium-expert-v2", + evaluator_env_num=8, + n_evaluator_episode=8, + ), + dataset=dict(env_id="halfcheetah-medium-expert-v2", ), + policy=dict( + cuda=True, + on_policy=False, + #load_path='./halfcheetah_medium_expert_v2_QGPO_seed0/ckpt/iteration_600000.pth.tar', + model=dict( + qgpo_critic=dict( + alpha=3, + q_alpha=1, + ), + device='cuda', + obs_dim=17, + action_dim=6, + ), + learn=dict( + learning_rate=1e-4, + batch_size=4096, + batch_size_q=256, + M=16, + diffusion_steps=15, + behavior_policy_stop_training_iter=600000, + energy_guided_policy_begin_training_iter=600000, + q_value_stop_training_iter=1100000, + ), + eval=dict( + guidance_scale=[0.0, 1.0, 2.0, 3.0, 5.0, 8.0, 10.0], + diffusion_steps=15, + evaluator=dict(eval_freq=50000, ), + ), + ), +) +main_config = EasyDict(main_config) + +create_config = dict( + env_manager=dict(type='base'), + policy=dict(type='qgpo', ), +) +create_config = EasyDict(create_config) diff --git a/dizoo/d4rl/config/halfcheetah_medium_expert_td3bc_config.py b/dizoo/d4rl/config/halfcheetah_medium_expert_td3bc_config.py old mode 100644 new mode 100755 index 8d25289131..ed99a2d3f0 --- a/dizoo/d4rl/config/halfcheetah_medium_expert_td3bc_config.py +++ b/dizoo/d4rl/config/halfcheetah_medium_expert_td3bc_config.py @@ -5,21 +5,22 @@ main_config = dict( exp_name='halfcheetah_medium_expert_td3-bc_seed0', env=dict( - env_id='halfcheetah-medium-expert-v0', + env_id='halfcheetah-medium-expert-v2', norm_obs=dict( - use_norm=True, + use_norm=True, offline_stats=dict(use_offline_stats=True, ), ), collector_env_num=1, evaluator_env_num=8, use_act_scale=True, n_evaluator_episode=8, - stop_value=6000, + stop_value=13000, ), policy=dict( + cuda=True, model=dict( - obs_shape=11, - action_shape=3, + obs_shape=17, + action_shape=6, ), learn=dict( train_epoch=30000, diff --git a/dizoo/d4rl/config/halfcheetah_medium_iql_config.py b/dizoo/d4rl/config/halfcheetah_medium_iql_config.py new file mode 100644 index 0000000000..440525a320 --- /dev/null +++ b/dizoo/d4rl/config/halfcheetah_medium_iql_config.py @@ -0,0 +1,53 @@ +# You can conduct Experiments on D4RL with this config file through the following command: +# cd ../entry && python d4rl_iql_main.py +from easydict import EasyDict + +main_config = dict( + exp_name="halfcheetah_medium_iql_seed0", + env=dict( + env_id='halfcheetah-medium-v2', + collector_env_num=1, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + stop_value=6000, + reward_norm="iql_locomotion", + ), + policy=dict( + cuda=True, + model=dict( + obs_shape=17, + action_shape=6, + ), + learn=dict( + data_path=None, + train_epoch=30000, + batch_size=4096, + learning_rate_q=3e-4, + learning_rate_policy=1e-4, + beta=0.05, + tau=0.7, + ), + collect=dict(data_type='d4rl', ), + eval=dict(evaluator=dict(eval_freq=5000, )), + other=dict(replay_buffer=dict(replay_buffer_size=2000000, ), ), + ), +) + +main_config = EasyDict(main_config) +main_config = main_config + +create_config = dict( + env=dict( + type='d4rl', + import_names=['dizoo.d4rl.envs.d4rl_env'], + ), + env_manager=dict(type='base'), + policy=dict( + type='iql', + import_names=['ding.policy.iql'], + ), + replay_buffer=dict(type='naive', ), +) +create_config = EasyDict(create_config) +create_config = create_config diff --git a/dizoo/d4rl/config/halfcheetah_medium_pd_config.py b/dizoo/d4rl/config/halfcheetah_medium_pd_config.py new file mode 100755 index 0000000000..ae4145f9a3 --- /dev/null +++ b/dizoo/d4rl/config/halfcheetah_medium_pd_config.py @@ -0,0 +1,99 @@ +from easydict import EasyDict + +main_config = dict( + exp_name="halfcheetah_medium_pd_seed0", + env=dict( + env_id='halfcheetah-medium-v2', + collector_env_num=1, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + returns_scale=1.0, + termination_penalty=-100, + max_path_length=1000, + use_padding=True, + include_returns=True, + normed=False, + stop_value=8000, + horizon=4, + obs_dim=17, + action_dim=6, + ), + policy=dict( + cuda=True, + model=dict( + diffuser_model='GaussianDiffusion', + diffuser_model_cfg=dict( + model='DiffusionUNet1d', + model_cfg=dict( + transition_dim=23, + dim=32, + dim_mults=[1, 4, 8], + returns_condition=False, + kernel_size=5, + attention=True, + ), + horizon=4, + obs_dim=17, + action_dim=6, + n_timesteps=20, + predict_epsilon=False, + loss_discount=1, + action_weight=10, + ), + value_model='ValueDiffusion', + value_model_cfg=dict( + model='TemporalValue', + model_cfg=dict( + horizon=4, + transition_dim=23, + dim=32, + dim_mults=[1, 4, 8], + kernel_size=5, + ), + horizon=4, + obs_dim=17, + action_dim=6, + n_timesteps=20, + predict_epsilon=True, + loss_discount=1, + ), + n_guide_steps=2, + scale=0.001, + t_stopgrad=4, + scale_grad_by_std=True, + ), + normalizer='GaussianNormalizer', + learn=dict( + data_path=None, + train_epoch=60000, + gradient_accumulate_every=2, + batch_size=32, + learning_rate=2e-4, + discount_factor=0.99, + plan_batch_size=64, + learner=dict(hook=dict(save_ckpt_after_iter=1000000000, )), + ), + collect=dict(data_type='diffuser_traj', ), + eval=dict( + evaluator=dict(eval_freq=500, ), + test_ret=0.9, + ), + other=dict(replay_buffer=dict(replay_buffer_size=2000000, ), ), + ), +) + +main_config = EasyDict(main_config) +main_config = main_config + +create_config = dict( + env=dict( + type='d4rl', + import_names=['dizoo.d4rl.envs.d4rl_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='pd', ), + replay_buffer=dict(type='naive', ), +) +create_config = EasyDict(create_config) +create_config = create_config diff --git a/dizoo/d4rl/config/halfcheetah_medium_replay_cql_config.py b/dizoo/d4rl/config/halfcheetah_medium_replay_cql_config.py old mode 100644 new mode 100755 index 010c7963d2..823e08d370 --- a/dizoo/d4rl/config/halfcheetah_medium_replay_cql_config.py +++ b/dizoo/d4rl/config/halfcheetah_medium_replay_cql_config.py @@ -5,7 +5,7 @@ main_config = dict( exp_name="halfcheetah_medium_replay_cql_seed0", env=dict( - env_id='halfcheetah-medium-replay-v0', + env_id='halfcheetah-medium-replay-v2', collector_env_num=1, evaluator_env_num=8, use_act_scale=True, @@ -15,8 +15,8 @@ policy=dict( cuda=True, model=dict( - obs_shape=11, - action_shape=3, + obs_shape=17, + action_shape=6, ), learn=dict( data_path=None, diff --git a/dizoo/d4rl/config/halfcheetah_medium_replay_dt_config.py b/dizoo/d4rl/config/halfcheetah_medium_replay_dt_config.py index 98260c4dd3..aa07e22280 100644 --- a/dizoo/d4rl/config/halfcheetah_medium_replay_dt_config.py +++ b/dizoo/d4rl/config/halfcheetah_medium_replay_dt_config.py @@ -2,37 +2,38 @@ from copy import deepcopy halfcheetah_dt_config = dict( - exp_name='halfcheetah_medium_replay_dt_seed0', + exp_name='dt_log/d4rl/halfcheetah/halfcheetah_medium_replay_dt_seed0', env=dict( env_id='HalfCheetah-v3', - norm_obs=dict(use_norm=False, ), - norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, use_act_scale=True, n_evaluator_episode=8, stop_value=6000, ), + dataset=dict( + env_type='mujoco', + rtg_scale=1000, + context_len=30, + data_dir_prefix='d4rl/halfcheetah_medium_replay-v2.pkl', + ), policy=dict( - stop_value=6000, cuda=True, + stop_value=6000, + state_mean=None, + state_std=None, + evaluator_env_num=8, env_name='HalfCheetah-v3', rtg_target=6000, # max target return to go max_eval_ep_len=1000, # max lenght of one episode - num_eval_ep=10, # num of evaluation episode - batch_size=64, wt_decay=1e-4, warmup_steps=10000, - num_updates_per_iter=100, context_len=20, - n_blocks=3, - embed_dim=128, - n_heads=1, - dropout_p=0.1, - log_dir='/home/wangzilin/research/dt/DI-engine/dizoo/d4rl/dt_data/halfcheetah_medium_replay_dt_log', + weight_decay=0.1, + clip_grad_norm_p=0.25, model=dict( - state_dim=17, - act_dim=6, + state_dim=11, + act_dim=3, n_blocks=3, h_dim=128, context_len=20, @@ -40,26 +41,13 @@ drop_p=0.1, continuous=True, ), - discount_factor=0.999, - nstep=3, - learn=dict( - dataset_path='/mnt/lustre/wangzilin/d4rl_data/halfcheetah-medium-replay-v2.pkl', - learning_rate=0.0001, - target_update_freq=100, - kappa=1.0, - min_q_weight=4.0, - ), - collect=dict(unroll_len=1, ), - eval=dict(evaluator=dict(evalu_freq=100, ), ), - other=dict( - eps=dict( - type='exp', - start=0.95, - end=0.1, - decay=10000, - ), - replay_buffer=dict(replay_buffer_size=1000, ), + batch_size=64, + learning_rate=1e-4, + collect=dict( + data_type='d4rl_trajectory', + unroll_len=1, ), + eval=dict(evaluator=dict(eval_freq=100, ), ), ), ) diff --git a/dizoo/d4rl/config/halfcheetah_medium_replay_iql_config.py b/dizoo/d4rl/config/halfcheetah_medium_replay_iql_config.py new file mode 100644 index 0000000000..0974735b72 --- /dev/null +++ b/dizoo/d4rl/config/halfcheetah_medium_replay_iql_config.py @@ -0,0 +1,53 @@ +# You can conduct Experiments on D4RL with this config file through the following command: +# cd ../entry && python d4rl_iql_main.py +from easydict import EasyDict + +main_config = dict( + exp_name="halfcheetah_medium_replay_iql_seed0", + env=dict( + env_id='halfcheetah-medium-replay-v2', + collector_env_num=1, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + stop_value=6000, + reward_norm="iql_locomotion", + ), + policy=dict( + cuda=True, + model=dict( + obs_shape=17, + action_shape=6, + ), + learn=dict( + data_path=None, + train_epoch=30000, + batch_size=4096, + learning_rate_q=3e-4, + learning_rate_policy=1e-4, + beta=0.05, + tau=0.7, + ), + collect=dict(data_type='d4rl', ), + eval=dict(evaluator=dict(eval_freq=5000, )), + other=dict(replay_buffer=dict(replay_buffer_size=2000000, ), ), + ), +) + +main_config = EasyDict(main_config) +main_config = main_config + +create_config = dict( + env=dict( + type='d4rl', + import_names=['dizoo.d4rl.envs.d4rl_env'], + ), + env_manager=dict(type='base'), + policy=dict( + type='iql', + import_names=['ding.policy.iql'], + ), + replay_buffer=dict(type='naive', ), +) +create_config = EasyDict(create_config) +create_config = create_config diff --git a/dizoo/d4rl/config/halfcheetah_medium_replay_td3bc_config.py b/dizoo/d4rl/config/halfcheetah_medium_replay_td3bc_config.py old mode 100644 new mode 100755 index 3561f320fb..22cf7ff544 --- a/dizoo/d4rl/config/halfcheetah_medium_replay_td3bc_config.py +++ b/dizoo/d4rl/config/halfcheetah_medium_replay_td3bc_config.py @@ -5,9 +5,9 @@ main_config = dict( exp_name='halfcheetah_medium_replay_td3-bc_seed0', env=dict( - env_id='halfcheetah-medium-replay-v0', + env_id='halfcheetah-medium-replay-v2', norm_obs=dict( - use_norm=True, + use_norm=True, offline_stats=dict(use_offline_stats=True, ), ), collector_env_num=1, @@ -17,9 +17,10 @@ stop_value=6000, ), policy=dict( + cuda=True, model=dict( - obs_shape=11, - action_shape=3, + obs_shape=17, + action_shape=6, ), learn=dict( train_epoch=30000, diff --git a/dizoo/d4rl/config/halfcheetah_medium_td3bc_config.py b/dizoo/d4rl/config/halfcheetah_medium_td3bc_config.py old mode 100644 new mode 100755 index ef6e2d3f40..541588e261 --- a/dizoo/d4rl/config/halfcheetah_medium_td3bc_config.py +++ b/dizoo/d4rl/config/halfcheetah_medium_td3bc_config.py @@ -5,21 +5,22 @@ main_config = dict( exp_name='halfcheetah_medium_td3-bc_seed0', env=dict( - env_id='halfcheetah-medium-v0', + env_id='halfcheetah-medium-v2', norm_obs=dict( - use_norm=True, + use_norm=True, offline_stats=dict(use_offline_stats=True, ), ), collector_env_num=1, evaluator_env_num=8, use_act_scale=True, n_evaluator_episode=8, - stop_value=6000, + stop_value=7600, ), policy=dict( + cuda=True, model=dict( - obs_shape=11, - action_shape=3, + obs_shape=17, + action_shape=6, ), learn=dict( train_epoch=30000, diff --git a/dizoo/d4rl/config/halfcheetah_random_cql_config.py b/dizoo/d4rl/config/halfcheetah_random_cql_config.py old mode 100644 new mode 100755 index bb4b321d80..58ead98dc5 --- a/dizoo/d4rl/config/halfcheetah_random_cql_config.py +++ b/dizoo/d4rl/config/halfcheetah_random_cql_config.py @@ -5,7 +5,7 @@ main_config = dict( exp_name="halfcheetah_expert_cql_seed0", env=dict( - env_id='halfcheetah-expert-v0', + env_id='halfcheetah-expert-v2', collector_env_num=1, evaluator_env_num=8, use_act_scale=True, @@ -15,8 +15,8 @@ policy=dict( cuda=True, model=dict( - obs_shape=11, - action_shape=3, + obs_shape=17, + action_shape=6, ), learn=dict( data_path=None, diff --git a/dizoo/d4rl/config/halfcheetah_random_td3bc_config.py b/dizoo/d4rl/config/halfcheetah_random_td3bc_config.py old mode 100644 new mode 100755 index dbe94d1a24..85c03478bc --- a/dizoo/d4rl/config/halfcheetah_random_td3bc_config.py +++ b/dizoo/d4rl/config/halfcheetah_random_td3bc_config.py @@ -5,9 +5,9 @@ main_config = dict( exp_name='halfcheetah_random_td3-bc_seed0', env=dict( - env_id='halfcheetah-random-v0', + env_id='halfcheetah-random-v2', norm_obs=dict( - use_norm=True, + use_norm=True, offline_stats=dict(use_offline_stats=True, ), ), collector_env_num=1, @@ -17,9 +17,10 @@ stop_value=6000, ), policy=dict( + cuda=True, model=dict( - obs_shape=11, - action_shape=3, + obs_shape=17, + action_shape=6, ), learn=dict( train_epoch=30000, diff --git a/dizoo/d4rl/config/hopper_expert_dt_config.py b/dizoo/d4rl/config/hopper_expert_dt_config.py index 592d9411df..26387afde5 100644 --- a/dizoo/d4rl/config/hopper_expert_dt_config.py +++ b/dizoo/d4rl/config/hopper_expert_dt_config.py @@ -2,34 +2,35 @@ from copy import deepcopy hopper_dt_config = dict( - exp_name='hopper_expert_dt_seed0', + exp_name='dt_log/d4rl/hopper/hopper_expert_dt_seed0', env=dict( env_id='Hopper-v3', - norm_obs=dict(use_norm=False, ), - norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, use_act_scale=True, n_evaluator_episode=8, - stop_value=6000, + stop_value=3600, + ), + dataset=dict( + env_type='mujoco', + rtg_scale=1000, + context_len=20, + data_dir_prefix='d4rl/hopper_expert-v2.pkl', ), policy=dict( - stop_value=6000, cuda=True, + stop_value=3600, + state_mean=None, + state_std=None, + evaluator_env_num=8, env_name='Hopper-v3', - rtg_target=6000, # max target return to go + rtg_target=3600, # max target return to go max_eval_ep_len=1000, # max lenght of one episode - num_eval_ep=10, # num of evaluation episode - batch_size=64, wt_decay=1e-4, warmup_steps=10000, - num_updates_per_iter=100, context_len=20, - n_blocks=3, - embed_dim=128, - n_heads=1, - dropout_p=0.1, - log_dir='/home/wangzilin/research/dt/DI-engine/dizoo/d4rl/dt_data/hopper_expert_dt_log', + weight_decay=0.1, + clip_grad_norm_p=0.25, model=dict( state_dim=11, act_dim=3, @@ -40,26 +41,13 @@ drop_p=0.1, continuous=True, ), - discount_factor=0.999, - nstep=3, - learn=dict( - dataset_path='/mnt/lustre/wangzilin/d4rl_data/hopper-expert-v2.pkl', - learning_rate=0.0001, - target_update_freq=100, - kappa=1.0, - min_q_weight=4.0, - ), - collect=dict(unroll_len=1, ), - eval=dict(evaluator=dict(evalu_freq=100, ), ), - other=dict( - eps=dict( - type='exp', - start=0.95, - end=0.1, - decay=10000, - ), - replay_buffer=dict(replay_buffer_size=1000, ), + batch_size=64, + learning_rate=1e-4, + collect=dict( + data_type='d4rl_trajectory', + unroll_len=1, ), + eval=dict(evaluator=dict(eval_freq=100, ), ), ), ) diff --git a/dizoo/d4rl/config/hopper_expert_td3bc_config.py b/dizoo/d4rl/config/hopper_expert_td3bc_config.py index b0874a0018..0e35474c48 100644 --- a/dizoo/d4rl/config/hopper_expert_td3bc_config.py +++ b/dizoo/d4rl/config/hopper_expert_td3bc_config.py @@ -5,11 +5,7 @@ main_config = dict( exp_name='hopper_expert_td3-bc_seed0', env=dict( - env_id='hopper-expert-v0', - norm_obs=dict( - use_norm=True, - offline_stats=dict(use_offline_stats=True, ), - ), + env_id='hopper-expert-v2', collector_env_num=1, evaluator_env_num=8, use_act_scale=True, @@ -17,6 +13,7 @@ stop_value=6000, ), policy=dict( + cuda=True, model=dict( obs_shape=11, action_shape=3, diff --git a/dizoo/d4rl/config/hopper_medium_bcq_config.py b/dizoo/d4rl/config/hopper_medium_bcq_config.py new file mode 100755 index 0000000000..06282d1680 --- /dev/null +++ b/dizoo/d4rl/config/hopper_medium_bcq_config.py @@ -0,0 +1,55 @@ +from easydict import EasyDict + +main_config = dict( + exp_name="hopper_medium_bcq_seed0_43_v0", + env=dict( + env_id='hopper-medium-v0', + collector_env_num=1, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + stop_value=3500, + ), + policy=dict( + cuda=True, + model=dict( + obs_shape=11, + action_shape=3, + actor_head_hidden_size=[400, 300], + critic_head_hidden_size=[400, 300], + phi=0.05, + ), + learn=dict( + data_path=None, + train_epoch=30000, + batch_size=100, + learning_rate_q=3e-3, + learning_rate_policy=3e-3, + learning_rate_alpha=3e-3, + lmbda=0.75, + learner=dict(hook=dict(save_ckpt_after_iter=1000000000, )), + ), + collect=dict(data_type='d4rl', ), + eval=dict(evaluator=dict(eval_freq=500, )), + other=dict(replay_buffer=dict(replay_buffer_size=2000000, ), ), + ), + seed=123, +) + +main_config = EasyDict(main_config) +main_config = main_config + +create_config = dict( + env=dict( + type='d4rl', + import_names=['dizoo.d4rl.envs.d4rl_env'], + ), + env_manager=dict(type='base'), + policy=dict( + type='bcq', + import_names=['ding.policy.bcq'], + ), + replay_buffer=dict(type='naive', ), +) +create_config = EasyDict(create_config) +create_config = create_config diff --git a/dizoo/d4rl/config/hopper_medium_dt_config.py b/dizoo/d4rl/config/hopper_medium_dt_config.py index ae3778a1d8..a5c389e67e 100644 --- a/dizoo/d4rl/config/hopper_medium_dt_config.py +++ b/dizoo/d4rl/config/hopper_medium_dt_config.py @@ -2,34 +2,35 @@ from copy import deepcopy hopper_dt_config = dict( - exp_name='hopper_medium_dt_seed0', + exp_name='dt_log/d4rl/hopper/hopper_medium_dt_seed0', env=dict( env_id='Hopper-v3', - norm_obs=dict(use_norm=False, ), - norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, use_act_scale=True, n_evaluator_episode=8, - stop_value=6000, + stop_value=3600, + ), + dataset=dict( + env_type='mujoco', + rtg_scale=1000, + context_len=20, + data_dir_prefix='d4rl/hopper_medium_expert-v2.pkl', ), policy=dict( - stop_value=6000, cuda=True, + stop_value=3600, + state_mean=None, + state_std=None, + evaluator_env_num=8, env_name='Hopper-v3', - rtg_target=6000, # max target return to go + rtg_target=3600, # max target return to go max_eval_ep_len=1000, # max lenght of one episode - num_eval_ep=10, # num of evaluation episode - batch_size=64, wt_decay=1e-4, warmup_steps=10000, - num_updates_per_iter=100, context_len=20, - n_blocks=3, - embed_dim=128, - n_heads=1, - dropout_p=0.1, - log_dir='/home/wangzilin/research/dt/DI-engine/dizoo/d4rl/dt_data/hopper_medium_dt_log', + weight_decay=0.1, + clip_grad_norm_p=0.25, model=dict( state_dim=11, act_dim=3, @@ -40,26 +41,13 @@ drop_p=0.1, continuous=True, ), - discount_factor=0.999, - nstep=3, - learn=dict( - dataset_path='/mnt/lustre/wangzilin/d4rl_data/hopper-medium-v2.pkl', - learning_rate=0.0001, - target_update_freq=100, - kappa=1.0, - min_q_weight=4.0, - ), - collect=dict(unroll_len=1, ), - eval=dict(evaluator=dict(evalu_freq=100, ), ), - other=dict( - eps=dict( - type='exp', - start=0.95, - end=0.1, - decay=10000, - ), - replay_buffer=dict(replay_buffer_size=1000, ), + batch_size=64, + learning_rate=1e-4, + collect=dict( + data_type='d4rl_trajectory', + unroll_len=1, ), + eval=dict(evaluator=dict(eval_freq=1000, ), ), ), ) diff --git a/dizoo/d4rl/config/hopper_medium_edac_config.py b/dizoo/d4rl/config/hopper_medium_edac_config.py new file mode 100755 index 0000000000..f14fad350f --- /dev/null +++ b/dizoo/d4rl/config/hopper_medium_edac_config.py @@ -0,0 +1,58 @@ +from easydict import EasyDict + +main_config = dict( + exp_name="hopper_medium_edac_seed0", + env=dict( + env_id='hopper-medium-v2', + collector_env_num=1, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + stop_value=3700, + ), + policy=dict( + cuda=True, + model=dict( + obs_shape=11, + action_shape=3, + ensemble_num=50, + actor_head_hidden_size=256, + actor_head_layer_num=3, + critic_head_hidden_size=256, + critic_head_layer_num=3, + ), + learn=dict( + data_path=None, + train_epoch=30000, + batch_size=256, + learning_rate_q=3e-4, + learning_rate_policy=3e-4, + learning_rate_alpha=3e-4, + alpha=1, + auto_alpha=True, + eta=1.0, + with_q_entropy=False, + learner=dict(hook=dict(save_ckpt_after_iter=100000, )), + ), + collect=dict(data_type='d4rl', ), + eval=dict(evaluator=dict(eval_freq=500, )), + ), + seed=0, +) + +main_config = EasyDict(main_config) +main_config = main_config + +create_config = dict( + env=dict( + type='d4rl', + import_names=['dizoo.d4rl.envs.d4rl_env'], + ), + env_manager=dict(type='base'), + policy=dict( + type='edac', + import_names=['ding.policy.edac'], + ), +) +create_config = EasyDict(create_config) +create_config = create_config diff --git a/dizoo/d4rl/config/hopper_medium_expert_bc_config.py b/dizoo/d4rl/config/hopper_medium_expert_bc_config.py index e04bd28069..348361dd2d 100644 --- a/dizoo/d4rl/config/hopper_medium_expert_bc_config.py +++ b/dizoo/d4rl/config/hopper_medium_expert_bc_config.py @@ -8,7 +8,7 @@ env=dict( env_id='hopper-medium-expert-v0', norm_obs=dict( - use_norm=True, + use_norm=True, offline_stats=dict(use_offline_stats=True, ), ), evaluator_env_num=8, @@ -38,7 +38,7 @@ data_type='d4rl', data_path=None, ), - eval=dict(evaluator=dict(eval_freq=-1,)), + eval=dict(evaluator=dict(eval_freq=-1, )), ), ) main_config = EasyDict(main_config) @@ -48,7 +48,7 @@ type='d4rl', import_names=['dizoo.d4rl.envs.d4rl_env'], ), - env_manager=dict(type='base',), + env_manager=dict(type='base', ), policy=dict( type='bc', import_names=['ding.policy.bc'], diff --git a/dizoo/d4rl/config/hopper_medium_expert_bcq_config.py b/dizoo/d4rl/config/hopper_medium_expert_bcq_config.py new file mode 100755 index 0000000000..ac48ee4847 --- /dev/null +++ b/dizoo/d4rl/config/hopper_medium_expert_bcq_config.py @@ -0,0 +1,55 @@ +from easydict import EasyDict + +main_config = dict( + exp_name="hopper_medium_expert_bcq_seed0", + env=dict( + env_id='hopper-medium-expert-v0', + collector_env_num=1, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + stop_value=3800, + ), + policy=dict( + cuda=True, + model=dict( + obs_shape=11, + action_shape=3, + actor_head_hidden_size=[400, 300], + critic_head_hidden_size=[400, 300], + phi=0.05, + ), + learn=dict( + data_path=None, + train_epoch=30000, + batch_size=100, + learning_rate_q=3e-3, + learning_rate_policy=3e-3, + learning_rate_alpha=3e-3, + lmbda=0.75, + learner=dict(hook=dict(save_ckpt_after_iter=1000000000, )), + ), + collect=dict(data_type='d4rl', ), + eval=dict(evaluator=dict(eval_freq=500, )), + other=dict(replay_buffer=dict(replay_buffer_size=2000000, ), ), + ), + seed=123, +) + +main_config = EasyDict(main_config) +main_config = main_config + +create_config = dict( + env=dict( + type='d4rl', + import_names=['dizoo.d4rl.envs.d4rl_env'], + ), + env_manager=dict(type='base'), + policy=dict( + type='bcq', + import_names=['ding.policy.bcq'], + ), + replay_buffer=dict(type='naive', ), +) +create_config = EasyDict(create_config) +create_config = create_config diff --git a/dizoo/d4rl/config/hopper_medium_expert_dt_config.py b/dizoo/d4rl/config/hopper_medium_expert_dt_config.py index cab3b10353..5934590bf1 100644 --- a/dizoo/d4rl/config/hopper_medium_expert_dt_config.py +++ b/dizoo/d4rl/config/hopper_medium_expert_dt_config.py @@ -2,34 +2,35 @@ from copy import deepcopy hopper_dt_config = dict( - exp_name='hopper_medium_expert_dt_seed0', + exp_name='dt_log/d4rl/hopper/hopper_medium_expert_dt', env=dict( env_id='Hopper-v3', - norm_obs=dict(use_norm=False, ), - norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, use_act_scale=True, n_evaluator_episode=8, - stop_value=6000, + stop_value=3600, + ), + dataset=dict( + env_type='mujoco', + rtg_scale=1000, + context_len=20, + data_dir_prefix='d4rl/hopper_medium_expert.pkl', ), policy=dict( - stop_value=6000, cuda=True, + stop_value=3600, + state_mean=None, + state_std=None, + evaluator_env_num=8, env_name='Hopper-v3', - rtg_target=6000, # max target return to go + rtg_target=3600, # max target return to go max_eval_ep_len=1000, # max lenght of one episode - num_eval_ep=10, # num of evaluation episode - batch_size=64, wt_decay=1e-4, warmup_steps=10000, - num_updates_per_iter=100, context_len=20, - n_blocks=3, - embed_dim=128, - n_heads=1, - dropout_p=0.1, - log_dir='/home/wangzilin/research/dt/DI-engine/dizoo/d4rl/dt_data/hopper_medium_expert_dt_log', + weight_decay=0.1, + clip_grad_norm_p=0.25, model=dict( state_dim=11, act_dim=3, @@ -40,26 +41,13 @@ drop_p=0.1, continuous=True, ), - discount_factor=0.999, - nstep=3, - learn=dict( - dataset_path='/mnt/lustre/wangzilin/d4rl_data/hopper-medium-expert-v2.pkl', - learning_rate=0.0001, - target_update_freq=100, - kappa=1.0, - min_q_weight=4.0, - ), - collect=dict(unroll_len=1, ), - eval=dict(evaluator=dict(evalu_freq=100, ), ), - other=dict( - eps=dict( - type='exp', - start=0.95, - end=0.1, - decay=10000, - ), - replay_buffer=dict(replay_buffer_size=1000, ), + batch_size=64, + learning_rate=1e-4, + collect=dict( + data_type='d4rl_trajectory', + unroll_len=1, ), + eval=dict(evaluator=dict(eval_freq=100, ), ), ), ) diff --git a/dizoo/d4rl/config/hopper_medium_expert_edac_config.py b/dizoo/d4rl/config/hopper_medium_expert_edac_config.py new file mode 100755 index 0000000000..5bbc5b375d --- /dev/null +++ b/dizoo/d4rl/config/hopper_medium_expert_edac_config.py @@ -0,0 +1,58 @@ +from easydict import EasyDict + +main_config = dict( + exp_name="hopper_medium_expert_edac_seed0", + env=dict( + env_id='hopper-medium-expert-v2', + collector_env_num=1, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + stop_value=5000, + ), + policy=dict( + cuda=True, + model=dict( + obs_shape=11, + action_shape=3, + ensemble_num=50, + actor_head_hidden_size=256, + actor_head_layer_num=3, + critic_head_hidden_size=256, + critic_head_layer_num=3, + ), + learn=dict( + data_path=None, + train_epoch=30000, + batch_size=256, + learning_rate_q=3e-4, + learning_rate_policy=3e-4, + learning_rate_alpha=3e-4, + alpha=1, + auto_alpha=False, + eta=1.0, + with_q_entropy=False, + learner=dict(hook=dict(save_ckpt_after_iter=100000, )), + ), + collect=dict(data_type='d4rl', ), + eval=dict(evaluator=dict(eval_freq=500, )), + ), + seed=0, +) + +main_config = EasyDict(main_config) +main_config = main_config + +create_config = dict( + env=dict( + type='d4rl', + import_names=['dizoo.d4rl.envs.d4rl_env'], + ), + env_manager=dict(type='base'), + policy=dict( + type='edac', + import_names=['ding.policy.edac'], + ), +) +create_config = EasyDict(create_config) +create_config = create_config diff --git a/dizoo/d4rl/config/hopper_medium_expert_ibc_ar_config.py b/dizoo/d4rl/config/hopper_medium_expert_ibc_ar_config.py index 061b8b53a6..5d1090dc77 100644 --- a/dizoo/d4rl/config/hopper_medium_expert_ibc_ar_config.py +++ b/dizoo/d4rl/config/hopper_medium_expert_ibc_ar_config.py @@ -8,7 +8,7 @@ env=dict( env_id='hopper-medium-expert-v0', norm_obs=dict( - use_norm=True, + use_norm=True, offline_stats=dict(use_offline_stats=True, ), ), evaluator_env_num=8, @@ -18,23 +18,19 @@ ), policy=dict( cuda=cuda, - model=dict( - obs_shape=11, - action_shape=3, - stochastic_optim=dict(type='ardfo',) - ), + model=dict(obs_shape=11, action_shape=3, stochastic_optim=dict(type='ardfo', )), learn=dict( multi_gpu=multi_gpu, train_epoch=15, batch_size=256, - optim=dict(learning_rate=1e-5,), + optim=dict(learning_rate=1e-5, ), learner=dict(hook=dict(log_show_after_iter=1000)), ), collect=dict( data_type='d4rl', data_path=None, ), - eval=dict(evaluator=dict(eval_freq=-1,)), + eval=dict(evaluator=dict(eval_freq=-1, )), ), ) main_config = EasyDict(main_config) @@ -44,7 +40,7 @@ type='d4rl', import_names=['dizoo.d4rl.envs.d4rl_env'], ), - env_manager=dict(type='base',), + env_manager=dict(type='base', ), policy=dict( type='ibc', import_names=['ding.policy.ibc'], diff --git a/dizoo/d4rl/config/hopper_medium_expert_ibc_config.py b/dizoo/d4rl/config/hopper_medium_expert_ibc_config.py index e7a72984b6..0f040970e6 100644 --- a/dizoo/d4rl/config/hopper_medium_expert_ibc_config.py +++ b/dizoo/d4rl/config/hopper_medium_expert_ibc_config.py @@ -8,7 +8,7 @@ env=dict( env_id='hopper-medium-expert-v0', norm_obs=dict( - use_norm=True, + use_norm=True, offline_stats=dict(use_offline_stats=True, ), ), evaluator_env_num=8, @@ -18,23 +18,19 @@ ), policy=dict( cuda=cuda, - model=dict( - obs_shape=11, - action_shape=3, - stochastic_optim=dict(type='dfo',) - ), + model=dict(obs_shape=11, action_shape=3, stochastic_optim=dict(type='dfo', )), learn=dict( multi_gpu=multi_gpu, train_epoch=15, batch_size=256, - optim=dict(learning_rate=1e-5,), + optim=dict(learning_rate=1e-5, ), learner=dict(hook=dict(log_show_after_iter=1000)), ), collect=dict( data_type='d4rl', data_path=None, ), - eval=dict(evaluator=dict(eval_freq=-1,)), + eval=dict(evaluator=dict(eval_freq=-1, )), ), ) main_config = EasyDict(main_config) @@ -44,7 +40,7 @@ type='d4rl', import_names=['dizoo.d4rl.envs.d4rl_env'], ), - env_manager=dict(type='base',), + env_manager=dict(type='base', ), policy=dict( type='ibc', import_names=['ding.policy.ibc'], diff --git a/dizoo/d4rl/config/hopper_medium_expert_ibc_mcmc_config.py b/dizoo/d4rl/config/hopper_medium_expert_ibc_mcmc_config.py index e5f6f3dbb1..478e0c5d44 100644 --- a/dizoo/d4rl/config/hopper_medium_expert_ibc_mcmc_config.py +++ b/dizoo/d4rl/config/hopper_medium_expert_ibc_mcmc_config.py @@ -8,7 +8,7 @@ env=dict( env_id='hopper-medium-expert-v0', norm_obs=dict( - use_norm=True, + use_norm=True, offline_stats=dict(use_offline_stats=True, ), ), evaluator_env_num=8, @@ -18,23 +18,19 @@ ), policy=dict( cuda=cuda, - model=dict( - obs_shape=11, - action_shape=3, - stochastic_optim=dict(type='mcmc',) - ), + model=dict(obs_shape=11, action_shape=3, stochastic_optim=dict(type='mcmc', )), learn=dict( multi_gpu=multi_gpu, train_epoch=15, batch_size=256, - optim=dict(learning_rate=1e-5,), + optim=dict(learning_rate=1e-5, ), learner=dict(hook=dict(log_show_after_iter=1000)), ), collect=dict( data_type='d4rl', data_path=None, ), - eval=dict(evaluator=dict(eval_freq=-1,)), + eval=dict(evaluator=dict(eval_freq=-1, )), ), ) main_config = EasyDict(main_config) @@ -44,7 +40,7 @@ type='d4rl', import_names=['dizoo.d4rl.envs.d4rl_env'], ), - env_manager=dict(type='base',), + env_manager=dict(type='base', ), policy=dict( type='ibc', import_names=['ding.policy.ibc'], diff --git a/dizoo/d4rl/config/hopper_medium_expert_iql_config.py b/dizoo/d4rl/config/hopper_medium_expert_iql_config.py new file mode 100644 index 0000000000..2eebce2771 --- /dev/null +++ b/dizoo/d4rl/config/hopper_medium_expert_iql_config.py @@ -0,0 +1,53 @@ +# You can conduct Experiments on D4RL with this config file through the following command: +# cd ../entry && python d4rl_iql_main.py +from easydict import EasyDict + +main_config = dict( + exp_name="hopper_medium_expert_iql_seed0", + env=dict( + env_id='hopper-medium-expert-v2', + collector_env_num=1, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + stop_value=6000, + reward_norm="iql_locomotion", + ), + policy=dict( + cuda=True, + model=dict( + obs_shape=11, + action_shape=3, + ), + learn=dict( + data_path=None, + train_epoch=30000, + batch_size=4096, + learning_rate_q=3e-4, + learning_rate_policy=1e-4, + beta=0.05, + tau=0.7, + ), + collect=dict(data_type='d4rl', ), + eval=dict(evaluator=dict(eval_freq=5000, )), + other=dict(replay_buffer=dict(replay_buffer_size=2000000, ), ), + ), +) + +main_config = EasyDict(main_config) +main_config = main_config + +create_config = dict( + env=dict( + type='d4rl', + import_names=['dizoo.d4rl.envs.d4rl_env'], + ), + env_manager=dict(type='base'), + policy=dict( + type='iql', + import_names=['ding.policy.iql'], + ), + replay_buffer=dict(type='naive', ), +) +create_config = EasyDict(create_config) +create_config = create_config diff --git a/dizoo/d4rl/config/hopper_medium_expert_pd_config.py b/dizoo/d4rl/config/hopper_medium_expert_pd_config.py new file mode 100755 index 0000000000..d7eeeeca5b --- /dev/null +++ b/dizoo/d4rl/config/hopper_medium_expert_pd_config.py @@ -0,0 +1,99 @@ +from easydict import EasyDict + +main_config = dict( + exp_name="hopper_medium_expert_pd_seed0", + env=dict( + env_id='hopper-medium-expert-v2', + collector_env_num=1, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + returns_scale=1.0, + termination_penalty=-100, + max_path_length=1000, + use_padding=True, + include_returns=True, + normed=False, + stop_value=8000, + horizon=32, + obs_dim=11, + action_dim=3, + ), + policy=dict( + cuda=True, + model=dict( + diffuser_model='GaussianDiffusion', + diffuser_model_cfg=dict( + model='DiffusionUNet1d', + model_cfg=dict( + transition_dim=14, + dim=32, + dim_mults=[1, 2, 4, 8], + returns_condition=False, + kernel_size=5, + attention=False, + ), + horizon=32, + obs_dim=11, + action_dim=3, + n_timesteps=20, + predict_epsilon=False, + loss_discount=1, + action_weight=10, + ), + value_model='ValueDiffusion', + value_model_cfg=dict( + model='TemporalValue', + model_cfg=dict( + horizon=32, + transition_dim=14, + dim=32, + dim_mults=[1, 2, 4, 8], + kernel_size=5, + ), + horizon=32, + obs_dim=11, + action_dim=3, + n_timesteps=20, + predict_epsilon=True, + loss_discount=1, + ), + n_guide_steps=2, + scale=0.0001, + t_stopgrad=4, + scale_grad_by_std=True, + ), + normalizer='GaussianNormalizer', + learn=dict( + data_path=None, + train_epoch=60000, + gradient_accumulate_every=2, + batch_size=32, + learning_rate=2e-4, + discount_factor=0.99, + plan_batch_size=64, + learner=dict(hook=dict(save_ckpt_after_iter=1000000000, )), + ), + collect=dict(data_type='diffuser_traj', ), + eval=dict( + evaluator=dict(eval_freq=500, ), + test_ret=0.9, + ), + other=dict(replay_buffer=dict(replay_buffer_size=2000000, ), ), + ), +) + +main_config = EasyDict(main_config) +main_config = main_config + +create_config = dict( + env=dict( + type='d4rl', + import_names=['dizoo.d4rl.envs.d4rl_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='pd', ), + replay_buffer=dict(type='naive', ), +) +create_config = EasyDict(create_config) +create_config = create_config diff --git a/dizoo/d4rl/config/hopper_medium_expert_qgpo_config.py b/dizoo/d4rl/config/hopper_medium_expert_qgpo_config.py new file mode 100644 index 0000000000..d47656afa2 --- /dev/null +++ b/dizoo/d4rl/config/hopper_medium_expert_qgpo_config.py @@ -0,0 +1,47 @@ +from easydict import EasyDict + +main_config = dict( + exp_name='hopper_medium_expert_v2_QGPO_seed0', + env=dict( + env_id="hopper-medium-expert-v2", + evaluator_env_num=8, + n_evaluator_episode=8, + ), + dataset=dict(env_id="hopper-medium-expert-v2", ), + policy=dict( + cuda=True, + on_policy=False, + #load_path='./hopper_medium_expert_v2_QGPO_seed0/ckpt/iteration_600000.pth.tar', + model=dict( + qgpo_critic=dict( + alpha=3, + q_alpha=1, + ), + device='cuda', + obs_dim=11, + action_dim=3, + ), + learn=dict( + learning_rate=1e-4, + batch_size=4096, + batch_size_q=256, + M=16, + diffusion_steps=15, + behavior_policy_stop_training_iter=600000, + energy_guided_policy_begin_training_iter=600000, + q_value_stop_training_iter=1100000, + ), + eval=dict( + guidance_scale=[0.0, 1.0, 2.0, 3.0, 5.0, 8.0, 10.0], + diffusion_steps=15, + evaluator=dict(eval_freq=50000, ), + ), + ), +) +main_config = EasyDict(main_config) + +create_config = dict( + env_manager=dict(type='base'), + policy=dict(type='qgpo', ), +) +create_config = EasyDict(create_config) diff --git a/dizoo/d4rl/config/hopper_medium_expert_td3bc_config.py b/dizoo/d4rl/config/hopper_medium_expert_td3bc_config.py index 19531debad..b51ed523fa 100644 --- a/dizoo/d4rl/config/hopper_medium_expert_td3bc_config.py +++ b/dizoo/d4rl/config/hopper_medium_expert_td3bc_config.py @@ -5,11 +5,7 @@ main_config = dict( exp_name='hopper_medium_expert_td3-bc_seed0', env=dict( - env_id='hopper-medium-expert-v0', - norm_obs=dict( - use_norm=True, - offline_stats=dict(use_offline_stats=True, ), - ), + env_id='hopper-medium-expert-v2', collector_env_num=1, evaluator_env_num=8, use_act_scale=True, @@ -17,6 +13,7 @@ stop_value=6000, ), policy=dict( + cuda=True, model=dict( obs_shape=11, action_shape=3, diff --git a/dizoo/d4rl/config/hopper_medium_iql_config.py b/dizoo/d4rl/config/hopper_medium_iql_config.py new file mode 100644 index 0000000000..61dbb5fac3 --- /dev/null +++ b/dizoo/d4rl/config/hopper_medium_iql_config.py @@ -0,0 +1,53 @@ +# You can conduct Experiments on D4RL with this config file through the following command: +# cd ../entry && python d4rl_iql_main.py +from easydict import EasyDict + +main_config = dict( + exp_name="hopper_medium_iql_seed0", + env=dict( + env_id='hopper-medium-v2', + collector_env_num=1, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + stop_value=6000, + reward_norm="iql_locomotion", + ), + policy=dict( + cuda=True, + model=dict( + obs_shape=11, + action_shape=3, + ), + learn=dict( + data_path=None, + train_epoch=30000, + batch_size=4096, + learning_rate_q=3e-4, + learning_rate_policy=1e-4, + beta=0.05, + tau=0.7, + ), + collect=dict(data_type='d4rl', ), + eval=dict(evaluator=dict(eval_freq=5000, )), + other=dict(replay_buffer=dict(replay_buffer_size=2000000, ), ), + ), +) + +main_config = EasyDict(main_config) +main_config = main_config + +create_config = dict( + env=dict( + type='d4rl', + import_names=['dizoo.d4rl.envs.d4rl_env'], + ), + env_manager=dict(type='base'), + policy=dict( + type='iql', + import_names=['ding.policy.iql'], + ), + replay_buffer=dict(type='naive', ), +) +create_config = EasyDict(create_config) +create_config = create_config diff --git a/dizoo/d4rl/config/hopper_medium_pd_config.py b/dizoo/d4rl/config/hopper_medium_pd_config.py new file mode 100755 index 0000000000..47f1c36ce0 --- /dev/null +++ b/dizoo/d4rl/config/hopper_medium_pd_config.py @@ -0,0 +1,99 @@ +from easydict import EasyDict + +main_config = dict( + exp_name="hopper_medium_pd_seed0", + env=dict( + env_id='hopper-medium-v2', + collector_env_num=1, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + returns_scale=1.0, + termination_penalty=-100, + max_path_length=1000, + use_padding=True, + include_returns=True, + normed=False, + stop_value=8000, + horizon=32, + obs_dim=11, + action_dim=3, + ), + policy=dict( + cuda=True, + model=dict( + diffuser_model='GaussianDiffusion', + diffuser_model_cfg=dict( + model='DiffusionUNet1d', + model_cfg=dict( + transition_dim=14, + dim=32, + dim_mults=[1, 2, 4, 8], + returns_condition=False, + kernel_size=5, + attention=False, + ), + horizon=32, + obs_dim=11, + action_dim=3, + n_timesteps=20, + predict_epsilon=False, + loss_discount=1, + action_weight=10, + ), + value_model='ValueDiffusion', + value_model_cfg=dict( + model='TemporalValue', + model_cfg=dict( + horizon=32, + transition_dim=14, + dim=32, + dim_mults=[1, 2, 4, 8], + kernel_size=5, + ), + horizon=32, + obs_dim=11, + action_dim=3, + n_timesteps=20, + predict_epsilon=True, + loss_discount=1, + ), + n_guide_steps=2, + scale=0.1, + t_stopgrad=2, + scale_grad_by_std=True, + ), + normalizer='GaussianNormalizer', + learn=dict( + data_path=None, + train_epoch=60000, + gradient_accumulate_every=2, + batch_size=32, + learning_rate=2e-4, + discount_factor=0.99, + plan_batch_size=64, + learner=dict(hook=dict(save_ckpt_after_iter=1000000000, )), + ), + collect=dict(data_type='diffuser_traj', ), + eval=dict( + evaluator=dict(eval_freq=500, ), + test_ret=0.9, + ), + other=dict(replay_buffer=dict(replay_buffer_size=2000000, ), ), + ), +) + +main_config = EasyDict(main_config) +main_config = main_config + +create_config = dict( + env=dict( + type='d4rl', + import_names=['dizoo.d4rl.envs.d4rl_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='pd', ), + replay_buffer=dict(type='naive', ), +) +create_config = EasyDict(create_config) +create_config = create_config diff --git a/dizoo/d4rl/config/hopper_medium_replay_dt_config.py b/dizoo/d4rl/config/hopper_medium_replay_dt_config.py index fd238c915e..a2615ba1b9 100644 --- a/dizoo/d4rl/config/hopper_medium_replay_dt_config.py +++ b/dizoo/d4rl/config/hopper_medium_replay_dt_config.py @@ -2,34 +2,35 @@ from copy import deepcopy hopper_dt_config = dict( - exp_name='hopper_medium_replay_dt_seed0', + exp_name='dt_log/d4rl/hopper/hopper_medium_replay_dt_seed0', env=dict( env_id='Hopper-v3', - norm_obs=dict(use_norm=False, ), - norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, use_act_scale=True, n_evaluator_episode=8, - stop_value=6000, + stop_value=3600, + ), + dataset=dict( + env_type='mujoco', + rtg_scale=1000, + context_len=30, + data_dir_prefix='d4rl/hopper_medium_replay-v2.pkl', ), policy=dict( - stop_value=6000, cuda=True, + stop_value=3600, + state_mean=None, + state_std=None, + evaluator_env_num=8, env_name='Hopper-v3', - rtg_target=6000, # max target return to go + rtg_target=3600, # max target return to go max_eval_ep_len=1000, # max lenght of one episode - num_eval_ep=10, # num of evaluation episode - batch_size=64, wt_decay=1e-4, warmup_steps=10000, - num_updates_per_iter=100, context_len=20, - n_blocks=3, - embed_dim=128, - n_heads=1, - dropout_p=0.1, - log_dir='/home/wangzilin/research/dt/DI-engine/dizoo/d4rl/dt_data/hopper_medium_replay_dt_log', + weight_decay=0.1, + clip_grad_norm_p=0.25, model=dict( state_dim=11, act_dim=3, @@ -40,26 +41,13 @@ drop_p=0.1, continuous=True, ), - discount_factor=0.999, - nstep=3, - learn=dict( - dataset_path='/mnt/lustre/wangzilin/d4rl_data/hopper-medium-replay-v2.pkl', - learning_rate=0.0001, - target_update_freq=100, - kappa=1.0, - min_q_weight=4.0, - ), - collect=dict(unroll_len=1, ), - eval=dict(evaluator=dict(evalu_freq=100, ), ), - other=dict( - eps=dict( - type='exp', - start=0.95, - end=0.1, - decay=10000, - ), - replay_buffer=dict(replay_buffer_size=1000, ), + batch_size=64, + learning_rate=1e-4, + collect=dict( + data_type='d4rl_trajectory', + unroll_len=1, ), + eval=dict(evaluator=dict(eval_freq=100, ), ), ), ) diff --git a/dizoo/d4rl/config/hopper_medium_replay_iql_config.py b/dizoo/d4rl/config/hopper_medium_replay_iql_config.py new file mode 100644 index 0000000000..df96a84aea --- /dev/null +++ b/dizoo/d4rl/config/hopper_medium_replay_iql_config.py @@ -0,0 +1,53 @@ +# You can conduct Experiments on D4RL with this config file through the following command: +# cd ../entry && python d4rl_iql_main.py +from easydict import EasyDict + +main_config = dict( + exp_name="hopper_medium_replay_iql_seed0", + env=dict( + env_id='hopper-medium-replay-v2', + collector_env_num=1, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + stop_value=6000, + reward_norm="iql_locomotion", + ), + policy=dict( + cuda=True, + model=dict( + obs_shape=11, + action_shape=3, + ), + learn=dict( + data_path=None, + train_epoch=30000, + batch_size=4096, + learning_rate_q=3e-4, + learning_rate_policy=1e-4, + beta=0.05, + tau=0.7, + ), + collect=dict(data_type='d4rl', ), + eval=dict(evaluator=dict(eval_freq=5000, )), + other=dict(replay_buffer=dict(replay_buffer_size=2000000, ), ), + ), +) + +main_config = EasyDict(main_config) +main_config = main_config + +create_config = dict( + env=dict( + type='d4rl', + import_names=['dizoo.d4rl.envs.d4rl_env'], + ), + env_manager=dict(type='base'), + policy=dict( + type='iql', + import_names=['ding.policy.iql'], + ), + replay_buffer=dict(type='naive', ), +) +create_config = EasyDict(create_config) +create_config = create_config diff --git a/dizoo/d4rl/config/hopper_medium_replay_td3bc_config.py b/dizoo/d4rl/config/hopper_medium_replay_td3bc_config.py index 8f754781db..6ca9cdef06 100644 --- a/dizoo/d4rl/config/hopper_medium_replay_td3bc_config.py +++ b/dizoo/d4rl/config/hopper_medium_replay_td3bc_config.py @@ -5,11 +5,7 @@ main_config = dict( exp_name='hopper_medium_replay_td3-bc_seed0', env=dict( - env_id='hopper-medium-replay-v0', - norm_obs=dict( - use_norm=True, - offline_stats=dict(use_offline_stats=True, ), - ), + env_id='hopper-medium-replay-v2', collector_env_num=1, evaluator_env_num=8, use_act_scale=True, @@ -17,6 +13,7 @@ stop_value=6000, ), policy=dict( + cuda=True, model=dict( obs_shape=11, action_shape=3, diff --git a/dizoo/d4rl/config/hopper_medium_td3bc_config.py b/dizoo/d4rl/config/hopper_medium_td3bc_config.py index cbf5fcce19..fd318545e6 100644 --- a/dizoo/d4rl/config/hopper_medium_td3bc_config.py +++ b/dizoo/d4rl/config/hopper_medium_td3bc_config.py @@ -5,11 +5,7 @@ main_config = dict( exp_name='hopper_medium_td3-bc_seed0', env=dict( - env_id='hopper-medium-v0', - norm_obs=dict( - use_norm=True, - offline_stats=dict(use_offline_stats=True, ), - ), + env_id='hopper-medium-v2', collector_env_num=1, evaluator_env_num=8, use_act_scale=True, @@ -17,6 +13,7 @@ stop_value=6000, ), policy=dict( + cuda=True, model=dict( obs_shape=11, action_shape=3, diff --git a/dizoo/d4rl/config/hopper_random_td3bc_config.py b/dizoo/d4rl/config/hopper_random_td3bc_config.py index 8cf796b5fb..abc1f3ab60 100644 --- a/dizoo/d4rl/config/hopper_random_td3bc_config.py +++ b/dizoo/d4rl/config/hopper_random_td3bc_config.py @@ -7,7 +7,7 @@ env=dict( env_id='hopper-random-v0', norm_obs=dict( - use_norm=True, + use_norm=True, offline_stats=dict(use_offline_stats=True, ), ), collector_env_num=1, @@ -17,6 +17,7 @@ stop_value=6000, ), policy=dict( + cuda=True, model=dict( obs_shape=11, action_shape=3, diff --git a/dizoo/d4rl/config/kitchen_complete_bc_config.py b/dizoo/d4rl/config/kitchen_complete_bc_config.py index 7160885da3..413696993d 100644 --- a/dizoo/d4rl/config/kitchen_complete_bc_config.py +++ b/dizoo/d4rl/config/kitchen_complete_bc_config.py @@ -8,7 +8,7 @@ env=dict( env_id='kitchen-complete-v0', norm_obs=dict( - use_norm=True, + use_norm=True, offline_stats=dict(use_offline_stats=True, ), ), evaluator_env_num=8, @@ -19,7 +19,7 @@ policy=dict( cuda=cuda, continuous=True, - loss_type='mse_loss', + loss_type='mse_loss', model=dict( obs_shape=60, action_shape=9, @@ -38,7 +38,7 @@ data_type='d4rl', data_path=None, ), - eval=dict(evaluator=dict(eval_freq=1000,)), + eval=dict(evaluator=dict(eval_freq=1000, )), ), ) main_config = EasyDict(main_config) @@ -48,7 +48,7 @@ type='d4rl', import_names=['dizoo.d4rl.envs.d4rl_env'], ), - env_manager=dict(type='base',), + env_manager=dict(type='base', ), policy=dict( type='bc', import_names=['ding.policy.bc'], diff --git a/dizoo/d4rl/config/kitchen_complete_ibc_ar_config.py b/dizoo/d4rl/config/kitchen_complete_ibc_ar_config.py index 403dc52eff..bbb7198af0 100644 --- a/dizoo/d4rl/config/kitchen_complete_ibc_ar_config.py +++ b/dizoo/d4rl/config/kitchen_complete_ibc_ar_config.py @@ -8,7 +8,7 @@ env=dict( env_id='kitchen-complete-v0', norm_obs=dict( - use_norm=True, + use_norm=True, offline_stats=dict(use_offline_stats=True, ), ), evaluator_env_num=8, @@ -18,23 +18,19 @@ ), policy=dict( cuda=cuda, - model=dict( - obs_shape=60, - action_shape=9, - stochastic_optim=dict(type='ardfo',) - ), + model=dict(obs_shape=60, action_shape=9, stochastic_optim=dict(type='ardfo', )), learn=dict( multi_gpu=multi_gpu, train_epoch=1000, batch_size=256, - optim=dict(learning_rate=1e-5,), + optim=dict(learning_rate=1e-5, ), learner=dict(hook=dict(log_show_after_iter=100)), ), collect=dict( data_type='d4rl', data_path=None, ), - eval=dict(evaluator=dict(eval_freq=1000,)), + eval=dict(evaluator=dict(eval_freq=1000, )), ), ) main_config = EasyDict(main_config) @@ -44,7 +40,7 @@ type='d4rl', import_names=['dizoo.d4rl.envs.d4rl_env'], ), - env_manager=dict(type='base',), + env_manager=dict(type='base', ), policy=dict( type='ibc', import_names=['ding.policy.ibc'], diff --git a/dizoo/d4rl/config/kitchen_complete_ibc_config.py b/dizoo/d4rl/config/kitchen_complete_ibc_config.py index 5c02f04a81..1606cb7792 100644 --- a/dizoo/d4rl/config/kitchen_complete_ibc_config.py +++ b/dizoo/d4rl/config/kitchen_complete_ibc_config.py @@ -8,7 +8,7 @@ env=dict( env_id='kitchen-complete-v0', norm_obs=dict( - use_norm=True, + use_norm=True, offline_stats=dict(use_offline_stats=True, ), ), evaluator_env_num=8, @@ -18,23 +18,19 @@ ), policy=dict( cuda=cuda, - model=dict( - obs_shape=60, - action_shape=9, - stochastic_optim=dict(type='dfo',) - ), + model=dict(obs_shape=60, action_shape=9, stochastic_optim=dict(type='dfo', )), learn=dict( multi_gpu=multi_gpu, train_epoch=1000, batch_size=256, - optim=dict(learning_rate=1e-5,), + optim=dict(learning_rate=1e-5, ), learner=dict(hook=dict(log_show_after_iter=100)), ), collect=dict( data_type='d4rl', data_path=None, ), - eval=dict(evaluator=dict(eval_freq=1000,)), + eval=dict(evaluator=dict(eval_freq=1000, )), ), ) main_config = EasyDict(main_config) @@ -44,7 +40,7 @@ type='d4rl', import_names=['dizoo.d4rl.envs.d4rl_env'], ), - env_manager=dict(type='base',), + env_manager=dict(type='base', ), policy=dict( type='ibc', import_names=['ding.policy.ibc'], diff --git a/dizoo/d4rl/config/kitchen_complete_ibc_mcmc_config.py b/dizoo/d4rl/config/kitchen_complete_ibc_mcmc_config.py index d93c5eb737..14924d5257 100644 --- a/dizoo/d4rl/config/kitchen_complete_ibc_mcmc_config.py +++ b/dizoo/d4rl/config/kitchen_complete_ibc_mcmc_config.py @@ -8,7 +8,7 @@ env=dict( env_id='kitchen-complete-v0', norm_obs=dict( - use_norm=True, + use_norm=True, offline_stats=dict(use_offline_stats=True, ), ), evaluator_env_num=8, @@ -18,23 +18,19 @@ ), policy=dict( cuda=cuda, - model=dict( - obs_shape=60, - action_shape=9, - stochastic_optim=dict(type='mcmc',) - ), + model=dict(obs_shape=60, action_shape=9, stochastic_optim=dict(type='mcmc', )), learn=dict( multi_gpu=multi_gpu, train_epoch=1000, batch_size=256, - optim=dict(learning_rate=1e-5,), + optim=dict(learning_rate=1e-5, ), learner=dict(hook=dict(log_show_after_iter=100)), ), collect=dict( data_type='d4rl', data_path=None, ), - eval=dict(evaluator=dict(eval_freq=1000,)), + eval=dict(evaluator=dict(eval_freq=1000, )), ), ) main_config = EasyDict(main_config) @@ -44,7 +40,7 @@ type='d4rl', import_names=['dizoo.d4rl.envs.d4rl_env'], ), - env_manager=dict(type='base',), + env_manager=dict(type='base', ), policy=dict( type='ibc', import_names=['ding.policy.ibc'], diff --git a/dizoo/d4rl/config/maze2d_large_pd_config.py b/dizoo/d4rl/config/maze2d_large_pd_config.py new file mode 100755 index 0000000000..a722488b89 --- /dev/null +++ b/dizoo/d4rl/config/maze2d_large_pd_config.py @@ -0,0 +1,82 @@ +from easydict import EasyDict + +main_config = dict( + exp_name="maze2d_large_pd_seed0", + env=dict( + env_id='maze2d-large-v1', + collector_env_num=1, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + returns_scale=1.0, + termination_penalty=None, + max_path_length=40000, + use_padding=False, + include_returns=False, + normed=False, + stop_value=500, + horizon=384, + obs_dim=4, + action_dim=2, + ), + policy=dict( + cuda=True, + model=dict( + diffuser_model='GaussianDiffusion', + diffuser_model_cfg=dict( + model='DiffusionUNet1d', + model_cfg=dict( + transition_dim=6, + dim=32, + dim_mults=[1, 4, 8], + returns_condition=False, + kernel_size=5, + attention=False, + ), + horizon=384, + obs_dim=4, + action_dim=2, + n_timesteps=256, + predict_epsilon=False, + loss_discount=1, + clip_denoised=True, + action_weight=1, + ), + value_model=None, + value_model_cfg=None, + ), + normalizer='LimitsNormalizer', + learn=dict( + data_path=None, + train_epoch=60000, + gradient_accumulate_every=2, + batch_size=32, + learning_rate=2e-4, + discount_factor=0.99, + plan_batch_size=1, + include_returns=False, + learner=dict(hook=dict(save_ckpt_after_iter=1000000000, )), + ), + collect=dict(data_type='diffuser_traj', ), + eval=dict( + evaluator=dict(eval_freq=500, ), + test_ret=0.9, + ), + other=dict(replay_buffer=dict(replay_buffer_size=2000000, ), ), + ), +) + +main_config = EasyDict(main_config) +main_config = main_config + +create_config = dict( + env=dict( + type='d4rl', + import_names=['dizoo.d4rl.envs.d4rl_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='pd', ), + replay_buffer=dict(type='naive', ), +) +create_config = EasyDict(create_config) +create_config = create_config diff --git a/dizoo/d4rl/config/maze2d_medium_pd_config.py b/dizoo/d4rl/config/maze2d_medium_pd_config.py new file mode 100755 index 0000000000..2a3d43d443 --- /dev/null +++ b/dizoo/d4rl/config/maze2d_medium_pd_config.py @@ -0,0 +1,82 @@ +from easydict import EasyDict + +main_config = dict( + exp_name="maze2d_medium_pd_seed0", + env=dict( + env_id='maze2d-medium-v1', + collector_env_num=1, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + returns_scale=1.0, + termination_penalty=None, + max_path_length=40000, + use_padding=False, + include_returns=False, + normed=False, + stop_value=357, + horizon=256, + obs_dim=4, + action_dim=2, + ), + policy=dict( + cuda=True, + model=dict( + diffuser_model='GaussianDiffusion', + diffuser_model_cfg=dict( + model='DiffusionUNet1d', + model_cfg=dict( + transition_dim=6, + dim=32, + dim_mults=[1, 4, 8], + returns_condition=False, + kernel_size=5, + attention=False, + ), + horizon=256, + obs_dim=4, + action_dim=2, + n_timesteps=256, + predict_epsilon=False, + loss_discount=1, + clip_denoised=True, + action_weight=1, + ), + value_model=None, + value_model_cfg=None, + ), + normalizer='LimitsNormalizer', + learn=dict( + data_path=None, + train_epoch=60000, + gradient_accumulate_every=2, + batch_size=32, + learning_rate=2e-4, + discount_factor=0.99, + plan_batch_size=1, + include_returns=False, + learner=dict(hook=dict(save_ckpt_after_iter=1000000000, )), + ), + collect=dict(data_type='diffuser_traj', ), + eval=dict( + evaluator=dict(eval_freq=500, ), + test_ret=0.9, + ), + other=dict(replay_buffer=dict(replay_buffer_size=2000000, ), ), + ), +) + +main_config = EasyDict(main_config) +main_config = main_config + +create_config = dict( + env=dict( + type='d4rl', + import_names=['dizoo.d4rl.envs.d4rl_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='pd', ), + replay_buffer=dict(type='naive', ), +) +create_config = EasyDict(create_config) +create_config = create_config diff --git a/dizoo/d4rl/config/maze2d_umaze_pd_config.py b/dizoo/d4rl/config/maze2d_umaze_pd_config.py new file mode 100755 index 0000000000..a10fef5844 --- /dev/null +++ b/dizoo/d4rl/config/maze2d_umaze_pd_config.py @@ -0,0 +1,82 @@ +from easydict import EasyDict + +main_config = dict( + exp_name="maze2d_umaze_pd_seed0", + env=dict( + env_id='maze2d-umaze-v1', + collector_env_num=1, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + returns_scale=1.0, + termination_penalty=None, + max_path_length=40000, + use_padding=False, + include_returns=False, + normed=False, + stop_value=190, + horizon=128, + obs_dim=4, + action_dim=2, + ), + policy=dict( + cuda=True, + model=dict( + diffuser_model='GaussianDiffusion', + diffuser_model_cfg=dict( + model='DiffusionUNet1d', + model_cfg=dict( + transition_dim=6, + dim=32, + dim_mults=[1, 4, 8], + returns_condition=False, + kernel_size=5, + attention=False, + ), + horizon=128, + obs_dim=4, + action_dim=2, + n_timesteps=64, + predict_epsilon=False, + loss_discount=1, + clip_denoised=True, + action_weight=1, + ), + value_model=None, + value_model_cfg=None, + ), + normalizer='LimitsNormalizer', + learn=dict( + data_path=None, + train_epoch=60000, + gradient_accumulate_every=2, + batch_size=32, + learning_rate=2e-4, + discount_factor=0.99, + plan_batch_size=1, + include_returns=False, + learner=dict(hook=dict(save_ckpt_after_iter=1000000000, )), + ), + collect=dict(data_type='diffuser_traj', ), + eval=dict( + evaluator=dict(eval_freq=500, ), + test_ret=0.9, + ), + other=dict(replay_buffer=dict(replay_buffer_size=2000000, ), ), + ), +) + +main_config = EasyDict(main_config) +main_config = main_config + +create_config = dict( + env=dict( + type='d4rl', + import_names=['dizoo.d4rl.envs.d4rl_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='pd', ), + replay_buffer=dict(type='naive', ), +) +create_config = EasyDict(create_config) +create_config = create_config diff --git a/dizoo/d4rl/config/pen_human_bc_config.py b/dizoo/d4rl/config/pen_human_bc_config.py index 6779ffd934..215b706ffc 100644 --- a/dizoo/d4rl/config/pen_human_bc_config.py +++ b/dizoo/d4rl/config/pen_human_bc_config.py @@ -8,7 +8,7 @@ env=dict( env_id='pen-human-v0', norm_obs=dict( - use_norm=True, + use_norm=True, offline_stats=dict(use_offline_stats=True, ), ), evaluator_env_num=8, @@ -38,7 +38,7 @@ data_type='d4rl', data_path=None, ), - eval=dict(evaluator=dict(eval_freq=1000,)), + eval=dict(evaluator=dict(eval_freq=1000, )), ), ) main_config = EasyDict(main_config) @@ -48,7 +48,7 @@ type='d4rl', import_names=['dizoo.d4rl.envs.d4rl_env'], ), - env_manager=dict(type='base',), + env_manager=dict(type='base', ), policy=dict( type='bc', import_names=['ding.policy.bc'], diff --git a/dizoo/d4rl/config/pen_human_ibc_ar_config.py b/dizoo/d4rl/config/pen_human_ibc_ar_config.py index b75e3b9f11..4f59733fd5 100644 --- a/dizoo/d4rl/config/pen_human_ibc_ar_config.py +++ b/dizoo/d4rl/config/pen_human_ibc_ar_config.py @@ -8,7 +8,7 @@ env=dict( env_id='pen-human-v0', norm_obs=dict( - use_norm=True, + use_norm=True, offline_stats=dict(use_offline_stats=True, ), ), evaluator_env_num=8, @@ -19,24 +19,20 @@ policy=dict( cuda=cuda, model=dict( - obs_shape=45, - action_shape=24, - hidden_size=128, - hidden_layer_num=4, - stochastic_optim=dict(type='ardfo',) + obs_shape=45, action_shape=24, hidden_size=128, hidden_layer_num=4, stochastic_optim=dict(type='ardfo', ) ), learn=dict( multi_gpu=multi_gpu, train_epoch=1000, batch_size=256, - optim=dict(learning_rate=1e-5,), + optim=dict(learning_rate=1e-5, ), learner=dict(hook=dict(log_show_after_iter=100)), ), collect=dict( data_type='d4rl', data_path=None, ), - eval=dict(evaluator=dict(eval_freq=1000,)), + eval=dict(evaluator=dict(eval_freq=1000, )), ), ) main_config = EasyDict(main_config) @@ -46,7 +42,7 @@ type='d4rl', import_names=['dizoo.d4rl.envs.d4rl_env'], ), - env_manager=dict(type='base',), + env_manager=dict(type='base', ), policy=dict( type='ibc', import_names=['ding.policy.ibc'], diff --git a/dizoo/d4rl/config/pen_human_ibc_config.py b/dizoo/d4rl/config/pen_human_ibc_config.py index 207487d921..9ed4f6d17b 100644 --- a/dizoo/d4rl/config/pen_human_ibc_config.py +++ b/dizoo/d4rl/config/pen_human_ibc_config.py @@ -8,7 +8,7 @@ env=dict( env_id='pen-human-v0', norm_obs=dict( - use_norm=True, + use_norm=True, offline_stats=dict(use_offline_stats=True, ), ), evaluator_env_num=8, @@ -18,23 +18,19 @@ ), policy=dict( cuda=cuda, - model=dict( - obs_shape=45, - action_shape=24, - stochastic_optim=dict(type='dfo',) - ), + model=dict(obs_shape=45, action_shape=24, stochastic_optim=dict(type='dfo', )), learn=dict( multi_gpu=multi_gpu, train_epoch=1000, batch_size=256, - optim=dict(learning_rate=1e-5,), + optim=dict(learning_rate=1e-5, ), learner=dict(hook=dict(log_show_after_iter=100)), ), collect=dict( data_type='d4rl', data_path=None, ), - eval=dict(evaluator=dict(eval_freq=1000,)), + eval=dict(evaluator=dict(eval_freq=1000, )), ), ) main_config = EasyDict(main_config) @@ -44,7 +40,7 @@ type='d4rl', import_names=['dizoo.d4rl.envs.d4rl_env'], ), - env_manager=dict(type='base',), + env_manager=dict(type='base', ), policy=dict( type='ibc', import_names=['ding.policy.ibc'], diff --git a/dizoo/d4rl/config/pen_human_ibc_mcmc_config.py b/dizoo/d4rl/config/pen_human_ibc_mcmc_config.py index cee0f631fd..4dd6b37f90 100644 --- a/dizoo/d4rl/config/pen_human_ibc_mcmc_config.py +++ b/dizoo/d4rl/config/pen_human_ibc_mcmc_config.py @@ -8,7 +8,7 @@ env=dict( env_id='pen-human-v0', norm_obs=dict( - use_norm=True, + use_norm=True, offline_stats=dict(use_offline_stats=True, ), ), evaluator_env_num=8, @@ -18,23 +18,19 @@ ), policy=dict( cuda=cuda, - model=dict( - obs_shape=45, - action_shape=24, - stochastic_optim=dict(type='mcmc',) - ), + model=dict(obs_shape=45, action_shape=24, stochastic_optim=dict(type='mcmc', )), learn=dict( multi_gpu=multi_gpu, train_epoch=1000, batch_size=256, - optim=dict(learning_rate=1e-5,), + optim=dict(learning_rate=1e-5, ), learner=dict(hook=dict(log_show_after_iter=100)), ), collect=dict( data_type='d4rl', data_path=None, ), - eval=dict(evaluator=dict(eval_freq=1000,)), + eval=dict(evaluator=dict(eval_freq=1000, )), ), ) main_config = EasyDict(main_config) @@ -44,7 +40,7 @@ type='d4rl', import_names=['dizoo.d4rl.envs.d4rl_env'], ), - env_manager=dict(type='base',), + env_manager=dict(type='base', ), policy=dict( type='ibc', import_names=['ding.policy.ibc'], diff --git a/dizoo/d4rl/config/walker2d_expert_cql_config.py b/dizoo/d4rl/config/walker2d_expert_cql_config.py old mode 100644 new mode 100755 index 271dbca013..346dd1a1f2 --- a/dizoo/d4rl/config/walker2d_expert_cql_config.py +++ b/dizoo/d4rl/config/walker2d_expert_cql_config.py @@ -5,7 +5,7 @@ main_config = dict( exp_name="walker2d_expert_cql_seed0", env=dict( - env_id='walker2d-expert-v0', + env_id='walker2d-expert-v2', collector_env_num=1, evaluator_env_num=8, use_act_scale=True, @@ -15,8 +15,8 @@ policy=dict( cuda=True, model=dict( - obs_shape=11, - action_shape=3, + obs_shape=17, + action_shape=6, ), learn=dict( data_path=None, diff --git a/dizoo/d4rl/config/walker2d_expert_dt_config.py b/dizoo/d4rl/config/walker2d_expert_dt_config.py index d50670dad7..3658f8ce03 100644 --- a/dizoo/d4rl/config/walker2d_expert_dt_config.py +++ b/dizoo/d4rl/config/walker2d_expert_dt_config.py @@ -1,38 +1,39 @@ from easydict import EasyDict from copy import deepcopy -walker2d_dt_config = dict( - exp_name='walker2d_expert_dt_seed0', +walk2d_dt_config = dict( + exp_name='dt_log/d4rl/walk2d/walk2d_expert_dt_seed0', env=dict( - env_id='Walker2d-v3', - norm_obs=dict(use_norm=False, ), - norm_reward=dict(use_norm=False, ), + env_id='Walk2d-v3', collector_env_num=1, evaluator_env_num=8, use_act_scale=True, n_evaluator_episode=8, - stop_value=6000, + stop_value=5000, + ), + dataset=dict( + env_type='mujoco', + rtg_scale=1000, + context_len=30, + data_dir_prefix='d4rl/walk2d_expert-v2.pkl', ), policy=dict( - stop_value=6000, cuda=True, - env_name='Walker2d-v3', - rtg_target=6000, # max target return to go + stop_value=5000, + state_mean=None, + state_std=None, + evaluator_env_num=8, + env_name='Walk2d-v3', + rtg_target=5000, # max target return to go max_eval_ep_len=1000, # max lenght of one episode - num_eval_ep=10, # num of evaluation episode - batch_size=64, wt_decay=1e-4, warmup_steps=10000, - num_updates_per_iter=100, context_len=20, - n_blocks=3, - embed_dim=128, - n_heads=1, - dropout_p=0.1, - log_dir='/home/wangzilin/research/dt/DI-engine/dizoo/d4rl/dt_data/walker2d_expert_dt_log', + weight_decay=0.1, + clip_grad_norm_p=0.25, model=dict( - state_dim=17, - act_dim=6, + state_dim=11, + act_dim=3, n_blocks=3, h_dim=128, context_len=20, @@ -40,32 +41,19 @@ drop_p=0.1, continuous=True, ), - discount_factor=0.999, - nstep=3, - learn=dict( - dataset_path='/mnt/lustre/wangzilin/d4rl_data/walker2d-expert-v2.pkl', - learning_rate=0.0001, - target_update_freq=100, - kappa=1.0, - min_q_weight=4.0, - ), - collect=dict(unroll_len=1, ), - eval=dict(evaluator=dict(evalu_freq=100, ), ), - other=dict( - eps=dict( - type='exp', - start=0.95, - end=0.1, - decay=10000, - ), - replay_buffer=dict(replay_buffer_size=1000, ), + batch_size=64, + learning_rate=1e-4, + collect=dict( + data_type='d4rl_trajectory', + unroll_len=1, ), + eval=dict(evaluator=dict(eval_freq=100, ), ), ), ) -walker2d_dt_config = EasyDict(walker2d_dt_config) -main_config = walker2d_dt_config -walker2d_dt_create_config = dict( +walk2d_dt_config = EasyDict(walk2d_dt_config) +main_config = walk2d_dt_config +walk2d_dt_create_config = dict( env=dict( type='mujoco', import_names=['dizoo.mujoco.envs.mujoco_env'], @@ -73,8 +61,8 @@ env_manager=dict(type='subprocess'), policy=dict(type='dt'), ) -walker2d_dt_create_config = EasyDict(walker2d_dt_create_config) -create_config = walker2d_dt_create_config +walk2d_dt_create_config = EasyDict(walk2d_dt_create_config) +create_config = walk2d_dt_create_config if __name__ == "__main__": from ding.entry import serial_pipeline_dt diff --git a/dizoo/d4rl/config/walker2d_expert_td3bc_config.py b/dizoo/d4rl/config/walker2d_expert_td3bc_config.py old mode 100644 new mode 100755 index c12d58b230..3cae080e92 --- a/dizoo/d4rl/config/walker2d_expert_td3bc_config.py +++ b/dizoo/d4rl/config/walker2d_expert_td3bc_config.py @@ -5,9 +5,9 @@ main_config = dict( exp_name='halfcheetah_expert_td3-bc_seed0', env=dict( - env_id='halfcheetah-expert-v0', + env_id='halfcheetah-expert-v2', norm_obs=dict( - use_norm=True, + use_norm=True, offline_stats=dict(use_offline_stats=True, ), ), collector_env_num=1, @@ -17,9 +17,10 @@ stop_value=6000, ), policy=dict( + cuda=True, model=dict( - obs_shape=11, - action_shape=3, + obs_shape=17, + action_shape=6, ), learn=dict( train_epoch=30000, diff --git a/dizoo/d4rl/config/walker2d_medium_cql_config.py b/dizoo/d4rl/config/walker2d_medium_cql_config.py old mode 100644 new mode 100755 index 0ea01a8689..afacebae0b --- a/dizoo/d4rl/config/walker2d_medium_cql_config.py +++ b/dizoo/d4rl/config/walker2d_medium_cql_config.py @@ -5,7 +5,7 @@ main_config = dict( exp_name="walker2d_medium_cql_seed0", env=dict( - env_id='walker2d-medium-v0', + env_id='walker2d-medium-v2', collector_env_num=1, evaluator_env_num=8, use_act_scale=True, @@ -15,8 +15,8 @@ policy=dict( cuda=True, model=dict( - obs_shape=11, - action_shape=3, + obs_shape=17, + action_shape=6, ), learn=dict( data_path=None, diff --git a/dizoo/d4rl/config/walker2d_medium_dt_config.py b/dizoo/d4rl/config/walker2d_medium_dt_config.py index e3b741a129..b8d88699ce 100644 --- a/dizoo/d4rl/config/walker2d_medium_dt_config.py +++ b/dizoo/d4rl/config/walker2d_medium_dt_config.py @@ -1,35 +1,36 @@ from easydict import EasyDict from copy import deepcopy -walker2d_dt_config = dict( - exp_name='walker2d_medium_dt_seed0', +walk2d_dt_config = dict( + exp_name='dt_log/d4rl/walk2d/walk2d_medium_dt', env=dict( env_id='Walker2d-v3', - norm_obs=dict(use_norm=False, ), - norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, use_act_scale=True, n_evaluator_episode=8, - stop_value=6000, + stop_value=5000, + ), + dataset=dict( + env_type='mujoco', + rtg_scale=1000, + context_len=20, + data_dir_prefix='d4rl/walker2d_medium-v2.pkl', ), policy=dict( - stop_value=6000, cuda=True, + stop_value=5000, + state_mean=None, + state_std=None, + evaluator_env_num=8, env_name='Walker2d-v3', - rtg_target=6000, # max target return to go + rtg_target=5000, # max target return to go max_eval_ep_len=1000, # max lenght of one episode - num_eval_ep=10, # num of evaluation episode - batch_size=64, wt_decay=1e-4, warmup_steps=10000, - num_updates_per_iter=100, context_len=20, - n_blocks=3, - embed_dim=128, - n_heads=1, - dropout_p=0.1, - log_dir='/home/wangzilin/research/dt/DI-engine/dizoo/d4rl/dt_data/walker2d_medium_dt_log', + weight_decay=0.1, + clip_grad_norm_p=0.25, model=dict( state_dim=17, act_dim=6, @@ -40,32 +41,19 @@ drop_p=0.1, continuous=True, ), - discount_factor=0.999, - nstep=3, - learn=dict( - dataset_path='/mnt/lustre/wangzilin/d4rl_data/walker2d-medium-v2.pkl', - learning_rate=0.0001, - target_update_freq=100, - kappa=1.0, - min_q_weight=4.0, - ), - collect=dict(unroll_len=1, ), - eval=dict(evaluator=dict(evalu_freq=100, ), ), - other=dict( - eps=dict( - type='exp', - start=0.95, - end=0.1, - decay=10000, - ), - replay_buffer=dict(replay_buffer_size=1000, ), + batch_size=64, + learning_rate=1e-4, + collect=dict( + data_type='d4rl_trajectory', + unroll_len=1, ), + eval=dict(evaluator=dict(eval_freq=100, ), ), ), ) -walker2d_dt_config = EasyDict(walker2d_dt_config) -main_config = walker2d_dt_config -walker2d_dt_create_config = dict( +walk2d_dt_config = EasyDict(walk2d_dt_config) +main_config = walk2d_dt_config +walk2d_dt_create_config = dict( env=dict( type='mujoco', import_names=['dizoo.mujoco.envs.mujoco_env'], @@ -73,8 +61,8 @@ env_manager=dict(type='subprocess'), policy=dict(type='dt'), ) -walker2d_dt_create_config = EasyDict(walker2d_dt_create_config) -create_config = walker2d_dt_create_config +walk2d_dt_create_config = EasyDict(walk2d_dt_create_config) +create_config = walk2d_dt_create_config if __name__ == "__main__": from ding.entry import serial_pipeline_dt diff --git a/dizoo/d4rl/config/walker2d_medium_expert_cql_config.py b/dizoo/d4rl/config/walker2d_medium_expert_cql_config.py old mode 100644 new mode 100755 index b15461ed5b..f05d15c346 --- a/dizoo/d4rl/config/walker2d_medium_expert_cql_config.py +++ b/dizoo/d4rl/config/walker2d_medium_expert_cql_config.py @@ -5,7 +5,7 @@ main_config = dict( exp_name="walker2d_medium_expert_cql_seed0", env=dict( - env_id='walker2d-medium-expert-v0', + env_id='walker2d-medium-expert-v2', collector_env_num=1, evaluator_env_num=8, use_act_scale=True, @@ -15,8 +15,8 @@ policy=dict( cuda=True, model=dict( - obs_shape=11, - action_shape=3, + obs_shape=17, + action_shape=6, ), learn=dict( data_path=None, diff --git a/dizoo/d4rl/config/walker2d_medium_expert_dt_config.py b/dizoo/d4rl/config/walker2d_medium_expert_dt_config.py index deb6051d10..225d00c2e3 100644 --- a/dizoo/d4rl/config/walker2d_medium_expert_dt_config.py +++ b/dizoo/d4rl/config/walker2d_medium_expert_dt_config.py @@ -1,38 +1,39 @@ from easydict import EasyDict from copy import deepcopy -walker2d_dt_config = dict( - exp_name='walker2d_medium_expert_dt_seed0', +walk2d_dt_config = dict( + exp_name='dt_log/d4rl/walk2d/walk2d_medium_expert_dt_seed0', env=dict( - env_id='Walker2d-v3', - norm_obs=dict(use_norm=False, ), - norm_reward=dict(use_norm=False, ), + env_id='Walk2d-v3', collector_env_num=1, evaluator_env_num=8, use_act_scale=True, n_evaluator_episode=8, - stop_value=6000, + stop_value=5000, + ), + dataset=dict( + env_type='mujoco', + rtg_scale=1000, + context_len=30, + data_dir_prefix='d4rl/walk2d_medium_expert-v2.pkl', ), policy=dict( - stop_value=6000, cuda=True, - env_name='Walker2d-v3', - rtg_target=6000, # max target return to go + stop_value=5000, + state_mean=None, + state_std=None, + evaluator_env_num=8, + env_name='Walk2d-v3', + rtg_target=5000, # max target return to go max_eval_ep_len=1000, # max lenght of one episode - num_eval_ep=10, # num of evaluation episode - batch_size=64, wt_decay=1e-4, warmup_steps=10000, - num_updates_per_iter=100, context_len=20, - n_blocks=3, - embed_dim=128, - n_heads=1, - dropout_p=0.1, - log_dir='/home/wangzilin/research/dt/DI-engine/dizoo/d4rl/dt_data/walker2d_medium_expert_dt_log', + weight_decay=0.1, + clip_grad_norm_p=0.25, model=dict( - state_dim=17, - act_dim=6, + state_dim=11, + act_dim=3, n_blocks=3, h_dim=128, context_len=20, @@ -40,32 +41,19 @@ drop_p=0.1, continuous=True, ), - discount_factor=0.999, - nstep=3, - learn=dict( - dataset_path='/mnt/lustre/wangzilin/d4rl_data/walker2d-medium-expert-v2.pkl', - learning_rate=0.0001, - target_update_freq=100, - kappa=1.0, - min_q_weight=4.0, - ), - collect=dict(unroll_len=1, ), - eval=dict(evaluator=dict(evalu_freq=100, ), ), - other=dict( - eps=dict( - type='exp', - start=0.95, - end=0.1, - decay=10000, - ), - replay_buffer=dict(replay_buffer_size=1000, ), + batch_size=64, + learning_rate=1e-4, + collect=dict( + data_type='d4rl_trajectory', + unroll_len=1, ), + eval=dict(evaluator=dict(eval_freq=100, ), ), ), ) -walker2d_dt_config = EasyDict(walker2d_dt_config) -main_config = walker2d_dt_config -walker2d_dt_create_config = dict( +walk2d_dt_config = EasyDict(walk2d_dt_config) +main_config = walk2d_dt_config +walk2d_dt_create_config = dict( env=dict( type='mujoco', import_names=['dizoo.mujoco.envs.mujoco_env'], @@ -73,8 +61,8 @@ env_manager=dict(type='subprocess'), policy=dict(type='dt'), ) -walker2d_dt_create_config = EasyDict(walker2d_dt_create_config) -create_config = walker2d_dt_create_config +walk2d_dt_create_config = EasyDict(walk2d_dt_create_config) +create_config = walk2d_dt_create_config if __name__ == "__main__": from ding.entry import serial_pipeline_dt diff --git a/dizoo/d4rl/config/walker2d_medium_expert_iql_config.py b/dizoo/d4rl/config/walker2d_medium_expert_iql_config.py new file mode 100644 index 0000000000..a6b4c7d1ac --- /dev/null +++ b/dizoo/d4rl/config/walker2d_medium_expert_iql_config.py @@ -0,0 +1,53 @@ +# You can conduct Experiments on D4RL with this config file through the following command: +# cd ../entry && python d4rl_iql_main.py +from easydict import EasyDict + +main_config = dict( + exp_name="walker2d_medium_expert_iql_seed0", + env=dict( + env_id='walker2d-medium-v2', + collector_env_num=1, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + stop_value=6000, + reward_norm="iql_locomotion", + ), + policy=dict( + cuda=True, + model=dict( + obs_shape=17, + action_shape=6, + ), + learn=dict( + data_path=None, + train_epoch=30000, + batch_size=4096, + learning_rate_q=3e-4, + learning_rate_policy=1e-4, + beta=0.05, + tau=0.7, + ), + collect=dict(data_type='d4rl', ), + eval=dict(evaluator=dict(eval_freq=5000, )), + other=dict(replay_buffer=dict(replay_buffer_size=2000000, ), ), + ), +) + +main_config = EasyDict(main_config) +main_config = main_config + +create_config = dict( + env=dict( + type='d4rl', + import_names=['dizoo.d4rl.envs.d4rl_env'], + ), + env_manager=dict(type='base'), + policy=dict( + type='iql', + import_names=['ding.policy.iql'], + ), + replay_buffer=dict(type='naive', ), +) +create_config = EasyDict(create_config) +create_config = create_config diff --git a/dizoo/d4rl/config/walker2d_medium_expert_pd_config.py b/dizoo/d4rl/config/walker2d_medium_expert_pd_config.py new file mode 100755 index 0000000000..06a80588aa --- /dev/null +++ b/dizoo/d4rl/config/walker2d_medium_expert_pd_config.py @@ -0,0 +1,99 @@ +from easydict import EasyDict + +main_config = dict( + exp_name="walker2d_medium_expert_pd_seed0", + env=dict( + env_id='walker2d-medium-expert-v2', + collector_env_num=1, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + returns_scale=1.0, + termination_penalty=-100, + max_path_length=1000, + use_padding=True, + include_returns=True, + normed=False, + stop_value=8000, + horizon=32, + obs_dim=17, + action_dim=6, + ), + policy=dict( + cuda=True, + model=dict( + diffuser_model='GaussianDiffusion', + diffuser_model_cfg=dict( + model='DiffusionUNet1d', + model_cfg=dict( + transition_dim=23, + dim=32, + dim_mults=[1, 2, 4, 8], + returns_condition=False, + kernel_size=5, + attention=False, + ), + horizon=32, + obs_dim=17, + action_dim=6, + n_timesteps=20, + predict_epsilon=False, + loss_discount=1, + action_weight=10, + ), + value_model='ValueDiffusion', + value_model_cfg=dict( + model='TemporalValue', + model_cfg=dict( + horizon=32, + transition_dim=23, + dim=32, + dim_mults=[1, 2, 4, 8], + kernel_size=5, + ), + horizon=32, + obs_dim=17, + action_dim=6, + n_timesteps=20, + predict_epsilon=True, + loss_discount=1, + ), + n_guide_steps=2, + scale=0.1, + t_stopgrad=2, + scale_grad_by_std=True, + ), + normalizer='GaussianNormalizer', + learn=dict( + data_path=None, + train_epoch=60000, + gradient_accumulate_every=2, + batch_size=32, + learning_rate=2e-4, + discount_factor=0.99, + plan_batch_size=64, + learner=dict(hook=dict(save_ckpt_after_iter=1000000000, )), + ), + collect=dict(data_type='diffuser_traj', ), + eval=dict( + evaluator=dict(eval_freq=500, ), + test_ret=0.9, + ), + other=dict(replay_buffer=dict(replay_buffer_size=2000000, ), ), + ), +) + +main_config = EasyDict(main_config) +main_config = main_config + +create_config = dict( + env=dict( + type='d4rl', + import_names=['dizoo.d4rl.envs.d4rl_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='pd', ), + replay_buffer=dict(type='naive', ), +) +create_config = EasyDict(create_config) +create_config = create_config diff --git a/dizoo/d4rl/config/walker2d_medium_expert_qgpo_config.py b/dizoo/d4rl/config/walker2d_medium_expert_qgpo_config.py new file mode 100644 index 0000000000..46c98f8dd1 --- /dev/null +++ b/dizoo/d4rl/config/walker2d_medium_expert_qgpo_config.py @@ -0,0 +1,47 @@ +from easydict import EasyDict + +main_config = dict( + exp_name='walker2d_medium_expert_v2_QGPO_seed0', + env=dict( + env_id="walker2d-medium-expert-v2", + evaluator_env_num=8, + n_evaluator_episode=8, + ), + dataset=dict(env_id="walker2d-medium-expert-v2", ), + policy=dict( + cuda=True, + on_policy=False, + #load_path='./walker2d_medium_expert_v2_QGPO_seed0/ckpt/iteration_600000.pth.tar', + model=dict( + qgpo_critic=dict( + alpha=3, + q_alpha=1, + ), + device='cuda', + obs_dim=17, + action_dim=6, + ), + learn=dict( + learning_rate=1e-4, + batch_size=4096, + batch_size_q=256, + M=16, + diffusion_steps=15, + behavior_policy_stop_training_iter=600000, + energy_guided_policy_begin_training_iter=600000, + q_value_stop_training_iter=1100000, + ), + eval=dict( + guidance_scale=[0.0, 1.0, 2.0, 3.0, 5.0, 8.0, 10.0], + diffusion_steps=15, + evaluator=dict(eval_freq=50000, ), + ), + ), +) +main_config = EasyDict(main_config) + +create_config = dict( + env_manager=dict(type='base'), + policy=dict(type='qgpo', ), +) +create_config = EasyDict(create_config) diff --git a/dizoo/d4rl/config/walker2d_medium_expert_td3bc_config.py b/dizoo/d4rl/config/walker2d_medium_expert_td3bc_config.py old mode 100644 new mode 100755 index 2aed878dd8..2473191de5 --- a/dizoo/d4rl/config/walker2d_medium_expert_td3bc_config.py +++ b/dizoo/d4rl/config/walker2d_medium_expert_td3bc_config.py @@ -5,9 +5,9 @@ main_config = dict( exp_name='walker2d_medium_expert_td3-bc_seed0', env=dict( - env_id='walker2d-medium-expert-v0', + env_id='walker2d-medium-expert-v2', norm_obs=dict( - use_norm=True, + use_norm=True, offline_stats=dict(use_offline_stats=True, ), ), collector_env_num=1, @@ -17,9 +17,10 @@ stop_value=6000, ), policy=dict( + cuda=True, model=dict( - obs_shape=11, - action_shape=3, + obs_shape=17, + action_shape=6, ), learn=dict( train_epoch=30000, diff --git a/dizoo/d4rl/config/walker2d_medium_iql_config.py b/dizoo/d4rl/config/walker2d_medium_iql_config.py new file mode 100644 index 0000000000..7e6756dcc6 --- /dev/null +++ b/dizoo/d4rl/config/walker2d_medium_iql_config.py @@ -0,0 +1,53 @@ +# You can conduct Experiments on D4RL with this config file through the following command: +# cd ../entry && python d4rl_iql_main.py +from easydict import EasyDict + +main_config = dict( + exp_name="walker2d_medium_iql_seed0", + env=dict( + env_id='walker2d-medium-v2', + collector_env_num=1, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + stop_value=6000, + reward_norm="iql_locomotion", + ), + policy=dict( + cuda=True, + model=dict( + obs_shape=17, + action_shape=6, + ), + learn=dict( + data_path=None, + train_epoch=30000, + batch_size=4096, + learning_rate_q=3e-4, + learning_rate_policy=1e-4, + beta=0.05, + tau=0.7, + ), + collect=dict(data_type='d4rl', ), + eval=dict(evaluator=dict(eval_freq=5000, )), + other=dict(replay_buffer=dict(replay_buffer_size=2000000, ), ), + ), +) + +main_config = EasyDict(main_config) +main_config = main_config + +create_config = dict( + env=dict( + type='d4rl', + import_names=['dizoo.d4rl.envs.d4rl_env'], + ), + env_manager=dict(type='base'), + policy=dict( + type='iql', + import_names=['ding.policy.iql'], + ), + replay_buffer=dict(type='naive', ), +) +create_config = EasyDict(create_config) +create_config = create_config diff --git a/dizoo/d4rl/config/walker2d_medium_pd_config.py b/dizoo/d4rl/config/walker2d_medium_pd_config.py new file mode 100755 index 0000000000..099364a763 --- /dev/null +++ b/dizoo/d4rl/config/walker2d_medium_pd_config.py @@ -0,0 +1,99 @@ +from easydict import EasyDict + +main_config = dict( + exp_name="walker2d_medium_pd_seed0", + env=dict( + env_id='walker2d-medium-v2', + collector_env_num=1, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + returns_scale=1.0, + termination_penalty=-100, + max_path_length=1000, + use_padding=True, + include_returns=True, + normed=False, + stop_value=8000, + horizon=32, + obs_dim=17, + action_dim=6, + ), + policy=dict( + cuda=True, + model=dict( + diffuser_model='GaussianDiffusion', + diffuser_model_cfg=dict( + model='DiffusionUNet1d', + model_cfg=dict( + transition_dim=23, + dim=32, + dim_mults=[1, 2, 4, 8], + returns_condition=False, + kernel_size=5, + attention=False, + ), + horizon=32, + obs_dim=17, + action_dim=6, + n_timesteps=20, + predict_epsilon=False, + loss_discount=1, + action_weight=10, + ), + value_model='ValueDiffusion', + value_model_cfg=dict( + model='TemporalValue', + model_cfg=dict( + horizon=32, + transition_dim=23, + dim=32, + dim_mults=[1, 2, 4, 8], + kernel_size=5, + ), + horizon=32, + obs_dim=17, + action_dim=6, + n_timesteps=20, + predict_epsilon=True, + loss_discount=1, + ), + n_guide_steps=2, + scale=0.1, + t_stopgrad=2, + scale_grad_by_std=True, + ), + normalizer='GaussianNormalizer', + learn=dict( + data_path=None, + train_epoch=60000, + gradient_accumulate_every=2, + batch_size=32, + learning_rate=2e-4, + discount_factor=0.99, + plan_batch_size=64, + learner=dict(hook=dict(save_ckpt_after_iter=1000000000, )), + ), + collect=dict(data_type='diffuser_traj', ), + eval=dict( + evaluator=dict(eval_freq=500, ), + test_ret=0.9, + ), + other=dict(replay_buffer=dict(replay_buffer_size=2000000, ), ), + ), +) + +main_config = EasyDict(main_config) +main_config = main_config + +create_config = dict( + env=dict( + type='d4rl', + import_names=['dizoo.d4rl.envs.d4rl_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='pd', ), + replay_buffer=dict(type='naive', ), +) +create_config = EasyDict(create_config) +create_config = create_config diff --git a/dizoo/d4rl/config/walker2d_medium_replay_cql_config.py b/dizoo/d4rl/config/walker2d_medium_replay_cql_config.py old mode 100644 new mode 100755 index 5dc454d158..23437423b6 --- a/dizoo/d4rl/config/walker2d_medium_replay_cql_config.py +++ b/dizoo/d4rl/config/walker2d_medium_replay_cql_config.py @@ -5,7 +5,7 @@ main_config = dict( exp_name="walker2d_medium_replay_cql_seed0", env=dict( - env_id='walker2d-medium-replay-v0', + env_id='walker2d-medium-replay-v2', collector_env_num=1, evaluator_env_num=8, use_act_scale=True, @@ -15,8 +15,8 @@ policy=dict( cuda=True, model=dict( - obs_shape=11, - action_shape=3, + obs_shape=17, + action_shape=6, ), learn=dict( data_path=None, diff --git a/dizoo/d4rl/config/walker2d_medium_replay_dt_config.py b/dizoo/d4rl/config/walker2d_medium_replay_dt_config.py index eb471c5370..b96375b242 100644 --- a/dizoo/d4rl/config/walker2d_medium_replay_dt_config.py +++ b/dizoo/d4rl/config/walker2d_medium_replay_dt_config.py @@ -1,38 +1,39 @@ from easydict import EasyDict from copy import deepcopy -walker2d_dt_config = dict( - exp_name='walker2d_medium_replay_dt_seed0', +walk2d_dt_config = dict( + exp_name='dt_log/d4rl/walk2d/walk2d_medium_replay_dt_seed0', env=dict( - env_id='Walker2d-v3', - norm_obs=dict(use_norm=False, ), - norm_reward=dict(use_norm=False, ), + env_id='Walk2d-v3', collector_env_num=1, evaluator_env_num=8, use_act_scale=True, n_evaluator_episode=8, - stop_value=6000, + stop_value=5000, + ), + dataset=dict( + env_type='mujoco', + rtg_scale=1000, + context_len=30, + data_dir_prefix='d4rl/walk2d_medium_replay-v2.pkl', ), policy=dict( - stop_value=6000, cuda=True, - env_name='Walker2d-v3', - rtg_target=6000, # max target return to go + stop_value=5000, + state_mean=None, + state_std=None, + evaluator_env_num=8, + env_name='Walk2d-v3', + rtg_target=5000, # max target return to go max_eval_ep_len=1000, # max lenght of one episode - num_eval_ep=10, # num of evaluation episode - batch_size=64, wt_decay=1e-4, warmup_steps=10000, - num_updates_per_iter=100, context_len=20, - n_blocks=3, - embed_dim=128, - n_heads=1, - dropout_p=0.1, - log_dir='/home/wangzilin/research/dt/DI-engine/dizoo/d4rl/dt_data/walker2d_medium_replay_dt_log', + weight_decay=0.1, + clip_grad_norm_p=0.25, model=dict( - state_dim=17, - act_dim=6, + state_dim=11, + act_dim=3, n_blocks=3, h_dim=128, context_len=20, @@ -40,32 +41,19 @@ drop_p=0.1, continuous=True, ), - discount_factor=0.999, - nstep=3, - learn=dict( - dataset_path='/mnt/lustre/wangzilin/d4rl_data/walker2d-medium-replay-v2.pkl', - learning_rate=0.0001, - target_update_freq=100, - kappa=1.0, - min_q_weight=4.0, - ), - collect=dict(unroll_len=1, ), - eval=dict(evaluator=dict(evalu_freq=100, ), ), - other=dict( - eps=dict( - type='exp', - start=0.95, - end=0.1, - decay=10000, - ), - replay_buffer=dict(replay_buffer_size=1000, ), + batch_size=64, + learning_rate=1e-4, + collect=dict( + data_type='d4rl_trajectory', + unroll_len=1, ), + eval=dict(evaluator=dict(eval_freq=100, ), ), ), ) -walker2d_dt_config = EasyDict(walker2d_dt_config) -main_config = walker2d_dt_config -walker2d_dt_create_config = dict( +walk2d_dt_config = EasyDict(walk2d_dt_config) +main_config = walk2d_dt_config +walk2d_dt_create_config = dict( env=dict( type='mujoco', import_names=['dizoo.mujoco.envs.mujoco_env'], @@ -73,8 +61,8 @@ env_manager=dict(type='subprocess'), policy=dict(type='dt'), ) -walker2d_dt_create_config = EasyDict(walker2d_dt_create_config) -create_config = walker2d_dt_create_config +walk2d_dt_create_config = EasyDict(walk2d_dt_create_config) +create_config = walk2d_dt_create_config if __name__ == "__main__": from ding.entry import serial_pipeline_dt diff --git a/dizoo/d4rl/config/walker2d_medium_replay_iql_config.py b/dizoo/d4rl/config/walker2d_medium_replay_iql_config.py new file mode 100644 index 0000000000..5b43225448 --- /dev/null +++ b/dizoo/d4rl/config/walker2d_medium_replay_iql_config.py @@ -0,0 +1,53 @@ +# You can conduct Experiments on D4RL with this config file through the following command: +# cd ../entry && python d4rl_iql_main.py +from easydict import EasyDict + +main_config = dict( + exp_name="walker2d_medium_replay_iql_seed0", + env=dict( + env_id='walker2d-medium-v2', + collector_env_num=1, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + stop_value=6000, + reward_norm="iql_locomotion", + ), + policy=dict( + cuda=True, + model=dict( + obs_shape=17, + action_shape=6, + ), + learn=dict( + data_path=None, + train_epoch=30000, + batch_size=4096, + learning_rate_q=3e-4, + learning_rate_policy=1e-4, + beta=0.05, + tau=0.7, + ), + collect=dict(data_type='d4rl', ), + eval=dict(evaluator=dict(eval_freq=5000, )), + other=dict(replay_buffer=dict(replay_buffer_size=2000000, ), ), + ), +) + +main_config = EasyDict(main_config) +main_config = main_config + +create_config = dict( + env=dict( + type='d4rl', + import_names=['dizoo.d4rl.envs.d4rl_env'], + ), + env_manager=dict(type='base'), + policy=dict( + type='iql', + import_names=['ding.policy.iql'], + ), + replay_buffer=dict(type='naive', ), +) +create_config = EasyDict(create_config) +create_config = create_config diff --git a/dizoo/d4rl/config/walker2d_medium_replay_td3bc_config.py b/dizoo/d4rl/config/walker2d_medium_replay_td3bc_config.py old mode 100644 new mode 100755 index 67cc95a1c2..0d13a8d7e2 --- a/dizoo/d4rl/config/walker2d_medium_replay_td3bc_config.py +++ b/dizoo/d4rl/config/walker2d_medium_replay_td3bc_config.py @@ -5,9 +5,9 @@ main_config = dict( exp_name='walker2d_medium_replay_td3-bc_seed0', env=dict( - env_id='walker2d-medium-replay-v0', + env_id='walker2d-medium-replay-v2', norm_obs=dict( - use_norm=True, + use_norm=True, offline_stats=dict(use_offline_stats=True, ), ), collector_env_num=1, @@ -17,9 +17,10 @@ stop_value=6000, ), policy=dict( + cuda=True, model=dict( - obs_shape=11, - action_shape=3, + obs_shape=17, + action_shape=6, ), learn=dict( train_epoch=30000, diff --git a/dizoo/d4rl/config/walker2d_medium_td3bc_config.py b/dizoo/d4rl/config/walker2d_medium_td3bc_config.py old mode 100644 new mode 100755 index dc76b5c012..d496087c96 --- a/dizoo/d4rl/config/walker2d_medium_td3bc_config.py +++ b/dizoo/d4rl/config/walker2d_medium_td3bc_config.py @@ -5,9 +5,9 @@ main_config = dict( exp_name='walker2d_medium_td3-bc_seed0', env=dict( - env_id='walker2d-medium-v0', + env_id='walker2d-medium-v2', norm_obs=dict( - use_norm=True, + use_norm=True, offline_stats=dict(use_offline_stats=True, ), ), collector_env_num=1, @@ -17,9 +17,10 @@ stop_value=6000, ), policy=dict( + cuda=True, model=dict( - obs_shape=11, - action_shape=3, + obs_shape=17, + action_shape=6, ), learn=dict( train_epoch=30000, diff --git a/dizoo/d4rl/config/walker2d_random_cql_config.py b/dizoo/d4rl/config/walker2d_random_cql_config.py old mode 100644 new mode 100755 index 271dbca013..346dd1a1f2 --- a/dizoo/d4rl/config/walker2d_random_cql_config.py +++ b/dizoo/d4rl/config/walker2d_random_cql_config.py @@ -5,7 +5,7 @@ main_config = dict( exp_name="walker2d_expert_cql_seed0", env=dict( - env_id='walker2d-expert-v0', + env_id='walker2d-expert-v2', collector_env_num=1, evaluator_env_num=8, use_act_scale=True, @@ -15,8 +15,8 @@ policy=dict( cuda=True, model=dict( - obs_shape=11, - action_shape=3, + obs_shape=17, + action_shape=6, ), learn=dict( data_path=None, diff --git a/dizoo/d4rl/config/walker2d_random_td3bc_config.py b/dizoo/d4rl/config/walker2d_random_td3bc_config.py old mode 100644 new mode 100755 index f252c14dbd..a38e1e8662 --- a/dizoo/d4rl/config/walker2d_random_td3bc_config.py +++ b/dizoo/d4rl/config/walker2d_random_td3bc_config.py @@ -5,9 +5,9 @@ main_config = dict( exp_name='walker2d_random_td3-bc_seed0', env=dict( - env_id='walker2d-random-v0', + env_id='walker2d-random-v2', norm_obs=dict( - use_norm=True, + use_norm=True, offline_stats=dict(use_offline_stats=True, ), ), collector_env_num=1, @@ -17,9 +17,10 @@ stop_value=6000, ), policy=dict( + cuda=True, model=dict( - obs_shape=11, - action_shape=3, + obs_shape=17, + action_shape=6, ), learn=dict( train_epoch=30000, diff --git a/dizoo/d4rl/dt_data/data/plot_figure_data.py b/dizoo/d4rl/dt_data/data/plot_figure_data.py deleted file mode 100644 index 5554f52fcb..0000000000 --- a/dizoo/d4rl/dt_data/data/plot_figure_data.py +++ /dev/null @@ -1,53 +0,0 @@ -import pandas as pd -import matplotlib.pyplot as plt -import seaborn as sns -import numpy as np -import argparse - -parser = argparse.ArgumentParser() -parser.add_argument('--hopper_ours_csv_path', type=str, default='dt_hopper_log/dt_Hopper-v3_log_22-05-20-03-47-07.csv') -parser.add_argument('--hopper_origin_csv_path', type=str, default='dt_hopper_log/dt_Hopper-v3_official.csv') -parser.add_argument( - '--halfcheetah_ours_csv_path', type=str, default='dt_halfcheetah_log/dt_HalfCheetah-v3_log_22-05-20-06-21-45.csv' -) -parser.add_argument( - '--halfcheetah_origin_csv_path', type=str, default='dt_halfcheetah_log/dt_HalfCheetah-v3_official.csv' -) - -parser.add_argument('--save_path', '-sp', type=str, default='./result.png') -args = parser.parse_args() - -sns.set_theme(style='darkgrid') -df_hopper_ours = pd.read_csv(args.hopper_ours_csv_path, header=None) -df_hopper_ours.columns = ['time_elapsed', 'total_updates', 'mean_action_loss', 'eval_avg_reward', 'eval_avg_ep_len'] -df_hopper_ours.drop(columns=['time_elapsed', 'mean_action_loss', 'eval_avg_ep_len']) - -ax = df_hopper_ours.plot(x='total_updates', y='eval_avg_reward', label='hopper_ours') - -df_half_ours = pd.read_csv(args.halfcheetah_ours_csv_path, header=None) -df_half_ours.columns = ['time_elapsed', 'total_updates', 'mean_action_loss', 'eval_avg_reward', 'eval_avg_ep_len'] -df_half_ours.drop(columns=['time_elapsed', 'mean_action_loss', 'eval_avg_ep_len']) - -ax = df_half_ours.plot(x='total_updates', y='eval_avg_reward', label='halfcheetah_ours', ax=ax) - -df_half_official = pd.read_csv(args.halfcheetah_origin_csv_path, header=None).iloc[:98] -df_half_official.columns = [ - 'time_elapsed', 'total_updates', 'mean_action_loss', 'eval_avg_reward', 'eval_avg_ep_len', 'd4rl_score' -] -df_half_official.drop(columns=['time_elapsed', 'mean_action_loss', 'eval_avg_ep_len', 'd4rl_score']) - -ax = df_half_official.plot(x='total_updates', y='eval_avg_reward', label='halfcheetah_official', ax=ax) - -df_hopper_official = pd.read_csv(args.hopper_origin_csv_path, header=None).iloc[:98] -df_hopper_official.columns = [ - 'time_elapsed', 'total_updates', 'mean_action_loss', 'eval_avg_reward', 'eval_avg_ep_len', 'd4rl_score' -] -df_hopper_official.drop(columns=['time_elapsed', 'mean_action_loss', 'eval_avg_ep_len', 'd4rl_score']) - -ax = df_hopper_official.plot(x='total_updates', y='eval_avg_reward', label='hopper_official', ax=ax) - -plt.xlabel('updates') -plt.ylabel('avg reward') -plt.title('avg reward along update times in evaluation') -plt.legend() -plt.savefig(args.save_path) diff --git a/dizoo/d4rl/dt_data/download_data.py b/dizoo/d4rl/dt_data/download_data.py deleted file mode 100644 index d2d6b9ef94..0000000000 --- a/dizoo/d4rl/dt_data/download_data.py +++ /dev/null @@ -1,70 +0,0 @@ -import os -import gym -import numpy as np - -import collections -import pickle - -import d4rl - - -def download_d4rl_data(): - datasets = [] - - data_dir = '/mnt/lustre/wangzilin/d4rl_data' - - print(data_dir) - - if not os.path.exists(data_dir): - os.makedirs(data_dir) - - for env_name in ['halfcheetah', 'walker2d', 'hopper']: - for dataset_type in ['random', 'medium', 'expert', 'medium-replay', 'medium-expert']: - - name = f'{env_name}-{dataset_type}-v2' - pkl_file_path = os.path.join(data_dir, name) - - print("processing: ", name) - - env = gym.make(name) - dataset = env.get_dataset() - - N = dataset['rewards'].shape[0] - data_ = collections.defaultdict(list) - - use_timeouts = False - if 'timeouts' in dataset: - use_timeouts = True - - episode_step = 0 - paths = [] - for i in range(N): - done_bool = bool(dataset['terminals'][i]) - if use_timeouts: - final_timestep = dataset['timeouts'][i] - else: - final_timestep = (episode_step == 1000 - 1) - for k in ['observations', 'next_observations', 'actions', 'rewards', 'terminals']: - data_[k].append(dataset[k][i]) - if done_bool or final_timestep: - episode_step = 0 - episode_data = {} - for k in data_: - episode_data[k] = np.array(data_[k]) - paths.append(episode_data) - data_ = collections.defaultdict(list) - episode_step += 1 - - returns = np.array([np.sum(p['rewards']) for p in paths]) - num_samples = np.sum([p['rewards'].shape[0] for p in paths]) - print(f'Number of samples collected: {num_samples}') - print( - f'Trajectory returns: mean = {np.mean(returns)}, std = {np.std(returns)}, max = {np.max(returns)}, min = {np.min(returns)}' - ) - - with open(f'{pkl_file_path}.pkl', 'wb') as f: - pickle.dump(paths, f) - - -if __name__ == "__main__": - download_d4rl_data() diff --git a/dizoo/d4rl/entry/d4rl_bcq_main.py b/dizoo/d4rl/entry/d4rl_bcq_main.py new file mode 100755 index 0000000000..099f6e025b --- /dev/null +++ b/dizoo/d4rl/entry/d4rl_bcq_main.py @@ -0,0 +1,21 @@ +from ding.entry import serial_pipeline_offline +from ding.config import read_config +from pathlib import Path + + +def train(args): + # launch from anywhere + config = Path(__file__).absolute().parent.parent / 'config' / args.config + config = read_config(str(config)) + config[0].exp_name = config[0].exp_name.replace('0', str(args.seed)) + serial_pipeline_offline(config, seed=args.seed) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument('--seed', '-s', type=int, default=0) + parser.add_argument('--config', '-c', type=str, default='halfcheetah_medium_bcq_config.py') + args = parser.parse_args() + train(args) diff --git a/dizoo/d4rl/entry/d4rl_cql_main.py b/dizoo/d4rl/entry/d4rl_cql_main.py index 9315a3644d..7a8934a90a 100644 --- a/dizoo/d4rl/entry/d4rl_cql_main.py +++ b/dizoo/d4rl/entry/d4rl_cql_main.py @@ -5,7 +5,7 @@ def train(args): # launch from anywhere - config = Path(__file__).absolute().parent.parent / 'config' / args.config + config = Path(__file__).absolute().parent.parent / 'config' / args.config config = read_config(str(config)) config[0].exp_name = config[0].exp_name.replace('0', str(args.seed)) serial_pipeline_offline(config, seed=args.seed) diff --git a/dizoo/d4rl/entry/d4rl_dt_main.py b/dizoo/d4rl/entry/d4rl_dt_main.py deleted file mode 100644 index 1cbd7ecbe7..0000000000 --- a/dizoo/d4rl/entry/d4rl_dt_main.py +++ /dev/null @@ -1,46 +0,0 @@ -from ding.entry import serial_pipeline_dt -from ding.config import read_config -import os -from copy import deepcopy - - -def train(args): - config_dir = '/mnt/lustre/wangzilin/research/dt/DI-engine/dizoo/d4rl/config' - if config_dir in args.config: - main_config, create_config = read_config(args.config) - else: - main_config, create_config = read_config(os.path.join(config_dir, args.config)) - - result_root = '/mnt/lustre/wangzilin/research/dt/dt/v2' - - itms = args.config.split('_') - env = itms[0] - sub_env = '_'.join(itms[1:3]) - if '_dt' in sub_env: - sub_env = sub_env[:-3] - print(env, sub_env) - - for seed in [0, 1, 2]: - - args.seed = seed - - log_dir = '_'.join([env, sub_env, str(args.seed)]) - - main_config.policy.log_dir = os.path.join(result_root, env, sub_env, log_dir) - - exp_name = '_'.join([env, sub_env, 'seed', str(args.seed)]) - - main_config.exp_name = exp_name - config = deepcopy([main_config, create_config]) - serial_pipeline_dt(config, seed=args.seed, max_train_iter=3000) - - -if __name__ == "__main__": - import argparse - from d4rl import set_dataset_path - - parser = argparse.ArgumentParser() - parser.add_argument('--seed', '-s', type=int, default=0) - parser.add_argument('--config', '-c', type=str, default='hopper_medium_expert_dt_config.py') - args = parser.parse_args() - train(args) diff --git a/dizoo/d4rl/entry/d4rl_dt_mujoco.py b/dizoo/d4rl/entry/d4rl_dt_mujoco.py new file mode 100644 index 0000000000..b6bf93e2c5 --- /dev/null +++ b/dizoo/d4rl/entry/d4rl_dt_mujoco.py @@ -0,0 +1,48 @@ +import gym +import torch +import numpy as np +from ditk import logging +from ding.model.template.decision_transformer import DecisionTransformer +from ding.policy import DTPolicy +from ding.envs import BaseEnvManagerV2 +from ding.envs.env_wrappers.env_wrappers import AllinObsWrapper +from ding.data import create_dataset +from ding.config import compile_config +from ding.framework import task, ding_init +from ding.framework.context import OfflineRLContext +from ding.framework.middleware import interaction_evaluator, trainer, CkptSaver, offline_data_fetcher_from_mem, offline_logger, termination_checker +from ding.utils import set_pkg_seed +from dizoo.d4rl.envs import D4RLEnv +from dizoo.d4rl.config.hopper_medium_dt_config import main_config, create_config + + +def main(): + # If you don't have offline data, you need to prepare if first and set the data_path in config + # For demostration, we also can train a RL policy (e.g. SAC) and collect some data + logging.getLogger().setLevel(logging.INFO) + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + ding_init(cfg) + with task.start(async_mode=False, ctx=OfflineRLContext()): + evaluator_env = BaseEnvManagerV2( + env_fn=[lambda: AllinObsWrapper(D4RLEnv(cfg.env)) for _ in range(cfg.env.evaluator_env_num)], + cfg=cfg.env.manager + ) + + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + + dataset = create_dataset(cfg) + # env_data_stats = dataset.get_d4rl_dataset_stats(cfg.policy.dataset_name) + cfg.policy.state_mean, cfg.policy.state_std = dataset.get_state_stats() + model = DecisionTransformer(**cfg.policy.model) + policy = DTPolicy(cfg.policy, model=model) + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) + task.use(offline_data_fetcher_from_mem(cfg, dataset)) + task.use(trainer(cfg, policy.learn_mode)) + task.use(termination_checker(max_train_iter=5e4)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=1000)) + task.use(offline_logger()) + task.run() + + +if __name__ == "__main__": + main() diff --git a/dizoo/d4rl/entry/d4rl_edac_main.py b/dizoo/d4rl/entry/d4rl_edac_main.py new file mode 100755 index 0000000000..b6710836cb --- /dev/null +++ b/dizoo/d4rl/entry/d4rl_edac_main.py @@ -0,0 +1,21 @@ +from ding.entry import serial_pipeline_offline +from ding.config import read_config +from pathlib import Path + + +def train(args): + # launch from anywhere + config = Path(__file__).absolute().parent.parent / 'config' / args.config + config = read_config(str(config)) + config[0].exp_name = config[0].exp_name.replace('0', str(args.seed)) + serial_pipeline_offline(config, seed=args.seed) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument('--seed', '-s', type=int, default=0) + parser.add_argument('--config', '-c', type=str, default='halfcheetah_medium_edac_config.py') + args = parser.parse_args() + train(args) diff --git a/dizoo/d4rl/entry/d4rl_ibc_main.py b/dizoo/d4rl/entry/d4rl_ibc_main.py index 1c0f57627e..a112916f6c 100644 --- a/dizoo/d4rl/entry/d4rl_ibc_main.py +++ b/dizoo/d4rl/entry/d4rl_ibc_main.py @@ -6,22 +6,25 @@ import torch import torch.multiprocessing as mp + def offline_worker(rank, config, args): dist_init(rank=rank, world_size=torch.cuda.device_count()) serial_pipeline_offline(config, seed=args.seed) + def train(args): # launch from anywhere - config = Path(__file__).absolute().parent.parent / 'config' / args.config + config = Path(__file__).absolute().parent.parent / 'config' / args.config config = read_config(str(config)) config[0].exp_name = config[0].exp_name.replace('0', str(args.seed)) - if not config[0].policy.learn.multi_gpu: + if not config[0].policy.multi_gpu: serial_pipeline_offline(config, seed=args.seed) else: os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = "29600" mp.spawn(offline_worker, nprocs=torch.cuda.device_count(), args=(config, args)) + if __name__ == "__main__": import argparse diff --git a/dizoo/d4rl/entry/d4rl_iql_main.py b/dizoo/d4rl/entry/d4rl_iql_main.py new file mode 100644 index 0000000000..ded097ee42 --- /dev/null +++ b/dizoo/d4rl/entry/d4rl_iql_main.py @@ -0,0 +1,21 @@ +from ding.entry import serial_pipeline_offline +from ding.config import read_config +from pathlib import Path + + +def train(args): + # launch from anywhere + config = Path(__file__).absolute().parent.parent / 'config' / args.config + config = read_config(str(config)) + config[0].exp_name = config[0].exp_name.replace('0', str(args.seed)) + serial_pipeline_offline(config, seed=args.seed) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument('--seed', '-s', type=int, default=10) + parser.add_argument('--config', '-c', type=str, default='halfcheetah_medium_iql_config.py') + args = parser.parse_args() + train(args) diff --git a/dizoo/d4rl/entry/d4rl_pd_main.py b/dizoo/d4rl/entry/d4rl_pd_main.py new file mode 100755 index 0000000000..0d355a77ea --- /dev/null +++ b/dizoo/d4rl/entry/d4rl_pd_main.py @@ -0,0 +1,21 @@ +from ding.entry import serial_pipeline_offline +from ding.config import read_config +from pathlib import Path + + +def train(args): + # launch from anywhere + config = Path(__file__).absolute().parent.parent / 'config' / args.config + config = read_config(str(config)) + config[0].exp_name = config[0].exp_name.replace('0', str(args.seed)) + serial_pipeline_offline(config, seed=args.seed) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument('--seed', '-s', type=int, default=10) + parser.add_argument('--config', '-c', type=str, default='halfcheetah_medium_pd_config.py') + args = parser.parse_args() + train(args) diff --git a/dizoo/d4rl/entry/d4rl_td3_bc_main.py b/dizoo/d4rl/entry/d4rl_td3_bc_main.py index bdf945978f..b25bf904a5 100644 --- a/dizoo/d4rl/entry/d4rl_td3_bc_main.py +++ b/dizoo/d4rl/entry/d4rl_td3_bc_main.py @@ -5,7 +5,7 @@ def train(args): # launch from anywhere - config = Path(__file__).absolute().parent.parent / 'config' / args.config + config = Path(__file__).absolute().parent.parent / 'config' / args.config config = read_config(str(config)) config[0].exp_name = config[0].exp_name.replace('0', str(args.seed)) serial_pipeline_offline(config, seed=args.seed) diff --git a/dizoo/d4rl/envs/d4rl_env.py b/dizoo/d4rl/envs/d4rl_env.py old mode 100644 new mode 100755 index 29100ab693..d14bdafaa5 --- a/dizoo/d4rl/envs/d4rl_env.py +++ b/dizoo/d4rl/envs/d4rl_env.py @@ -2,6 +2,9 @@ import copy import numpy as np import gym +import matplotlib.pyplot as plt +import einops +import imageio from easydict import EasyDict from ding.envs import BaseEnv, BaseEnvTimestep @@ -11,6 +14,40 @@ from .d4rl_wrappers import wrap_d4rl from ding.utils import ENV_REGISTRY +MAZE_BOUNDS = {'maze2d-umaze-v1': (0, 5, 0, 5), 'maze2d-medium-v1': (0, 8, 0, 8), 'maze2d-large-v1': (0, 9, 0, 12)} + + +def plot2img(fig, remove_margins=True): + # https://stackoverflow.com/a/35362787/2912349 + # https://stackoverflow.com/a/54334430/2912349 + + from matplotlib.backends.backend_agg import FigureCanvasAgg + + if remove_margins: + fig.subplots_adjust(left=0, bottom=0, right=1, top=1, wspace=0, hspace=0) + + canvas = FigureCanvasAgg(fig) + canvas.draw() + img_as_string, (width, height) = canvas.print_to_buffer() + return np.fromstring(img_as_string, dtype='uint8').reshape((height, width, 4)) + + +def zipsafe(*args): + length = len(args[0]) + assert all([len(a) == length for a in args]) + return zip(*args) + + +def zipkw(*args, **kwargs): + nargs = len(args) + keys = kwargs.keys() + vals = [kwargs[k] for k in keys] + zipped = zipsafe(*args, *vals) + for items in zipped: + zipped_args = items[:nargs] + zipped_kwargs = {k: v for k, v in zipsafe(keys, items[nargs:])} + yield zipped_args, zipped_kwargs + @ENV_REGISTRY.register('d4rl') class D4RLEnv(BaseEnv): @@ -19,12 +56,19 @@ def __init__(self, cfg: dict) -> None: self._cfg = cfg self._use_act_scale = cfg.use_act_scale self._init_flag = False + if 'maze' in self._cfg.env_id: + self.observations = [] + self._extent = (0, 1, 1, 0) def reset(self) -> np.ndarray: if not self._init_flag: self._env = self._make_env(only_info=False) self._env.observation_space.dtype = np.float32 # To unify the format of envs in DI-engine self._observation_space = self._env.observation_space + if 'maze' in self._cfg.env_id: + new_low = np.tile(self._observation_space.low, 2) + new_high = np.tile(self._observation_space.high, 2) + self._observation_space = gym.spaces.Box(low=new_low, high=new_high) self._action_space = self._env.action_space self._reward_space = gym.spaces.Box( low=self._env.reward_range[0], high=self._env.reward_range[1], shape=(1, ), dtype=np.float32 @@ -35,9 +79,15 @@ def reset(self) -> np.ndarray: self._env.seed(self._seed + np_seed) elif hasattr(self, '_seed'): self._env.seed(self._seed) + if 'maze' in self._cfg.env_id: + target = self._env.get_target() + self.target_obs = np.array([*target, 0, 0]) obs = self._env.reset() + if 'maze' in self._cfg.env_id: + self.observations.append(obs) + obs = np.hstack((obs, self.target_obs)) obs = to_ndarray(obs).astype('float32') - self._final_eval_reward = 0. + self._eval_episode_return = 0. return obs def close(self) -> None: @@ -56,17 +106,69 @@ def step(self, action: Union[np.ndarray, list]) -> BaseEnvTimestep: action_range = {'min': self.action_space.low[0], 'max': self.action_space.high[0], 'dtype': np.float32} action = affine_transform(action, min_val=action_range['min'], max_val=action_range['max']) obs, rew, done, info = self._env.step(action) - self._final_eval_reward += rew + self._eval_episode_return += rew + if 'maze' in self._cfg.env_id: + self.observations.append(obs) + obs = np.hstack([obs, self.target_obs]) obs = to_ndarray(obs).astype('float32') rew = to_ndarray([rew]) # wrapped to be transfered to a array with shape (1,) if done: - info['final_eval_reward'] = self._final_eval_reward + info['eval_episode_return'] = self._eval_episode_return + # self.composite('/mnt/PD/render/rollout.png',self.observations,ncol=1) return BaseEnvTimestep(obs, rew, done, info) + def renders(self, observations, conditions=None, title=None): + bounds = MAZE_BOUNDS[self._cfg.env_id] + + observations = observations + .5 + if len(bounds) == 2: + _, scale = bounds + observations /= scale + elif len(bounds) == 4: + _, iscale, _, jscale = bounds + observations[:, 0] /= iscale + observations[:, 1] /= jscale + else: + raise RuntimeError(f'Unrecognized bounds for {self._cfg.env_id}: {bounds}') + + if conditions is not None: + conditions /= scale + + plt.clf() + fig = plt.gcf() + fig.set_size_inches(5, 5) + plt.imshow(self._background * .5, extent=self._extent, cmap=plt.cm.binary, vmin=0, vmax=1) + + path_length = len(observations) + colors = plt.cm.jet(np.linspace(0, 1, path_length)) + plt.plot(observations[:, 1], observations[:, 0], c='black', zorder=10) + plt.scatter(observations[:, 1], observations[:, 0], c=colors, zorder=20) + plt.axis('off') + plt.title(title) + img = plot2img(fig, remove_margins=self._remove_margins) + return img + + def composite(self, savepath, paths, ncol=5, **kwargs): + assert len(paths) % ncol == 0, 'Number of paths must be divisible by number of columns' + + images = [] + for path, kw in zipkw(paths, **kwargs): + img = self.renders(*path, **kw) + images.append(img) + images = np.stack(images, axis=0) + + nrow = len(images) // ncol + images = einops.rearrange(images, '(nrow ncol) H W C -> (nrow H) (ncol W) C', nrow=nrow, ncol=ncol) + imageio.imsave(savepath, images) + print(f'Saved {len(paths)} samples to: {savepath}') + def _make_env(self, only_info=False): return wrap_d4rl( self._cfg.env_id, - norm_obs=self._cfg.get('norm_obs', EasyDict(use_norm=False, offline_stats=dict(use_offline_stats=False, )),), + norm_obs=self._cfg.get( + 'norm_obs', + EasyDict(use_norm=False, offline_stats=dict(use_offline_stats=False, )), + ), norm_reward=self._cfg.get('norm_reward', EasyDict(use_norm=False, )), only_info=only_info ) diff --git a/dizoo/dmc2gym/config/cartpole_balance/cartpole_balance_dreamer_config.py b/dizoo/dmc2gym/config/cartpole_balance/cartpole_balance_dreamer_config.py new file mode 100644 index 0000000000..623cfaacf1 --- /dev/null +++ b/dizoo/dmc2gym/config/cartpole_balance/cartpole_balance_dreamer_config.py @@ -0,0 +1,95 @@ +from easydict import EasyDict + +from ding.entry import serial_pipeline_dreamer + +cuda = False + +cartpole_balance_dreamer_config = dict( + exp_name='dmc2gym_cartpole_balance_dreamer', + env=dict( + env_id='dmc2gym_cartpole_balance', + domain_name='cartpole', + task_name='balance', + frame_skip=1, + warp_frame=True, + scale=True, + clip_rewards=False, + action_repeat=2, + frame_stack=1, + from_pixels=True, + resize=64, + collector_env_num=1, + evaluator_env_num=1, + n_evaluator_episode=1, + stop_value=1000, # 1000 + ), + policy=dict( + cuda=cuda, + # it is better to put random_collect_size in policy.other + random_collect_size=2500, + model=dict( + obs_shape=(3, 64, 64), + action_shape=1, + actor_dist='normal', + ), + learn=dict( + lambda_=0.95, + learning_rate=3e-5, + batch_size=16, + batch_length=64, + imag_sample=True, + discount=0.997, + reward_EMA=True, + ), + collect=dict( + n_sample=1, + unroll_len=1, + action_size=1, # has to be specified + collect_dyn_sample=True, + ), + command=dict(), + eval=dict(evaluator=dict(eval_freq=5000, )), + other=dict( + # environment buffer + replay_buffer=dict(replay_buffer_size=500000, periodic_thruput_seconds=60), + ), + ), + world_model=dict( + pretrain=100, + train_freq=2, + cuda=cuda, + model=dict( + state_size=(3, 64, 64), # has to be specified + obs_type='RGB', + action_size=1, # has to be specified + action_type='continuous', + reward_size=1, + batch_size=16, + ), + ), +) + +cartpole_balance_dreamer_config = EasyDict(cartpole_balance_dreamer_config) + +cartpole_balance_create_config = dict( + env=dict( + type='dmc2gym', + import_names=['dizoo.dmc2gym.envs.dmc2gym_env'], + ), + env_manager=dict(type='base'), + policy=dict( + type='dreamer', + import_names=['ding.policy.mbpolicy.dreamer'], + ), + replay_buffer=dict(type='sequence', ), + world_model=dict( + type='dreamer', + import_names=['ding.world_model.dreamer'], + ), +) +cartpole_balance_create_config = EasyDict(cartpole_balance_create_config) + +if __name__ == '__main__': + serial_pipeline_dreamer( + (cartpole_balance_dreamer_config, cartpole_balance_create_config), seed=0, max_env_step=500000 + ) diff --git a/dizoo/dmc2gym/config/cheetah_run/cheetah_run_dreamer_config.py b/dizoo/dmc2gym/config/cheetah_run/cheetah_run_dreamer_config.py new file mode 100644 index 0000000000..22b6ae911b --- /dev/null +++ b/dizoo/dmc2gym/config/cheetah_run/cheetah_run_dreamer_config.py @@ -0,0 +1,93 @@ +from easydict import EasyDict + +from ding.entry import serial_pipeline_dreamer + +cuda = False + +cheetah_run_dreamer_config = dict( + exp_name='dmc2gym_cheetah_run_dreamer', + env=dict( + env_id='dmc2gym_cheetah_run', + domain_name='cheetah', + task_name='run', + frame_skip=1, + warp_frame=True, + scale=True, + clip_rewards=False, + action_repeat=2, + frame_stack=1, + from_pixels=True, + resize=64, + collector_env_num=1, + evaluator_env_num=1, + n_evaluator_episode=1, + stop_value=1000, # 1000 + ), + policy=dict( + cuda=cuda, + # it is better to put random_collect_size in policy.other + random_collect_size=2500, + model=dict( + obs_shape=(3, 64, 64), + action_shape=6, + actor_dist='normal', + ), + learn=dict( + lambda_=0.95, + learning_rate=3e-5, + batch_size=16, + batch_length=64, + imag_sample=True, + discount=0.997, + reward_EMA=True, + ), + collect=dict( + n_sample=1, + unroll_len=1, + action_size=6, # has to be specified + collect_dyn_sample=True, + ), + command=dict(), + eval=dict(evaluator=dict(eval_freq=5000, )), + other=dict( + # environment buffer + replay_buffer=dict(replay_buffer_size=500000, periodic_thruput_seconds=60), + ), + ), + world_model=dict( + pretrain=100, + train_freq=2, + cuda=cuda, + model=dict( + state_size=(3, 64, 64), # has to be specified + obs_type='RGB', + action_size=6, # has to be specified + action_type='continuous', + reward_size=1, + batch_size=16, + ), + ), +) + +cheetah_run_dreamer_config = EasyDict(cheetah_run_dreamer_config) + +cheetah_run_create_config = dict( + env=dict( + type='dmc2gym', + import_names=['dizoo.dmc2gym.envs.dmc2gym_env'], + ), + env_manager=dict(type='base'), + policy=dict( + type='dreamer', + import_names=['ding.policy.mbpolicy.dreamer'], + ), + replay_buffer=dict(type='sequence', ), + world_model=dict( + type='dreamer', + import_names=['ding.world_model.dreamer'], + ), +) +cheetah_run_create_config = EasyDict(cheetah_run_create_config) + +if __name__ == '__main__': + serial_pipeline_dreamer((cheetah_run_dreamer_config, cheetah_run_create_config), seed=0, max_env_step=500000) diff --git a/dizoo/dmc2gym/config/dmc2gym_dreamer_config.py b/dizoo/dmc2gym/config/dmc2gym_dreamer_config.py new file mode 100644 index 0000000000..de8e09e3d8 --- /dev/null +++ b/dizoo/dmc2gym/config/dmc2gym_dreamer_config.py @@ -0,0 +1,93 @@ +from easydict import EasyDict + +from ding.entry import serial_pipeline_dreamer + +cuda = False + +cartpole_balance_dreamer_config = dict( + exp_name='dmc2gym_cartpole_balance_dreamer', + env=dict( + env_id='dmc2gym_cartpole_balance', + domain_name='cartpole', + task_name='balance', + frame_skip=1, + warp_frame=True, + scale=True, + clip_rewards=False, + action_repeat=2, + frame_stack=1, + from_pixels=True, + resize=64, + collector_env_num=1, + evaluator_env_num=1, + n_evaluator_episode=1, + stop_value=1000, # 1000 + ), + policy=dict( + cuda=cuda, + # it is better to put random_collect_size in policy.other + random_collect_size=2500, + model=dict( + obs_shape=(3, 64, 64), + action_shape=1, + actor_dist='normal', + ), + learn=dict( + lambda_=0.95, + learning_rate=3e-5, + batch_size=16, + batch_length=64, + imag_sample=True, + discount=0.997, + reward_EMA=True, + ), + collect=dict( + n_sample=1, + unroll_len=1, + action_size=1, # has to be specified + collect_dyn_sample=True, + ), + command=dict(), + eval=dict(evaluator=dict(eval_freq=5000, )), + other=dict( + # environment buffer + replay_buffer=dict(replay_buffer_size=500000, periodic_thruput_seconds=60), + ), + ), + world_model=dict( + pretrain=100, + train_freq=2, + cuda=cuda, + model=dict( + state_size=(3, 64, 64), # has to be specified + action_size=1, # has to be specified + reward_size=1, + batch_size=16, + ), + ), +) + +cartpole_balance_dreamer_config = EasyDict(cartpole_balance_dreamer_config) + +cartpole_balance_create_config = dict( + env=dict( + type='dmc2gym', + import_names=['dizoo.dmc2gym.envs.dmc2gym_env'], + ), + env_manager=dict(type='base'), + policy=dict( + type='dreamer', + import_names=['ding.policy.mbpolicy.dreamer'], + ), + replay_buffer=dict(type='sequence', ), + world_model=dict( + type='dreamer', + import_names=['ding.world_model.dreamer'], + ), +) +cartpole_balance_create_config = EasyDict(cartpole_balance_create_config) + +if __name__ == '__main__': + serial_pipeline_dreamer( + (cartpole_balance_dreamer_config, cartpole_balance_create_config), seed=0, max_env_step=1000000 + ) diff --git a/dizoo/dmc2gym/config/dmc2gym_ppo_config.py b/dizoo/dmc2gym/config/dmc2gym_ppo_config.py new file mode 100644 index 0000000000..207b398e63 --- /dev/null +++ b/dizoo/dmc2gym/config/dmc2gym_ppo_config.py @@ -0,0 +1,63 @@ +from easydict import EasyDict + +cartpole_balance_ppo_config = dict( + exp_name='dmc2gym_cartpole_balance_ppo', + env=dict( + env_id='dmc2gym_cartpole_balance', + domain_name='cartpole', + task_name='balance', + from_pixels=False, + norm_obs=dict(use_norm=False, ), + norm_reward=dict(use_norm=False, ), + collector_env_num=1, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + stop_value=1000, + ), + policy=dict( + cuda=True, + recompute_adv=True, + action_space='discrete', + model=dict( + obs_shape=5, + action_shape=1, + action_space='discrete', + encoder_hidden_size_list=[64, 64, 128], + critic_head_hidden_size=128, + actor_head_hidden_size=128, + ), + learn=dict( + epoch_per_collect=2, + batch_size=64, + learning_rate=0.001, + value_weight=0.5, + entropy_weight=0.01, + clip_ratio=0.2, + learner=dict(hook=dict(save_ckpt_after_iter=100)), + ), + collect=dict( + n_sample=256, + unroll_len=1, + discount_factor=0.9, + gae_lambda=0.95, + ), + other=dict(replay_buffer=dict(replay_buffer_size=10000, ), ), + ) +) +cartpole_balance_ppo_config = EasyDict(cartpole_balance_ppo_config) +main_config = cartpole_balance_ppo_config + +cartpole_balance_create_config = dict( + env=dict( + type='dmc2gym', + import_names=['dizoo.dmc2gym.envs.dmc2gym_env'], + ), + env_manager=dict(type='base'), + policy=dict(type='ppo'), + replay_buffer=dict(type='naive', ), +) +cartpole_balance_create_config = EasyDict(cartpole_balance_create_config) +create_config = cartpole_balance_create_config + +# To use this config, you can enter dizoo/dmc2gym/entry to call dmc2gym_onppo_main.py diff --git a/dizoo/dmc2gym/config/dmc2gym_sac_pixel_config.py b/dizoo/dmc2gym/config/dmc2gym_sac_pixel_config.py new file mode 100644 index 0000000000..88c4e6a9d7 --- /dev/null +++ b/dizoo/dmc2gym/config/dmc2gym_sac_pixel_config.py @@ -0,0 +1,79 @@ +from easydict import EasyDict +import os +# os.environ['MUJOCO_GL']="egl" +dmc2gym_sac_config = dict( + exp_name='dmc2gym_sac_pixel_seed0', + env=dict( + env_id='dmc2gym-v0', + domain_name="cartpole", + task_name="swingup", + frame_skip=4, + warp_frame=True, + scale=True, + clip_rewards=False, + frame_stack=3, + from_pixels=True, # pixel obs + channels_first=False, # obs shape (height, width, 3) + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=1e6, + manager=dict(shared_memory=False, ), + ), + policy=dict( + model_type='pixel', + cuda=True, + random_collect_size=10000, + model=dict( + obs_shape=(3, 84, 84), + action_shape=1, + twin_critic=True, + encoder_hidden_size_list=[32, 32, 32], + actor_head_hidden_size=1024, + critic_head_hidden_size=1024, + share_encoder=True, + ), + learn=dict( + ignore_done=True, + update_per_collect=1, + batch_size=128, + learning_rate_q=1e-3, + learning_rate_policy=1e-3, + learning_rate_alpha=3e-4, + target_theta=0.005, + discount_factor=0.99, + alpha=0.2, + reparameterization=True, + auto_alpha=True, + ), + collect=dict( + n_sample=1, + unroll_len=1, + ), + eval=dict(), + other=dict(replay_buffer=dict(replay_buffer_size=100000, ), ), + ), +) + +dmc2gym_sac_config = EasyDict(dmc2gym_sac_config) +main_config = dmc2gym_sac_config + +dmc2gym_sac_create_config = dict( + env=dict( + type='dmc2gym', + import_names=['dizoo.dmc2gym.envs.dmc2gym_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict( + type='sac', + import_names=['ding.policy.sac'], + ), + replay_buffer=dict(type='naive', ), +) +dmc2gym_sac_create_config = EasyDict(dmc2gym_sac_create_config) +create_config = dmc2gym_sac_create_config + +if __name__ == "__main__": + # or you can enter `ding -m serial -c ant_sac_config.py -s 0 --env-step 1e7` + from ding.entry import serial_pipeline + serial_pipeline((main_config, create_config), seed=0) diff --git a/dizoo/dmc2gym/config/dmc2gym_sac_state_config.py b/dizoo/dmc2gym/config/dmc2gym_sac_state_config.py new file mode 100644 index 0000000000..840629423b --- /dev/null +++ b/dizoo/dmc2gym/config/dmc2gym_sac_state_config.py @@ -0,0 +1,69 @@ +from easydict import EasyDict + +dmc2gym_sac_config = dict( + exp_name='dmc2gym_sac_state_seed0', + env=dict( + env_id='dmc2gym-v0', + domain_name="cartpole", + task_name="swingup", + frame_skip=8, + frame_stack=1, + from_pixels=False, # state obs + channels_first=False, # obs shape (height, width, 3) + collector_env_num=16, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=1e6, + manager=dict(shared_memory=False, ), + ), + policy=dict( + model_type='state', + cuda=True, + random_collect_size=10000, + model=dict( + obs_shape=5, + action_shape=1, + twin_critic=True, + action_space='reparameterization', + actor_head_hidden_size=256, + critic_head_hidden_size=256, + ), + learn=dict( + ignore_done=True, + update_per_collect=1, + batch_size=256, + learning_rate_q=1e-3, + learning_rate_policy=1e-3, + learning_rate_alpha=3e-4, + target_theta=0.005, + discount_factor=0.99, + alpha=0.2, + reparameterization=True, + auto_alpha=True, + ), + collect=dict( + n_sample=1, + unroll_len=1, + ), + eval=dict(), + other=dict(replay_buffer=dict(replay_buffer_size=1000000, ), ), + ), +) + +dmc2gym_sac_config = EasyDict(dmc2gym_sac_config) +main_config = dmc2gym_sac_config + +dmc2gym_sac_create_config = dict( + env=dict( + type='dmc2gym', + import_names=['dizoo.dmc2gym.envs.dmc2gym_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict( + type='sac', + import_names=['ding.policy.sac'], + ), + replay_buffer=dict(type='naive', ), +) +dmc2gym_sac_create_config = EasyDict(dmc2gym_sac_create_config) +create_config = dmc2gym_sac_create_config diff --git a/dizoo/dmc2gym/config/walker_walk/walker_walk_dreamer_config.py b/dizoo/dmc2gym/config/walker_walk/walker_walk_dreamer_config.py new file mode 100644 index 0000000000..da7f9e0edb --- /dev/null +++ b/dizoo/dmc2gym/config/walker_walk/walker_walk_dreamer_config.py @@ -0,0 +1,92 @@ +from easydict import EasyDict + +from ding.entry import serial_pipeline_dreamer + +cuda = False + +walker_walk_dreamer_config = dict( + exp_name='dmc2gym_walker_walk_dreamer', + env=dict( + env_id='dmc2gym_walker_walk', + domain_name='walker', + task_name='walk', + frame_skip=1, + warp_frame=True, + scale=True, + clip_rewards=False, + action_repeat=2, + frame_stack=1, + from_pixels=True, + resize=64, + collector_env_num=1, + evaluator_env_num=1, + n_evaluator_episode=1, + stop_value=1000, # 1000 + ), + policy=dict( + cuda=cuda, + # it is better to put random_collect_size in policy.other + random_collect_size=2500, + model=dict( + action_shape=6, + actor_dist='normal', + ), + learn=dict( + lambda_=0.95, + learning_rate=3e-5, + batch_size=16, + batch_length=64, + imag_sample=True, + discount=0.997, + reward_EMA=True, + ), + collect=dict( + n_sample=1, + unroll_len=1, + action_size=6, # has to be specified + collect_dyn_sample=True, + ), + command=dict(), + eval=dict(evaluator=dict(eval_freq=5000, )), + other=dict( + # environment buffer + replay_buffer=dict(replay_buffer_size=500000, periodic_thruput_seconds=60), + ), + ), + world_model=dict( + pretrain=100, + train_freq=2, + cuda=cuda, + model=dict( + state_size=(3, 64, 64), # has to be specified + obs_type='RGB', + action_size=6, # has to be specified + action_type='continuous', + reward_size=1, + batch_size=16, + ), + ), +) + +walker_walk_dreamer_config = EasyDict(walker_walk_dreamer_config) + +walker_walk_create_config = dict( + env=dict( + type='dmc2gym', + import_names=['dizoo.dmc2gym.envs.dmc2gym_env'], + ), + env_manager=dict(type='base'), + policy=dict( + type='dreamer', + import_names=['ding.policy.mbpolicy.dreamer'], + ), + replay_buffer=dict(type='sequence', ), + world_model=dict( + type='dreamer', + import_names=['ding.world_model.dreamer'], + ), +) +walker_walk_create_config = EasyDict(walker_walk_create_config) + +if __name__ == '__main__': + serial_pipeline_dreamer((walker_walk_dreamer_config, walker_walk_create_config), seed=0, max_env_step=500000) diff --git a/dizoo/dmc2gym/dmc2gym_cheetah.png b/dizoo/dmc2gym/dmc2gym_cheetah.png new file mode 100644 index 0000000000..4fa1c9d9b5 Binary files /dev/null and b/dizoo/dmc2gym/dmc2gym_cheetah.png differ diff --git a/dizoo/dmc2gym/entry/dmc2gym_onppo_main.py b/dizoo/dmc2gym/entry/dmc2gym_onppo_main.py new file mode 100644 index 0000000000..412a46577a --- /dev/null +++ b/dizoo/dmc2gym/entry/dmc2gym_onppo_main.py @@ -0,0 +1,124 @@ +import os +from easydict import EasyDict +from functools import partial +from tensorboardX import SummaryWriter +import dmc2gym + +from ding.worker import BaseLearner, SampleSerialCollector, InteractionSerialEvaluator +from ding.model import VAC +from ding.policy import PPOPolicy +from ding.envs import DingEnvWrapper, EvalEpisodeReturnWrapper, BaseEnvManager +from ding.config import compile_config +from ding.utils import set_pkg_seed +from dizoo.dmc2gym.config.dmc2gym_ppo_config import cartpole_balance_ppo_config +from dizoo.dmc2gym.envs.dmc2gym_env import * + + +class Dmc2GymWrapper(gym.Wrapper): + + def __init__(self, env, cfg): + super().__init__(env) + cfg = EasyDict(cfg) + self._cfg = cfg + + env_info = dmc2gym_env_info[cfg.domain_name][cfg.task_name] + + self._observation_space = env_info["observation_space"]( + from_pixels=self._cfg["from_pixels"], + height=self._cfg["height"], + width=self._cfg["width"], + channels_first=self._cfg["channels_first"] + ) + self._action_space = env_info["action_space"] + self._reward_space = env_info["reward_space"](self._cfg["frame_skip"]) + + def _process_obs(self, obs): + if self._cfg["from_pixels"]: + obs = to_ndarray(obs).astype(np.uint8) + else: + obs = to_ndarray(obs).astype(np.float32) + return obs + + def step(self, action): + action = np.array([action]).astype('float32') + obs, reward, done, info = self.env.step(action) + return self._process_obs(obs), reward, done, info + + def reset(self): + obs = self.env.reset() + return self._process_obs(obs) + + +def wrapped_dmc2gym_env(cfg): + default_cfg = { + "frame_skip": 3, + "from_pixels": True, + "visualize_reward": False, + "height": 100, + "width": 100, + "channels_first": True, + } + default_cfg.update(cfg) + + return DingEnvWrapper( + dmc2gym.make( + domain_name=default_cfg["domain_name"], + task_name=default_cfg["task_name"], + seed=1, + visualize_reward=default_cfg["visualize_reward"], + from_pixels=default_cfg["from_pixels"], + height=default_cfg["height"], + width=default_cfg["width"], + frame_skip=default_cfg["frame_skip"] + ), + cfg={ + 'env_wrapper': [ + lambda env: Dmc2GymWrapper(env, default_cfg), + lambda env: EvalEpisodeReturnWrapper(env), + ] + } + ) + + +def main(cfg, seed=0, max_env_step=int(1e10), max_train_iter=int(1e10)): + cfg = compile_config( + cfg, BaseEnvManager, PPOPolicy, BaseLearner, SampleSerialCollector, InteractionSerialEvaluator, save_cfg=True + ) + collector_env_num, evaluator_env_num = cfg.env.collector_env_num, cfg.env.evaluator_env_num + collector_env = BaseEnvManager( + env_fn=[partial(wrapped_dmc2gym_env, cfg=cartpole_balance_ppo_config.env) for _ in range(collector_env_num)], + cfg=cfg.env.manager + ) + evaluator_env = BaseEnvManager( + env_fn=[partial(wrapped_dmc2gym_env, cfg=cartpole_balance_ppo_config.env) for _ in range(evaluator_env_num)], + cfg=cfg.env.manager + ) + + collector_env.seed(seed) + evaluator_env.seed(seed, dynamic_seed=False) + set_pkg_seed(seed, use_cuda=cfg.policy.cuda) + + model = VAC(**cfg.policy.model) + policy = PPOPolicy(cfg.policy, model=model) + tb_logger = SummaryWriter(os.path.join('./{}/log/'.format(cfg.exp_name), 'serial')) + learner = BaseLearner(cfg.policy.learn.learner, policy.learn_mode, tb_logger, exp_name=cfg.exp_name) + collector = SampleSerialCollector( + cfg.policy.collect.collector, collector_env, policy.collect_mode, tb_logger, exp_name=cfg.exp_name + ) + evaluator = InteractionSerialEvaluator( + cfg.policy.eval.evaluator, evaluator_env, policy.eval_mode, tb_logger, exp_name=cfg.exp_name + ) + + while True: + if evaluator.should_eval(learner.train_iter): + stop, reward = evaluator.eval(learner.save_checkpoint, learner.train_iter, collector.envstep) + if stop: + break + new_data = collector.collect(train_iter=learner.train_iter) + learner.train(new_data, collector.envstep) + if collector.envstep >= max_env_step or learner.train_iter >= max_train_iter: + break + + +if __name__ == '__main__': + main(cartpole_balance_ppo_config) diff --git a/dizoo/dmc2gym/entry/dmc2gym_sac_pixel_main.py b/dizoo/dmc2gym/entry/dmc2gym_sac_pixel_main.py new file mode 100644 index 0000000000..e03fcc2f05 --- /dev/null +++ b/dizoo/dmc2gym/entry/dmc2gym_sac_pixel_main.py @@ -0,0 +1,89 @@ +from tensorboardX import SummaryWriter +from ditk import logging +import os +import numpy as np +from ding.model.template.qac import ContinuousQAC +from ding.policy import SACPolicy +from ding.envs import BaseEnvManagerV2 +from ding.data import DequeBuffer +from ding.config import compile_config +from ding.framework import task +from ding.framework.context import OnlineRLContext +from ding.framework.middleware import data_pusher, StepCollector, interaction_evaluator, \ + CkptSaver, OffPolicyLearner, termination_checker +from ding.utils import set_pkg_seed +from dizoo.dmc2gym.envs.dmc2gym_env import DMC2GymEnv +from dizoo.dmc2gym.config.dmc2gym_sac_pixel_config import main_config, create_config + + +def main(): + logging.getLogger().setLevel(logging.INFO) + main_config.exp_name = 'dmc2gym_sac_pixel_seed0' + main_config.policy.cuda = True + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + + num_seed = 1 + for seed_i in range(num_seed): + tb_logger = SummaryWriter(os.path.join('./{}/log/'.format(cfg.exp_name), 'seed' + str(seed_i))) + + with task.start(async_mode=False, ctx=OnlineRLContext()): + collector_env = BaseEnvManagerV2( + env_fn=[lambda: DMC2GymEnv(cfg.env) for _ in range(cfg.env.collector_env_num)], cfg=cfg.env.manager + ) + evaluator_env = BaseEnvManagerV2( + env_fn=[lambda: DMC2GymEnv(cfg.env) for _ in range(cfg.env.evaluator_env_num)], cfg=cfg.env.manager + ) + + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + + model = ContinuousQAC(**cfg.policy.model) + logging.info(model) + buffer_ = DequeBuffer(size=cfg.policy.other.replay_buffer.replay_buffer_size) + policy = SACPolicy(cfg.policy, model=model) + + def _add_scalar(ctx): + if ctx.eval_value != -np.inf: + tb_logger.add_scalar('evaluator_step/reward', ctx.eval_value, global_step=ctx.env_step) + collector_rewards = [ctx.trajectories[i]['reward'] for i in range(len(ctx.trajectories))] + collector_mean_reward = sum(collector_rewards) / len(ctx.trajectories) + # collector_max_reward = max(collector_rewards) + # collector_min_reward = min(collector_rewards) + tb_logger.add_scalar('collecter_step/mean_reward', collector_mean_reward, global_step=ctx.env_step) + # tb_logger.add_scalar('collecter_step/max_reward', collector_max_reward, global_step= ctx.env_step) + # tb_logger.add_scalar('collecter_step/min_reward', collector_min_reward, global_step= ctx.env_step) + tb_logger.add_scalar( + 'collecter_step/avg_env_step_per_episode', + ctx.env_step / ctx.env_episode, + global_step=ctx.env_step + ) + + def _add_train_scalar(ctx): + len_train = len(ctx.train_output) + cur_lr_q_avg = sum([ctx.train_output[i]['cur_lr_q'] for i in range(len_train)]) / len_train + cur_lr_p_avg = sum([ctx.train_output[i]['cur_lr_p'] for i in range(len_train)]) / len_train + critic_loss_avg = sum([ctx.train_output[i]['critic_loss'] for i in range(len_train)]) / len_train + policy_loss_avg = sum([ctx.train_output[i]['policy_loss'] for i in range(len_train)]) / len_train + total_loss_avg = sum([ctx.train_output[i]['total_loss'] for i in range(len_train)]) / len_train + tb_logger.add_scalar('learner_step/cur_lr_q_avg', cur_lr_q_avg, global_step=ctx.env_step) + tb_logger.add_scalar('learner_step/cur_lr_p_avg', cur_lr_p_avg, global_step=ctx.env_step) + tb_logger.add_scalar('learner_step/critic_loss_avg', critic_loss_avg, global_step=ctx.env_step) + tb_logger.add_scalar('learner_step/policy_loss_avg', policy_loss_avg, global_step=ctx.env_step) + tb_logger.add_scalar('learner_step/total_loss_avg', total_loss_avg, global_step=ctx.env_step) + + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) + task.use( + StepCollector( + cfg, policy.collect_mode, collector_env, random_collect_size=cfg.policy.random_collect_size + ) + ) + task.use(_add_scalar) + task.use(data_pusher(cfg, buffer_)) + task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) + task.use(_add_train_scalar) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=int(1e5))) + task.use(termination_checker(max_env_step=int(5e6))) + task.run() + + +if __name__ == "__main__": + main() diff --git a/dizoo/dmc2gym/entry/dmc2gym_sac_state_main.py b/dizoo/dmc2gym/entry/dmc2gym_sac_state_main.py new file mode 100644 index 0000000000..6f2393071c --- /dev/null +++ b/dizoo/dmc2gym/entry/dmc2gym_sac_state_main.py @@ -0,0 +1,88 @@ +from ditk import logging +from ding.model import ContinuousQAC +from ding.policy import SACPolicy +from ding.envs import BaseEnvManagerV2 +from ding.data import DequeBuffer +from ding.config import compile_config +from ding.framework import task +from ding.framework.context import OnlineRLContext +from ding.framework.middleware import data_pusher, StepCollector, interaction_evaluator, \ + CkptSaver, OffPolicyLearner, termination_checker +from ding.utils import set_pkg_seed +from dizoo.dmc2gym.envs.dmc2gym_env import DMC2GymEnv +from dizoo.dmc2gym.config.dmc2gym_sac_state_config import main_config, create_config +import numpy as np +from tensorboardX import SummaryWriter +import os + + +def main(): + logging.getLogger().setLevel(logging.INFO) + main_config.exp_name = 'dmc2gym_sac_state_nseed_5M' + main_config.policy.cuda = True + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + + num_seed = 4 + for seed_i in range(num_seed): + tb_logger = SummaryWriter(os.path.join('./{}/log/'.format(cfg.exp_name), 'seed' + str(seed_i))) + + with task.start(async_mode=False, ctx=OnlineRLContext()): + collector_env = BaseEnvManagerV2( + env_fn=[lambda: DMC2GymEnv(cfg.env) for _ in range(cfg.env.collector_env_num)], cfg=cfg.env.manager + ) + evaluator_env = BaseEnvManagerV2( + env_fn=[lambda: DMC2GymEnv(cfg.env) for _ in range(cfg.env.evaluator_env_num)], cfg=cfg.env.manager + ) + + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + + model = ContinuousQAC(**cfg.policy.model) + buffer_ = DequeBuffer(size=cfg.policy.other.replay_buffer.replay_buffer_size) + policy = SACPolicy(cfg.policy, model=model) + + def _add_scalar(ctx): + if ctx.eval_value != -np.inf: + tb_logger.add_scalar('evaluator_step/reward', ctx.eval_value, global_step=ctx.env_step) + collector_rewards = [ctx.trajectories[i]['reward'] for i in range(len(ctx.trajectories))] + collector_mean_reward = sum(collector_rewards) / len(ctx.trajectories) + # collector_max_reward = max(collector_rewards) + # collector_min_reward = min(collector_rewards) + tb_logger.add_scalar('collecter_step/mean_reward', collector_mean_reward, global_step=ctx.env_step) + # tb_logger.add_scalar('collecter_step/max_reward', collector_max_reward, global_step= ctx.env_step) + # tb_logger.add_scalar('collecter_step/min_reward', collector_min_reward, global_step= ctx.env_step) + tb_logger.add_scalar( + 'collecter_step/avg_env_step_per_episode', + ctx.env_step / ctx.env_episode, + global_step=ctx.env_step + ) + + def _add_train_scalar(ctx): + len_train = len(ctx.train_output) + cur_lr_q_avg = sum([ctx.train_output[i]['cur_lr_q'] for i in range(len_train)]) / len_train + cur_lr_p_avg = sum([ctx.train_output[i]['cur_lr_p'] for i in range(len_train)]) / len_train + critic_loss_avg = sum([ctx.train_output[i]['critic_loss'] for i in range(len_train)]) / len_train + policy_loss_avg = sum([ctx.train_output[i]['policy_loss'] for i in range(len_train)]) / len_train + total_loss_avg = sum([ctx.train_output[i]['total_loss'] for i in range(len_train)]) / len_train + tb_logger.add_scalar('learner_step/cur_lr_q_avg', cur_lr_q_avg, global_step=ctx.env_step) + tb_logger.add_scalar('learner_step/cur_lr_p_avg', cur_lr_p_avg, global_step=ctx.env_step) + tb_logger.add_scalar('learner_step/critic_loss_avg', critic_loss_avg, global_step=ctx.env_step) + tb_logger.add_scalar('learner_step/policy_loss_avg', policy_loss_avg, global_step=ctx.env_step) + tb_logger.add_scalar('learner_step/total_loss_avg', total_loss_avg, global_step=ctx.env_step) + + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) + task.use( + StepCollector( + cfg, policy.collect_mode, collector_env, random_collect_size=cfg.policy.random_collect_size + ) + ) + task.use(_add_scalar) + task.use(data_pusher(cfg, buffer_)) + task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) + task.use(_add_train_scalar) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=int(1e5))) + task.use(termination_checker(max_env_step=int(5e6))) + task.run() + + +if __name__ == "__main__": + main() diff --git a/dizoo/dmc2gym/envs/dmc2gym_env.py b/dizoo/dmc2gym/envs/dmc2gym_env.py index b05b091554..8b2ffc519c 100644 --- a/dizoo/dmc2gym/envs/dmc2gym_env.py +++ b/dizoo/dmc2gym/envs/dmc2gym_env.py @@ -3,14 +3,16 @@ from gym.spaces import Box import numpy as np from ding.envs import BaseEnv, BaseEnvTimestep +from ding.envs.common.common_function import affine_transform from ding.torch_utils import to_ndarray from ding.utils import ENV_REGISTRY import dmc2gym +from ding.envs import WarpFrameWrapper, ScaledFloatFrameWrapper, ClipRewardWrapper, ActionRepeatWrapper, FrameStackWrapper def dmc2gym_observation_space(dim, minimum=-np.inf, maximum=np.inf, dtype=np.float32) -> Callable: - def observation_space(from_pixels=True, height=100, width=100, channels_first=True) -> Box: + def observation_space(from_pixels=True, height=84, width=84, channels_first=True) -> Box: if from_pixels: shape = [3, height, width] if channels_first else [height, width, 3] return Box(low=0, high=255, shape=shape, dtype=np.uint8) @@ -40,6 +42,9 @@ def reward_space(frame_skip=1) -> Box: return reward_space +""" +default observation, state, action, reward space for dmc2gym env +""" dmc2gym_env_info = { "ball_in_cup": { "catch": { @@ -106,13 +111,20 @@ def __init__(self, cfg: dict = {}) -> None: assert cfg.task_name in dmc2gym_env_info[ cfg.domain_name], '{}/{}'.format(cfg.task_name, dmc2gym_env_info[cfg.domain_name].keys()) + # default config for dmc2gym env self._cfg = { - "frame_skip": 3, + "frame_skip": 4, + 'warp_frame': False, + 'scale': False, + 'clip_rewards': False, + 'action_repeat': 1, + "frame_stack": 3, "from_pixels": True, "visualize_reward": False, - "height": 100, - "width": 100, + "height": 84, + "width": 84, "channels_first": True, + "resize": 84, } self._cfg.update(cfg) @@ -141,9 +153,26 @@ def reset(self) -> np.ndarray: from_pixels=self._cfg["from_pixels"], height=self._cfg["height"], width=self._cfg["width"], - frame_skip=self._cfg["frame_skip"] + frame_skip=self._cfg["frame_skip"], + channels_first=self._cfg["channels_first"], ) + # optional env wrapper + if self._cfg['warp_frame']: + self._env = WarpFrameWrapper(self._env, size=self._cfg['resize']) + if self._cfg['scale']: + self._env = ScaledFloatFrameWrapper(self._env) + if self._cfg['clip_rewards']: + self._env = ClipRewardWrapper(self._env) + if self._cfg['action_repeat']: + self._env = ActionRepeatWrapper(self._env, self._cfg['action_repeat']) + if self._cfg['frame_stack'] > 1: + self._env = FrameStackWrapper(self._env, self._cfg['frame_stack']) + + # set the obs, action space of wrapped env + self._observation_space = self._env.observation_space + self._action_space = self._env.action_space + if self._replay_path is not None: if gym.version.VERSION > '0.22.0': self._env.metadata.update({'render_modes': ["rgb_array"]}) @@ -165,13 +194,10 @@ def reset(self) -> np.ndarray: elif hasattr(self, '_seed'): self._env.seed(self._seed) - self._final_eval_reward = 0 + self._eval_episode_return = 0 obs = self._env.reset() - if self._cfg["from_pixels"]: - obs = to_ndarray(obs).astype(np.uint8) - else: - obs = to_ndarray(obs).astype(np.float32) + obs = to_ndarray(obs).astype(np.float32) return obs @@ -187,17 +213,15 @@ def seed(self, seed: int, dynamic_seed: bool = True) -> None: def step(self, action: np.ndarray) -> BaseEnvTimestep: action = action.astype('float32') + action = affine_transform(action, min_val=self._env.action_space.low, max_val=self._env.action_space.high) obs, rew, done, info = self._env.step(action) - self._final_eval_reward += rew + self._eval_episode_return += rew if done: - info['final_eval_reward'] = self._final_eval_reward + info['eval_episode_return'] = self._eval_episode_return - if self._cfg["from_pixels"]: - obs = to_ndarray(obs).astype(np.uint8) - else: - obs = to_ndarray(obs).astype(np.float32) + obs = to_ndarray(obs).astype(np.float32) + rew = to_ndarray([rew]).astype(np.float32) # wrapped to be transferred to a array with shape (1,) - rew = to_ndarray([rew]).astype(np.float32) # wrapped to be transfered to a array with shape (1,) return BaseEnvTimestep(obs, rew, done, info) def enable_save_replay(self, replay_path: Optional[str] = None) -> None: @@ -222,4 +246,4 @@ def reward_space(self) -> gym.spaces.Space: return self._reward_space def __repr__(self) -> str: - return "DI-engine Deepmind Control Suite to gym Env: " + self._cfg["domain_name"] + ":" + self._cfg["task_name"] + return "DI-engine DeepMind Control Suite to gym Env: " + self._cfg["domain_name"] + ":" + self._cfg["task_name"] diff --git a/dizoo/evogym/__init__.py b/dizoo/evogym/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/dizoo/evogym/config/bridgewalker_ddpg_config.py b/dizoo/evogym/config/bridgewalker_ddpg_config.py new file mode 100644 index 0000000000..b712a0020e --- /dev/null +++ b/dizoo/evogym/config/bridgewalker_ddpg_config.py @@ -0,0 +1,69 @@ +from easydict import EasyDict + +bridgewalker_ddpg_config = dict( + exp_name='evogym_bridgewalker_ddpg_seed0', + env=dict( + env_id='BridgeWalker-v0', + robot='speed_bot', + robot_dir='../envs', + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=10, + manager=dict(shared_memory=True, ), + # The path to save the game replay + # replay_path='./evogym_walker_ddpg_seed0/video', + ), + policy=dict( + cuda=True, + # load_path="./evogym_walker_ddpg_seed0/ckpt/ckpt_best.pth.tar", + random_collect_size=1000, + model=dict( + obs_shape=59, + action_shape=10, + twin_critic=False, + actor_head_hidden_size=256, + critic_head_hidden_size=256, + action_space='regression', + ), + learn=dict( + update_per_collect=1, + batch_size=256, + learning_rate_actor=1e-3, + learning_rate_critic=1e-3, + ignore_done=False, + target_theta=0.005, + discount_factor=0.99, # discount_factor: 0.97-0.99 + actor_update_freq=1, + noise=False, + ), + collect=dict( + n_sample=1, + unroll_len=1, + noise_sigma=0.1, + ), + other=dict(replay_buffer=dict(replay_buffer_size=1000000, ), ), + ) +) +bridgewalker_ddpg_config = EasyDict(bridgewalker_ddpg_config) +main_config = bridgewalker_ddpg_config + +bridgewalker_ddpg_create_config = dict( + env=dict( + type='evogym', + import_names=['dizoo.evogym.envs.evogym_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict( + type='ddpg', + import_names=['ding.policy.ddpg'], + ), + replay_buffer=dict(type='naive', ), +) +bridgewalker_ddpg_create_config = EasyDict(bridgewalker_ddpg_create_config) +create_config = bridgewalker_ddpg_create_config + +if __name__ == "__main__": + # or you can enter `ding -m serial -c evogym_bridgewalker_ddpg_config.py -s 0 --env-step 1e7` + from ding.entry import serial_pipeline + serial_pipeline((main_config, create_config), seed=0) diff --git a/dizoo/evogym/config/carrier_ppo_config.py b/dizoo/evogym/config/carrier_ppo_config.py new file mode 100644 index 0000000000..16a3919392 --- /dev/null +++ b/dizoo/evogym/config/carrier_ppo_config.py @@ -0,0 +1,65 @@ +from easydict import EasyDict + +carry_ppo_config = dict( + exp_name='evogym_carrier_ppo_seed1', + env=dict( + env_id='Carrier-v0', + robot='carry_bot', + robot_dir='./dizoo/evogym/envs', + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=10, + manager=dict(shared_memory=True, ), + # The path to save the game replay + # replay_path='./evogym_carry_ppo_seed0/video', + ), + policy=dict( + cuda=True, + recompute_adv=True, + # load_path="./evogym_carry_ppo_seed0/ckpt/ckpt_best.pth.tar", + model=dict( + obs_shape=70, + action_shape=12, + action_space='continuous', + ), + action_space='continuous', + learn=dict( + epoch_per_collect=10, + batch_size=256, + learning_rate=3e-3, + value_weight=0.5, + entropy_weight=0.01, + clip_ratio=0.2, + adv_norm=True, + value_norm=True, + ), + collect=dict( + n_sample=2048, + gae_lambda=0.97, + ), + eval=dict(evaluator=dict(eval_freq=5000, )), + ) +) +carry_ppo_config = EasyDict(carry_ppo_config) +main_config = carry_ppo_config + +carry_ppo_create_config = dict( + env=dict( + type='evogym', + import_names=['dizoo.evogym.envs.evogym_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict( + type='ppo', + import_names=['ding.policy.ppo'], + ), + replay_buffer=dict(type='naive', ), +) +carry_ppo_create_config = EasyDict(carry_ppo_create_config) +create_config = carry_ppo_create_config + +if __name__ == "__main__": + # or you can enter `ding -m serial -c evogym_carry_ppo_config.py -s 0 --env-step 1e7` + from ding.entry import serial_pipeline_onpolicy + serial_pipeline_onpolicy((main_config, create_config), seed=0) diff --git a/dizoo/evogym/config/walker_ddpg_config.py b/dizoo/evogym/config/walker_ddpg_config.py new file mode 100644 index 0000000000..58d1c65048 --- /dev/null +++ b/dizoo/evogym/config/walker_ddpg_config.py @@ -0,0 +1,69 @@ +from easydict import EasyDict + +walker_ddpg_config = dict( + exp_name='evogym_walker_ddpg_seed0', + env=dict( + env_id='Walker-v0', + robot='speed_bot', + robot_dir='./dizoo/evogym/envs', + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=10, + manager=dict(shared_memory=True, ), + # The path to save the game replay + # replay_path='./evogym_walker_ddpg_seed0/video', + ), + policy=dict( + cuda=True, + # load_path="./evogym_walker_ddpg_seed0/ckpt/ckpt_best.pth.tar", + random_collect_size=1000, + model=dict( + obs_shape=58, + action_shape=10, + twin_critic=False, + actor_head_hidden_size=256, + critic_head_hidden_size=256, + action_space='regression', + ), + learn=dict( + update_per_collect=1, + batch_size=256, + learning_rate_actor=1e-3, + learning_rate_critic=1e-3, + ignore_done=False, + target_theta=0.005, + discount_factor=0.99, # discount_factor: 0.97-0.99 + actor_update_freq=1, + noise=False, + ), + collect=dict( + n_sample=1, + unroll_len=1, + noise_sigma=0.1, + ), + other=dict(replay_buffer=dict(replay_buffer_size=1000000, ), ), + ) +) +walker_ddpg_config = EasyDict(walker_ddpg_config) +main_config = walker_ddpg_config + +walker_ddpg_create_config = dict( + env=dict( + type='evogym', + import_names=['dizoo.evogym.envs.evogym_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict( + type='ddpg', + import_names=['ding.policy.ddpg'], + ), + replay_buffer=dict(type='naive', ), +) +walker_ddpg_create_config = EasyDict(walker_ddpg_create_config) +create_config = walker_ddpg_create_config + +if __name__ == "__main__": + # or you can enter `ding -m serial -c evogym_walker_ddpg_config.py -s 0 --env-step 1e7` + from ding.entry import serial_pipeline + serial_pipeline((main_config, create_config), seed=0) diff --git a/dizoo/evogym/config/walker_ppo_config.py b/dizoo/evogym/config/walker_ppo_config.py new file mode 100644 index 0000000000..2f1939a3d0 --- /dev/null +++ b/dizoo/evogym/config/walker_ppo_config.py @@ -0,0 +1,65 @@ +from easydict import EasyDict + +walker_ppo_config = dict( + exp_name='evogym_walker_ppo_seed0', + env=dict( + env_id='Walker-v0', + robot='speed_bot', + robot_dir='./dizoo/evogym/envs', + collector_env_num=1, + evaluator_env_num=1, + n_evaluator_episode=1, + stop_value=10, + manager=dict(shared_memory=True, ), + # The path to save the game replay + # replay_path='./evogym_walker_ppo_seed0/video', + ), + policy=dict( + cuda=True, + recompute_adv=True, + # load_path="./evogym_walker_ppo_seed0/ckpt/ckpt_best.pth.tar", + model=dict( + obs_shape=58, + action_shape=10, + action_space='continuous', + ), + action_space='continuous', + learn=dict( + epoch_per_collect=10, + batch_size=256, + learning_rate=3e-4, + value_weight=0.5, + entropy_weight=0.0, + clip_ratio=0.2, + adv_norm=True, + value_norm=True, + ), + collect=dict( + n_sample=2048, + gae_lambda=0.97, + ), + eval=dict(evaluator=dict(eval_freq=5000, )), + ) +) +walker_ppo_config = EasyDict(walker_ppo_config) +main_config = walker_ppo_config + +walker_ppo_create_config = dict( + env=dict( + type='evogym', + import_names=['dizoo.evogym.envs.evogym_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict( + type='ppo', + import_names=['ding.policy.ppo'], + ), + replay_buffer=dict(type='naive', ), +) +walker_ppo_create_config = EasyDict(walker_ppo_create_config) +create_config = walker_ppo_create_config + +if __name__ == "__main__": + # or you can enter `ding -m serial -c evogym_walker_ppo_config.py -s 0 --env-step 1e7` + from ding.entry import serial_pipeline_onpolicy + serial_pipeline_onpolicy((main_config, create_config), seed=0) diff --git a/dizoo/evogym/entry/walker_ppo_eval.py b/dizoo/evogym/entry/walker_ppo_eval.py new file mode 100644 index 0000000000..7c1e706092 --- /dev/null +++ b/dizoo/evogym/entry/walker_ppo_eval.py @@ -0,0 +1,57 @@ +import os +import gym +import torch +from tensorboardX import SummaryWriter +from easydict import EasyDict +from functools import partial + +from ding.config import compile_config +from ding.worker import BaseLearner, SampleSerialCollector, InteractionSerialEvaluator, AdvancedReplayBuffer +from ding.envs import BaseEnvManager +from ding.envs import get_vec_env_setting, create_env_manager +from ding.policy import PPOPolicy +from ding.utils import set_pkg_seed + +from dizoo.evogym.config.walker_ppo_config import main_config, create_config + + +def main(cfg, create_cfg, seed=0): + cfg = compile_config( + cfg, + BaseEnvManager, + PPOPolicy, + BaseLearner, + SampleSerialCollector, + InteractionSerialEvaluator, + AdvancedReplayBuffer, + create_cfg=create_cfg, + save_cfg=True + ) + + create_cfg.policy.type = create_cfg.policy.type + '_command' + env_fn = None + cfg = compile_config(cfg, seed=seed, env=env_fn, auto=True, create_cfg=create_cfg, save_cfg=True) + # Create main components: env, policy + env_fn, collector_env_cfg, evaluator_env_cfg = get_vec_env_setting(cfg.env) + evaluator_env = create_env_manager(cfg.env.manager, [partial(env_fn, cfg=c) for c in evaluator_env_cfg]) + + evaluator_env.enable_save_replay(cfg.env.replay_path) + + # Set random seed for all package and instance + evaluator_env.seed(seed, dynamic_seed=False) + set_pkg_seed(seed, use_cuda=cfg.policy.cuda) + + # Set up RL Policy + policy = PPOPolicy(cfg.policy) + policy.eval_mode.load_state_dict(torch.load(cfg.policy.load_path, map_location='cpu')) + + # evaluate + tb_logger = SummaryWriter(os.path.join('./{}/log/'.format(cfg.exp_name), 'serial')) + evaluator = InteractionSerialEvaluator( + cfg.policy.eval.evaluator, evaluator_env, policy.eval_mode, tb_logger, exp_name=cfg.exp_name + ) + evaluator.eval() + + +if __name__ == "__main__": + main(main_config, create_config, seed=0) diff --git a/dizoo/evogym/envs/__init__.py b/dizoo/evogym/envs/__init__.py new file mode 100644 index 0000000000..6ee6f41343 --- /dev/null +++ b/dizoo/evogym/envs/__init__.py @@ -0,0 +1 @@ +from .evogym_env import EvoGymEnv diff --git a/dizoo/evogym/envs/evogym_env.py b/dizoo/evogym/envs/evogym_env.py new file mode 100644 index 0000000000..e0d8d59cdf --- /dev/null +++ b/dizoo/evogym/envs/evogym_env.py @@ -0,0 +1,178 @@ +from typing import Any, Union, List, Optional +import os +import time +import copy +import numpy as np +import gym +from easydict import EasyDict + +from ding.envs import BaseEnv, BaseEnvTimestep, EvalEpisodeReturnWrapper +from ding.envs.common.common_function import affine_transform +from ding.torch_utils import to_ndarray, to_list +from ding.utils import ENV_REGISTRY + +import evogym.envs +from evogym import WorldObject, sample_robot +from evogym.sim import EvoSim + + +@ENV_REGISTRY.register('evogym') +class EvoGymEnv(BaseEnv): + + @classmethod + def default_config(cls: type) -> EasyDict: + cfg = EasyDict(copy.deepcopy(cls.config)) + cfg.cfg_type = cls.__name__ + 'Dict' + return cfg + + config = dict( + env_id='Walker-v0', + robot='speed_bot', # refer to 'world data' for more robots configurations + robot_h=5, # only used for random robots + robot_w=5, # only used for random robots + robot_pd=None, # only used for random robots, probability distributions of randomly generated components) + robot_dir="" # only used for defined robots, path to the robot config, env/world_data/my_bot.json + ) + + def __init__(self, cfg: dict) -> None: + self._cfg = cfg + self._init_flag = False + self._replay_path = None + if 'robot_dir' not in self._cfg.keys(): + self._cfg = '../' + + def reset(self) -> np.ndarray: + if not self._init_flag: + self._env = self._make_env() + self._env.observation_space.dtype = np.float32 # To unify the format of envs in DI-engine + self._observation_space = self._env.observation_space + self.num_actuators = self._env.get_actuator_indices('robot').size + # by default actions space is double (float64), create a new space with type of type float (float32) + self._action_space = gym.spaces.Box(low=0.6, high=1.6, shape=(self.num_actuators, ), dtype=np.float32) + self._reward_space = gym.spaces.Box( + low=self._env.reward_range[0], high=self._env.reward_range[1], shape=(1, ), dtype=np.float32 + ) + self._init_flag = True + if hasattr(self, '_seed') and hasattr(self, '_dynamic_seed') and self._dynamic_seed: + np_seed = 100 * np.random.randint(1, 1000) + self._env.seed(self._seed + np_seed) + elif hasattr(self, '_seed'): + self._env.seed(self._seed) + if self._replay_path is not None: + gym.logger.set_level(gym.logger.DEBUG) + # make render mode compatible with gym + if gym.version.VERSION > '0.22.0': + self._env.metadata.update({'render_modes': ["rgb_array"]}) + else: + self._env.metadata.update({'render.modes': ["rgb_array"]}) + self._env = gym.wrappers.RecordVideo( + self._env, + video_folder=self._replay_path, + episode_trigger=lambda episode_id: True, + name_prefix='rl-video-{}-{}'.format(id(self), time.time()) + ) + obs = self._env.reset() + obs = to_ndarray(obs).astype('float32') + return obs + + def close(self) -> None: + if self._init_flag: + self._env.close() + self._init_flag = False + + def seed(self, seed: int, dynamic_seed: bool = True) -> None: + self._seed = seed + self._dynamic_seed = dynamic_seed + np.random.seed(self._seed) + + def step(self, action: Union[np.ndarray, list]) -> BaseEnvTimestep: + action = to_ndarray(action).astype(np.float32) + obs, rew, done, info = self._env.step(action) + obs = to_ndarray(obs).astype(np.float32) + rew = to_ndarray([rew]).astype(np.float32) + return BaseEnvTimestep(obs, rew, done, info) + + def _make_env(self): + # robot configuration can be read from file or created randomly + if self._cfg.robot in [None, 'random']: + h, w = 5, 5 + pd = None + if 'robot_h' in self._cfg.keys(): + assert self._cfg.robot_h > 0 + h = self._cfg.robot_h + if 'robot_w' in self._cfg.keys(): + assert self._cfg.robot_w > 0 + w = self._cfg.robot_w + if 'robot_pd' in self._cfg.keys(): + assert isinstance(self._cfg.robot_pd, np.ndarray) + assert self._cfg.robot_w > 0 + pd = self._cfg.robot_pd + structure = sample_robot((h, w), pd) + else: + structure = self.read_robot_from_file(self._cfg.robot, self._cfg.robot_dir) + env = gym.make(self._cfg.env_id, body=structure[0]) + env = EvalEpisodeReturnWrapper(env) + return env + + def enable_save_replay(self, replay_path: Optional[str] = None) -> None: + if replay_path is None: + replay_path = './video' + self._replay_path = replay_path + + def random_action(self) -> np.ndarray: + return self.action_space.sample() + + def __repr__(self) -> str: + return "DI-engine EvoGym Env({})".format(self._cfg.env_id) + + @staticmethod + def create_collector_env_cfg(cfg: dict) -> List[dict]: + collector_cfg = copy.deepcopy(cfg) + collector_env_num = collector_cfg.pop('collector_env_num', 1) + return [collector_cfg for _ in range(collector_env_num)] + + @staticmethod + def create_evaluator_env_cfg(cfg: dict) -> List[dict]: + evaluator_cfg = copy.deepcopy(cfg) + evaluator_env_num = evaluator_cfg.pop('evaluator_env_num', 1) + return [evaluator_cfg for _ in range(evaluator_env_num)] + + @property + def observation_space(self) -> gym.spaces.Space: + return self._observation_space + + @property + def action_space(self) -> gym.spaces.Space: + return self._action_space + + @property + def reward_space(self) -> gym.spaces.Space: + return self._reward_space + + @staticmethod + def read_robot_from_file(file_name, root_dir='../'): + possible_paths = [ + os.path.join(file_name), + os.path.join(f'{file_name}.npz'), + os.path.join(f'{file_name}.json'), + os.path.join(root_dir, 'world_data', file_name), + os.path.join(root_dir, 'world_data', f'{file_name}.npz'), + os.path.join(root_dir, 'world_data', f'{file_name}.json'), + ] + + best_path = None + for path in possible_paths: + if os.path.exists(path): + best_path = path + break + + if best_path.endswith('json'): + robot_object = WorldObject.from_json(best_path) + return (robot_object.get_structure(), robot_object.get_connections()) + if best_path.endswith('npz'): + structure_data = np.load(best_path) + structure = [] + for key, value in structure_data.items(): + structure.append(value) + return tuple(structure) + return None diff --git a/dizoo/evogym/envs/test/test_evogym_env.py b/dizoo/evogym/envs/test/test_evogym_env.py new file mode 100644 index 0000000000..3b0f17410c --- /dev/null +++ b/dizoo/evogym/envs/test/test_evogym_env.py @@ -0,0 +1,37 @@ +import pytest +import numpy as np +from easydict import EasyDict + +from ding.utils import set_pkg_seed +from dizoo.evogym.envs import EvoGymEnv + + +@pytest.mark.envtest +@pytest.mark.parametrize('robot', ['speed_bot', 'random']) +def test_evogym_env_eval_episode_return(robot): + set_pkg_seed(1234, use_cuda=False) + env = EvoGymEnv(EasyDict({'env_id': 'Walker-v0', 'robot': robot, 'robot_dir': '../'})) + env.seed(1234) + env.reset() + action_dim = env.action_space.shape + eval_episode_return = np.array([0.], dtype=np.float32) + if robot == 'speed_bot': + assert env.observation_space.shape == (58, ) + assert action_dim == (10, ) + + while True: + action = np.random.random(size=action_dim) + timestep = env.step(action) + eval_episode_return += timestep.reward + print("{}(dtype: {})".format(timestep.reward, timestep.reward.dtype)) + if timestep.done: + print( + "{}({}), {}({})".format( + timestep.info['eval_episode_return'], type(timestep.info['eval_episode_return']), + eval_episode_return, type(eval_episode_return) + ) + ) + # timestep.reward and the cumulative reward in wrapper EvalEpisodeReturn are not the same. + assert abs(timestep.info['eval_episode_return'].item() - eval_episode_return.item()) / \ + abs(timestep.info['eval_episode_return'].item()) < 1e-5 + break diff --git a/dizoo/evogym/envs/test/visualize_simple_env.py b/dizoo/evogym/envs/test/visualize_simple_env.py new file mode 100644 index 0000000000..2203209fbe --- /dev/null +++ b/dizoo/evogym/envs/test/visualize_simple_env.py @@ -0,0 +1,33 @@ +import gym +from evogym import sample_robot +from gym.wrappers import Monitor + +# import envs from the envs folder and register them +import evogym.envs +from dizoo.evogym.envs.viewer import DingEvoViewer +from evogym.sim import EvoSim + +if __name__ == '__main__': + gym.logger.set_level(gym.logger.DEBUG) + # create a random robot + body, connections = sample_robot((5, 5)) + + # make the SimpleWalkingEnv using gym.make and with the robot information + #env = EvoGymEnv(EasyDict({'env_id': 'Walker-v0', 'robot': 'speed_bot', 'robot_dir': '../'})) + #env.enable_save_replay('video') + + env = gym.make('Walker-v0', body=body) + env.default_viewer = DingEvoViewer(EvoSim(env.world)) + env = Monitor(env, './video', force=True) + env.__class__.render = env.default_viewer.render + env.metadata['render.modes'] = 'rgb_array' + + env.reset() + # step the environment for 200 iterations + for i in range(100): + action = env.action_space.sample() + ob, reward, done, info = env.step(action) + x = env.render() + if done: + env.reset() + env.close() diff --git a/dizoo/evogym/envs/world_data/carry_bot.json b/dizoo/evogym/envs/world_data/carry_bot.json new file mode 100644 index 0000000000..bd9af791d9 --- /dev/null +++ b/dizoo/evogym/envs/world_data/carry_bot.json @@ -0,0 +1,142 @@ +{ + "grid_width": 5, + "grid_height": 5, + "objects": { + "new_object_1": { + "indices": [ + 20, + 15, + 16, + 17, + 18, + 19, + 10, + 11, + 12, + 13, + 14, + 5, + 6, + 8, + 9, + 0, + 1, + 3, + 4 + ], + "types": [ + 1, + 3, + 3, + 3, + 3, + 3, + 1, + 3, + 3, + 3, + 1, + 2, + 4, + 4, + 2, + 2, + 4, + 4, + 2 + ], + "neighbors": { + "20": [ + 15 + ], + "15": [ + 10, + 16, + 20 + ], + "16": [ + 15, + 11, + 17 + ], + "17": [ + 16, + 12, + 18 + ], + "18": [ + 17, + 13, + 19 + ], + "19": [ + 18, + 14 + ], + "10": [ + 5, + 11, + 15 + ], + "11": [ + 10, + 6, + 12, + 16 + ], + "12": [ + 11, + 13, + 17 + ], + "13": [ + 12, + 14, + 8, + 18 + ], + "14": [ + 9, + 13, + 19 + ], + "5": [ + 10, + 0, + 6 + ], + "6": [ + 5, + 1, + 11 + ], + "8": [ + 9, + 3, + 13 + ], + "9": [ + 14, + 4, + 8 + ], + "0": [ + 5, + 1 + ], + "1": [ + 0, + 6 + ], + "3": [ + 4, + 8 + ], + "4": [ + 9, + 3 + ] + } + } + } +} \ No newline at end of file diff --git a/dizoo/evogym/envs/world_data/simple_evironment.json b/dizoo/evogym/envs/world_data/simple_evironment.json new file mode 100644 index 0000000000..b820344789 --- /dev/null +++ b/dizoo/evogym/envs/world_data/simple_evironment.json @@ -0,0 +1,178 @@ +{ + "grid_width": 20, + "grid_height": 10, + "objects": { + "box": { + "indices": [ + 164, + 165, + 166, + 144, + 145, + 146 + ], + "types": [ + 2, + 2, + 2, + 2, + 2, + 2 + ], + "neighbors": { + "164": [ + 144, + 165 + ], + "165": [ + 164, + 166, + 145 + ], + "166": [ + 146, + 165 + ], + "144": [ + 164, + 145 + ], + "145": [ + 144, + 146, + 165 + ], + "146": [ + 145, + 166 + ] + } + }, + "ground": { + "indices": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11, + 12, + 13, + 14, + 15, + 16, + 17, + 18, + 19 + ], + "types": [ + 5, + 5, + 5, + 5, + 5, + 5, + 5, + 5, + 5, + 5, + 5, + 5, + 5, + 5, + 5, + 5, + 5, + 5, + 5, + 5 + ], + "neighbors": { + "0": [ + 1 + ], + "1": [ + 0, + 2 + ], + "2": [ + 1, + 3 + ], + "3": [ + 2, + 4 + ], + "4": [ + 3, + 5 + ], + "5": [ + 4, + 6 + ], + "6": [ + 5, + 7 + ], + "7": [ + 6, + 8 + ], + "8": [ + 7, + 9 + ], + "9": [ + 8, + 10 + ], + "10": [ + 9, + 11 + ], + "11": [ + 10, + 12 + ], + "12": [ + 11, + 13 + ], + "13": [ + 12, + 14 + ], + "14": [ + 13, + 15 + ], + "15": [ + 14, + 16 + ], + "16": [ + 15, + 17 + ], + "17": [ + 16, + 18 + ], + "18": [ + 17, + 19 + ], + "19": [ + 18 + ] + } + } + } +} \ No newline at end of file diff --git a/dizoo/evogym/envs/world_data/speed_bot.json b/dizoo/evogym/envs/world_data/speed_bot.json new file mode 100644 index 0000000000..9f7694774a --- /dev/null +++ b/dizoo/evogym/envs/world_data/speed_bot.json @@ -0,0 +1,120 @@ +{ + "grid_width": 5, + "grid_height": 5, + "objects": { + "new_object_1": { + "indices": [ + 15, + 16, + 17, + 18, + 19, + 10, + 11, + 12, + 13, + 14, + 5, + 6, + 8, + 9, + 0, + 4 + ], + "types": [ + 2, + 3, + 3, + 3, + 1, + 1, + 3, + 3, + 3, + 1, + 3, + 1, + 1, + 3, + 3, + 3 + ], + "neighbors": { + "15": [ + 10, + 16 + ], + "16": [ + 15, + 11, + 17 + ], + "17": [ + 16, + 12, + 18 + ], + "18": [ + 17, + 13, + 19 + ], + "19": [ + 18, + 14 + ], + "10": [ + 5, + 11, + 15 + ], + "11": [ + 10, + 6, + 12, + 16 + ], + "12": [ + 11, + 13, + 17 + ], + "13": [ + 12, + 14, + 8, + 18 + ], + "14": [ + 9, + 13, + 19 + ], + "5": [ + 10, + 0, + 6 + ], + "6": [ + 5, + 11 + ], + "8": [ + 9, + 13 + ], + "9": [ + 14, + 8, + 4 + ], + "0": [ + 5 + ], + "4": [ + 9 + ] + } + } + } +} \ No newline at end of file diff --git a/dizoo/evogym/evogym.gif b/dizoo/evogym/evogym.gif new file mode 100644 index 0000000000..d50b0d9874 Binary files /dev/null and b/dizoo/evogym/evogym.gif differ diff --git a/dizoo/frozen_lake/FrozenLake.gif b/dizoo/frozen_lake/FrozenLake.gif new file mode 100644 index 0000000000..db46a98e39 Binary files /dev/null and b/dizoo/frozen_lake/FrozenLake.gif differ diff --git a/dizoo/frozen_lake/__init__.py b/dizoo/frozen_lake/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/dizoo/frozen_lake/config/__init__.py b/dizoo/frozen_lake/config/__init__.py new file mode 100644 index 0000000000..9bec16a088 --- /dev/null +++ b/dizoo/frozen_lake/config/__init__.py @@ -0,0 +1 @@ +from .frozen_lake_dqn_config import main_config, create_config diff --git a/dizoo/frozen_lake/config/frozen_lake_dqn_config.py b/dizoo/frozen_lake/config/frozen_lake_dqn_config.py new file mode 100644 index 0000000000..84fe0de199 --- /dev/null +++ b/dizoo/frozen_lake/config/frozen_lake_dqn_config.py @@ -0,0 +1,64 @@ +from easydict import EasyDict + +frozen_lake_dqn_config = dict( + exp_name='frozen_lake_seed0', + env=dict( + collector_env_num=8, + evaluator_env_num=5, + n_evaluator_episode=10, + env_id='FrozenLake-v1', + desc=None, + map_name="4x4", + is_slippery=False, + save_replay_gif=False, + ), + policy=dict( + cuda=True, + load_path='frozen_lake_seed0/ckpt/ckpt_best.pth.tar', + model=dict( + obs_shape=16, + action_shape=4, + encoder_hidden_size_list=[128, 128, 64], + dueling=True, + ), + nstep=3, + discount_factor=0.97, + learn=dict( + update_per_collect=5, + batch_size=256, + learning_rate=0.001, + ), + collect=dict(n_sample=10), + eval=dict(evaluator=dict(eval_freq=40, )), + other=dict( + eps=dict( + type='exp', + start=0.8, + end=0.1, + decay=10000, + ), + replay_buffer=dict(replay_buffer_size=20000, ), + ), + ), +) + +frozen_lake_dqn_config = EasyDict(frozen_lake_dqn_config) +main_config = frozen_lake_dqn_config + +frozen_lake_dqn_create_config = dict( + env=dict( + type='frozen_lake', + import_names=['dizoo.frozen_lake.envs.frozen_lake_env'], + ), + env_manager=dict(type='base'), + policy=dict(type='dqn'), + replay_buffer=dict(type='deque', import_names=['ding.data.buffer.deque_buffer_wrapper']), +) + +frozen_lake_dqn_create_config = EasyDict(frozen_lake_dqn_create_config) +create_config = frozen_lake_dqn_create_config + +if __name__ == "__main__": + # or you can enter `ding -m serial -c frozen_lake_dqn_config.py -s 0` + from ding.entry import serial_pipeline + serial_pipeline((main_config, create_config), max_env_step=5000, seed=0) diff --git a/dizoo/frozen_lake/envs/__init__.py b/dizoo/frozen_lake/envs/__init__.py new file mode 100644 index 0000000000..dfec345139 --- /dev/null +++ b/dizoo/frozen_lake/envs/__init__.py @@ -0,0 +1 @@ +from .frozen_lake_env import FrozenLakeEnv diff --git a/dizoo/frozen_lake/envs/frozen_lake_env.py b/dizoo/frozen_lake/envs/frozen_lake_env.py new file mode 100644 index 0000000000..72f179077a --- /dev/null +++ b/dizoo/frozen_lake/envs/frozen_lake_env.py @@ -0,0 +1,144 @@ +from typing import Any, Dict, List, Optional +import imageio +import os +import gymnasium as gymn +import numpy as np +from ding.envs import BaseEnv, BaseEnvTimestep +from ding.torch_utils import to_ndarray +from ding.utils import ENV_REGISTRY + + +@ENV_REGISTRY.register('frozen_lake') +class FrozenLakeEnv(BaseEnv): + + def __init__(self, cfg) -> None: + self._cfg = cfg + assert self._cfg.env_id == "FrozenLake-v1", "yout name is not FrozernLake_v1" + self._init_flag = False + self._save_replay_bool = False + self._save_replay_count = 0 + self._init_flag = False + self._frames = [] + self._replay_path = False + + def reset(self) -> np.ndarray: + if not self._init_flag: + if not self._cfg.desc: #specify maps non-preloaded maps + self._env = gymn.make( + self._cfg.env_id, + desc=self._cfg.desc, + map_name=self._cfg.map_name, + is_slippery=self._cfg.is_slippery, + render_mode="rgb_array" + ) + self._observation_space = self._env.observation_space + self._action_space = self._env.action_space + self._reward_space = gymn.spaces.Box( + low=self._env.reward_range[0], high=self._env.reward_range[1], shape=(1, ), dtype=np.float32 + ) + self._init_flag = True + self._eval_episode_return = 0 + if hasattr(self, '_seed') and hasattr(self, '_dynamic_seed') and self._dynamic_seed: + np_seed = 100 * np.random.randint(1, 1000) + self._env_seed = self._seed + np_seed + elif hasattr(self, '_seed'): + self._env_seed = self._seed + if hasattr(self, '_seed'): + obs, info = self._env.reset(seed=self._env_seed) + else: + obs, info = self._env.reset() + obs = np.eye(16, dtype=np.float32)[obs - 1] + return obs + + def close(self) -> None: + if self._init_flag: + self._env.close() + self._init_flag = False + + def seed(self, seed: int, dynamic_seed: bool = True) -> None: + self._seed = seed + self._dynamic_seed = dynamic_seed + np.random.seed(self._seed) + + def step(self, action: Dict) -> BaseEnvTimestep: + obs, rew, terminated, truncated, info = self._env.step(action[0]) + self._eval_episode_return += rew + obs = np.eye(16, dtype=np.float32)[obs - 1] + rew = to_ndarray([rew]) + if self._save_replay_bool: + picture = self._env.render() + self._frames.append(picture) + if terminated or truncated: + done = True + else: + done = False + if done: + info['eval_episode_return'] = self._eval_episode_return + if self._save_replay_bool: + assert self._replay_path is not None, "your should have a path" + path = os.path.join( + self._replay_path, '{}_episode_{}.gif'.format(self._cfg.env_id, self._save_replay_count) + ) + self.frames_to_gif(self._frames, path) + self._frames = [] + self._save_replay_count += 1 + rew = rew.astype(np.float32) + return BaseEnvTimestep(obs, rew, done, info) + + def random_action(self) -> Dict: + raw_action = self._env.action_space.sample() + my_type = type(self._env.action_space) + return [raw_action] + + def __repr__(self) -> str: + return "DI-engine Frozen Lake Env" + + @property + def observation_space(self) -> gymn.spaces.Space: + return self._observation_space + + @property + def action_space(self) -> gymn.spaces.Space: + return self._action_space + + @property + def reward_space(self) -> gymn.spaces.Space: + return self._reward_space + + def enable_save_replay(self, replay_path: Optional[str] = None) -> None: + if replay_path is None: + replay_path = './video' + self._replay_path = replay_path + self._save_replay_bool = True + self._save_replay_count = 0 + self._frames = [] + + @staticmethod + def frames_to_gif(frames: List[imageio.core.util.Array], gif_path: str, duration: float = 0.1) -> None: + """ + Convert a list of frames into a GIF. + Args: + - frames (List[imageio.core.util.Array]): A list of frames, each frame is an image. + - gif_path (str): The path to save the GIF file. + - duration (float): Duration between each frame in the GIF (seconds). + + Returns: + None, the GIF file is saved directly to the specified path. + """ + # Save all frames as temporary image files + temp_image_files = [] + for i, frame in enumerate(frames): + temp_image_file = f"frame_{i}.png" # Temporary file name + imageio.imwrite(temp_image_file, frame) # Save the frame as a PNG file + temp_image_files.append(temp_image_file) + + # Use imageio to convert temporary image files to GIF + with imageio.get_writer(gif_path, mode='I', duration=duration) as writer: + for temp_image_file in temp_image_files: + image = imageio.imread(temp_image_file) + writer.append_data(image) + + # Clean up temporary image files + for temp_image_file in temp_image_files: + os.remove(temp_image_file) + print(f"GIF saved as {gif_path}") diff --git a/dizoo/frozen_lake/envs/test_frozen_lake_env.py b/dizoo/frozen_lake/envs/test_frozen_lake_env.py new file mode 100644 index 0000000000..c313a264e0 --- /dev/null +++ b/dizoo/frozen_lake/envs/test_frozen_lake_env.py @@ -0,0 +1,44 @@ +import numpy as np +import pytest +from dizoo.frozen_lake.envs import FrozenLakeEnv +from easydict import EasyDict + + +@pytest.mark.envtest +class TestGymHybridEnv: + + def test_my_lake(self): + env = FrozenLakeEnv( + EasyDict({ + 'env_id': 'FrozenLake-v1', + 'desc': None, + 'map_name': "4x4", + 'is_slippery': False, + }) + ) + for _ in range(5): + env.seed(314, dynamic_seed=False) + assert env._seed == 314 + obs = env.reset() + assert obs.shape == ( + 16, + ), "Considering the one-hot encoding format, your observation should have a dimensionality of 16." + for i in range(10): + env.enable_save_replay("./video") + # Both ``env.random_action()``, and utilizing ``np.random`` as well as action space, + # can generate legal random action. + if i < 5: + random_action = np.array([env.action_space.sample()]) + else: + random_action = env.random_action() + timestep = env.step(random_action) + print(timestep) + assert isinstance(timestep.obs, np.ndarray) + assert isinstance(timestep.done, bool) + assert timestep.obs.shape == (16, ) + assert timestep.reward.shape == (1, ) + assert timestep.reward >= env.reward_space.low + assert timestep.reward <= env.reward_space.high + + print(env.observation_space, env.action_space, env.reward_space) + env.close() diff --git a/dizoo/gfootball/README.md b/dizoo/gfootball/README.md index 67e097796a..5b66855f20 100644 --- a/dizoo/gfootball/README.md +++ b/dizoo/gfootball/README.md @@ -9,39 +9,39 @@ ├── README.md ├── __init__.py ├── config -│   ├── gfootball_counter_mappo_config.py -│   ├── gfootball_counter_masac_config.py -│   ├── gfootball_keeper_mappo_config.py -│   └── gfootball_keeper_masac_config.py +│ ├── gfootball_counter_mappo_config.py +│ ├── gfootball_counter_masac_config.py +│ ├── gfootball_keeper_mappo_config.py +│ └── gfootball_keeper_masac_config.py ├── entry -│   ├── __init__.py -│   ├── gfootball_bc_config.py -│   ├── gfootball_bc_kaggle5th_main.py -│   ├── gfootball_bc_rule_lt0_main.py -│   ├── gfootball_bc_rule_main.py -│   ├── gfootball_dqn_config.py -│   └── parallel -│   ├── show_dataset.py -│   ├── test_accuracy.py +│ ├── __init__.py +│ ├── gfootball_bc_config.py +│ ├── gfootball_bc_kaggle5th_main.py +│ ├── gfootball_bc_rule_lt0_main.py +│ ├── gfootball_bc_rule_main.py +│ ├── gfootball_dqn_config.py +│ └── parallel +│ ├── show_dataset.py +│ ├── test_accuracy.py ├── envs -│   ├── __init__.py -│   ├── action -│   ├── fake_dataset.py -│   ├── gfootball_academy_env.py -│   ├── gfootball_env.py -│   ├── gfootballsp_env.py -│   ├── obs -│   ├── reward -│   └── tests +│ ├── __init__.py +│ ├── action +│ ├── fake_dataset.py +│ ├── gfootball_academy_env.py +│ ├── gfootball_env.py +│ ├── gfootballsp_env.py +│ ├── obs +│ ├── reward +│ └── tests ├── gfootball.gif ├── model -│   ├── __init__.py -│   ├── bots -│   ├── conv1d -│   └── q_network +│ ├── __init__.py +│ ├── bots +│ ├── conv1d +│ └── q_network ├── policy -│   ├── __init__.py -│   └── ppo_lstm.py +│ ├── __init__.py +│ └── ppo_lstm.py └── replay.py ``` @@ -127,4 +127,3 @@ https://drive.google.com/file/d/1n-_bF63IQ49b-p0nEZt_NPTL-dmNkoKs/view?usp=shari ### Self Play PPO (work in progress) 使用self-play的PPO算法进行训练的入口，使用DI-engine提供的league模块和PPO算法。具体请见`dizoo/gfootball/entry/parallel/gfootball_ppo_parallel_config.py`入口。 - diff --git a/dizoo/gfootball/config/gfootball_counter_mappo_config.py b/dizoo/gfootball/config/gfootball_counter_mappo_config.py index 92ed3c8841..59c74bf5bd 100644 --- a/dizoo/gfootball/config/gfootball_counter_mappo_config.py +++ b/dizoo/gfootball/config/gfootball_counter_mappo_config.py @@ -35,8 +35,6 @@ action_shape=19, ), learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, epoch_per_collect=10, batch_size=3200, learning_rate=5e-4, diff --git a/dizoo/gfootball/config/gfootball_counter_masac_config.py b/dizoo/gfootball/config/gfootball_counter_masac_config.py index 3abff56713..b30acbca6a 100644 --- a/dizoo/gfootball/config/gfootball_counter_masac_config.py +++ b/dizoo/gfootball/config/gfootball_counter_masac_config.py @@ -76,7 +76,7 @@ import_names=['dizoo.gfootball.envs.gfootball_academy_env'], ), env_manager=dict(type='subprocess'), - policy=dict(type='sac_discrete', ), + policy=dict(type='discrete_sac'), ) gfootball_keeper_masac_default_create_config = EasyDict(gfootball_keeper_masac_default_create_config) create_config = gfootball_keeper_masac_default_create_config diff --git a/dizoo/gfootball/config/gfootball_keeper_mappo_config.py b/dizoo/gfootball/config/gfootball_keeper_mappo_config.py index e6907eb267..5623d6c801 100644 --- a/dizoo/gfootball/config/gfootball_keeper_mappo_config.py +++ b/dizoo/gfootball/config/gfootball_keeper_mappo_config.py @@ -32,8 +32,6 @@ action_shape=19, ), learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, epoch_per_collect=10, batch_size=3200, learning_rate=5e-4, diff --git a/dizoo/gfootball/config/gfootball_keeper_masac_config.py b/dizoo/gfootball/config/gfootball_keeper_masac_config.py index b0cd015386..a89c6dc460 100644 --- a/dizoo/gfootball/config/gfootball_keeper_masac_config.py +++ b/dizoo/gfootball/config/gfootball_keeper_masac_config.py @@ -74,7 +74,7 @@ import_names=['dizoo.gfootball.envs.gfootball_academy_env'], ), env_manager=dict(type='subprocess'), - policy=dict(type='sac_discrete', ), + policy=dict(type='discrete_sac'), ) gfootball_keeper_masac_default_create_config = EasyDict(gfootball_keeper_masac_default_create_config) create_config = gfootball_keeper_masac_default_create_config diff --git a/dizoo/gfootball/envs/gfootball_academy_env.py b/dizoo/gfootball/envs/gfootball_academy_env.py index e0fd084b9d..183ac41d25 100644 --- a/dizoo/gfootball/envs/gfootball_academy_env.py +++ b/dizoo/gfootball/envs/gfootball_academy_env.py @@ -223,7 +223,7 @@ def reset(self): self._env.seed(self._seed + np_seed) elif hasattr(self, '_seed'): self._env.seed(self._seed) - self._final_eval_reward = 0 + self._eval_episode_return = 0 return obs @@ -272,10 +272,10 @@ def step(self, actions): If done=True and sum(rewards)<=0 the reward is 1. If done=True and sum(rewards)>0 the reward is 100. """ - infos['final_eval_reward'] = infos['score_reward'] # TODO(pu) + infos['eval_episode_return'] = infos['score_reward'] # TODO(pu) return BaseEnvTimestep(obs, np.array(-int(done)).astype(np.float32), done, infos) else: - infos['final_eval_reward'] = infos['score_reward'] + infos['eval_episode_return'] = infos['score_reward'] return BaseEnvTimestep(obs, np.array(100).astype(np.float32), done, infos) def get_obs(self): diff --git a/dizoo/gfootball/envs/gfootball_env.py b/dizoo/gfootball/envs/gfootball_env.py index dcbaa85c12..49209d21eb 100644 --- a/dizoo/gfootball/envs/gfootball_env.py +++ b/dizoo/gfootball/envs/gfootball_env.py @@ -126,7 +126,7 @@ def step(self, action: np.array) -> 'GfootballEnv.timestep': info = {'cum_reward': self._reward_helper.cum_reward} if self._is_done: - info['final_eval_reward'] = to_ndarray(self._reward_helper.cum_reward) + info['eval_episode_return'] = to_ndarray(self._reward_helper.cum_reward) if self._save_replay_gif: path = os.path.join( self._replay_path, '{}_episode_{}.gif'.format(self.env_name, self._save_replay_gif_count) diff --git a/dizoo/gfootball/envs/gfootballsp_env.py b/dizoo/gfootball/envs/gfootballsp_env.py index 00860506b7..7c3131b9db 100644 --- a/dizoo/gfootball/envs/gfootballsp_env.py +++ b/dizoo/gfootball/envs/gfootballsp_env.py @@ -56,9 +56,9 @@ def _make_env(self): self._env.seed(self._seed) self._launch_env_flag = True if self.is_evaluator: - self._final_eval_reward = [0, 0] + self._eval_episode_return = [0, 0] else: - self._final_eval_reward = [0, 0] + self._eval_episode_return = [0, 0] def reset(self) -> np.ndarray: if not self._launch_env_flag: @@ -96,22 +96,22 @@ def step(self, action) -> 'GfootballEnv.timestep': obs = to_ndarray(self._encoder.encode(raw_obs)) rew = [rew, rew] obs = [obs, obs] - self._final_eval_reward[0] += raw_rew - self._final_eval_reward[1] += raw_rew + self._eval_episode_return[0] += raw_rew + self._eval_episode_return[1] += raw_rew else: rew = GfootballEnv.calc_reward(raw_rew[0], self._prev_obs, raw_obs[0]) rew_oppo = GfootballEnv.calc_reward(raw_rew[1], self._prev_obs, raw_obs[1]) rew = [rew, rew_oppo] obs = [to_ndarray(self._encoder.encode(raw_obs[0])), to_ndarray(self._encoder.encode(raw_obs[1]))] - self._final_eval_reward[0] += raw_rew[0] - self._final_eval_reward[1] += raw_rew[1] + self._eval_episode_return[0] += raw_rew[0] + self._eval_episode_return[1] += raw_rew[1] if done: if self.is_evaluator: - info['final_eval_reward'] = self._final_eval_reward + info['eval_episode_return'] = self._eval_episode_return else: - info[0]['final_eval_reward'] = self._final_eval_reward[0] - info[1]['final_eval_reward'] = self._final_eval_reward[1] + info[0]['eval_episode_return'] = self._eval_episode_return[0] + info[1]['eval_episode_return'] = self._eval_episode_return[1] return BaseEnvTimestep(obs, rew, done, info) diff --git a/dizoo/gfootball/policy/ppo_lstm.py b/dizoo/gfootball/policy/ppo_lstm.py index 8e6c4c7b31..4380f1261c 100644 --- a/dizoo/gfootball/policy/ppo_lstm.py +++ b/dizoo/gfootball/policy/ppo_lstm.py @@ -36,8 +36,6 @@ class PPOPolicy(Policy): nstep_return=False, nstep=3, learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, # How many updates(iterations) to train after collector's one collection. # Bigger "update_per_collect" means bigger off-policy. # collect data -> update policy-> collect data -> ... diff --git a/dizoo/gym_anytrading/config/stocks_dqn_config.py b/dizoo/gym_anytrading/config/stocks_dqn_config.py index c16ab0a5a5..c05a1f5974 100644 --- a/dizoo/gym_anytrading/config/stocks_dqn_config.py +++ b/dizoo/gym_anytrading/config/stocks_dqn_config.py @@ -78,13 +78,11 @@ import_names=['dizoo.gym_anytrading.envs.stocks_env'], ), env_manager=dict(type='base'), - policy=dict( - type='dqn', - ), + policy=dict(type='dqn', ), evaluator=dict( type='trading_interaction', import_names=['dizoo.gym_anytrading.worker'], - ), + ), ) stocks_dqn_create_config = EasyDict(stocks_dqn_create_config) create_config = stocks_dqn_create_config diff --git a/dizoo/gym_anytrading/envs/README.md b/dizoo/gym_anytrading/envs/README.md index f569c10de3..0be3fe219b 100644 --- a/dizoo/gym_anytrading/envs/README.md +++ b/dizoo/gym_anytrading/envs/README.md @@ -57,9 +57,9 @@ If profit or loss occurs, it means that one of the following two cycles in state ### Current Profit Calculation -According to the above defination, we can easily know that the formula of accumulative profit is: +According to the above definition, we can easily know that the formula of accumulative profit is: -$\prod_{buying\ long}(r_{curr}/r_{pre}\ *\ cost) + \prod_{short\ selling}((2-r_{curr}/r_{pre})\ *\ cost)$ +$\prod_{buying\ long}(r_{curr}/r_{pre}\ *\ cost) * \prod_{short\ selling}((2-r_{curr}/r_{pre})\ *\ cost)$ ### Reward Function @@ -78,7 +78,7 @@ Comparing the objective function ($\mathbb{E}_{\tau}\sum\ r$) in reinforcement l so that maximize $\mathbb{E}_{\tau} \sum r$ is equivalent to maximize $\mathbb{E}_{\tau}[\prod_{buying\ long}(r_{curr}/r_{pre}\ *\ cost) + \prod_{short\ selling}((2-r_{curr}/r_{pre})\ *\ cost)]$ -The experimental results show that such a defination is better than the original gym-anytrading accumulated reward function :$\sum(r_{curr} - r_{pre})$. +The experimental results show that such a definition is better than the original gym-anytrading accumulated reward function :$\sum(r_{curr} - r_{pre})$. ### Render Function As you see, you can use `render` method to plot the position and profit at one episode. diff --git a/dizoo/gym_anytrading/envs/statemachine.png b/dizoo/gym_anytrading/envs/statemachine.png index 2c355939f1..bbc3d39d89 100644 Binary files a/dizoo/gym_anytrading/envs/statemachine.png and b/dizoo/gym_anytrading/envs/statemachine.png differ diff --git a/dizoo/gym_anytrading/envs/stocks_env.py b/dizoo/gym_anytrading/envs/stocks_env.py index a71ae3f59b..8d34caa182 100644 --- a/dizoo/gym_anytrading/envs/stocks_env.py +++ b/dizoo/gym_anytrading/envs/stocks_env.py @@ -69,7 +69,7 @@ def _process_data(self, start_idx: int = None) -> Any: # validate index if start_idx is None: if self.train_range == None or self.test_range == None: - self.start_idx = np.random.randint(self.window_size, len(self.df) - self._cfg.eps_length) + self.start_idx = np.random.randint(self.window_size - 1, len(self.df) - self._cfg.eps_length) elif self._env_id[-1] == 'e': boundary = int(len(self.df) * (1 + self.test_range)) assert len(self.df) - self._cfg.eps_length > boundary + self.window_size,\ diff --git a/dizoo/gym_anytrading/envs/trading_env.py b/dizoo/gym_anytrading/envs/trading_env.py index dfce59e317..a29c5b2215 100644 --- a/dizoo/gym_anytrading/envs/trading_env.py +++ b/dizoo/gym_anytrading/envs/trading_env.py @@ -185,7 +185,7 @@ def step(self, action: np.ndarray) -> BaseEnvTimestep: if self._env_id[-1] == 'e' and self.cnt % self.plot_freq == 0: self.render() info['max_possible_profit'] = np.log(self.max_possible_profit()) - info['final_eval_reward'] = self._total_reward + info['eval_episode_return'] = self._total_reward step_reward = to_ndarray([step_reward]).astype(np.float32) return BaseEnvTimestep(observation, step_reward, self._done, info) @@ -252,7 +252,7 @@ def create_collector_env_cfg(cfg: dict) -> List[dict]: collector_env_num = cfg.pop('collector_env_num') collector_env_cfg = [copy.deepcopy(cfg) for _ in range(collector_env_num)] for i in range(collector_env_num): - collector_env_cfg[i]['env_id'] += ('-' + str(i) + 'e') + collector_env_cfg[i]['env_id'] += ('-' + str(i) + 'c') return collector_env_cfg # override diff --git a/dizoo/gym_anytrading/worker/trading_serial_evaluator.py b/dizoo/gym_anytrading/worker/trading_serial_evaluator.py index 66da63bce8..4287d48423 100644 --- a/dizoo/gym_anytrading/worker/trading_serial_evaluator.py +++ b/dizoo/gym_anytrading/worker/trading_serial_evaluator.py @@ -6,7 +6,7 @@ from ding.envs import BaseEnvManager from ding.worker import VectorEvalMonitor, InteractionSerialEvaluator -from ding.torch_utils import to_tensor, to_ndarray +from ding.torch_utils import to_tensor, to_ndarray, to_item from ding.utils import SERIAL_EVALUATOR_REGISTRY, import_module @@ -30,6 +30,7 @@ class TradingSerialEvaluator(InteractionSerialEvaluator): ), type='trading_interaction', ) + def __init__( self, cfg: dict, @@ -65,7 +66,7 @@ def eval( - n_episode (:obj:`int`): Number of evaluation episodes. Returns: - stop_flag (:obj:`bool`): Whether this training program can be ended. - - return_info (:obj:`dict`): Current evaluation return information. + - episode_info (:obj:`dict`): Current evaluation return information. ''' if n_episode is None: @@ -73,7 +74,6 @@ def eval( assert n_episode is not None, "please indicate eval n_episode" envstep_count = 0 info = {} - return_info = [] eval_monitor = TradingEvalMonitor(self._env.env_num, n_episode) self._env.reset() self._policy.reset() @@ -103,11 +103,9 @@ def eval( if t.done: # Env reset is done by env_manager automatically. self._policy.reset([env_id]) - reward = t.info['final_eval_reward'] - if 'episode_info' in t.info: - eval_monitor.update_info(env_id, t.info['episode_info']) + reward = t.info['eval_episode_return'] + eval_monitor.update_info(env_id, t.info) eval_monitor.update_reward(env_id, reward) - return_info.append(t.info) #========== only used by anytrading ======= if 'max_possible_profit' in t.info: @@ -122,7 +120,7 @@ def eval( ) envstep_count += 1 duration = self._timer.value - episode_reward = eval_monitor.get_episode_reward() + episode_return = eval_monitor.get_episode_return() info = { 'train_iter': train_iter, 'ckpt_name': 'iteration_{}.pth.tar'.format(train_iter), @@ -132,11 +130,11 @@ def eval( 'evaluate_time': duration, 'avg_envstep_per_sec': envstep_count / duration, 'avg_time_per_episode': n_episode / duration, - 'reward_mean': np.mean(episode_reward), - 'reward_std': np.std(episode_reward), - 'reward_max': np.max(episode_reward), - 'reward_min': np.min(episode_reward), - # 'each_reward': episode_reward, + 'reward_mean': np.mean(episode_return), + 'reward_std': np.std(episode_return), + 'reward_max': np.max(episode_return), + 'reward_min': np.min(episode_return), + # 'each_reward': episode_return, } episode_info = eval_monitor.get_episode_info() if episode_info is not None: @@ -172,19 +170,20 @@ def eval( from ding.utils import fps self._tb_logger.add_video(video_title, videos, render_iter, fps(self._env)) - eval_reward = np.mean(episode_reward) - if eval_reward > self._max_eval_reward: + episode_return = np.mean(episode_return) + if episode_return > self._max_episode_return: if save_ckpt_fn: save_ckpt_fn('ckpt_best.pth.tar') - self._max_eval_reward = eval_reward - stop_flag = eval_reward >= self._stop_value and train_iter > 0 + self._max_episode_return = episode_return + stop_flag = episode_return >= self._stop_value and train_iter > 0 if stop_flag: self._logger.info( "[DI-engine serial pipeline] " + - "Current eval_reward: {} is greater than stop_value: {}".format(eval_reward, self._stop_value) + + "Current episode_return: {} is greater than stop_value: {}".format(episode_return, self._stop_value) + ", so your RL agent is converged, you can refer to 'log/evaluator/evaluator_logger.txt' for details." ) - return stop_flag, return_info + episode_info = to_item(episode_info) + return stop_flag, episode_info class TradingEvalMonitor(VectorEvalMonitor): @@ -193,7 +192,7 @@ class TradingEvalMonitor(VectorEvalMonitor): Inherit VectorEvalMonitor for trading env. Add func update_max_profit and get_max_episode_profit in order to log the max_profit for every episode. Interfaces: - Besides (__init__, is_finished, update_info, update_reward, get_episode_reward,\ + Besides (__init__, is_finished, update_info, update_reward, get_episode_return,\ get_latest_reward, get_current_episode, get_episode_info), there are\ (update_max_profit, get_max_episode_profit). """ diff --git a/dizoo/gym_hybrid/config/gym_hybrid_ddpg_config.py b/dizoo/gym_hybrid/config/gym_hybrid_ddpg_config.py index 9c00f408ee..854f5f3939 100644 --- a/dizoo/gym_hybrid/config/gym_hybrid_ddpg_config.py +++ b/dizoo/gym_hybrid/config/gym_hybrid_ddpg_config.py @@ -10,13 +10,10 @@ env_id='Moving-v0', # ['Sliding-v0', 'Moving-v0'] n_evaluator_episode=5, stop_value=1.8, - # The path to save the game replay - # replay_path='gym_hybrid_ddpg_seed0/video', ), policy=dict( cuda=True, priority=False, - load_path="./gym_hybrid_ddpg_seed0/ckpt/ckpt_best.pth.tar", random_collect_size=0, # hybrid action space not support random collect now action_space='hybrid', model=dict( @@ -71,4 +68,4 @@ if __name__ == "__main__": # or you can enter `ding -m serial -c gym_hybrid_ddpg_config.py -s 0` from ding.entry import serial_pipeline - serial_pipeline([main_config, create_config], seed=0) + serial_pipeline([main_config, create_config], seed=0, max_env_step=int(1e7)) diff --git a/dizoo/gym_hybrid/config/gym_hybrid_hppo_config.py b/dizoo/gym_hybrid/config/gym_hybrid_hppo_config.py index edf043eb62..2011972e19 100644 --- a/dizoo/gym_hybrid/config/gym_hybrid_hppo_config.py +++ b/dizoo/gym_hybrid/config/gym_hybrid_hppo_config.py @@ -13,7 +13,6 @@ ), policy=dict( cuda=True, - priority=False, action_space='hybrid', recompute_adv=True, model=dict( @@ -32,14 +31,12 @@ epoch_per_collect=10, batch_size=320, learning_rate=3e-4, - value_weight=0.5, - entropy_weight=0.03, - clip_ratio=0.2, + entropy_weight=0.5, adv_norm=True, value_norm=True, ), collect=dict( - n_sample=int(3200), + n_sample=3200, discount_factor=0.99, gae_lambda=0.95, collector=dict(collect_print_freq=1000, ), @@ -64,4 +61,4 @@ if __name__ == "__main__": # or you can enter `ding -m serial -c gym_hybrid_hppo_config.py -s 0` from ding.entry import serial_pipeline_onpolicy - serial_pipeline_onpolicy([main_config, create_config], seed=0) + serial_pipeline_onpolicy([main_config, create_config], seed=0, max_env_step=int(1e7)) diff --git a/dizoo/gym_hybrid/config/gym_hybrid_mpdqn_config.py b/dizoo/gym_hybrid/config/gym_hybrid_mpdqn_config.py index 6791349877..dd2bf32343 100644 --- a/dizoo/gym_hybrid/config/gym_hybrid_mpdqn_config.py +++ b/dizoo/gym_hybrid/config/gym_hybrid_mpdqn_config.py @@ -13,9 +13,6 @@ ), policy=dict( cuda=True, - priority=False, - # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. - priority_IS_weight=False, discount_factor=0.99, nstep=1, model=dict( @@ -28,11 +25,6 @@ action_mask=[[1, 0], [0, 1], [0, 0]], ), learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, - # How many updates(iterations) to train after collector's one collection. - # Bigger "update_per_collect" means bigger off-policy. - # collect data -> update policy-> collect data -> ... update_per_collect=500, # 10~500 batch_size=320, learning_rate_dis=3e-4, @@ -83,4 +75,4 @@ if __name__ == "__main__": # or you can enter `ding -m serial -c gym_hybrid_mpdqn_config.py -s 0` from ding.entry import serial_pipeline - serial_pipeline([main_config, create_config], seed=0) + serial_pipeline([main_config, create_config], seed=0, max_env_step=int(1e7)) diff --git a/dizoo/gym_hybrid/config/gym_hybrid_pdqn_config.py b/dizoo/gym_hybrid/config/gym_hybrid_pdqn_config.py index 817177adcf..79d28b07bd 100644 --- a/dizoo/gym_hybrid/config/gym_hybrid_pdqn_config.py +++ b/dizoo/gym_hybrid/config/gym_hybrid_pdqn_config.py @@ -13,9 +13,6 @@ ), policy=dict( cuda=True, - priority=False, - # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. - priority_IS_weight=False, discount_factor=0.99, nstep=1, model=dict( @@ -26,11 +23,6 @@ ), ), learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, - # How many updates(iterations) to train after collector's one collection. - # Bigger "update_per_collect" means bigger off-policy. - # collect data -> update policy-> collect data -> ... update_per_collect=500, # 10~500 batch_size=320, learning_rate_dis=3e-4, @@ -81,4 +73,4 @@ if __name__ == "__main__": # or you can enter `ding -m serial -c gym_hybrid_pdqn_config.py -s 0` from ding.entry import serial_pipeline - serial_pipeline([main_config, create_config], seed=0) + serial_pipeline([main_config, create_config], seed=0, max_env_step=int(1e7)) diff --git a/dizoo/gym_hybrid/entry/gym_hybrid_ddpg_eval.py b/dizoo/gym_hybrid/entry/gym_hybrid_ddpg_eval.py index ef30169a48..388928a61b 100644 --- a/dizoo/gym_hybrid/entry/gym_hybrid_ddpg_eval.py +++ b/dizoo/gym_hybrid/entry/gym_hybrid_ddpg_eval.py @@ -8,33 +8,23 @@ from ding.config import compile_config from ding.worker import BaseLearner, SampleSerialCollector, InteractionSerialEvaluator, AdvancedReplayBuffer from ding.envs import BaseEnvManager, DingEnvWrapper -from ding.envs import get_vec_env_setting, create_env_manager +from ding.envs import get_vec_env_setting from ding.policy import DDPGPolicy -from ding.model import QAC +from ding.model import ContinuousQAC from ding.utils import set_pkg_seed from ding.rl_utils import get_epsilon_greedy_fn from dizoo.gym_hybrid.config.gym_hybrid_ddpg_config import gym_hybrid_ddpg_config, gym_hybrid_ddpg_create_config def main(main_cfg, create_cfg, seed=0): - cfg = compile_config( - main_cfg, - BaseEnvManager, - DDPGPolicy, - BaseLearner, - SampleSerialCollector, - InteractionSerialEvaluator, - AdvancedReplayBuffer, - create_cfg=create_cfg, - save_cfg=True - ) - - create_cfg.policy.type = create_cfg.policy.type + '_command' - env_fn = None - cfg = compile_config(cfg, seed=seed, env=env_fn, auto=True, create_cfg=create_cfg, save_cfg=True) + # Specify evaluation arguments + main_cfg.policy.load_path = './ckpt_best.pth.tar' + main_cfg.env.replay_path = './' + main_cfg.env.evaluator_env_num = 1 # only 1 env for save replay + cfg = compile_config(main_cfg, seed=seed, auto=True, create_cfg=create_cfg, save_cfg=True) # Create main components: env, policy env_fn, collector_env_cfg, evaluator_env_cfg = get_vec_env_setting(cfg.env) - evaluator_env = create_env_manager(cfg.env.manager, [partial(env_fn, cfg=c) for c in evaluator_env_cfg]) + evaluator_env = BaseEnvManager([partial(env_fn, cfg=c) for c in evaluator_env_cfg], cfg.env.manager) evaluator_env.enable_save_replay(cfg.env.replay_path) # switch save replay interface @@ -43,7 +33,7 @@ def main(main_cfg, create_cfg, seed=0): set_pkg_seed(seed, use_cuda=cfg.policy.cuda) # Set up RL Policy - model = QAC(**cfg.policy.model) + model = ContinuousQAC(**cfg.policy.model) policy = DDPGPolicy(cfg.policy, model=model) policy.eval_mode.load_state_dict(torch.load(cfg.policy.load_path, map_location='cpu')) diff --git a/dizoo/gym_hybrid/entry/gym_hybrid_ddpg_main.py b/dizoo/gym_hybrid/entry/gym_hybrid_ddpg_main.py index 8b2cce8656..aaa8753d6b 100644 --- a/dizoo/gym_hybrid/entry/gym_hybrid_ddpg_main.py +++ b/dizoo/gym_hybrid/entry/gym_hybrid_ddpg_main.py @@ -7,7 +7,7 @@ from ding.worker import BaseLearner, SampleSerialCollector, InteractionSerialEvaluator, AdvancedReplayBuffer from ding.envs import BaseEnvManager from ding.policy import DDPGPolicy -from ding.model import QAC +from ding.model import ContinuousQAC from ding.utils import set_pkg_seed from ding.rl_utils import get_epsilon_greedy_fn from dizoo.gym_hybrid.envs.gym_hybrid_env import GymHybridEnv @@ -43,7 +43,7 @@ def main(cfg, seed=0): set_pkg_seed(seed, use_cuda=cfg.policy.cuda) # Set up RL Policy - model = QAC(**cfg.policy.model) + model = ContinuousQAC(**cfg.policy.model) policy = DDPGPolicy(cfg.policy, model=model) # Set up collection, training and evaluation utilities diff --git a/dizoo/gym_hybrid/envs/README.md b/dizoo/gym_hybrid/envs/README.md new file mode 100644 index 0000000000..a89b9c113e --- /dev/null +++ b/dizoo/gym_hybrid/envs/README.md @@ -0,0 +1,21 @@ +# Modified gym-hybrid + +The gym-hybrid directory is modified from https://github.com/thomashirtz/gym-hybrid. +We add the HardMove environment additionally. (Please refer to https://arxiv.org/abs/2109.05490 Section 5.1 for details about HardMove env.) + +Specifically, the modified gym-hybrid contains the following three types of environments: + +- Moving-v0 +- Sliding-v0 +- HardMove-v0 + +### Install Guide + +```bash +cd DI-engine/dizoo/gym_hybrid/envs/gym-hybrid +pip install -e . +``` + +## Acknowledgement + +https://github.com/thomashirtz/gym-hybrid \ No newline at end of file diff --git a/dizoo/gym_hybrid/envs/gym-hybrid/README.md b/dizoo/gym_hybrid/envs/gym-hybrid/README.md new file mode 100644 index 0000000000..153b07db96 --- /dev/null +++ b/dizoo/gym_hybrid/envs/gym-hybrid/README.md @@ -0,0 +1,147 @@ +# gym-hybrid + +Repository containing a collection of environment for reinforcement learning task possessing discrete-continuous hybrid action space. + +## "Sliding-v0" and "Moving-v0" + + + +"Moving-v0" and "Sliding-v0" are sandbox environments for parameterized action-space algorithms. The goal of the agent is to stop inside a target area. + +The field is a square with a side length of 2. The target area is a circle with radius 0.1. There is three discrete actions: turn, accelerate, and break. In addition to the action, there is 2 possible complementary parameters: acceleration and rotation. + +The episode terminates if one of the three condition is filled: +* the agent stop inside the target area, +* the agent leaves the field, +* the step count is higher than the limit (set by default at 200). + +The moving environment doesn't take into account the conservation of inertia, while the sliding environment does. `Sliding-v0` is therefore more realistic than `Moving-v0`. + +All the parameters, actions, states and rewards are the same between the two environments. Only the underlying physics changes. + +### State +The [state](https://github.com/thomashirtz/gym-hybrid/blob/fee4bf5de2dc1dd0d2a5431498124b2c071a2344/gym_hybrid/environments.py#L126) is constituted of a list of 10 elements. The environment related values are: the current step divided by the maximum step, and the position of the target (x and y). The player related values are the position (x and y), the speed, the direction (cosine and sine), the distance related to the target, and an indicator that becomes 1 if the player is inside the target zone. +```python +state = [ + agent.x, + agent.y, + agent.speed, + np.cos(agent.theta), + np.sin(agent.theta), + target.x, + target.y, + distance, + 0 if distance > target_radius else 1, + current_step / max_step +] +``` + +### Reward +The [reward](https://github.com/thomashirtz/gym-hybrid/blob/fee4bf5de2dc1dd0d2a5431498124b2c071a2344/gym_hybrid/environments.py#L141) is the distance of the agent from the target of the last step minus the current distance. There is a penalty (set by default at a low value) to incentivize the learning algorithm to score as quickly as possible. A bonus reward of one is added if the player achieve to stop inside the target area. A malus of one is applied if the step count exceed the limit or if the player leaves the field. + +### Actions + +**The action ids are:** +1. Accelerate +2. Turn +3. Break + +**The parameters are:** +1. Acceleration value +2. Rotation value + +**There is two distinct way to format an action:** + +Action with all the parameters (convenient if the model output all the parameters): +```python +action = (action_id, [acceleration_value, rotation_value]) +``` +Example of a valid actions: +```python +action = (0, [0.1, 0.4]) +action = (1, [0.0, 0.2]) +action = (2, [0.1, 0.3]) +``` +Note: Only the parameter related to the action chosen will be used. + +Action with only the parameter related to the action id (convenient for algorithms that output only the parameter +of the chosen action, since it doesn't require to pad the action): +```python +action = (0, [acceleration_value]) +action = (1, [rotation_value]) +action = (2, []) +``` +Example of valid actions: +```python +action = (0, [0.1]) +action = (1, [0.2]) +action = (2, []) +``` +### Basics +Make and initialize an environment: +```python +import gym +import gym_parametrized + +sliding_env = gym.make('Sliding-v0') +sliding_env.reset() + +moving_env = gym.make('Moving-v0') +moving_env.reset() +``` + +Get the action space and the observation space: +```python +ACTION_SPACE = env.action_space[0].n +PARAMETERS_SPACE = env.action_space[1].shape[0] +OBSERVATION_SPACE = env.observation_space.shape[0] +``` + +Run a random agent: +```python +done = False +while not done: + state, reward, done, info = env.step(env.action_space.sample()) + print(f'State: {state} Reward: {reward} Done: {done}') +``` +### Parameters +The parameter that can be modified during the initialization are: +* `seed` (default = None) +* `max_turn`, angle in radi that can be achieved in one step (default = np.pi/2) +* `max_acceleration`, acceleration that can be achieved in one step (if the input parameter is 1) (default = 0.5) +* `delta_t`, time step of one step (default = 0.005) +* `max_step`, limit of the number of step before the end of an environment (default = 200) +* `penalty`, value substracted to the reward each step to incentivise the agent to finish the environment quicker (default = 0.001) + +Initialization with custom parameters: +```python +env = gym.make( + 'Moving-v0', + seed=0, + max_turn=1, + max_acceleration=1.0, + delta_t=0.001, + max_step=500, + penalty=0.01 +) +``` + +### Render & Recording +Two testing files are avalaible to show users how to render and record the environment: +* [Python file example for recording](tests/moving_record.py) +* [Python file example for rendering](tests/moving_render.py) + +## Disclaimer +Even though the mechanics of the environment are done, maybe the hyperparameters will need some further adjustments. + +## Reference +This environment is described in several papers such as: +[Parametrized Deep Q-Networks Learning, Xiong et al., 2018](https://arxiv.org/pdf/1810.06394.pdf) +[Hybrid Actor-Critic Reinforcement Learning in Parameterized Action Space, Fan et al., 2019](https://arxiv.org/pdf/1903.01344.pdf) + +## Installation + +Direct Installation from github using pip by running this command: +```shell +pip install git+https://github.com/thomashirtz/gym-hybrid#egg=gym-hybrid +``` diff --git a/dizoo/gym_hybrid/envs/gym-hybrid/gym_hybrid/__init__.py b/dizoo/gym_hybrid/envs/gym-hybrid/gym_hybrid/__init__.py new file mode 100644 index 0000000000..89cb5d7764 --- /dev/null +++ b/dizoo/gym_hybrid/envs/gym-hybrid/gym_hybrid/__init__.py @@ -0,0 +1,17 @@ +from gym.envs.registration import register +from gym_hybrid.environments import MovingEnv +from gym_hybrid.environments import SlidingEnv +from gym_hybrid.environments import HardMoveEnv + +register( + id='Moving-v0', + entry_point='gym_hybrid:MovingEnv', +) +register( + id='Sliding-v0', + entry_point='gym_hybrid:SlidingEnv', +) +register( + id='HardMove-v0', + entry_point='gym_hybrid:HardMoveEnv', +) diff --git a/dizoo/gym_hybrid/envs/gym-hybrid/gym_hybrid/agents.py b/dizoo/gym_hybrid/envs/gym-hybrid/gym_hybrid/agents.py new file mode 100644 index 0000000000..4b8669ee98 --- /dev/null +++ b/dizoo/gym_hybrid/envs/gym-hybrid/gym_hybrid/agents.py @@ -0,0 +1,117 @@ +from itertools import product + +import numpy as np + + +class BaseAgent: + + def __init__(self, break_value: float, delta_t: float): + self.x = None + self.y = None + self.phi = None # angle of the velocity vector + self.theta = None # direction of the agent + self.speed = None + self.delta_t = delta_t + self.break_value = break_value + + def accelerate(self, value: float) -> None: + raise NotImplementedError + + def break_(self) -> None: + raise NotImplementedError + + def turn(self, value: float) -> None: + raise NotImplementedError + + def reset(self, x: float, y: float, direction: float) -> None: + self.x = x + self.y = y + self.speed = 0 + self.theta = direction + + def _step(self) -> None: + angle = self.theta if self.phi is None else self.phi + self.x += self.delta_t * self.speed * np.cos(angle) + self.y += self.delta_t * self.speed * np.sin(angle) + + +class MovingAgent(BaseAgent): + + def __init__(self, break_value: float, delta_t: float): + super(MovingAgent, self).__init__(break_value, delta_t) + + def accelerate(self, value: float) -> None: + self.speed += value + self._step() + + def break_(self) -> None: + self.speed = 0 if self.speed < self.break_value else self.speed - self.break_value + self._step() + + def turn(self, value: float) -> None: + self.theta = (self.theta + value) % (2 * np.pi) + self._step() + + +class SlidingAgent(BaseAgent): + + def __init__(self, break_value: float, delta_t: float): + super(SlidingAgent, self).__init__(break_value, delta_t) + self.phi = 0 + + def accelerate(self, value: float) -> None: + # Adding two polar vectors: https://math.stackexchange.com/a/1365938/849658 + # phi_1, r_1 = self.theta, value # the direction of the agent and the magnitude induced by the action + # phi_2, r_2 = self.phi, self.speed # the direction of the velocity vector and its magnitude + speed = np.sqrt(value ** 2 + self.speed ** 2 + 2 * value * self.speed * np.cos(self.phi - self.theta)) + angle = self.theta + np.arctan2( + self.speed * np.sin(self.phi - self.theta), value + self.speed * np.cos(self.phi - self.theta) + ) + self.speed = speed + self.phi = angle + self._step() + + def break_(self) -> None: + self.speed = 0 if self.speed < self.break_value else self.speed - self.break_value + self.phi = self.theta if self.speed == 0 else self.phi # not sure it is needed + self._step() + + def turn(self, value: float) -> None: + self.theta = (self.theta + value) % (2 * np.pi) + self._step() + + +class HardMoveAgent(BaseAgent): + + def __init__(self, break_value: float, delta_t: float, num_actuators: int = 4): + super(HardMoveAgent, self).__init__(break_value, delta_t) + self.phi = 0 + self.num_actuators = num_actuators + # NOTE: meta_to_mask + self.K = 2 ** self.num_actuators + self.meta_to_mask = list(product(*[list(range(2)) for _ in range(self.num_actuators)])) + + def accelerate(self, value: float) -> None: + pass + + def break_(self) -> None: + pass + + def turn(self, value: float) -> None: + pass + + def move(self, move_direction_meta: int, move_distances: list) -> None: + move_directions_mask = self.meta_to_mask[int(move_direction_meta)] + self.move_vector = np.array( + [ + move_directions_mask[i] * move_distances[i] * + np.array([np.cos(i * 2 * np.pi / self.num_actuators), + np.sin(i * 2 * np.pi / self.num_actuators)]) for i in range(len(move_distances)) + ] + ).sum(0) + self._step() + self.theta = np.arctan(self.y / self.x) # direction of the agent, in radian + + def _step(self) -> None: + self.x = self.x + self.move_vector[0] + self.y = self.y + self.move_vector[1] diff --git a/dizoo/gym_hybrid/envs/gym-hybrid/gym_hybrid/bg.jpg b/dizoo/gym_hybrid/envs/gym-hybrid/gym_hybrid/bg.jpg new file mode 100644 index 0000000000..40a6332bc5 Binary files /dev/null and b/dizoo/gym_hybrid/envs/gym-hybrid/gym_hybrid/bg.jpg differ diff --git a/dizoo/gym_hybrid/envs/gym-hybrid/gym_hybrid/environments.py b/dizoo/gym_hybrid/envs/gym-hybrid/gym_hybrid/environments.py new file mode 100644 index 0000000000..9716bc4484 --- /dev/null +++ b/dizoo/gym_hybrid/envs/gym-hybrid/gym_hybrid/environments.py @@ -0,0 +1,405 @@ +from collections import namedtuple +from typing import Optional +from typing import Tuple + +import gym +import numpy as np +import cv2 +import os +from gym import spaces +from gym.utils import seeding + +# gym.logger.set_level(40) # noqa + +from .agents import BaseAgent, MovingAgent, SlidingAgent, HardMoveAgent + +# Action Id +ACCELERATE = 0 +TURN = 1 +BREAK = 2 + +Target = namedtuple('Target', ['x', 'y', 'radius']) + + +class Action: + """" + Action class to store and standardize the action for the environment. + """ + + def __init__(self, id_: int, parameters: list): + """" + Initialization of an action. + + Args: + id_: The id of the selected action. + parameters: The parameters of an action. + """ + self.id = id_ + self.parameters = parameters + + @property + def parameter(self) -> float: + """" + Property method to return the parameter related to the action selected. + + Returns: + The parameter related to this action_id + """ + if len(self.parameters) == 2: + return self.parameters[self.id] + else: + return self.parameters[0] + + +class BaseEnv(gym.Env): + """" + Gym environment parent class. + """ + + def __init__( + self, + seed: Optional[int] = None, + max_turn: float = np.pi / 2, + max_acceleration: float = 0.5, + delta_t: float = 0.005, + max_step: int = 200, + penalty: float = 0.001, + break_value: float = 0.1, + ): + """Initialization of the gym environment. + + Args: + seed (int): Seed used to get reproducible results. + max_turn (float): Maximum turn during one step (in radian). + max_acceleration (float): Maximum acceleration during one step. + delta_t (float): Time duration of one step. + max_step (int): Maximum number of steps in one episode. + penalty (float): Score penalty given at the agent every step. + break_value (float): Break value when performing break action. + """ + # Agent Parameters + self.max_turn = max_turn + self.max_acceleration = max_acceleration + self.break_value = break_value + + # Environment Parameters + self.delta_t = delta_t + self.max_step = max_step + self.field_size = 1.0 + self.target_radius = 0.1 + self.penalty = penalty + + # Initialization + self.seed(seed) + self.target = None + self.viewer = None + self.current_step = None + self.agent = BaseAgent(break_value=break_value, delta_t=delta_t) + + parameters_min = np.array([0, -1]) + parameters_max = np.array([1, +1]) + + self.action_space = spaces.Tuple((spaces.Discrete(3), spaces.Box(parameters_min, parameters_max))) + self.observation_space = spaces.Box(np.ones(10), -np.ones(10)) + dirname = os.path.dirname(__file__) + self.bg = cv2.imread(os.path.join(dirname, 'bg.jpg')) + self.bg = cv2.cvtColor(self.bg, cv2.COLOR_BGR2RGB) + self.bg = cv2.resize(self.bg, (800, 800)) + self.target_img = cv2.imread(os.path.join(dirname, 'target.png'), cv2.IMREAD_UNCHANGED) + self.target_img = cv2.resize(self.target_img, (60, 60)) + + def seed(self, seed: Optional[int] = None) -> list: + self.np_random, seed = seeding.np_random(seed) # noqa + return [seed] + + def reset(self) -> list: + self.current_step = 0 + + limit = self.field_size - self.target_radius + low = [-limit, -limit, self.target_radius] + high = [limit, limit, self.target_radius] + self.target = Target(*self.np_random.uniform(low, high)) + + low = [-self.field_size, -self.field_size, 0] + high = [self.field_size, self.field_size, 2 * np.pi] + self.agent.reset(*self.np_random.uniform(low, high)) + + return self.get_state() + + def step(self, raw_action: Tuple[int, list]) -> Tuple[list, float, bool, dict]: + action = Action(*raw_action) + last_distance = self.distance + self.current_step += 1 + + if action.id == TURN: + rotation = self.max_turn * max(min(action.parameter, 1), -1) + self.agent.turn(rotation) + elif action.id == ACCELERATE: + acceleration = self.max_acceleration * max(min(action.parameter, 1), 0) + self.agent.accelerate(acceleration) + elif action.id == BREAK: + self.agent.break_() + + if self.distance < self.target_radius and self.agent.speed == 0: + reward = self.get_reward(last_distance, True) + done = True + elif abs(self.agent.x) > self.field_size or abs(self.agent.y + ) > self.field_size or self.current_step > self.max_step: + reward = -1 + done = True + else: + reward = self.get_reward(last_distance) + done = False + + return self.get_state(), reward, done, {} + + def get_state(self) -> list: + state = [ + self.agent.x, self.agent.y, self.agent.speed, + np.cos(self.agent.theta), + np.sin(self.agent.theta), self.target.x, self.target.y, self.distance, + 0 if self.distance > self.target_radius else 1, self.current_step / self.max_step + ] + return state + + def get_reward(self, last_distance: float, goal: bool = False) -> float: + return last_distance - self.distance - self.penalty + (1 if goal else 0) + + @property + def distance(self) -> float: + return self.get_distance(self.agent.x, self.agent.y, self.target.x, self.target.y) + + @staticmethod + def get_distance(x1: float, y1: float, x2: float, y2: float) -> float: + return np.sqrt(((x1 - x2) ** 2) + ((y1 - y2) ** 2)).item() + + def render(self, mode='human'): + screen_width = 400 + screen_height = 400 + unit_x = screen_width / 2 + unit_y = screen_height / 2 + agent_radius = 0.05 + + if self.viewer is None: + from gym.envs.classic_control import rendering + self.viewer = rendering.Viewer(screen_width, screen_height) + + agent = rendering.make_circle(unit_x * agent_radius) + self.agent_trans = rendering.Transform( + translation=(unit_x * (1 + self.agent.x), unit_y * (1 + self.agent.y)) + ) # noqa + agent.add_attr(self.agent_trans) + agent.set_color(0.1, 0.3, 0.9) + self.viewer.add_geom(agent) + + t, r, m = 0.1 * unit_x, 0.04 * unit_y, 0.06 * unit_x + arrow = rendering.FilledPolygon([(t, 0), (m, r), (m, -r)]) + self.arrow_trans = rendering.Transform(rotation=self.agent.theta) # noqa + arrow.add_attr(self.arrow_trans) + arrow.add_attr(self.agent_trans) + arrow.set_color(0, 0, 0) + self.viewer.add_geom(arrow) + + target = rendering.make_circle(unit_x * self.target_radius, filled=False) + target_trans = rendering.Transform(translation=(unit_x * (1 + self.target.x), unit_y * (1 + self.target.y))) + target.add_attr(target_trans) + target.set_color(0, 0.6, 0) + self.viewer.add_geom(target) + + self.arrow_trans.set_rotation(self.agent.theta) + self.agent_trans.set_translation(unit_x * (1 + self.agent.x), unit_y * (1 + self.agent.y)) + + ret = self.viewer.render(return_rgb_array=mode == 'rgb_array') + # add background + ret = np.where(ret == 255, self.bg, ret) + # add target logo + # # x, y = int(unit_x * (1 + self.target.x)), int(unit_y * (1 - self.target.y)) + # # x, y = x - 20, y + 25 # seed0 + # target_area = ret[x:x+60, y:y+60] + # rgb_img = cv2.cvtColor(self.target_img[..., :3], cv2.COLOR_BGR2RGB) + # target_area = np.where(self.target_img[..., -1:] == 0, target_area, rgb_img) + # ret[x:x+60, y:y+60] = target_area + # add frame + frames = np.array([60, 60, 30]).reshape(1, 1, -1) + ret[:6] = frames + ret[:, :6] = frames + ret[-6:] = frames + ret[:, -6:] = frames + return ret + + def close(self): + if self.viewer: + self.viewer.close() + self.viewer = None + + +class MovingEnv(BaseEnv): + + def __init__( + self, + seed: int = None, + max_turn: float = np.pi / 2, + max_acceleration: float = 0.5, + delta_t: float = 0.005, + max_step: int = 200, + penalty: float = 0.001, + break_value: float = 0.1, + ): + super(MovingEnv, self).__init__( + seed=seed, + max_turn=max_turn, + max_acceleration=max_acceleration, + delta_t=delta_t, + max_step=max_step, + penalty=penalty, + break_value=break_value, + ) + + self.agent = MovingAgent( + break_value=break_value, + delta_t=delta_t, + ) + + +class SlidingEnv(BaseEnv): + + def __init__( + self, + seed: int = None, + max_turn: float = np.pi / 2, + max_acceleration: float = 0.5, + delta_t: float = 0.005, + max_step: int = 200, + penalty: float = 0.001, + break_value: float = 0.1 + ): + super(SlidingEnv, self).__init__( + seed=seed, + max_turn=max_turn, + max_acceleration=max_acceleration, + delta_t=delta_t, + max_step=max_step, + penalty=penalty, + break_value=break_value + ) + + self.agent = SlidingAgent(break_value=break_value, delta_t=delta_t) + + +class HardMoveEnv(gym.Env): + """" + HardMove environment. Please refer to https://arxiv.org/abs/2109.05490 for details. + """ + + def __init__( + self, + num_actuators: int = 4, + seed: Optional[int] = None, + max_turn: float = np.pi / 2, + max_acceleration: float = 0.5, + delta_t: float = 0.005, + max_step: int = 25, + penalty: float = 0.001, + break_value: float = 0.1, + ): + """Initialization of the gym environment. + + Args: + seed (int): Seed used to get reproducible results. + max_turn (float): Maximum turn during one step (in radian). + max_acceleration (float): Maximum acceleration during one step. + delta_t (float): Time duration of one step. + max_step (int): Maximum number of steps in one episode. + penalty (float): Score penalty given at the agent every step. + break_value (float): Break value when performing break action. + """ + # Agent Parameters + self.num_actuators = num_actuators + self.max_turn = max_turn + self.max_acceleration = max_acceleration + self.break_value = break_value + + # Environment Parameters + self.delta_t = delta_t + self.max_step = max_step + self.field_size = 1.0 + self.target_radius = 0.1 + self.penalty = penalty + + # Initialization + self.seed(seed) + self.target = None + self.viewer = None + self.current_step = None + self.agent = HardMoveAgent(break_value=break_value, delta_t=delta_t, num_actuators=self.num_actuators) + + parameters_min = np.array([-1 for i in range(self.num_actuators)]) + parameters_max = np.array([+1 for i in range(self.num_actuators)]) + + self.action_space = spaces.Tuple( + (spaces.Discrete(int(2 ** self.num_actuators)), spaces.Box(parameters_min, parameters_max)) + ) + self.observation_space = spaces.Box(np.ones(10), -np.ones(10)) + + def seed(self, seed: Optional[int] = None) -> list: + self.np_random, seed = seeding.np_random(seed) # noqa + return [seed] + + def reset(self) -> list: + self.current_step = 0 + + limit = self.field_size - self.target_radius + low = [-limit, -limit, self.target_radius] + high = [limit, limit, self.target_radius] + self.target = Target(*self.np_random.uniform(low, high)) + + low = [-self.field_size, -self.field_size, 0] + high = [self.field_size, self.field_size, 2 * np.pi] + self.agent.reset(*self.np_random.uniform(low, high)) + + return self.get_state() + + def step(self, raw_action: Tuple[int, list]) -> Tuple[list, float, bool, dict]: + move_direction_meta = raw_action[0] # shape (1,) in {2**n} + move_distances = raw_action[1] # shape (2**n,) + last_distance = self.distance + self.current_step += 1 + + self.agent.move(move_direction_meta, move_distances) + if self.distance < self.target_radius: + reward = self.get_reward(last_distance, True) + done = True + elif abs(self.agent.x) > self.field_size or abs(self.agent.y + ) > self.field_size or self.current_step > self.max_step: + reward = -1 + done = True + else: + reward = self.get_reward(last_distance) + done = False + + return self.get_state(), reward, done, {} + + def get_state(self) -> list: + state = [ + self.agent.x, self.agent.y, self.agent.speed, + np.cos(self.agent.theta), + np.sin(self.agent.theta), self.target.x, self.target.y, self.distance, + 0 if self.distance > self.target_radius else 1, self.current_step / self.max_step + ] + return state + + def get_reward(self, last_distance: float, goal: bool = False) -> float: + return last_distance - self.distance - self.penalty + (1 if goal else 0) + + @property + def distance(self) -> float: + return self.get_distance(self.agent.x, self.agent.y, self.target.x, self.target.y) + + @staticmethod + def get_distance(x1: float, y1: float, x2: float, y2: float) -> float: + return np.sqrt(((x1 - x2) ** 2) + ((y1 - y2) ** 2)).item() + + def close(self): + if self.viewer: + self.viewer.close() + self.viewer = None diff --git a/dizoo/gym_hybrid/envs/gym-hybrid/gym_hybrid/target.png b/dizoo/gym_hybrid/envs/gym-hybrid/gym_hybrid/target.png new file mode 100644 index 0000000000..66bd053adf Binary files /dev/null and b/dizoo/gym_hybrid/envs/gym-hybrid/gym_hybrid/target.png differ diff --git a/dizoo/gym_hybrid/envs/gym-hybrid/setup.py b/dizoo/gym_hybrid/envs/gym-hybrid/setup.py new file mode 100644 index 0000000000..248ccb4535 --- /dev/null +++ b/dizoo/gym_hybrid/envs/gym-hybrid/setup.py @@ -0,0 +1,8 @@ +from setuptools import setup + +setup( + name='gym_hybrid', + version='0.0.2', # original gym_hybrid version='0.0.1' + packages=['gym_hybrid'], + install_requires=['gym', 'numpy'], +) diff --git a/dizoo/gym_hybrid/envs/gym-hybrid/tests/hardmove.py b/dizoo/gym_hybrid/envs/gym-hybrid/tests/hardmove.py new file mode 100644 index 0000000000..bde6b6eb8f --- /dev/null +++ b/dizoo/gym_hybrid/envs/gym-hybrid/tests/hardmove.py @@ -0,0 +1,17 @@ +import time +import gym +import gym_hybrid + +if __name__ == '__main__': + env = gym.make('HardMove-v0') + env.reset() + + ACTION_SPACE = env.action_space[0].n + PARAMETERS_SPACE = env.action_space[1].shape[0] + OBSERVATION_SPACE = env.observation_space.shape[0] + + done = False + while not done: + state, reward, done, info = env.step(env.action_space.sample()) + print(f'State: {state} Reward: {reward} Done: {done}') + time.sleep(0.1) diff --git a/dizoo/gym_hybrid/envs/gym-hybrid/tests/moving.py b/dizoo/gym_hybrid/envs/gym-hybrid/tests/moving.py new file mode 100644 index 0000000000..52315decd9 --- /dev/null +++ b/dizoo/gym_hybrid/envs/gym-hybrid/tests/moving.py @@ -0,0 +1,17 @@ +import time +import gym +import gym_hybrid + +if __name__ == '__main__': + env = gym.make('Moving-v0') + env.reset() + + ACTION_SPACE = env.action_space[0].n + PARAMETERS_SPACE = env.action_space[1].shape[0] + OBSERVATION_SPACE = env.observation_space.shape[0] + + done = False + while not done: + state, reward, done, info = env.step(env.action_space.sample()) + print(f'State: {state} Reward: {reward} Done: {done}') + time.sleep(0.1) diff --git a/dizoo/gym_hybrid/envs/gym-hybrid/tests/record.py b/dizoo/gym_hybrid/envs/gym-hybrid/tests/record.py new file mode 100644 index 0000000000..d97eaa13b2 --- /dev/null +++ b/dizoo/gym_hybrid/envs/gym-hybrid/tests/record.py @@ -0,0 +1,14 @@ +import gym +import gym_hybrid + +if __name__ == '__main__': + env = gym.make('Sliding-v0') + env = gym.wrappers.Monitor(env, "./video", force=True) + env.metadata["render.modes"] = ["human", "rgb_array"] + env.reset() + + done = False + while not done: + _, _, done, _ = env.step(env.action_space.sample()) + + env.close() diff --git a/dizoo/gym_hybrid/envs/gym-hybrid/tests/render.py b/dizoo/gym_hybrid/envs/gym-hybrid/tests/render.py new file mode 100644 index 0000000000..a382525fc4 --- /dev/null +++ b/dizoo/gym_hybrid/envs/gym-hybrid/tests/render.py @@ -0,0 +1,16 @@ +import time +import gym +import gym_hybrid + +if __name__ == '__main__': + env = gym.make('Sliding-v0') + env.reset() + + done = False + while not done: + _, _, done, _ = env.step(env.action_space.sample()) + env.render() + time.sleep(0.1) + + time.sleep(1) + env.close() diff --git a/dizoo/gym_hybrid/envs/gym-hybrid/tests/sliding.py b/dizoo/gym_hybrid/envs/gym-hybrid/tests/sliding.py new file mode 100644 index 0000000000..4a44dc0332 --- /dev/null +++ b/dizoo/gym_hybrid/envs/gym-hybrid/tests/sliding.py @@ -0,0 +1,17 @@ +import time +import gym +import gym_hybrid + +if __name__ == '__main__': + env = gym.make('Sliding-v0') + env.reset() + + ACTION_SPACE = env.action_space[0].n + PARAMETERS_SPACE = env.action_space[1].shape[0] + OBSERVATION_SPACE = env.observation_space.shape[0] + + done = False + while not done: + state, reward, done, info = env.step(env.action_space.sample()) + print(f'State: {state} Reward: {reward} Done: {done}') + time.sleep(0.1) diff --git a/dizoo/gym_hybrid/envs/gym_hybrid_env.py b/dizoo/gym_hybrid/envs/gym_hybrid_env.py index 73925944c2..9f02925d1a 100644 --- a/dizoo/gym_hybrid/envs/gym_hybrid_env.py +++ b/dizoo/gym_hybrid/envs/gym_hybrid_env.py @@ -1,10 +1,14 @@ -from typing import Any, List, Dict, Union, Optional -import time +import copy +import os +from typing import Dict, Optional + import gym import gym_hybrid -import copy +import matplotlib.pyplot as plt import numpy as np from easydict import EasyDict +from matplotlib import animation + from ding.envs import BaseEnv, BaseEnvTimestep from ding.envs.common import affine_transform from ding.torch_utils import to_ndarray @@ -13,33 +17,35 @@ @ENV_REGISTRY.register('gym_hybrid') class GymHybridEnv(BaseEnv): - default_env_id = ['Sliding-v0', 'Moving-v0'] + default_env_id = ['Sliding-v0', 'Moving-v0', 'HardMove-v0'] + + @classmethod + def default_config(cls: type) -> EasyDict: + cfg = EasyDict(copy.deepcopy(cls.config)) + cfg.cfg_type = cls.__name__ + 'Dict' + return cfg + + config = dict( + env_id='Moving-v0', + act_scale=True, + ) def __init__(self, cfg: EasyDict) -> None: self._cfg = cfg self._env_id = cfg.env_id assert self._env_id in self.default_env_id self._act_scale = cfg.act_scale - self._init_flag = False self._replay_path = None + self._save_replay = False + self._save_replay_count = 0 + self._init_flag = False def reset(self) -> np.ndarray: if not self._init_flag: - self._env = gym.make(self._env_id) - if self._replay_path is not None: - if gym.version.VERSION > '0.22.0': - # Gym removed classic control rendering to support using pygame instead. - # And thus, gym hybrid currently do not support rendering. - self._env.metadata["render_modes"] = ["human", "rgb_array"] - else: - self._env = gym.wrappers.RecordVideo( - self._env, - video_folder=self._replay_path, - episode_trigger=lambda episode_id: True, - name_prefix='rl-video-{}'.format(id(self)) - ) - self._env.metadata["render.modes"] = ["human", "rgb_array"] - + if self._env_id == 'HardMove-v0': + self._env = gym.make(self._env_id, num_actuators=self._cfg.num_actuators) + else: + self._env = gym.make(self._env_id) self._observation_space = self._env.observation_space self._action_space = self._env.action_space self._reward_space = gym.spaces.Box( @@ -51,7 +57,7 @@ def reset(self) -> np.ndarray: self._env.seed(self._seed + np_seed) elif hasattr(self, '_seed'): self._env.seed(self._seed) - self._final_eval_reward = 0 + self._eval_episode_return = 0 obs = self._env.reset() obs = to_ndarray(obs).astype(np.float32) return obs @@ -68,16 +74,21 @@ def seed(self, seed: int, dynamic_seed: bool = True) -> None: def step(self, action: Dict) -> BaseEnvTimestep: if self._act_scale: - # acceleration_value. - action['action_args'][0] = affine_transform(action['action_args'][0], min_val=0, max_val=1) - # rotation_value. Following line can be omitted, because in the affine_transform function, - # we have already done the clip(-1,1) operation - action['action_args'][1] = affine_transform(action['action_args'][1], min_val=-1, max_val=1) - action = [action['action_type'], action['action_args']] + if self._env_id == 'HardMove-v0': + action = [ + action['action_type'], [affine_transform(i, min_val=-1, max_val=1) for i in action['action_args']] + ] + else: + # acceleration_value. + action['action_args'][0] = affine_transform(action['action_args'][0], min_val=0, max_val=1) + # rotation_value. Following line can be omitted, because in the affine_transform function, + # we have already done the clip(-1,1) operation + action['action_args'][1] = affine_transform(action['action_args'][1], min_val=-1, max_val=1) + action = [action['action_type'], action['action_args']] + if self._save_replay: + self._frames.append(self._env.render(mode='rgb_array')) obs, rew, done, info = self._env.step(action) - self._final_eval_reward += rew - if done: - info['final_eval_reward'] = self._final_eval_reward + obs = to_ndarray(obs) if isinstance(obs, list): # corner case for i in range(len(obs)): @@ -87,10 +98,22 @@ def step(self, action: Dict) -> BaseEnvTimestep: assert isinstance(obs, np.ndarray) and obs.shape == (10, ) obs = obs.astype(np.float32) - rew = to_ndarray([rew]) # wrapped to be transfered to a numpy array with shape (1,) + rew = to_ndarray([rew]) # wrapped to be transferred to a numpy array with shape (1,) if isinstance(rew, list): rew = rew[0] assert isinstance(rew, np.ndarray) and rew.shape == (1, ) + self._eval_episode_return += rew.item() + if done: + info['eval_episode_return'] = self._eval_episode_return + if self._save_replay: + if self._env_id == 'HardMove-v0': + self._env_id = f'hardmove_n{self._cfg.num_actuators}' + path = os.path.join( + self._replay_path, '{}_episode_{}.gif'.format(self._env_id, self._save_replay_count) + ) + self.display_frames_as_gif(self._frames, path) + self._frames = [] + self._save_replay_count += 1 info['action_args_mask'] = np.array([[1, 0], [0, 1], [0, 0]]) return BaseEnvTimestep(obs, rew, done, info) @@ -105,11 +128,6 @@ def random_action(self) -> Dict: def __repr__(self) -> str: return "DI-engine gym hybrid Env" - def enable_save_replay(self, replay_path: Optional[str] = None) -> None: - if replay_path is None: - replay_path = './video' - self._replay_path = replay_path - @property def observation_space(self) -> gym.spaces.Space: return self._observation_space @@ -121,3 +139,22 @@ def action_space(self) -> gym.spaces.Space: @property def reward_space(self) -> gym.spaces.Space: return self._reward_space + + def enable_save_replay(self, replay_path: Optional[str] = None) -> None: + if replay_path is None: + replay_path = './video' + self._replay_path = replay_path + self._save_replay = True + self._save_replay_count = 0 + self._frames = [] + + @staticmethod + def display_frames_as_gif(frames: list, path: str) -> None: + patch = plt.imshow(frames[0]) + plt.axis('off') + + def animate(i): + patch.set_data(frames[i]) + + anim = animation.FuncAnimation(plt.gcf(), animate, frames=len(frames), interval=5) + anim.save(path, writer='imagemagick', fps=20) diff --git a/dizoo/gym_hybrid/envs/test_gym_hybrid_env.py b/dizoo/gym_hybrid/envs/test_gym_hybrid_env.py index e39c1c1dc1..896987f33f 100644 --- a/dizoo/gym_hybrid/envs/test_gym_hybrid_env.py +++ b/dizoo/gym_hybrid/envs/test_gym_hybrid_env.py @@ -1,15 +1,24 @@ -import pytest import numpy as np -from easydict import EasyDict - +import pytest from dizoo.gym_hybrid.envs import GymHybridEnv +from easydict import EasyDict @pytest.mark.envtest class TestGymHybridEnv: def test_naive(self): - env = GymHybridEnv(EasyDict({'env_id': 'Moving-v0', 'act_scale': False})) + env = GymHybridEnv( + EasyDict( + { + 'env_id': 'Moving-v0', + 'act_scale': False, + 'save_replay_gif': False, + 'replay_path_gif': None, + 'replay_path': None + } + ) + ) env.enable_save_replay('./video') env.seed(314, dynamic_seed=False) assert env._seed == 314 diff --git a/dizoo/gym_hybrid/moving_v0.gif b/dizoo/gym_hybrid/moving_v0.gif index 255fd21201..2823dcc8f6 100644 Binary files a/dizoo/gym_hybrid/moving_v0.gif and b/dizoo/gym_hybrid/moving_v0.gif differ diff --git a/dizoo/gym_pybullet_drones/__init__.py b/dizoo/gym_pybullet_drones/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/dizoo/gym_pybullet_drones/config/flythrugate_onppo_config.py b/dizoo/gym_pybullet_drones/config/flythrugate_onppo_config.py new file mode 100644 index 0000000000..ee7bc55dc1 --- /dev/null +++ b/dizoo/gym_pybullet_drones/config/flythrugate_onppo_config.py @@ -0,0 +1,61 @@ +from easydict import EasyDict + +flythrugate_ppo_config = dict( + exp_name='flythrugate_ppo_seed0', + env=dict( + manager=dict(shared_memory=False, reset_inplace=True), + env_id='flythrugate-aviary-v0', + norm_obs=dict(use_norm=False, ), + norm_reward=dict(use_norm=False, ), + collector_env_num=8, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + stop_value=0, + action_type="VEL", + ), + policy=dict( + cuda=True, + recompute_adv=True, + # load_path="./flythrugate_ppo_seed0/ckpt/ckpt_best.pth.tar", + model=dict( + obs_shape=12, + action_shape=4, + action_space='continuous', + ), + action_space='continuous', + learn=dict( + epoch_per_collect=10, + batch_size=64, + learning_rate=3e-4, + value_weight=0.5, + entropy_weight=0.0, + clip_ratio=0.2, + adv_norm=True, + value_norm=True, + ), + collect=dict( + n_sample=2048, + gae_lambda=0.97, + ), + eval=dict(evaluator=dict(eval_freq=5000, )), + ), +) +flythrugate_ppo_config = EasyDict(flythrugate_ppo_config) +main_config = flythrugate_ppo_config + +flythrugate_ppo_create_config = dict( + env=dict( + type='gym_pybullet_drones', + import_names=['dizoo.gym_pybullet_drones.envs.gym_pybullet_drones_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='ppo', ), +) +flythrugate_ppo_create_config = EasyDict(flythrugate_ppo_create_config) +create_config = flythrugate_ppo_create_config + +if __name__ == "__main__": + # or you can enter `ding -m serial_onpolicy -c flythrugate_ppo_config.py -s 0 --env-step 1e7` + from ding.entry import serial_pipeline_onpolicy + serial_pipeline_onpolicy((main_config, create_config), seed=0) diff --git a/dizoo/gym_pybullet_drones/config/takeoffaviary_onppo_config.py b/dizoo/gym_pybullet_drones/config/takeoffaviary_onppo_config.py new file mode 100644 index 0000000000..fe2f7bfa9c --- /dev/null +++ b/dizoo/gym_pybullet_drones/config/takeoffaviary_onppo_config.py @@ -0,0 +1,61 @@ +from easydict import EasyDict + +takeoffaviary_ppo_config = dict( + exp_name='takeoffaviary_ppo_seed0', + env=dict( + manager=dict(shared_memory=False, reset_inplace=True), + env_id='takeoff-aviary-v0', + norm_obs=dict(use_norm=False, ), + norm_reward=dict(use_norm=False, ), + collector_env_num=8, + evaluator_env_num=8, + use_act_scale=True, + n_evaluator_episode=8, + stop_value=0, + action_type="VEL", + ), + policy=dict( + cuda=True, + recompute_adv=True, + # load_path="./takeoffaviary_ppo_seed0/ckpt/ckpt_best.pth.tar", + model=dict( + obs_shape=12, + action_shape=4, + action_space='continuous', + ), + action_space='continuous', + learn=dict( + epoch_per_collect=10, #reduce + batch_size=64, + learning_rate=3e-4, #tune; pytorch lr scheduler + value_weight=0.5, + entropy_weight=0.0, #0.001 + clip_ratio=0.2, #0.1 + adv_norm=True, + value_norm=True, + ), + collect=dict( + n_sample=2048, + gae_lambda=0.97, + ), + eval=dict(evaluator=dict(eval_freq=5000, )), + ), +) +takeoffaviary_ppo_config = EasyDict(takeoffaviary_ppo_config) +main_config = takeoffaviary_ppo_config + +takeoffaviary_ppo_create_config = dict( + env=dict( + type='gym_pybullet_drones', + import_names=['dizoo.gym_pybullet_drones.envs.gym_pybullet_drones_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='ppo', ), +) +takeoffaviary_ppo_create_config = EasyDict(takeoffaviary_ppo_create_config) +create_config = takeoffaviary_ppo_create_config + +if __name__ == "__main__": + # or you can enter `ding -m serial_onpolicy -c takeoffaviary_ppo_config.py -s 0 --env-step 1e7` + from ding.entry import serial_pipeline_onpolicy + serial_pipeline_onpolicy((main_config, create_config), seed=0) diff --git a/dizoo/gym_pybullet_drones/entry/flythrugate_onppo_eval.py b/dizoo/gym_pybullet_drones/entry/flythrugate_onppo_eval.py new file mode 100644 index 0000000000..05337e8c24 --- /dev/null +++ b/dizoo/gym_pybullet_drones/entry/flythrugate_onppo_eval.py @@ -0,0 +1,55 @@ +import os +import gym +import torch +from tensorboardX import SummaryWriter +from easydict import EasyDict + +from ding.config import compile_config +from ding.worker import BaseLearner, SampleSerialCollector, InteractionSerialEvaluator, NaiveReplayBuffer +from ding.envs import BaseEnvManager, DingEnvWrapper +from ding.policy import PPOPolicy +from ding.model import VAC +from ding.utils import set_pkg_seed + +from dizoo.gym_pybullet_drones.envs.gym_pybullet_drones_env import GymPybulletDronesEnv +from dizoo.gym_pybullet_drones.config.flythrugate_onppo_config import flythrugate_ppo_config + + +def main(cfg, seed=0, max_iterations=int(1e10)): + cfg = compile_config( + cfg, + BaseEnvManager, + PPOPolicy, + BaseLearner, + SampleSerialCollector, + InteractionSerialEvaluator, + NaiveReplayBuffer, + save_cfg=True + ) + collector_env_num, evaluator_env_num = cfg.env.collector_env_num, cfg.env.evaluator_env_num + + info = cfg.env.manager + + cfg.env['record'] = True + cfg.env['gui'] = True + cfg.env['print_debug_info'] = True + cfg.env['plot_observation'] = True + + evaluator_env = BaseEnvManager( + env_fn=[lambda: GymPybulletDronesEnv(cfg.env) for _ in range(evaluator_env_num)], cfg=cfg.env.manager + ) + + evaluator_env.seed(seed, dynamic_seed=False) + set_pkg_seed(seed, use_cuda=cfg.policy.cuda) + + model = VAC(**cfg.policy.model) + policy = PPOPolicy(cfg.policy, model=model) + policy.eval_mode.load_state_dict(torch.load(cfg.policy.load_path, map_location='cpu')) + + tb_logger = SummaryWriter(os.path.join('./log/', 'serial')) + evaluator = InteractionSerialEvaluator(cfg.policy.eval.evaluator, evaluator_env, policy.eval_mode, tb_logger) + evaluator.eval() + + +if __name__ == "__main__": + main(flythrugate_ppo_config) diff --git a/dizoo/gym_pybullet_drones/entry/takeoffaviary_onppo_eval.py b/dizoo/gym_pybullet_drones/entry/takeoffaviary_onppo_eval.py new file mode 100644 index 0000000000..3ff48259cf --- /dev/null +++ b/dizoo/gym_pybullet_drones/entry/takeoffaviary_onppo_eval.py @@ -0,0 +1,53 @@ +import os +import gym +import torch +from tensorboardX import SummaryWriter +from easydict import EasyDict + +from ding.config import compile_config +from ding.worker import BaseLearner, SampleSerialCollector, InteractionSerialEvaluator, NaiveReplayBuffer +from ding.envs import BaseEnvManager, DingEnvWrapper +from ding.policy import PPOPolicy +from ding.model import VAC +from ding.utils import set_pkg_seed + +from dizoo.gym_pybullet_drones.envs.gym_pybullet_drones_env import GymPybulletDronesEnv +from dizoo.gym_pybullet_drones.config.takeoffaviary_onppo_config import takeoffaviary_ppo_config + + +def main(cfg, seed=0, max_iterations=int(1e10)): + cfg = compile_config( + cfg, + BaseEnvManager, + PPOPolicy, + BaseLearner, + SampleSerialCollector, + InteractionSerialEvaluator, + NaiveReplayBuffer, + save_cfg=True + ) + collector_env_num, evaluator_env_num = cfg.env.collector_env_num, cfg.env.evaluator_env_num + + cfg.env['record'] = True + cfg.env['gui'] = True + cfg.env['print_debug_info'] = True + cfg.env['plot_observation'] = True + + evaluator_env = BaseEnvManager( + env_fn=[lambda: GymPybulletDronesEnv(cfg.env) for _ in range(evaluator_env_num)], cfg=cfg.env.manager + ) + + evaluator_env.seed(seed, dynamic_seed=False) + set_pkg_seed(seed, use_cuda=cfg.policy.cuda) + + model = VAC(**cfg.policy.model) + policy = PPOPolicy(cfg.policy, model=model) + policy.eval_mode.load_state_dict(torch.load(cfg.policy.load_path, map_location='cpu')) + + tb_logger = SummaryWriter(os.path.join('./log/', 'serial')) + evaluator = InteractionSerialEvaluator(cfg.policy.eval.evaluator, evaluator_env, policy.eval_mode, tb_logger) + evaluator.eval() + + +if __name__ == "__main__": + main(takeoffaviary_ppo_config) diff --git a/dizoo/gym_pybullet_drones/envs/__init__.py b/dizoo/gym_pybullet_drones/envs/__init__.py new file mode 100644 index 0000000000..2a2fba6a46 --- /dev/null +++ b/dizoo/gym_pybullet_drones/envs/__init__.py @@ -0,0 +1 @@ +from .gym_pybullet_drones_env import GymPybulletDronesEnv diff --git a/dizoo/gym_pybullet_drones/envs/gym_pybullet_drones_env.py b/dizoo/gym_pybullet_drones/envs/gym_pybullet_drones_env.py new file mode 100644 index 0000000000..8b1ca6fcce --- /dev/null +++ b/dizoo/gym_pybullet_drones/envs/gym_pybullet_drones_env.py @@ -0,0 +1,270 @@ +from typing import Optional, Callable +import numpy as np +import copy +import gym +from gym.spaces import Box +import gym_pybullet_drones +from gym_pybullet_drones.utils.enums import DroneModel, Physics +from gym_pybullet_drones.envs.single_agent_rl.BaseSingleAgentAviary import ActionType, ObservationType +from gym_pybullet_drones.utils.Logger import Logger +from ding.envs import BaseEnv, BaseEnvTimestep +from ding.torch_utils import to_ndarray +from ding.utils import ENV_REGISTRY + +from easydict import EasyDict + + +def gym_pybullet_drones_observation_space(dim, minimum=-np.inf, maximum=np.inf, dtype=np.float32) -> Callable: + lower_bound = np.repeat(minimum, dim).astype(dtype) + upper_bound = np.repeat(maximum, dim).astype(dtype) + lower_bound[2] = 0.0 + return Box(lower_bound, upper_bound, dtype=dtype) + + +def drones_action_dim(type_of_action) -> int: + if type_of_action in [ActionType.RPM, ActionType.DYN, ActionType.VEL]: + return 4 + elif type_of_action == ActionType.PID: + return 3 + elif type_of_action == ActionType.TUN: + return 6 + elif type_of_action in [ActionType.ONE_D_DYN, ActionType.ONE_D_PID, ActionType.ONE_D_RPM]: + return 1 + else: + raise ValueError('Invalid action type.') + + +def gym_pybullet_drones_action_space(drone_num=1, minimum=-1, maximum=1, dtype=np.float32) -> Callable: + + def _gym_pybullet_drones_action_space(type_of_action) -> Box: + dim = drones_action_dim(type_of_action) + return Box( + np.repeat(minimum, dim * drone_num).astype(dtype), + np.repeat(maximum, dim * drone_num).astype(dtype), + dtype=dtype + ) + + return _gym_pybullet_drones_action_space + + +def gym_pybullet_drones_reward_space(minimum=-10000, maximum=0, dtype=np.float32) -> Callable: + return Box(np.repeat(minimum, 1).astype(dtype), np.repeat(maximum, 1).astype(dtype), dtype=dtype) + + +gym_pybullet_drones_env_info = { + "takeoff-aviary-v0": { + "observation_space": gym_pybullet_drones_observation_space(12, minimum=-1, maximum=1), + "action_space": gym_pybullet_drones_action_space(drone_num=1, minimum=-1, maximum=1), + "reward_space": gym_pybullet_drones_reward_space() + }, + "flythrugate-aviary-v0": { + "observation_space": gym_pybullet_drones_observation_space(12, minimum=-1, maximum=1), + "action_space": gym_pybullet_drones_action_space(drone_num=1, minimum=-1, maximum=1), + "reward_space": gym_pybullet_drones_reward_space() + }, +} + +action_type = { + "PID": ActionType.PID, + "DYN": ActionType.DYN, + "VEL": ActionType.VEL, + "RPM": ActionType.RPM, + "TUN": ActionType.TUN, + "ONE_D_DYN": ActionType.ONE_D_DYN, + "ONE_D_PID": ActionType.ONE_D_PID, + "ONE_D_RPM": ActionType.ONE_D_RPM, +} + + +@ENV_REGISTRY.register('gym_pybullet_drones') +class GymPybulletDronesEnv(BaseEnv): + """ + Gym_Pybullet_Drones Environment for training and simulating UAV drones in pybullet physical engine. + The tasks are registered in the standard of gym library. + url: 'https://github.com/utiasDSL/gym-pybullet-drones' + """ + + @classmethod + def default_config(cls: type) -> EasyDict: + cfg = EasyDict(copy.deepcopy(cls.config)) + cfg.cfg_type = cls.__name__ + 'Dict' + return cfg + + config = { + 'num_drones': 1, + 'print_debug_info': False, + 'output_folder': "./results", + 'plot_observation': False, + 'freq': 240, + 'aggregate_phy_steps': 1, + 'gui': False, + 'record': False, + "action_type": "RPM", + } + + def __init__(self, cfg: dict = {}) -> None: + self.raw_cfg = copy.deepcopy(cfg) + for k, v in self.default_config().items(): + if k not in cfg: + cfg[k] = v + + if cfg["num_drones"] == 1: + self.env_kwargs = { + 'drone_model': DroneModel.CF2X, + 'initial_xyzs': None, + 'initial_rpys': None, + 'physics': Physics.PYB, + 'freq': 240, + 'aggregate_phy_steps': 1, + 'gui': False, + 'record': False, + 'obs': ObservationType.KIN, + 'act': ActionType.RPM + } + else: + # TODO(zjow): develop envs that support multi drones. + self.env_kwargs = { + 'drone_model': DroneModel.CF2X, + 'num_drones': 2, + 'neighbourhood_radius': np.inf, + 'initial_xyzs': None, + 'initial_rpys': None, + 'physics': Physics.PYB, + 'freq': 240, + 'aggregate_phy_steps': 1, + 'gui': False, + 'record': False, + 'obs': ObservationType.KIN, + 'act': ActionType.RPM + } + + self._cfg = cfg + + for k, _ in self.env_kwargs.items(): + if k in cfg: + self.env_kwargs[k] = cfg[k] + + self.env_kwargs["act"] = action_type[cfg["action_type"]] + self.action_type = self.env_kwargs["act"] + + self._env_id = cfg.env_id + self._init_flag = False + self._replay_path = None + + self._observation_space = gym_pybullet_drones_env_info[cfg.env_id]["observation_space"] + self._action_space = gym_pybullet_drones_env_info[cfg.env_id]["action_space"](self.action_type) + self._action_dim = drones_action_dim(self.action_type) * self._cfg["num_drones"] + self._reward_space = gym_pybullet_drones_env_info[cfg.env_id]["reward_space"] + + self.env_step_count = 0 + + def reset(self) -> np.ndarray: + if not self._init_flag: + + self._env = gym.make(self._env_id, **self.env_kwargs) + + if self._cfg["plot_observation"]: + self.observation_logger = Logger( + logging_freq_hz=int(self._env.SIM_FREQ / self._env.AGGR_PHY_STEPS), + num_drones=1, + output_folder=self._cfg["output_folder"] + ) + + self._init_flag = True + + if hasattr(self, '_seed') and hasattr(self, '_dynamic_seed') and self._dynamic_seed: + np_seed = 100 * np.random.randint(1, 1000) + self._env.seed(self._seed + np_seed) + elif hasattr(self, '_seed'): + self._env.seed(self._seed) + + self._eval_episode_return = 0 + obs = self._env.reset() + obs = to_ndarray(obs).astype(np.float32) + self.env_step_count = 0 + if self._cfg["plot_observation"]: + self.observation_logger.log( + drone=0, + timestamp=self.env_step_count / self._env.SIM_FREQ, + state=np.hstack([obs[0:3], np.zeros(4), obs[3:15], + np.resize(np.zeros(self._action_dim), (4))]), + control=np.zeros(12) + ) + if self._cfg["print_debug_info"]: + if self.env_step_count % self._env.SIM_FREQ == 0: + self._env.render() + self.env_step_count += 1 + + return obs + + def close(self) -> None: + if self._init_flag: + self._env.close() + self._init_flag = False + + def seed(self, seed: int, dynamic_seed: bool = True) -> None: + self._seed = seed + self._dynamic_seed = dynamic_seed + np.random.seed(self._seed) + + def step(self, action: np.ndarray) -> BaseEnvTimestep: + # action = action.astype('float32') + obs, rew, done, info = self._env.step(action) + if self._cfg["plot_observation"]: + self.observation_logger.log( + drone=0, + timestamp=self.env_step_count / self._env.SIM_FREQ, + state=np.hstack([obs[0:3], np.zeros(4), obs[3:15], + np.resize(action, (4))]), + control=np.zeros(12) + ) + + if self._cfg["print_debug_info"]: + if self.env_step_count % self._env.SIM_FREQ == 0: + self._env.render() + self.env_step_count += 1 + self._eval_episode_return += rew + if done: + info['eval_episode_return'] = self._eval_episode_return + if self._cfg["print_debug_info"]: + self.plot_observation_curve() + + obs = to_ndarray(obs).astype(np.float32) + rew = to_ndarray([rew]).astype(np.float32) # wrapped to be transfered to a array with shape (1,) + return BaseEnvTimestep(obs, rew, done, info) + + def enable_save_replay(self, replay_path: Optional[str] = None) -> None: + if replay_path is None: + replay_path = './video' + self._replay_path = replay_path + + def random_action(self) -> np.ndarray: + return self.action_space.sample().astype(np.float32) + + @property + def observation_space(self) -> gym.spaces.Space: + if not self._init_flag: + return self._observation_space + else: + return self._env.observation_space + + @property + def action_space(self) -> gym.spaces.Space: + if not self._init_flag: + return self._action_space + else: + return self._env.action_space + + @property + def reward_space(self) -> gym.spaces.Space: + return self._reward_space + + def __repr__(self) -> str: + return "DI-engine gym_pybullet_drones Env: " + self._cfg["env_id"] + + def plot_observation_curve(self) -> None: + if self._cfg["plot_observation"]: + self.observation_logger.plot() + + def clone(self, caller: str) -> 'GymPybulletDronesEnv': + return GymPybulletDronesEnv(self.raw_cfg) diff --git a/dizoo/gym_pybullet_drones/envs/test_ding_env.py b/dizoo/gym_pybullet_drones/envs/test_ding_env.py new file mode 100644 index 0000000000..eae113ef79 --- /dev/null +++ b/dizoo/gym_pybullet_drones/envs/test_ding_env.py @@ -0,0 +1,32 @@ +import pytest +from easydict import EasyDict +import gym_pybullet_drones + +from ding.envs import BaseEnv, BaseEnvTimestep +from dizoo.gym_pybullet_drones.envs.gym_pybullet_drones_env import GymPybulletDronesEnv + + +@pytest.mark.envtest +class TestGymPybulletDronesEnv: + + def test_naive(self): + cfg = {"env_id": "takeoff-aviary-v0"} + cfg = EasyDict(cfg) + env = GymPybulletDronesEnv(cfg) + + env.reset() + done = False + while not done: + action = env.action_space.sample() + assert action.shape[0] == 4 + + for i in range(action.shape[0]): + assert action[i] >= env.action_space.low[i] and action[i] <= env.action_space.high[i] + + obs, reward, done, info = env.step(action) + + assert obs.shape[0] == 12 + for i in range(obs.shape[0]): + assert obs[i] >= env.observation_space.low[i] and obs[i] <= env.observation_space.high[i] + + assert reward >= env.reward_space.low and reward <= env.reward_space.high diff --git a/dizoo/gym_pybullet_drones/envs/test_ori_env.py b/dizoo/gym_pybullet_drones/envs/test_ori_env.py new file mode 100644 index 0000000000..bff4538ee9 --- /dev/null +++ b/dizoo/gym_pybullet_drones/envs/test_ori_env.py @@ -0,0 +1,27 @@ +import pytest +import gym +import numpy as np + +import gym_pybullet_drones + + +@pytest.mark.envtest +class TestGymPybulletDronesOriEnv: + + def test_naive(self): + env = gym.make("takeoff-aviary-v0") + env.reset() + done = False + while not done: + action = env.action_space.sample() + assert action.shape[0] == 4 + + for i in range(action.shape[0]): + assert action[i] >= env.action_space.low[i] and action[i] <= env.action_space.high[i] + + obs, reward, done, info = env.step(action) + assert obs.shape[0] == 12 + for i in range(obs.shape[0]): + assert obs[i] >= env.observation_space.low[i] and obs[i] <= env.observation_space.high[i] + + assert reward >= env.reward_space.low and reward <= env.reward_space.high diff --git a/dizoo/gym_pybullet_drones/gym_pybullet_drones.gif b/dizoo/gym_pybullet_drones/gym_pybullet_drones.gif new file mode 100644 index 0000000000..f732f7ce65 Binary files /dev/null and b/dizoo/gym_pybullet_drones/gym_pybullet_drones.gif differ diff --git a/dizoo/gym_soccer/config/gym_soccer_pdqn_config.py b/dizoo/gym_soccer/config/gym_soccer_pdqn_config.py index 8e797e9761..fc2e409184 100644 --- a/dizoo/gym_soccer/config/gym_soccer_pdqn_config.py +++ b/dizoo/gym_soccer/config/gym_soccer_pdqn_config.py @@ -26,11 +26,6 @@ ), ), learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, - # How many updates(iterations) to train after collector's one collection. - # Bigger "update_per_collect" means bigger off-policy. - # collect data -> update policy-> collect data -> ... update_per_collect=500, # 10 ~ 500 batch_size=320, learning_rate_dis=3e-4, diff --git a/dizoo/gym_soccer/envs/gym_soccer_env.py b/dizoo/gym_soccer/envs/gym_soccer_env.py index 2f4ac3ab8b..b0e759b495 100644 --- a/dizoo/gym_soccer/envs/gym_soccer_env.py +++ b/dizoo/gym_soccer/envs/gym_soccer_env.py @@ -29,7 +29,7 @@ def reset(self) -> np.array: if not self._init_flag: self._env = gym.make(self._env_id, replay_path=self._replay_path, port=self._cfg.port) # TODO self._init_flag = True - self._final_eval_reward = 0 + self._eval_episode_return = 0 obs = self._env.reset() obs = to_ndarray(obs).astype(np.float32) return obs @@ -45,9 +45,9 @@ def step(self, action: List) -> BaseEnvTimestep: action[5][0] = affine_transform(action[5][0], min_val=-180, max_val=180) obs, rew, done, info = self._env.step(action) - self._final_eval_reward += rew + self._eval_episode_return += rew if done: - info['final_eval_reward'] = self._final_eval_reward + info['eval_episode_return'] = self._eval_episode_return obs = to_ndarray(obs).astype(np.float32) # reward wrapped to be transfered to a numpy array with shape (1,) rew = to_ndarray([rew]) diff --git a/dizoo/gym_soccer/envs/test_gym_soccer_env.py b/dizoo/gym_soccer/envs/test_gym_soccer_env.py index 3360ce9e0b..50bedd89ac 100644 --- a/dizoo/gym_soccer/envs/test_gym_soccer_env.py +++ b/dizoo/gym_soccer/envs/test_gym_soccer_env.py @@ -28,7 +28,7 @@ def test_naive(self): if timestep.done: print('reset env') env.reset() - assert env._final_eval_reward == 0 + assert env._eval_episode_return == 0 print(env.info()) # env.replay_log("./video/20211019011053-base_left_0-vs-base_right_0.rcg") env.close() diff --git a/dizoo/image_classification/entry/imagenet_res18_config.py b/dizoo/image_classification/entry/imagenet_res18_config.py index 7081069365..bd4f473dd6 100644 --- a/dizoo/image_classification/entry/imagenet_res18_config.py +++ b/dizoo/image_classification/entry/imagenet_res18_config.py @@ -4,8 +4,8 @@ exp_name='imagenet_res18', policy=dict( cuda=True, + multi_gpu=True, learn=dict( - multi_gpu=True, bp_update_sync=True, train_epoch=200, batch_size=32, @@ -27,9 +27,7 @@ learn_data_path='/mnt/lustre/share/images/train', eval_data_path='/mnt/lustre/share/images/val', ), - eval=dict( - batch_size=32, evaluator=dict(eval_freq=1, multi_gpu=True, stop_value=dict(loss=0.5, acc1=75.0, acc5=95.0)) - ), + eval=dict(batch_size=32, evaluator=dict(eval_freq=1, stop_value=dict(loss=0.5, acc1=75.0, acc5=95.0))), ), env=dict(), ) diff --git a/dizoo/image_classification/entry/imagenet_res18_main.py b/dizoo/image_classification/entry/imagenet_res18_main.py index f99bcafdcd..7e49736ab2 100644 --- a/dizoo/image_classification/entry/imagenet_res18_main.py +++ b/dizoo/image_classification/entry/imagenet_res18_main.py @@ -115,7 +115,7 @@ def gt(self, metric1: dict, metric2: dict) -> bool: def main(cfg: dict, seed: int) -> None: cfg = compile_config(cfg, seed=seed, policy=ImageClassificationPolicy, evaluator=MetricSerialEvaluator) - if cfg.policy.learn.multi_gpu: + if cfg.policy.multi_gpu: rank, world_size = dist_init() else: rank, world_size = 0, 1 @@ -127,7 +127,7 @@ def main(cfg: dict, seed: int) -> None: policy = ImageClassificationPolicy(cfg.policy, model=model, enable_field=['learn', 'eval']) learn_dataset = ImageNetDataset(cfg.policy.collect.learn_data_path, is_training=True) eval_dataset = ImageNetDataset(cfg.policy.collect.eval_data_path, is_training=False) - if cfg.policy.learn.multi_gpu: + if cfg.policy.multi_gpu: learn_sampler = DistributedSampler(learn_dataset) eval_sampler = DistributedSampler(eval_dataset) else: diff --git a/dizoo/image_classification/policy/policy.py b/dizoo/image_classification/policy/policy.py index 8b4fc55fd3..e6eb9f60ba 100644 --- a/dizoo/image_classification/policy/policy.py +++ b/dizoo/image_classification/policy/policy.py @@ -56,7 +56,7 @@ def _forward_learn(self, data): backward_time = self._timer.value with self._timer: - if self._cfg.learn.multi_gpu: + if self._cfg.multi_gpu: self.sync_gradients(self._learn_model) sync_time = self._timer.value self._optimizer.step() diff --git a/dizoo/ising_env/__init__.py b/dizoo/ising_env/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/dizoo/ising_env/config/ising_mfq_config.py b/dizoo/ising_env/config/ising_mfq_config.py new file mode 100644 index 0000000000..bfe720c106 --- /dev/null +++ b/dizoo/ising_env/config/ising_mfq_config.py @@ -0,0 +1,68 @@ +from easydict import EasyDict +from ding.utils import set_pkg_seed + +obs_shape = 4 +action_shape = 2 +num_agents = 100 +dim_spin = 2 +agent_view_sight = 1 + +ising_mfq_config = dict( + exp_name='ising_mfq_seed0', + env=dict( + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + num_agents=num_agents, + dim_spin=dim_spin, + agent_view_sight=agent_view_sight, + manager=dict(shared_memory=False, ), + ), + policy=dict( + cuda=True, + priority=False, + model=dict( + obs_shape=obs_shape + action_shape, # for we will concat the pre_action_prob into obs + action_shape=action_shape, + encoder_hidden_size_list=[128, 128, 512], + init_bias=0, + ), + nstep=3, + discount_factor=0.99, + learn=dict( + update_per_collect=10, + batch_size=32, + learning_rate=0.0001, + target_update_freq=500, + ), + collect=dict(n_sample=96, ), + eval=dict(evaluator=dict(eval_freq=1000, )), + other=dict( + eps=dict( + type='exp', + start=1., + end=0.05, + decay=250000, + ), + replay_buffer=dict(replay_buffer_size=100000, ), + ), + ), +) +ising_mfq_config = EasyDict(ising_mfq_config) +main_config = ising_mfq_config +ising_mfq_create_config = dict( + env=dict( + type='ising_model', + import_names=['dizoo.ising_env.envs.ising_model_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='dqn'), +) +ising_mfq_create_config = EasyDict(ising_mfq_create_config) +create_config = ising_mfq_create_config + +if __name__ == '__main__': + # or you can enter `ding -m serial -c ising_mfq_config.py -s 0` + from ding.entry import serial_pipeline + seed = 1 + serial_pipeline((main_config, create_config), seed=seed, max_env_step=5e4) diff --git a/dizoo/ising_env/entry/ising_mfq_eval.py b/dizoo/ising_env/entry/ising_mfq_eval.py new file mode 100644 index 0000000000..f688a003a7 --- /dev/null +++ b/dizoo/ising_env/entry/ising_mfq_eval.py @@ -0,0 +1,15 @@ +from dizoo.ising_env.config.ising_mfq_config import main_config, create_config +from ding.entry import eval + + +def main(): + main_config.env.collector_env_num = 1 + main_config.env.evaluator_env_num = 1 + main_config.env.n_evaluator_episode = 1 + ckpt_path = './ckpt_best.pth.tar' + replay_path = './replay_videos' + eval((main_config, create_config), seed=1, load_path=ckpt_path, replay_path=replay_path) + + +if __name__ == "__main__": + main() diff --git a/dizoo/ising_env/envs/__init__.py b/dizoo/ising_env/envs/__init__.py new file mode 100644 index 0000000000..81cae0e551 --- /dev/null +++ b/dizoo/ising_env/envs/__init__.py @@ -0,0 +1 @@ +from .ising_model_env import IsingModelEnv diff --git a/dizoo/ising_env/envs/ising_model/Ising.py b/dizoo/ising_env/envs/ising_model/Ising.py new file mode 100644 index 0000000000..02b8e2d3f6 --- /dev/null +++ b/dizoo/ising_env/envs/ising_model/Ising.py @@ -0,0 +1,114 @@ +import numpy as np + +from dizoo.ising_env.envs.ising_model.multiagent.core import IsingWorld, IsingAgent + + +class Scenario(): + + def _calc_mask(self, agent, shape_size): + # compute the neighbour mask for each agent + if agent.view_sight == -1: + # fully observed + agent.spin_mask += 1 + elif agent.view_sight == 0: + # observe itself + agent.spin_mask[agent.state.id] = 1 + elif agent.view_sight > 0: + # observe neighbours + delta = list(range(-int(agent.view_sight), int(agent.view_sight) + 1, 1)) + delta.remove(0) # agent itself is not counted as neighbour of itself + for dt in delta: + row = agent.state.p_pos[0] + col = agent.state.p_pos[1] + row_dt = row + dt + col_dt = col + dt + if row_dt in range(0, shape_size): + agent.spin_mask[agent.state.id + shape_size * dt] = 1 + if col_dt in range(0, shape_size): + agent.spin_mask[agent.state.id + dt] = 1 + + # the graph is cyclic, most left and most right are neighbours + if agent.state.p_pos[0] < agent.view_sight: + tar = shape_size - (np.array(range(0, int(agent.view_sight - agent.state.p_pos[0]), 1)) + 1) + tar = tar * shape_size + agent.state.p_pos[1] + agent.spin_mask[tar] = [1] * len(tar) + + if agent.state.p_pos[1] < agent.view_sight: + tar = shape_size - (np.array(range(0, int(agent.view_sight - agent.state.p_pos[1]), 1)) + 1) + tar = agent.state.p_pos[0] * shape_size + tar + agent.spin_mask[tar] = [1] * len(tar) + + if agent.state.p_pos[0] >= shape_size - agent.view_sight: + tar = np.array(range(0, int(agent.view_sight - (shape_size - 1 - agent.state.p_pos[0])), 1)) + tar = tar * shape_size + agent.state.p_pos[1] + agent.spin_mask[tar] = [1] * len(tar) + + if agent.state.p_pos[1] >= shape_size - agent.view_sight: + tar = np.array(range(0, int(agent.view_sight - (shape_size - 1 - agent.state.p_pos[1])), 1)) + tar = agent.state.p_pos[0] * shape_size + tar + agent.spin_mask[tar] = [1] * len(tar) + + def make_world(self, num_agents=100, agent_view=1): + world = IsingWorld() + world.agent_view_sight = agent_view + world.dim_spin = 2 + world.dim_pos = 2 + world.n_agents = num_agents + world.shape_size = int(np.ceil(np.power(num_agents, 1.0 / world.dim_pos))) + world.global_state = np.zeros((world.shape_size, ) * world.dim_pos) + # assume 0 external magnetic field + world.field = np.zeros((world.shape_size, ) * world.dim_pos) + + world.agents = [IsingAgent(view_sight=world.agent_view_sight) for i in range(num_agents)] + + # make initial conditions + self.reset_world(world) + + return world + + def reset_world(self, world): + + world_mat = np.array( + range(np.power(world.shape_size, world.dim_pos))). \ + reshape((world.shape_size,) * world.dim_pos) + # init agent state and global state + for i, agent in enumerate(world.agents): + agent.name = 'agent %d' % i + agent.color = np.array([0.35, 0.35, 0.85]) + agent.state.id = i + agent.state.p_pos = np.where(world_mat == i) + agent.state.spin = np.random.choice(world.dim_spin) + agent.spin_mask = np.zeros(world.n_agents) + + assert world.dim_pos == 2, "cyclic neighbour only support 2D now" + self._calc_mask(agent, world.shape_size) + world.global_state[agent.state.p_pos] = agent.state.spin + + n_ups = np.count_nonzero(world.global_state.flatten()) + n_downs = world.n_agents - n_ups + world.order_param = abs(n_ups - n_downs) / (world.n_agents + 0.0) + + def reward(self, agent, world): + # turn the state into -1/1 for easy computing + world.global_state[np.where(world.global_state == 0)] = -1 + + mask_display = agent.spin_mask.reshape((int(np.sqrt(world.n_agents)), -1)) + + local_reward = - 0.5 * world.global_state[agent.state.p_pos] \ + * np.sum(world.global_state.flatten() * agent.spin_mask) + + world.global_state[np.where(world.global_state == -1)] = 0 + return -local_reward + + def observation(self, agent, world): + # get positions of all entities in this agent's reference frame + # agent state is updated in the world.step() function already + # update the changes of the world + + # return the neighbour state + return world.global_state.flatten()[np.where(agent.spin_mask == 1)] + + def done(self, agent, world): + if world.order_param == 1.0: + return True + return False diff --git a/dizoo/ising_env/envs/ising_model/__init__.py b/dizoo/ising_env/envs/ising_model/__init__.py new file mode 100644 index 0000000000..9a4b6f71af --- /dev/null +++ b/dizoo/ising_env/envs/ising_model/__init__.py @@ -0,0 +1,7 @@ +import imp +import os.path as osp + + +def load(name): + pathname = osp.join(osp.dirname(__file__), name) + return imp.load_source('', pathname) diff --git a/dizoo/ising_env/envs/ising_model/multiagent/__init__.py b/dizoo/ising_env/envs/ising_model/multiagent/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/dizoo/ising_env/envs/ising_model/multiagent/core.py b/dizoo/ising_env/envs/ising_model/multiagent/core.py new file mode 100644 index 0000000000..b0da45e3a4 --- /dev/null +++ b/dizoo/ising_env/envs/ising_model/multiagent/core.py @@ -0,0 +1,131 @@ +import numpy as np + + +class IsingEntityState(object): + + def __init__(self): + self.id = None + self.p_pos = None + + +class IsingAgentState(IsingEntityState): + + def __init__(self): + super(IsingAgentState, self).__init__() + # up or down + self.spin = None + + +class IsingAction(object): + + def __init__(self): + # action + self.a = None + + +# properties and state of physical world entity +class IsingEntity(object): + + def __init__(self): + # name + self.name = '' + # properties: + self.size = 0.050 + # entity can move / be pushed + self.movable = False + # color + self.color = None + # state: position and spin + self.state = IsingEntityState() + + +class IsingAgent(IsingEntity): + + def __init__(self, view_sight=1): + super(IsingAgent, self).__init__() + # agents are movable by default + self.movable = False + # -1: observe the whole state, 0: itself, 1: neighbour of 1 unit + self.view_sight = view_sight + self.spin_mask = None # the mask for who is neighbours + # state + self.state = IsingAgentState() + self.state.spin_range = [0, 1] + # action + self.action = IsingAction() + self.action.a_range = [0, 1] + # script behavior to execute + self.action_callback = None + + +# multi-agent world +class IsingWorld(object): + + def __init__(self): + # list of agents and entities (can change at execution-time!) + self.agents = [] + self.n_agents = 1 + self.agent_view_sight = 1 + # position dimensionality + self.dim_pos = 2 + # state dimension + self.dim_spin = 2 + # color dimensionality + self.dim_color = 3 + # world size + self.shape_size = 1 + # ising specific + self.global_state = None # log all spins + self.moment = 1 + self.field = None # external magnetic field + self.temperature = .1 # Temperature (in units of energy) + self.interaction = 1 # Interaction (ferromagnetic if positive, + # antiferromagnetic if negative) + self.order_param = 1.0 + self.order_param_delta = 0.01 # log the change of order parameter for "done" + self.n_up = 0 + self.n_down = 0 + + # return all entities in the world + @property + def entities(self): + return self.agents + + # return all agents controllable by external policies + @property + def policy_agents(self): + return [agent for agent in self.agents if agent.action_callback is None] + + # return all agents controlled by world scripts, no use for now + @property + def scripted_agents(self): + return [agent for agent in self.agents if agent.action_callback is not None] + + # update state of the world + def step(self): + + # set actions for scripted agents, no use for now + for agent in self.scripted_agents: + agent.action = agent.action_callback(agent, self) + + # update agent state, and to the global_state + for agent in self.agents: + self.update_agent_state(agent) + self.global_state[agent.state.p_pos] = agent.state.spin + + # update the world's order parameters + self.n_up = np.count_nonzero(self.global_state.flatten()) + self.n_down = self.n_agents - self.n_up + order_param_old = self.order_param + self.order_param = abs(self.n_up - self.n_down) / (self.n_agents + 0.0) + # self.order_param_delta = (self.order_param - order_param_old) / + # order_param_old + + def update_agent_state(self, agent): + if agent.action.a == 0: + # agent.state.spin = agent.state.spin + agent.state.spin = 0 + else: + # print(agent.name + " change spin") + # agent.state.spin = 1.0 - agent.state.spin + agent.state.spin = 1 diff --git a/dizoo/ising_env/envs/ising_model/multiagent/environment.py b/dizoo/ising_env/envs/ising_model/multiagent/environment.py new file mode 100644 index 0000000000..81970e9723 --- /dev/null +++ b/dizoo/ising_env/envs/ising_model/multiagent/environment.py @@ -0,0 +1,117 @@ +import gym +from gym import spaces +import numpy as np + + +class IsingMultiAgentEnv(gym.Env): + metadata = {'render.modes': ['human', 'rgb_array']} + + def __init__( + self, + world, + reset_callback=None, + reward_callback=None, + observation_callback=None, + info_callback=None, + done_callback=None + ): + + self.world = world + self.agents = self.world.policy_agents + # number of controllable agents + self.n = len(world.policy_agents) + assert self.n == len(world.agents) + # scenario callbacks + self.reset_callback = reset_callback + self.reward_callback = reward_callback + self.observation_callback = observation_callback + self.info_callback = info_callback + self.done_callback = done_callback + # environment parameters + self.discrete_action_space = True + + # if true, every agent has the same reward + self.shared_reward = False + self.time = 0 + + # configure spaces + self.action_space = [] + self.observation_space = [] + + if self.discrete_action_space: + self.action_space.append(spaces.Discrete(self.world.dim_spin)) + else: + raise NotImplementedError("only discrete action is allowed") + + # observation space, called self-defined scenario.observation + # define the size of the observation here + # use the global state + mask + # self.observation_space.append(spaces.MultiBinary(self.n * 2)) + self.observation_space.append(spaces.MultiBinary(4 * self.world.agent_view_sight)) + + def _step(self, action_n): + "descend from gym.env, env.step() actually calls this step" + obs_n = [] + reward_n = [] + done_n = [] + self.agents = self.world.policy_agents + + display_action = action_n.reshape((int(np.sqrt(self.n)), -1)) + + # set action for each agent + for i, agent in enumerate(self.agents): + self._set_action(action_n[i], agent) + + # update the agent's new state and global state + # world(in core.py) is a part of env, created though examples' make_world + self.world.step() + + # observation, reward, done function are implemented in different examples + for agent in self.agents: + obs_n.append(self._get_obs(agent)) + reward_n.append(self._get_reward(agent)) + done_n.append(self._get_done(agent)) + + # all agents get total reward in cooperative case + reward = np.sum(reward_n) + if self.shared_reward: + reward_n = [reward] * self.n + + return obs_n, reward_n, done_n, self.world.order_param, \ + self.world.n_up, self.world.n_down + + def _reset(self): + # reset world, init agent state and global state + self.reset_callback(self.world) + # record observations for each agent + obs_n = [] + self.agents = self.world.policy_agents + for agent in self.agents: + obs_n.append(self._get_obs(agent)) + return obs_n + + # get observation for a particular agent + def _get_obs(self, agent): + if self.observation_callback is None: + return np.zeros(0) + # call the scenario's observation here + return self.observation_callback(agent, self.world) + + # get dones for a particular agent + def _get_done(self, agent): + if self.done_callback is None: + return False + # call the scenario's done here + return self.done_callback(agent, self.world) + + # get reward for a particular agent + def _get_reward(self, agent): + if self.reward_callback is None: + return 0.0 + # call the scenario's reward here + return self.reward_callback(agent, self.world) + + # set env action for a particular agent + def _set_action(self, action, agent): + agent.action.a = 0 if action <= 0 else 1 + assert len(action) == 1, "action dimenion error!" diff --git a/dizoo/ising_env/envs/ising_model_env.py b/dizoo/ising_env/envs/ising_model_env.py new file mode 100644 index 0000000000..e3ccbe890a --- /dev/null +++ b/dizoo/ising_env/envs/ising_model_env.py @@ -0,0 +1,174 @@ +from typing import Any, Optional +import copy +import os +import sys +import numpy as np +from numpy import dtype +import gym +import matplotlib.pyplot as plt +import imageio +from ding.envs import BaseEnv, BaseEnvTimestep +from ding.torch_utils import to_ndarray, to_list +from ding.utils import ENV_REGISTRY +from dizoo.ising_env.envs.ising_model.multiagent.environment import IsingMultiAgentEnv +import dizoo.ising_env.envs.ising_model as ising_model_ + + +@ENV_REGISTRY.register('ising_model') +class IsingModelEnv(BaseEnv): + """ + Overview: + Ising Model Environment for Multi-Agent Reinforcement Learning according to the paper: \ + [Mean Field Multi-Agent Reinforcement Learning](https://arxiv.org/abs/1802.05438). \ + The environment is a grid of agents, each of which can be in one of two states: \ + spin up or spin down. The agents interact with their neighbors according to the Ising model, \ + and the goal is to maximize the global order parameter, which is the average spin of all agents. \ + Details of the environment can be found in the \ + [DI-engine-Doc](https://di-engine-docs.readthedocs.io/zh-cn/latest/13_envs/index.html). + Interface: + `__init__`, `reset`, `close`, `seed`, `step`, `random_action`, `num_agents`, \ + `observation_space`, `action_space`, `reward_space`. + """ + + def __init__(self, cfg: dict) -> None: + self._cfg = cfg + self._init_flag = False + self._action_space = gym.spaces.Discrete(cfg.dim_spin) # default 2 + self._observation_space = gym.spaces.MultiBinary(4 * cfg.agent_view_sight) + self._reward_space = gym.spaces.Box(low=float("-inf"), high=float("inf"), shape=(1, ), dtype=np.float32) + self._replay_path = None + + def calculate_action_prob(self, actions): + num_action = self._action_space.n + N = actions.shape[0] # agent_num + # Convert actions to one_hot encoding + one_hot_actions = np.eye(num_action)[actions.flatten()] + action_prob = np.zeros((N, num_action)) + + for i in range(N): + # Select only the one_hot actions of agents visible to agent i + visible_actions = one_hot_actions[self._env.agents[i].spin_mask == 1] + if visible_actions.size > 0: + # Calculate the average of the one_hot encoding for visible agents only + action_prob[i] = visible_actions.mean(axis=0) + else: + # If no visible agents, action_prob remains zero for agent i + action_prob[i] = np.zeros(num_action) + + return action_prob + + def reset(self) -> np.ndarray: + if hasattr(self, '_seed') and hasattr(self, '_dynamic_seed') and self._dynamic_seed: + np_seed = 100 * np.random.randint(1, 1000) + self._cfg.seed = self._seed + np_seed + elif hasattr(self, '_seed'): + self._cfg.seed = self._seed + if not self._init_flag: + # self._env = MujocoMulti(env_args=self._cfg) + ising_model = ising_model_.load('Ising.py').Scenario() + self._env = IsingMultiAgentEnv( + world=ising_model.make_world(num_agents=self._cfg.num_agents, agent_view=1), + reset_callback=ising_model.reset_world, + reward_callback=ising_model.reward, + observation_callback=ising_model.observation, + done_callback=ising_model.done + ) + self._init_flag = True + obs = self._env._reset() + obs = np.stack(obs) + self.pre_action = np.zeros(self._cfg.num_agents, dtype=np.int32) + # consider the last global state as pre action prob + pre_action_prob = self.calculate_action_prob(self._env.world.global_state.flatten().astype(int)) + obs = np.concatenate([obs, pre_action_prob], axis=1) + obs = to_ndarray(obs).astype(np.float32) + self._eval_episode_return = 0 + self.cur_step = 0 + if self._replay_path is not None: + self._frames = [] + return obs + + def close(self) -> None: + if self._init_flag: + self._env.close() + self._init_flag = False + + def seed(self, seed: int, dynamic_seed: bool = True) -> None: + self._seed = seed + self._dynamic_seed = dynamic_seed + np.random.seed(self._seed) + + def step(self, action: np.ndarray) -> BaseEnvTimestep: + action = to_ndarray(action) + if len(action.shape) == 1: + action = np.expand_dims(action, axis=1) + obs, rew, done, order_param, ups, downs = self._env._step(action) + info = {"order_param": order_param, "ups": ups, "downs": downs, 'pre_action': self.pre_action} + pre_action_prob = self.calculate_action_prob(self.pre_action) + self.pre_action = action + obs = np.stack(obs) + obs = np.concatenate([obs, pre_action_prob], axis=1) + obs = to_ndarray(obs).astype(np.float32) + rew = np.stack(rew) + self._eval_episode_return += np.sum(rew) + self.cur_step += 1 + + if self._replay_path is not None: + # transform the action to a 2D grid. e.g. (100,) -> (10, 10) + action_matrix = action.reshape((int(np.sqrt(self._cfg.num_agents)), -1)) + self._frames.append(self.render(action_matrix, info)) + + done = done[0] # dones are the same for all agents + if done: + info['eval_episode_return'] = self._eval_episode_return / self.cur_step + if self._replay_path is not None: + path = os.path.join(self._replay_path, '{}_episode.gif'.format(self._save_replay_count)) + self.display_frames_as_gif(self._frames, path) + self._save_replay_count += 1 + return BaseEnvTimestep(obs, rew, done, info) + + def random_action(self) -> np.ndarray: + random_action = self.action_space.sample() + random_action = to_ndarray([random_action], dtype=np.int64) + return random_action + + def enable_save_replay(self, replay_path: Optional[str] = None) -> None: + if replay_path is None: + replay_path = './video' + self._replay_path = replay_path + if not os.path.exists(replay_path): + os.makedirs(replay_path) + self._save_replay_count = 0 + + def render(self, action_matrix, info) -> None: + fig, ax = plt.subplots(figsize=(5, 5)) + ax.imshow(action_matrix, cmap='gray', vmin=0, vmax=1, interpolation='none') + ax.set_title(f"Step {self.cur_step}: Order={info['order_param']}, Up={info['ups']}, Down={info['downs']}") + # save the figure to buffer + fig.canvas.draw() + image = np.frombuffer(fig.canvas.tostring_rgb(), dtype='uint8') + image = image.reshape(fig.canvas.get_width_height()[::-1] + (3, )) + plt.close(fig) + return image + + @staticmethod + def display_frames_as_gif(frames: list, output_path: str) -> None: + imageio.mimsave(output_path, frames, duration=50) + + @property + def num_agents(self) -> Any: + return self._env.n + + @property + def observation_space(self) -> gym.spaces.Space: + return self._env.observation_space[0] + + @property + def action_space(self) -> gym.spaces.Space: + return self._env.action_space[0] + + @property + def reward_space(self) -> gym.spaces.Space: + return self._reward_space + + def __repr__(self) -> str: + return "DI-engine Ising Model Env({})".format(self._cfg.env_id) diff --git a/dizoo/ising_env/envs/test_ising_model_env.py b/dizoo/ising_env/envs/test_ising_model_env.py new file mode 100644 index 0000000000..7211be8657 --- /dev/null +++ b/dizoo/ising_env/envs/test_ising_model_env.py @@ -0,0 +1,39 @@ +import pytest +import numpy as np +from dizoo.ising_env.envs import IsingModelEnv +from easydict import EasyDict + +num_agents = 100 + + +@pytest.mark.envtest +class TestIsingModelEnv: + + def test_ising(self): + env = IsingModelEnv(EasyDict({'num_agents': num_agents, 'dim_spin': 2, 'agent_view_sight': 1})) + env.seed(314, dynamic_seed=False) + assert env._seed == 314 + obs = env.reset() + assert obs.shape == (num_agents, 4 + 2) + for _ in range(5): + env.reset() + np.random.seed(314) + print('=' * 60) + for i in range(10): + # Both ``env.random_action()``, and utilizing ``np.random`` as well as action space, + # can generate legal random action. + if i < 5: + random_action = np.random.randint(0, env.action_space.n, size=(num_agents, 1)) + else: + random_action = np.array([env.action_space.sample() for _ in range(num_agents)]) + random_action = np.expand_dims(random_action, axis=1) + timestep = env.step(random_action) + print('timestep', timestep, '\n') + assert isinstance(timestep.obs, np.ndarray) + assert isinstance(timestep.done, bool) + assert timestep.obs.shape == (num_agents, 4 + 2) + assert timestep.reward.shape == (num_agents, 1) + assert timestep.reward[0] >= env.reward_space.low + assert timestep.reward[0] <= env.reward_space.high + print(env.observation_space, env.action_space, env.reward_space) + env.close() diff --git a/dizoo/ising_env/ising_env.gif b/dizoo/ising_env/ising_env.gif new file mode 100644 index 0000000000..e8a9fdcddc Binary files /dev/null and b/dizoo/ising_env/ising_env.gif differ diff --git a/dizoo/league_demo/game_env.py b/dizoo/league_demo/game_env.py index 7005f5fee0..6c97d12616 100644 --- a/dizoo/league_demo/game_env.py +++ b/dizoo/league_demo/game_env.py @@ -62,10 +62,10 @@ def step(self, actions: List[int]) -> BaseEnvTimestep: dones = True, True infos = { 'result': results[0], - 'final_eval_reward': rewards[0] + 'eval_episode_return': rewards[0] }, { 'result': results[1], - 'final_eval_reward': rewards[1] + 'eval_episode_return': rewards[1] } return BaseEnvTimestep(observations, rewards, True, infos) diff --git a/dizoo/league_demo/league_demo_collector.py b/dizoo/league_demo/league_demo_collector.py index 7b2826270e..4afe72bbe8 100644 --- a/dizoo/league_demo/league_demo_collector.py +++ b/dizoo/league_demo/league_demo_collector.py @@ -163,7 +163,7 @@ def envstep(self) -> int: Overview: Print the total envstep count. Return: - - envstep (:obj:`int`): the total envstep count + - envstep (:obj:`int`): The total envstep count. """ return self._total_envstep_count @@ -278,8 +278,8 @@ def collect(self, if timestep.done: self._total_episode_count += 1 info = { - 'reward0': timestep.info[0]['final_eval_reward'], - 'reward1': timestep.info[1]['final_eval_reward'], + 'reward0': timestep.info[0]['eval_episode_return'], + 'reward1': timestep.info[1]['eval_episode_return'], 'time': self._env_info[env_id]['time'], 'step': self._env_info[env_id]['step'], 'probs0': probs0, @@ -312,8 +312,8 @@ def _output_log(self, train_iter: int) -> None: episode_count = len(self._episode_info) envstep_count = sum([d['step'] for d in self._episode_info]) duration = sum([d['time'] for d in self._episode_info]) - episode_reward0 = [d['reward0'] for d in self._episode_info] - episode_reward1 = [d['reward1'] for d in self._episode_info] + episode_return0 = [d['reward0'] for d in self._episode_info] + episode_return1 = [d['reward1'] for d in self._episode_info] probs0 = [d['probs0'] for d in self._episode_info] probs1 = [d['probs1'] for d in self._episode_info] self._total_duration += duration @@ -324,14 +324,14 @@ def _output_log(self, train_iter: int) -> None: 'avg_envstep_per_sec': envstep_count / duration, 'avg_episode_per_sec': episode_count / duration, 'collect_time': duration, - 'reward0_mean': np.mean(episode_reward0), - 'reward0_std': np.std(episode_reward0), - 'reward0_max': np.max(episode_reward0), - 'reward0_min': np.min(episode_reward0), - 'reward1_mean': np.mean(episode_reward1), - 'reward1_std': np.std(episode_reward1), - 'reward1_max': np.max(episode_reward1), - 'reward1_min': np.min(episode_reward1), + 'reward0_mean': np.mean(episode_return0), + 'reward0_std': np.std(episode_return0), + 'reward0_max': np.max(episode_return0), + 'reward0_min': np.min(episode_return0), + 'reward1_mean': np.mean(episode_return1), + 'reward1_std': np.std(episode_return1), + 'reward1_max': np.max(episode_return1), + 'reward1_min': np.min(episode_return1), 'total_envstep_count': self._total_envstep_count, 'total_episode_count': self._total_episode_count, 'total_duration': self._total_duration, diff --git a/dizoo/mario/mario_dqn_example.py b/dizoo/mario/mario_dqn_example.py index deab06858e..945ee33479 100644 --- a/dizoo/mario/mario_dqn_example.py +++ b/dizoo/mario/mario_dqn_example.py @@ -4,7 +4,7 @@ from ding.policy import DQNPolicy from ding.envs import DingEnvWrapper, SubprocessEnvManagerV2 from ding.envs.env_wrappers import MaxAndSkipWrapper, WarpFrameWrapper, ScaledFloatFrameWrapper, FrameStackWrapper, \ - FinalEvalRewardEnv + EvalEpisodeReturnWrapper, TimeLimitWrapper from ding.data import DequeBuffer from ding.config import compile_config from ding.framework import task @@ -26,7 +26,8 @@ def wrapped_mario_env(): lambda env: WarpFrameWrapper(env, size=84), lambda env: ScaledFloatFrameWrapper(env), lambda env: FrameStackWrapper(env, n_frames=4), - lambda env: FinalEvalRewardEnv(env), + lambda env: TimeLimitWrapper(env, max_limit=400), + lambda env: EvalEpisodeReturnWrapper(env), ] } ) @@ -57,7 +58,7 @@ def main(): task.use(nstep_reward_enhancer(cfg)) task.use(data_pusher(cfg, buffer_)) task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) - task.use(CkptSaver(cfg, policy, train_freq=1000)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=1000)) task.run() diff --git a/dizoo/mario/mario_dqn_main.py b/dizoo/mario/mario_dqn_main.py index 011cb8e5b3..7b7bfe1efb 100644 --- a/dizoo/mario/mario_dqn_main.py +++ b/dizoo/mario/mario_dqn_main.py @@ -7,7 +7,7 @@ from ding.worker import BaseLearner, SampleSerialCollector, InteractionSerialEvaluator, AdvancedReplayBuffer from ding.envs import SyncSubprocessEnvManager, DingEnvWrapper, BaseEnvManager from ding.envs.env_wrappers import MaxAndSkipWrapper, WarpFrameWrapper, ScaledFloatFrameWrapper, FrameStackWrapper, \ - FinalEvalRewardEnv + EvalEpisodeReturnWrapper from ding.policy import DQNPolicy from ding.model import DQN from ding.utils import set_pkg_seed @@ -26,7 +26,7 @@ def wrapped_mario_env(): lambda env: WarpFrameWrapper(env, size=84), lambda env: ScaledFloatFrameWrapper(env), lambda env: FrameStackWrapper(env, n_frames=4), - lambda env: FinalEvalRewardEnv(env), + lambda env: EvalEpisodeReturnWrapper(env), ] } ) @@ -94,9 +94,7 @@ def main(cfg, seed=0): break learner.train(train_data, collector.envstep) # evaluate - evaluator_env = BaseEnvManager( - env_fn=[wrapped_mario_env for _ in range(evaluator_env_num)], cfg=cfg.env.manager - ) + evaluator_env = BaseEnvManager(env_fn=[wrapped_mario_env for _ in range(evaluator_env_num)], cfg=cfg.env.manager) evaluator_env.enable_save_replay(cfg.env.replay_path) # switch save replay interface evaluator = InteractionSerialEvaluator( cfg.policy.eval.evaluator, evaluator_env, policy.eval_mode, tb_logger, exp_name=cfg.exp_name diff --git a/dizoo/maze/__init__.py b/dizoo/maze/__init__.py new file mode 100644 index 0000000000..3bfa255bd5 --- /dev/null +++ b/dizoo/maze/__init__.py @@ -0,0 +1,3 @@ +from gym.envs.registration import register + +register(id='Maze', entry_point='dizoo.maze.envs:Maze') diff --git a/dizoo/maze/config/maze_bc_config.py b/dizoo/maze/config/maze_bc_config.py new file mode 100644 index 0000000000..18c9d6ade8 --- /dev/null +++ b/dizoo/maze/config/maze_bc_config.py @@ -0,0 +1,56 @@ +from easydict import EasyDict + +maze_size = 16 +num_actions = 4 +maze_pc_config = dict( + exp_name="maze_bc_seed0", + env=dict( + collector_env_num=1, + evaluator_env_num=5, + n_evaluator_episode=5, + env_id='Maze', + size=maze_size, + wall_type='tunnel', + stop_value=1 + ), + policy=dict( + cuda=True, + maze_size=maze_size, + num_actions=num_actions, + max_bfs_steps=100, + model=dict( + obs_shape=[3, maze_size, maze_size], + action_shape=num_actions, + encoder_hidden_size_list=[ + 128, + 256, + 512, + 1024, + ], + strides=[1, 1, 1, 1] + ), + learn=dict( + # update_per_collect=4, + batch_size=256, + learning_rate=0.005, + train_epoch=5000, + optimizer='SGD', + ), + eval=dict(evaluator=dict(n_episode=5)), + collect=dict(), + ), +) +maze_pc_config = EasyDict(maze_pc_config) +main_config = maze_pc_config +maze_pc_create_config = dict( + env=dict( + type='maze', + import_names=['dizoo.maze.envs.maze_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='bc'), +) +maze_pc_create_config = EasyDict(maze_pc_create_config) +create_config = maze_pc_create_config + +# You can run `dizoo/maze/entry/maze_bc_main.py` to run this config. diff --git a/dizoo/maze/config/maze_pc_config.py b/dizoo/maze/config/maze_pc_config.py new file mode 100644 index 0000000000..2a5f40b278 --- /dev/null +++ b/dizoo/maze/config/maze_pc_config.py @@ -0,0 +1,57 @@ +from easydict import EasyDict + +maze_size = 16 +num_actions = 4 +maze_pc_config = dict( + exp_name="maze_pc_seed0", + train_seeds=5, + env=dict( + collector_env_num=8, + evaluator_env_num=5, + n_evaluator_episode=5, + env_id='Maze', + size=maze_size, + wall_type='tunnel', + stop_value=1, + ), + policy=dict( + cuda=True, + maze_size=maze_size, + num_actions=num_actions, + max_bfs_steps=100, + model=dict( + obs_shape=[8, maze_size, maze_size], + action_shape=num_actions, + encoder_hidden_size_list=[ + 128, + 256, + 512, + 1024, + ], + ), + learn=dict( + batch_size=32, + learning_rate=0.0005, + train_epoch=100, + optimizer='Adam', + ), + eval=dict(evaluator=dict(n_episode=5)), + collect=dict(), + ), +) +maze_pc_config = EasyDict(maze_pc_config) +main_config = maze_pc_config +maze_pc_create_config = dict( + env=dict( + type='maze', + import_names=['dizoo.maze.envs.maze_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='pc_bfs'), +) +maze_pc_create_config = EasyDict(maze_pc_create_config) +create_config = maze_pc_create_config + +if __name__ == '__main__': + from ding.entry import serial_pipeline_pc + serial_pipeline_pc([maze_pc_config, maze_pc_create_config], seed=0) diff --git a/dizoo/maze/entry/maze_bc_main.py b/dizoo/maze/entry/maze_bc_main.py new file mode 100644 index 0000000000..3a42d4e921 --- /dev/null +++ b/dizoo/maze/entry/maze_bc_main.py @@ -0,0 +1,200 @@ +from typing import Union, Optional, Tuple +import os +from functools import partial +from copy import deepcopy + +import easydict +import torch +import numpy as np +from tensorboardX import SummaryWriter +from torch.utils.data import DataLoader, Dataset + +from ding.envs import get_vec_env_setting, create_env_manager +from ding.worker import BaseLearner, InteractionSerialEvaluator +from ding.config import read_config, compile_config +from ding.policy import create_policy +from ding.utils import set_pkg_seed +from dizoo.maze.envs.maze_env import Maze + + +# BFS algorithm +def get_vi_sequence(env, observation): + """Returns [L, W, W] optimal actions.""" + xy = np.where(observation[Ellipsis, -1] == 1) + start_x, start_y = xy[0][0], xy[1][0] + target_location = env.target_location + nav_map = env.nav_map + current_points = [target_location] + chosen_actions = {target_location: 0} + visited_points = {target_location: True} + vi_sequence = [] + + vi_map = np.full((env.size, env.size), fill_value=env.n_action, dtype=np.int32) + + found_start = False + while current_points and not found_start: + next_points = [] + for point_x, point_y in current_points: + for (action, (next_point_x, next_point_y)) in [(0, (point_x - 1, point_y)), (1, (point_x, point_y - 1)), + (2, (point_x + 1, point_y)), (3, (point_x, point_y + 1))]: + + if (next_point_x, next_point_y) in visited_points: + continue + + if not (0 <= next_point_x < len(nav_map) and 0 <= next_point_y < len(nav_map[next_point_x])): + continue + + if nav_map[next_point_x][next_point_y] == 'x': + continue + + next_points.append((next_point_x, next_point_y)) + visited_points[(next_point_x, next_point_y)] = True + chosen_actions[(next_point_x, next_point_y)] = action + vi_map[next_point_x, next_point_y] = action + + if next_point_x == start_x and next_point_y == start_y: + found_start = True + vi_sequence.append(vi_map.copy()) + current_points = next_points + track_back = [] + if found_start: + cur_x, cur_y = start_x, start_y + while cur_x != target_location[0] or cur_y != target_location[1]: + act = vi_sequence[-1][cur_x, cur_y] + track_back.append((torch.FloatTensor(env.process_states([cur_x, cur_y], env.get_maze_map())), act)) + if act == 0: + cur_x += 1 + elif act == 1: + cur_y += 1 + elif act == 2: + cur_x -= 1 + elif act == 3: + cur_y -= 1 + + return np.array(vi_sequence), track_back + + +class BCDataset(Dataset): + + def __init__(self, all_data): + self._data = all_data + + def __getitem__(self, item): + return {'obs': self._data[item][0], 'action': self._data[item][1]} + + def __len__(self): + return len(self._data) + + +def load_bc_dataset(train_seeds=1, test_seeds=1, batch_size=32): + + def load_env(seed): + ccc = easydict.EasyDict({'size': 16}) + e = Maze(ccc) + e.seed(seed) + e.reset() + return e + + envs = [load_env(i) for i in range(train_seeds + test_seeds)] + data_train = [] + data_test = [] + + for idx, env in enumerate(envs): + if idx < train_seeds: + data = data_train + else: + data = data_test + + start_obs = env.process_states(env._get_obs(), env.get_maze_map()) + _, track_back = get_vi_sequence(env, start_obs) + + data += track_back + + train_data = BCDataset(data_train) + test_data = BCDataset(data_test) + + train_dataset = DataLoader(train_data, batch_size=batch_size, shuffle=True) + test_dataset = DataLoader(test_data, batch_size=batch_size, shuffle=True) + return train_dataset, test_dataset + + +def serial_pipeline_bc( + input_cfg: Union[str, Tuple[dict, dict]], + seed: int = 0, + model: Optional[torch.nn.Module] = None, + max_iter=int(1e6), +) -> Union['Policy', bool]: # noqa + r""" + Overview: + Serial pipeline entry of imitation learning. + Arguments: + - input_cfg (:obj:`Union[str, Tuple[dict, dict]]`): Config in dict type. \ + ``str`` type means config file path. \ + ``Tuple[dict, dict]`` type means [user_config, create_cfg]. + - seed (:obj:`int`): Random seed. + - data_path (:obj:`str`): Path of training data. + - model (:obj:`Optional[torch.nn.Module]`): Instance of torch.nn.Module. + Returns: + - policy (:obj:`Policy`): Converged policy. + - convergence (:obj:`bool`): whether il training is converged + """ + if isinstance(input_cfg, str): + cfg, create_cfg = read_config(input_cfg) + else: + cfg, create_cfg = deepcopy(input_cfg) + cfg = compile_config(cfg, seed=seed, auto=True, create_cfg=create_cfg) + + # Env, Policy + env_fn, _, evaluator_env_cfg = get_vec_env_setting(cfg.env) + evaluator_env = create_env_manager(cfg.env.manager, [partial(env_fn, cfg=c) for c in evaluator_env_cfg]) + # Random seed + evaluator_env.seed(cfg.seed, dynamic_seed=False) + set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) + policy = create_policy(cfg.policy, model=model, enable_field=['learn', 'eval']) + + # Main components + tb_logger = SummaryWriter(os.path.join('./{}/log/'.format(cfg.exp_name), 'serial')) + dataloader, test_dataloader = load_bc_dataset() + learner = BaseLearner(cfg.policy.learn.learner, policy.learn_mode, tb_logger, exp_name=cfg.exp_name) + evaluator = InteractionSerialEvaluator( + cfg.policy.eval.evaluator, evaluator_env, policy.eval_mode, tb_logger, exp_name=cfg.exp_name + ) + + # ========== + # Main loop + # ========== + learner.call_hook('before_run') + stop = False + iter_cnt = 0 + for epoch in range(cfg.policy.learn.train_epoch): + # Evaluate policy performance + loss_list = [] + for _, bat in enumerate(test_dataloader): + bat['action'] = bat['action'].long() + res = policy._forward_eval(bat['obs']) + res = torch.argmax(res['logit'], dim=1) + loss_list.append(torch.sum(res == bat['action'].squeeze(-1)).item() / bat['action'].shape[0]) + label = 'validation_acc' + tb_logger.add_scalar(label, sum(loss_list) / len(loss_list), iter_cnt) + for i, train_data in enumerate(dataloader): + if evaluator.should_eval(learner.train_iter): + stop, reward = evaluator.eval(learner.save_checkpoint, learner.train_iter) + if stop: + break + train_data['action'] = train_data['action'].long() + learner.train(train_data) + iter_cnt += 1 + if iter_cnt >= max_iter: + stop = True + break + if stop: + break + + learner.call_hook('after_run') + print('final reward is: {}'.format(reward)) + return policy, stop + + +if __name__ == '__main__': + from dizoo.maze.config.maze_bc_config import main_config, create_config + serial_pipeline_bc([main_config, create_config], seed=0) diff --git a/dizoo/maze/envs/__init__.py b/dizoo/maze/envs/__init__.py new file mode 100644 index 0000000000..ab42c5b39d --- /dev/null +++ b/dizoo/maze/envs/__init__.py @@ -0,0 +1 @@ +from .maze_env import Maze diff --git a/dizoo/maze/envs/maze_env.py b/dizoo/maze/envs/maze_env.py new file mode 100644 index 0000000000..f441b7d698 --- /dev/null +++ b/dizoo/maze/envs/maze_env.py @@ -0,0 +1,380 @@ +from typing import List + +import copy +import numpy as np +import gym +from gym import spaces +from gym.utils import seeding + +from ding.envs import BaseEnvTimestep +from ding.utils import ENV_REGISTRY + + +@ENV_REGISTRY.register('maze') +class Maze(gym.Env): + """ + Environment with random maze layouts. The ASCII representation of the mazes include the following objects: + - ``: empty + - `x`: wall + - `S`: the start location (optional) + - `T`: the target location. + """ + KEY_EMPTY = 0 + KEY_WALL = 1 + KEY_TARGET = 2 + KEY_START = 3 + ASCII_MAP = { + KEY_EMPTY: ' ', + KEY_WALL: 'x', + KEY_TARGET: 'T', + KEY_START: 'S', + } + + def __init__( + self, + cfg, + ): + self._size = cfg.size + self._init_flag = False + self._random_start = True + self._seed = None + self._step = 0 + + def reset(self): + self.active_init() + obs = self._get_obs() + self._step = 0 + return self.process_states(obs, self.get_maze_map()) + + def seed(self, seed: int, dynamic_seed: bool = True) -> None: + self._seed = seed + self._dynamic_seed = dynamic_seed + np.random.seed(self._seed) + + def active_init( + self, + tabular_obs=False, + reward_fn=lambda x, y, tx, ty: 1 if (x == tx and y == ty) else 0, + done_fn=lambda x, y, tx, ty: x == tx and y == ty + ): + self._maze = self.generate_maze(self.size, self._seed, 'tunnel') + self._num_maze_keys = len(Maze.ASCII_MAP.keys()) + nav_map = self.maze_to_ascii(self._maze) + self._map = nav_map + self._tabular_obs = tabular_obs + self._reward_fn = reward_fn + self._done_fn = done_fn + if self._reward_fn is None: + self._reward_fn = lambda x, y, tx, ty: float(x == tx and y == ty) + if self._done_fn is None: + self._done_fn = lambda x, y, tx, ty: False + + self._max_x = len(self._map) + if not self._max_x: + raise ValueError('Invalid map.') + self._max_y = len(self._map[0]) + if not all(len(m) == self._max_y for m in self._map): + raise ValueError('Invalid map.') + self._start_x, self._start_y = self._find_initial_point() + self._target_x, self._target_y = self._find_target_point() + self._x, self._y = self._start_x, self._start_y + + self._n_state = self._max_x * self._max_y + self._n_action = 4 + + if self._tabular_obs: + self.observation_space = spaces.Discrete(self._n_state) + else: + self.observation_space = spaces.Box(low=0.0, high=np.inf, shape=(16, 16, 3)) + + self.action_space = spaces.Discrete(self._n_action) + self.reward_space = spaces.Box(low=0, high=1, shape=(1, ), dtype=np.float32) + + def random_start(self): + init_x, init_y = self._x, self._y + while True: # Find empty grid cell. + self._x = self.np_random.integers(self._max_x) + self._y = self.np_random.integers(self._max_y) + if self._map[self._x][self._y] != 'x': + break + ret = copy.deepcopy(self.process_states(self._get_obs(), self.get_maze_map())) + self._x, self._y = init_x, init_y + return ret + + def close(self) -> None: + if self._init_flag: + self._env.close() + self._init_flag = False + + @property + def num_maze_keys(self): + return self._num_maze_keys + + @property + def size(self): + return self._size + + def process_states(self, observations, maze_maps): + """Returns [B, W, W, 3] binary values. Channels are (wall; goal; obs)""" + loc = np.eye(self._size * self._size, dtype=np.int64)[observations[0] * self._size + observations[1]] + loc = np.reshape(loc, [self._size, self._size]) + maze_maps = maze_maps.astype(np.int64) + + states = np.concatenate([maze_maps, loc[Ellipsis, None]], axis=-1, dtype=np.int64) + return states + + def get_maze_map(self, stacked=True): + if not stacked: + return self._maze.copy() + wall = self._maze.copy() + target_x, target_y = self.target_location + assert wall[target_x][target_y] == Maze.KEY_TARGET + wall[target_x][target_y] = 0 + target = np.zeros((self._size, self._size)) + target[target_x][target_y] = 1 + assert wall[self._start_x][self._start_y] == Maze.KEY_START + wall[self._start_x][self._start_y] = 0 + return np.stack([wall, target], axis=-1) + + def generate_maze(self, size, seed, wall_type): + rng, _ = seeding.np_random(seed) + maze = np.full((size, size), fill_value=Maze.KEY_EMPTY, dtype=int) + + if wall_type == 'none': + maze[[0, -1], :] = Maze.KEY_WALL + maze[:, [0, -1]] = Maze.KEY_WALL + elif wall_type == 'tunnel': + self.sample_wall(maze, rng) + elif wall_type.startswith('blocks:'): + maze[[0, -1], :] = Maze.KEY_WALL + maze[:, [0, -1]] = Maze.KEY_WALL + self.sample_blocks(maze, rng, int(wall_type.split(':')[-1])) + else: + raise ValueError('Unknown wall type: %s' % wall_type) + + loc_target = self.sample_location(maze, rng) + maze[loc_target] = Maze.KEY_TARGET + + loc_start = self.sample_location(maze, rng) + maze[loc_start] = Maze.KEY_START + self._start_x, self._start_y = loc_start + + return maze + + def sample_blocks(self, maze, rng, num_blocks): + """Sample single-block 'wall' or 'obstacles'.""" + for _ in range(num_blocks): + loc = self.sample_location(maze, rng) + maze[loc] = Maze.KEY_WALL + + def sample_wall( + self, maze, rng, shortcut_prob=0.1, inner_wall_thickness=1, outer_wall_thickness=1, corridor_thickness=2 + ): + room = maze + + # step 1: fill everything as wall + room[:] = Maze.KEY_WALL + + # step 2: prepare + # we move two pixels at a time, because the walls are also occupying pixels + delta = inner_wall_thickness + corridor_thickness + dx = [delta, -delta, 0, 0] + dy = [0, 0, delta, -delta] + + def get_loc_type(y, x): + # remember there is a outside wall of 1 pixel surrounding the room + if (y < outer_wall_thickness or y + corridor_thickness - 1 >= room.shape[0] - outer_wall_thickness): + return 'invalid' + if (x < outer_wall_thickness or x + corridor_thickness - 1 >= room.shape[1] - outer_wall_thickness): + return 'invalid' + # already visited + if room[y, x] == Maze.KEY_EMPTY: + return 'occupied' + return 'valid' + + def connect_pixel(y, x, ny, nx): + pixel = Maze.KEY_EMPTY + if ny == y: + room[y:y + corridor_thickness, min(x, nx):max(x, nx) + corridor_thickness] = pixel + else: + room[min(y, ny):max(y, ny) + corridor_thickness, x:x + corridor_thickness] = pixel + + def carve_passage_from(y, x): + room[y, x] = Maze.KEY_EMPTY + for direction in rng.permutation(len(dx)): + ny = y + dy[direction] + nx = x + dx[direction] + + loc_type = get_loc_type(ny, nx) + if loc_type == 'invalid': + continue + elif loc_type == 'valid': + connect_pixel(y, x, ny, nx) + # recursion + carve_passage_from(ny, nx) + else: + # occupied + # we create shortcut with some probability, this is because + # we do not want to restrict to only one feasible path. + if rng.random() < shortcut_prob: + connect_pixel(y, x, ny, nx) + + carve_passage_from(outer_wall_thickness, outer_wall_thickness) + + def sample_location(self, maze, rng): + for _ in range(1000): + x, y = rng.integers(low=1, high=self._size, size=2) + if maze[x, y] == Maze.KEY_EMPTY: + return x, y + raise ValueError('Cannot sample empty location, make maze bigger?') + + @staticmethod + def key_to_ascii(key): + if key in Maze.ASCII_MAP: + return Maze.ASCII_MAP[key] + assert (Maze.KEY_OBJ <= key < Maze.KEY_OBJ + Maze.MAX_OBJ_TYPES) + return chr(ord('1') + key - Maze.KEY_OBJ) + + def maze_to_ascii(self, maze): + return [[Maze.key_to_ascii(x) for x in row] for row in maze] + + def tabular_obs_action(self, status_obs, action, include_maze_layout=False): + tabular_obs = self.get_tabular_obs(status_obs) + multiplier = self._n_action + if include_maze_layout: + multiplier += self._num_maze_keys + return multiplier * tabular_obs + action + + @staticmethod + def create_collector_env_cfg(cfg: dict) -> List[dict]: + collector_env_num = cfg.pop('collector_env_num') + cfg = copy.deepcopy(cfg) + cfg.is_train = True + return [cfg for _ in range(collector_env_num)] + + @staticmethod + def create_evaluator_env_cfg(cfg: dict) -> List[dict]: + evaluator_env_num = cfg.pop('evaluator_env_num') + cfg = copy.deepcopy(cfg) + cfg.is_train = False + return [cfg for _ in range(evaluator_env_num)] + + @property + def nav_map(self): + return self._map + + @property + def n_state(self): + return self._n_state + + @property + def n_action(self): + return self._n_action + + @property + def target_location(self): + return self._target_x, self._target_y + + @property + def tabular_obs(self): + return self._tabular_obs + + def _find_initial_point(self): + for x in range(self._max_x): + for y in range(self._max_y): + if self._map[x][y] == 'S': + break + if self._map[x][y] == 'S': + break + else: + return None, None + + return x, y + + def _find_target_point(self): + for x in range(self._max_x): + for y in range(self._max_y): + if self._map[x][y] == 'T': + break + if self._map[x][y] == 'T': + break + else: + raise ValueError('Target point not found in map.') + + return x, y + + def _get_obs(self): + if self._tabular_obs: + return self._x * self._max_y + self._y + else: + return np.array([self._x, self._y]) + + def get_tabular_obs(self, status_obs): + return self._max_y * status_obs[..., 0] + status_obs[..., 1] + + def get_xy(self, state): + x = state / self._max_y + y = state % self._max_y + return x, y + + def step(self, action): + last_x, last_y = self._x, self._y + if action == 0: + if self._x < self._max_x - 1: + self._x += 1 + elif action == 1: + if self._y < self._max_y - 1: + self._y += 1 + elif action == 2: + if self._x > 0: + self._x -= 1 + elif action == 3: + if self._y > 0: + self._y -= 1 + + if self._map[self._x][self._y] == 'x': + self._x, self._y = last_x, last_y + self._step += 1 + reward = self._reward_fn(self._x, self._y, self._target_x, self._target_y) + done = self._done_fn(self._x, self._y, self._target_x, self._target_y) + info = {} + if self._step > 100: + done = True + if done: + info['final_eval_reward'] = reward + info['eval_episode_return'] = reward + return BaseEnvTimestep(self.process_states(self._get_obs(), self.get_maze_map()), reward, done, info) + + +def get_value_map(env): + """Returns [W, W, A] one-hot VI actions.""" + target_location = env.target_location + nav_map = env.nav_map + current_points = [target_location] + chosen_actions = {target_location: 0} + visited_points = {target_location: True} + + while current_points: + next_points = [] + for point_x, point_y in current_points: + for (action, (next_point_x, next_point_y)) in [(0, (point_x - 1, point_y)), (1, (point_x, point_y - 1)), + (2, (point_x + 1, point_y)), (3, (point_x, point_y + 1))]: + + if (next_point_x, next_point_y) in visited_points: + continue + + if not (0 <= next_point_x < len(nav_map) and 0 <= next_point_y < len(nav_map[next_point_x])): + continue + + if nav_map[next_point_x][next_point_y] == 'x': + continue + + next_points.append((next_point_x, next_point_y)) + visited_points[(next_point_x, next_point_y)] = True + chosen_actions[(next_point_x, next_point_y)] = action + current_points = next_points + + value_map = np.zeros([env.size, env.size, env.n_action]) + for (x, y), action in chosen_actions.items(): + value_map[x][y][action] = 1 + return value_map diff --git a/dizoo/maze/envs/test_maze_env.py b/dizoo/maze/envs/test_maze_env.py new file mode 100644 index 0000000000..b8350d46d3 --- /dev/null +++ b/dizoo/maze/envs/test_maze_env.py @@ -0,0 +1,28 @@ +import pytest +import os +import numpy as np +from dizoo.maze.envs.maze_env import Maze +from easydict import EasyDict +import copy + + +@pytest.mark.envtest +class TestMazeEnv: + + def test_maze(self): + env = Maze(EasyDict({'size': 16})) + env.seed(314) + assert env._seed == 314 + obs = env.reset() + assert obs.shape == (16, 16, 3) + min_val, max_val = 0, 3 + for i in range(100): + random_action = np.random.randint(min_val, max_val, size=(1, )) + timestep = env.step(random_action) + print(timestep) + print(timestep.obs.max()) + assert isinstance(timestep.obs, np.ndarray) + assert isinstance(timestep.done, bool) + if timestep.done: + env.reset() + env.close() diff --git a/dizoo/metadrive/__init__.py b/dizoo/metadrive/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/dizoo/metadrive/config/__init__.py b/dizoo/metadrive/config/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/dizoo/metadrive/config/metadrive_onppo_config.py b/dizoo/metadrive/config/metadrive_onppo_config.py new file mode 100644 index 0000000000..05585b8a58 --- /dev/null +++ b/dizoo/metadrive/config/metadrive_onppo_config.py @@ -0,0 +1,111 @@ +from easydict import EasyDict +from functools import partial +from tensorboardX import SummaryWriter +import metadrive +import gym +from ding.envs import BaseEnvManager, SyncSubprocessEnvManager +from ding.config import compile_config +from ding.model.template import ContinuousQAC, VAC +from ding.policy import PPOPolicy +from ding.worker import SampleSerialCollector, InteractionSerialEvaluator, BaseLearner +from dizoo.metadrive.env.drive_env import MetaDrivePPOOriginEnv +from dizoo.metadrive.env.drive_wrapper import DriveEnvWrapper + +metadrive_basic_config = dict( + exp_name='metadrive_onppo_seed0', + env=dict( + metadrive=dict( + use_render=False, + traffic_density=0.10, # Density of vehicles occupying the roads, range in [0,1] + map='XSOS', # Int or string: an easy way to fill map_config + horizon=4000, # Max step number + driving_reward=1.0, # Reward to encourage agent to move forward. + speed_reward=0.1, # Reward to encourage agent to drive at a high speed + use_lateral_reward=False, # reward for lane keeping + out_of_road_penalty=40.0, # Penalty to discourage driving out of road + crash_vehicle_penalty=40.0, # Penalty to discourage collision + decision_repeat=20, # Reciprocal of decision frequency + out_of_route_done=True, # Game over if driving out of road + ), + manager=dict( + shared_memory=False, + max_retry=2, + context='spawn', + ), + n_evaluator_episode=16, + stop_value=255, + collector_env_num=8, + evaluator_env_num=8, + ), + policy=dict( + cuda=True, + action_space='continuous', + model=dict( + obs_shape=[5, 84, 84], + action_shape=2, + action_space='continuous', + bound_type='tanh', + encoder_hidden_size_list=[128, 128, 64], + ), + learn=dict( + epoch_per_collect=10, + batch_size=64, + learning_rate=3e-4, + entropy_weight=0.001, + value_weight=0.5, + clip_ratio=0.02, + adv_norm=False, + value_norm=True, + grad_clip_value=10, + ), + collect=dict(n_sample=3000, ), + eval=dict(evaluator=dict(eval_freq=1000, ), ), + ), +) +main_config = EasyDict(metadrive_basic_config) + + +def wrapped_env(env_cfg, wrapper_cfg=None): + return DriveEnvWrapper(MetaDrivePPOOriginEnv(env_cfg), wrapper_cfg) + + +def main(cfg): + cfg = compile_config( + cfg, SyncSubprocessEnvManager, PPOPolicy, BaseLearner, SampleSerialCollector, InteractionSerialEvaluator + ) + collector_env_num, evaluator_env_num = cfg.env.collector_env_num, cfg.env.evaluator_env_num + collector_env = SyncSubprocessEnvManager( + env_fn=[partial(wrapped_env, cfg.env.metadrive) for _ in range(collector_env_num)], + cfg=cfg.env.manager, + ) + evaluator_env = SyncSubprocessEnvManager( + env_fn=[partial(wrapped_env, cfg.env.metadrive) for _ in range(evaluator_env_num)], + cfg=cfg.env.manager, + ) + model = VAC(**cfg.policy.model) + policy = PPOPolicy(cfg.policy, model=model) + tb_logger = SummaryWriter('./log/{}/'.format(cfg.exp_name)) + learner = BaseLearner(cfg.policy.learn.learner, policy.learn_mode, tb_logger, exp_name=cfg.exp_name) + collector = SampleSerialCollector( + cfg.policy.collect.collector, collector_env, policy.collect_mode, tb_logger, exp_name=cfg.exp_name + ) + evaluator = InteractionSerialEvaluator( + cfg.policy.eval.evaluator, evaluator_env, policy.eval_mode, tb_logger, exp_name=cfg.exp_name + ) + learner.call_hook('before_run') + while True: + if evaluator.should_eval(learner.train_iter): + stop, rate = evaluator.eval(learner.save_checkpoint, learner.train_iter, collector.envstep) + if stop: + break + # Sampling data from environments + new_data = collector.collect(cfg.policy.collect.n_sample, train_iter=learner.train_iter) + learner.train(new_data, collector.envstep) + learner.call_hook('after_run') + collector.close() + evaluator.close() + learner.close() + + +if __name__ == '__main__': + main(main_config) diff --git a/dizoo/metadrive/config/metadrive_onppo_eval_config.py b/dizoo/metadrive/config/metadrive_onppo_eval_config.py new file mode 100644 index 0000000000..c9dab89ed2 --- /dev/null +++ b/dizoo/metadrive/config/metadrive_onppo_eval_config.py @@ -0,0 +1,96 @@ +from easydict import EasyDict +from functools import partial +from tensorboardX import SummaryWriter +import torch +from ding.envs import BaseEnvManager, SyncSubprocessEnvManager +from ding.config import compile_config +from ding.model.template import VAC +from ding.policy import PPOPolicy +from ding.worker import SampleSerialCollector, InteractionSerialEvaluator, BaseLearner +from dizoo.metadrive.env.drive_env import MetaDrivePPOOriginEnv +from dizoo.metadrive.env.drive_wrapper import DriveEnvWrapper + +# Load the trained model from this direction, if None, it will initialize from scratch +model_dir = None +metadrive_basic_config = dict( + exp_name='metadrive_onppo_eval_seed0', + env=dict( + metadrive=dict( + use_render=True, + traffic_density=0.10, # Density of vehicles occupying the roads, range in [0,1] + map='XSOS', # Int or string: an easy way to fill map_config + horizon=4000, # Max step number + driving_reward=1.0, # Reward to encourage agent to move forward. + speed_reward=0.10, # Reward to encourage agent to drive at a high speed + use_lateral_reward=False, # reward for lane keeping + out_of_road_penalty=40.0, # Penalty to discourage driving out of road + crash_vehicle_penalty=40.0, # Penalty to discourage collision + decision_repeat=20, # Reciprocal of decision frequency + out_of_route_done=True, # Game over if driving out of road + show_bird_view=False, # Only used to evaluate, whether to draw five channels of bird-view image + ), + manager=dict( + shared_memory=False, + max_retry=2, + context='spawn', + ), + n_evaluator_episode=16, + stop_value=255, + collector_env_num=1, + evaluator_env_num=1, + ), + policy=dict( + cuda=True, + action_space='continuous', + model=dict( + obs_shape=[5, 84, 84], + action_shape=2, + action_space='continuous', + bound_type='tanh', + encoder_hidden_size_list=[128, 128, 64], + ), + learn=dict( + epoch_per_collect=10, + batch_size=64, + learning_rate=3e-4, + entropy_weight=0.001, + value_weight=0.5, + clip_ratio=0.02, + adv_norm=False, + value_norm=True, + grad_clip_value=10, + ), + collect=dict(n_sample=1000, ), + eval=dict(evaluator=dict(eval_freq=1000, ), ), + ), +) +main_config = EasyDict(metadrive_basic_config) + + +def wrapped_env(env_cfg, wrapper_cfg=None): + return DriveEnvWrapper(MetaDrivePPOOriginEnv(env_cfg), wrapper_cfg) + + +def main(cfg): + cfg = compile_config(cfg, BaseEnvManager, PPOPolicy, BaseLearner, SampleSerialCollector, InteractionSerialEvaluator) + evaluator_env_num = cfg.env.evaluator_env_num + show_bird_view = cfg.env.metadrive.show_bird_view + wrapper_cfg = {'show_bird_view': show_bird_view} + evaluator_env = BaseEnvManager( + env_fn=[partial(wrapped_env, cfg.env.metadrive, wrapper_cfg) for _ in range(evaluator_env_num)], + cfg=cfg.env.manager, + ) + model = VAC(**cfg.policy.model) + policy = PPOPolicy(cfg.policy, model=model) + if model_dir is not None: + policy._load_state_dict_collect(torch.load(model_dir, map_location='cpu')) + tb_logger = SummaryWriter('./log/{}/'.format(cfg.exp_name)) + evaluator = InteractionSerialEvaluator( + cfg.policy.eval.evaluator, evaluator_env, policy.eval_mode, tb_logger, exp_name=cfg.exp_name + ) + stop, rate = evaluator.eval() + evaluator.close() + + +if __name__ == '__main__': + main(main_config) diff --git a/dizoo/metadrive/env/__init__.py b/dizoo/metadrive/env/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/dizoo/metadrive/env/drive_env.py b/dizoo/metadrive/env/drive_env.py new file mode 100644 index 0000000000..87087b8f97 --- /dev/null +++ b/dizoo/metadrive/env/drive_env.py @@ -0,0 +1,364 @@ +import copy +import gym +import numpy as np +from ditk import logging +from typing import Union, Dict, AnyStr, Tuple, Optional +from gym.envs.registration import register +from metadrive.manager.traffic_manager import TrafficMode +from metadrive.obs.top_down_obs_multi_channel import TopDownMultiChannel +from metadrive.constants import RENDER_MODE_NONE, DEFAULT_AGENT, REPLAY_DONE, TerminationState +from metadrive.envs.base_env import BaseEnv +from metadrive.component.map.base_map import BaseMap +from metadrive.component.map.pg_map import parse_map_config, MapGenerateMethod +from metadrive.component.pgblock.first_block import FirstPGBlock +from metadrive.component.vehicle.base_vehicle import BaseVehicle +from metadrive.utils import Config, merge_dicts, get_np_random, clip +from metadrive.envs.base_env import BASE_DEFAULT_CONFIG +from metadrive.component.road_network import Road +from metadrive.component.algorithm.blocks_prob_dist import PGBlockDistConfig + +METADRIVE_DEFAULT_CONFIG = dict( + # ===== Generalization ===== + start_seed=0, + environment_num=10, + decision_repeat=20, + block_dist_config=PGBlockDistConfig, + + # ===== Map Config ===== + map=3, # int or string: an easy way to fill map_config + random_lane_width=False, + random_lane_num=False, + map_config={ + BaseMap.GENERATE_TYPE: MapGenerateMethod.BIG_BLOCK_NUM, + BaseMap.GENERATE_CONFIG: None, # it can be a file path / block num / block ID sequence + BaseMap.LANE_WIDTH: 3.5, + BaseMap.LANE_NUM: 3, + "exit_length": 50, + }, + + # ===== Traffic ===== + traffic_density=0.1, + need_inverse_traffic=False, + traffic_mode=TrafficMode.Trigger, # "Respawn", "Trigger" + random_traffic=False, # Traffic is randomized at default. + traffic_vehicle_config=dict( + show_navi_mark=False, + show_dest_mark=False, + enable_reverse=False, + show_lidar=False, + show_lane_line_detector=False, + show_side_detector=False, + ), + + # ===== Object ===== + accident_prob=0., # accident may happen on each block with this probability, except multi-exits block + + # ===== Others ===== + use_AI_protector=False, + save_level=0.5, + is_multi_agent=False, + vehicle_config=dict(spawn_lane_index=(FirstPGBlock.NODE_1, FirstPGBlock.NODE_2, 0)), + + # ===== Agent ===== + random_spawn_lane_index=True, + target_vehicle_configs={ + DEFAULT_AGENT: dict( + use_special_color=True, + spawn_lane_index=(FirstPGBlock.NODE_1, FirstPGBlock.NODE_2, 0), + ) + }, + + # ===== Reward Scheme ===== + # See: https://github.com/decisionforce/metadrive/issues/283 + success_reward=10.0, + out_of_road_penalty=5.0, + crash_vehicle_penalty=5.0, + crash_object_penalty=5.0, + driving_reward=1.0, + speed_reward=0.1, + use_lateral_reward=False, + + # ===== Cost Scheme ===== + crash_vehicle_cost=1.0, + crash_object_cost=1.0, + out_of_road_cost=1.0, + + # ===== Termination Scheme ===== + out_of_route_done=False, + on_screen=False, + show_bird_view=False, +) + + +class MetaDrivePPOOriginEnv(BaseEnv): + + @classmethod + def default_config(cls) -> "Config": + config = super(MetaDrivePPOOriginEnv, cls).default_config() + config.update(METADRIVE_DEFAULT_CONFIG) + config.register_type("map", str, int) + config["map_config"].register_type("config", None) + return config + + def __init__(self, config: dict = None): + self.raw_cfg = config + self.default_config_copy = Config(self.default_config(), unchangeable=True) + self.init_flag = False + + @property + def observation_space(self): + return gym.spaces.Box(0, 1, shape=(84, 84, 5), dtype=np.float32) + + @property + def action_space(self): + return gym.spaces.Box(-1, 1, shape=(2, ), dtype=np.float32) + + @property + def reward_space(self): + return gym.spaces.Box(-100, 100, shape=(1, ), dtype=np.float32) + + def seed(self, seed, dynamic_seed=False): + # TODO implement dynamic_seed mechanism + super().seed(seed) + + def reset(self): + if not self.init_flag: + super(MetaDrivePPOOriginEnv, self).__init__(self.raw_cfg) + self.start_seed = self.config["start_seed"] + self.env_num = self.config["environment_num"] + self.init_flag = True + obs = super().reset() + return obs + + def _merge_extra_config(self, config: Union[dict, "Config"]) -> "Config": + config = self.default_config().update(config, allow_add_new_key=False) + if config["vehicle_config"]["lidar"]["distance"] > 50: + config["max_distance"] = config["vehicle_config"]["lidar"]["distance"] + return config + + def _post_process_config(self, config): + config = super(MetaDrivePPOOriginEnv, self)._post_process_config(config) + if not config["rgb_clip"]: + logging.warning( + "You have set rgb_clip = False, which means the observation will be uint8 values in [0, 255]. " + "Please make sure you have parsed them later before feeding them to network!" + ) + config["map_config"] = parse_map_config( + easy_map_config=config["map"], new_map_config=config["map_config"], default_config=self.default_config_copy + ) + config["vehicle_config"]["rgb_clip"] = config["rgb_clip"] + config["vehicle_config"]["random_agent_model"] = config["random_agent_model"] + if config.get("gaussian_noise", 0) > 0: + assert config["vehicle_config"]["lidar"]["gaussian_noise"] == 0, "You already provide config!" + assert config["vehicle_config"]["side_detector"]["gaussian_noise"] == 0, "You already provide config!" + assert config["vehicle_config"]["lane_line_detector"]["gaussian_noise"] == 0, "You already provide config!" + config["vehicle_config"]["lidar"]["gaussian_noise"] = config["gaussian_noise"] + config["vehicle_config"]["side_detector"]["gaussian_noise"] = config["gaussian_noise"] + config["vehicle_config"]["lane_line_detector"]["gaussian_noise"] = config["gaussian_noise"] + if config.get("dropout_prob", 0) > 0: + assert config["vehicle_config"]["lidar"]["dropout_prob"] == 0, "You already provide config!" + assert config["vehicle_config"]["side_detector"]["dropout_prob"] == 0, "You already provide config!" + assert config["vehicle_config"]["lane_line_detector"]["dropout_prob"] == 0, "You already provide config!" + config["vehicle_config"]["lidar"]["dropout_prob"] = config["dropout_prob"] + config["vehicle_config"]["side_detector"]["dropout_prob"] = config["dropout_prob"] + config["vehicle_config"]["lane_line_detector"]["dropout_prob"] = config["dropout_prob"] + target_v_config = copy.deepcopy(config["vehicle_config"]) + if not config["is_multi_agent"]: + target_v_config.update(config["target_vehicle_configs"][DEFAULT_AGENT]) + config["target_vehicle_configs"][DEFAULT_AGENT] = target_v_config + return config + + def step(self, actions: Union[np.ndarray, Dict[AnyStr, np.ndarray]]): + actions = self._preprocess_actions(actions) + engine_info = self._step_simulator(actions) + o, r, d, i = self._get_step_return(actions, engine_info=engine_info) + return o, r, d, i + + def cost_function(self, vehicle_id: str): + vehicle = self.vehicles[vehicle_id] + step_info = dict() + step_info["cost"] = 0 + if self._is_out_of_road(vehicle): + step_info["cost"] = self.config["out_of_road_cost"] + elif vehicle.crash_vehicle: + step_info["cost"] = self.config["crash_vehicle_cost"] + elif vehicle.crash_object: + step_info["cost"] = self.config["crash_object_cost"] + return step_info['cost'], step_info + + def _is_out_of_road(self, vehicle): + ret = vehicle.on_yellow_continuous_line or vehicle.on_white_continuous_line or \ + (not vehicle.on_lane) or vehicle.crash_sidewalk + if self.config["out_of_route_done"]: + ret = ret or vehicle.out_of_route + return ret + + def done_function(self, vehicle_id: str): + vehicle = self.vehicles[vehicle_id] + done = False + done_info = { + TerminationState.CRASH_VEHICLE: False, + TerminationState.CRASH_OBJECT: False, + TerminationState.CRASH_BUILDING: False, + TerminationState.OUT_OF_ROAD: False, + TerminationState.SUCCESS: False, + TerminationState.MAX_STEP: False, + TerminationState.ENV_SEED: self.current_seed, + } + if self._is_arrive_destination(vehicle): + done = True + logging.info("Episode ended! Reason: arrive_dest.") + done_info[TerminationState.SUCCESS] = True + if self._is_out_of_road(vehicle): + done = True + logging.info("Episode ended! Reason: out_of_road.") + done_info[TerminationState.OUT_OF_ROAD] = True + if vehicle.crash_vehicle: + done = True + logging.info("Episode ended! Reason: crash vehicle ") + done_info[TerminationState.CRASH_VEHICLE] = True + if vehicle.crash_object: + done = True + done_info[TerminationState.CRASH_OBJECT] = True + logging.info("Episode ended! Reason: crash object ") + if vehicle.crash_building: + done = True + done_info[TerminationState.CRASH_BUILDING] = True + logging.info("Episode ended! Reason: crash building ") + if self.config["max_step_per_agent"] is not None and \ + self.episode_lengths[vehicle_id] >= self.config["max_step_per_agent"]: + done = True + done_info[TerminationState.MAX_STEP] = True + logging.info("Episode ended! Reason: max step ") + + if self.config["horizon"] is not None and \ + self.episode_lengths[vehicle_id] >= self.config["horizon"] and not self.is_multi_agent: + # single agent horizon has the same meaning as max_step_per_agent + done = True + done_info[TerminationState.MAX_STEP] = True + logging.info("Episode ended! Reason: max step ") + + done_info[TerminationState.CRASH] = ( + done_info[TerminationState.CRASH_VEHICLE] or done_info[TerminationState.CRASH_OBJECT] + or done_info[TerminationState.CRASH_BUILDING] + ) + return done, done_info + + def reward_function(self, vehicle_id: str): + """ + Override this func to get a new reward function + :param vehicle_id: id of BaseVehicle + :return: reward + """ + vehicle = self.vehicles[vehicle_id] + step_info = dict() + + # Reward for moving forward in current lane + if vehicle.lane in vehicle.navigation.current_ref_lanes: + current_lane = vehicle.lane + positive_road = 1 + else: + current_lane = vehicle.navigation.current_ref_lanes[0] + current_road = vehicle.navigation.current_road + positive_road = 1 if not current_road.is_negative_road() else -1 + long_last, _ = current_lane.local_coordinates(vehicle.last_position) + long_now, lateral_now = current_lane.local_coordinates(vehicle.position) + + # reward for lane keeping, without it vehicle can learn to overtake but fail to keep in lane + if self.config["use_lateral_reward"]: + lateral_factor = clip(1 - 2 * abs(lateral_now) / vehicle.navigation.get_current_lane_width(), 0.0, 1.0) + else: + lateral_factor = 1.0 + + reward = 0.0 + reward += self.config["driving_reward"] * (long_now - long_last) * lateral_factor * positive_road + reward += self.config["speed_reward"] * (vehicle.speed / vehicle.max_speed) * positive_road + + step_info["step_reward"] = reward + + if self._is_arrive_destination(vehicle): + reward = +self.config["success_reward"] + elif self._is_out_of_road(vehicle): + reward = -self.config["out_of_road_penalty"] + elif vehicle.crash_vehicle: + reward = -self.config["crash_vehicle_penalty"] + elif vehicle.crash_object: + reward = -self.config["crash_object_penalty"] + return reward, step_info + + def _get_reset_return(self): + ret = {} + self.engine.after_step() + for v_id, v in self.vehicles.items(): + self.observations[v_id].reset(self, v) + ret[v_id] = self.observations[v_id].observe(v) + return ret if self.is_multi_agent else self._wrap_as_single_agent(ret) + + def switch_to_third_person_view(self) -> (str, BaseVehicle): + if self.main_camera is None: + return + self.main_camera.reset() + if self.config["prefer_track_agent"] is not None and self.config["prefer_track_agent"] in self.vehicles.keys(): + new_v = self.vehicles[self.config["prefer_track_agent"]] + current_track_vehicle = new_v + else: + if self.main_camera.is_bird_view_camera(): + current_track_vehicle = self.current_track_vehicle + else: + vehicles = list(self.engine.agents.values()) + if len(vehicles) <= 1: + return + if self.current_track_vehicle in vehicles: + vehicles.remove(self.current_track_vehicle) + new_v = get_np_random().choice(vehicles) + current_track_vehicle = new_v + self.main_camera.track(current_track_vehicle) + return + + def switch_to_top_down_view(self): + self.main_camera.stop_track() + + def setup_engine(self): + super(MetaDrivePPOOriginEnv, self).setup_engine() + self.engine.accept("b", self.switch_to_top_down_view) + self.engine.accept("q", self.switch_to_third_person_view) + from metadrive.manager.traffic_manager import TrafficManager + from metadrive.manager.map_manager import MapManager + self.engine.register_manager("map_manager", MapManager()) + self.engine.register_manager("traffic_manager", TrafficManager()) + + def _is_arrive_destination(self, vehicle): + long, lat = vehicle.navigation.final_lane.local_coordinates(vehicle.position) + flag = (vehicle.navigation.final_lane.length - 5 < long < vehicle.navigation.final_lane.length + 5) and ( + vehicle.navigation.get_current_lane_width() / 2 >= lat >= + (0.5 - vehicle.navigation.get_current_lane_num()) * vehicle.navigation.get_current_lane_width() + ) + return flag + + def _reset_global_seed(self, force_seed=None): + """ + Current seed is set to force seed if force_seed is not None. + Otherwise, current seed is randomly generated. + """ + current_seed = force_seed if force_seed is not None else \ + get_np_random(self._DEBUG_RANDOM_SEED).randint(self.start_seed, self.start_seed + self.env_num) + self.seed(current_seed) + + def _get_observations(self): + return {DEFAULT_AGENT: self.get_single_observation(self.config["vehicle_config"])} + + def get_single_observation(self, _=None): + return TopDownMultiChannel( + self.config["vehicle_config"], + self.config["on_screen"], + self.config["rgb_clip"], + frame_stack=3, + post_stack=10, + frame_skip=1, + resolution=(84, 84), + max_distance=36, + ) + + def clone(self, caller: str): + cfg = copy.deepcopy(self.raw_cfg) + return MetaDrivePPOOriginEnv(cfg) diff --git a/dizoo/metadrive/env/drive_utils.py b/dizoo/metadrive/env/drive_utils.py new file mode 100644 index 0000000000..2009e5a52d --- /dev/null +++ b/dizoo/metadrive/env/drive_utils.py @@ -0,0 +1,121 @@ +from typing import Optional, List +from gym import utils +from abc import ABC, abstractmethod +from typing import Any, Dict, Optional +from easydict import EasyDict +from itertools import product +import gym +import copy +import numpy as np +import matplotlib.pyplot as plt +from ding.utils.default_helper import deep_merge_dicts + + +class AAA(): + + def __init__(self) -> None: + self.x = 0 + + +def deep_update( + original: dict, + new_dict: dict, + new_keys_allowed: bool = False, + whitelist: Optional[List[str]] = None, + override_all_if_type_changes: Optional[List[str]] = None +): + """ + Overview: + Updates original dict with values from new_dict recursively. + + .. note:: + + If new key is introduced in new_dict, then if new_keys_allowed is not + True, an error will be thrown. Further, for sub-dicts, if the key is + in the whitelist, then new subkeys can be introduced. + + Arguments: + - original (:obj:`dict`): Dictionary with default values. + - new_dict (:obj:`dict`): Dictionary with values to be updated + - new_keys_allowed (:obj:`bool`): Whether new keys are allowed. + - whitelist (Optional[List[str]]): List of keys that correspond to dict + values where new subkeys can be introduced. This is only at the top + level. + - override_all_if_type_changes(Optional[List[str]]): List of top level + keys with value=dict, for which we always simply override the + entire value (:obj:`dict`), if the "type" key in that value dict changes. + """ + whitelist = whitelist or [] + override_all_if_type_changes = override_all_if_type_changes or [] + for k, value in new_dict.items(): + if k not in original and not new_keys_allowed: + raise RuntimeError("Unknown config parameter `{}`. Base config have: {}.".format(k, original.keys())) + # Both original value and new one are dicts. + if isinstance(original.get(k), dict) and isinstance(value, dict): + # Check old type vs old one. If different, override entire value. + if k in override_all_if_type_changes and \ + "type" in value and "type" in original[k] and \ + value["type"] != original[k]["type"]: + original[k] = value + # Whitelisted key -> ok to add new subkeys. + elif k in whitelist: + deep_update(original[k], value, True) + # Non-whitelisted key. + else: + deep_update(original[k], value, new_keys_allowed) + # Original value not a dict OR new value not a dict: + # Override entire value. + else: + original[k] = value + return original + + +class BaseDriveEnv(gym.Env, utils.EzPickle): + config = dict() + + @abstractmethod + def __init__(self, cfg: Dict, **kwargs) -> None: + if 'cfg_type' not in cfg: + self._cfg = self.__class__.default_config() + self._cfg = deep_merge_dicts(self._cfg, cfg) + else: + self._cfg = cfg + utils.EzPickle.__init__(self) + + @abstractmethod + def step(self, action: Any) -> Any: + """ + Run one step of the environment and return the observation dict. + """ + raise NotImplementedError + + @abstractmethod + def reset(self, *args, **kwargs) -> Any: + """ + Reset current environment. + """ + raise NotImplementedError + + @abstractmethod + def close(self) -> None: + """ + Release all resources in environment and close. + """ + raise NotImplementedError + + @abstractmethod + def seed(self, seed: int) -> None: + """ + Set random seed. + """ + raise NotImplementedError + + @classmethod + def default_config(cls: type) -> EasyDict: + cfg = EasyDict(cls.config) + cfg.cfg_type = cls.__name__ + 'Config' + return copy.deepcopy(cfg) + + @abstractmethod + def __repr__(self) -> str: + raise NotImplementedError diff --git a/dizoo/metadrive/env/drive_wrapper.py b/dizoo/metadrive/env/drive_wrapper.py new file mode 100644 index 0000000000..9b1a1373fd --- /dev/null +++ b/dizoo/metadrive/env/drive_wrapper.py @@ -0,0 +1,149 @@ +from typing import Any, Dict, Optional +from easydict import EasyDict +import matplotlib.pyplot as plt +import gym +import copy +import numpy as np +from ding.envs.env.base_env import BaseEnvTimestep +from ding.torch_utils.data_helper import to_ndarray +from ding.utils.default_helper import deep_merge_dicts +from dizoo.metadrive.env.drive_utils import BaseDriveEnv + + +def draw_multi_channels_top_down_observation(obs, show_time=0.5): + num_channels = obs.shape[-1] + assert num_channels == 5 + channel_names = [ + "Road and navigation", "Ego now and previous pos", "Neighbor at step t", "Neighbor at step t-1", + "Neighbor at step t-2" + ] + fig, axs = plt.subplots(1, num_channels, figsize=(15, 4), dpi=80) + count = 0 + + def close_event(): + plt.close() + + timer = fig.canvas.new_timer(interval=show_time * 1000) + timer.add_callback(close_event) + for i, name in enumerate(channel_names): + count += 1 + ax = axs[i] + ax.imshow(obs[..., i], cmap="bone") + ax.set_xticks([]) + ax.set_yticks([]) + ax.set_title(name) + fig.suptitle("Multi-channels Top-down Observation") + timer.start() + plt.show() + plt.close() + + +class DriveEnvWrapper(gym.Wrapper): + """ + Overview: + Environment wrapper to make ``gym.Env`` align with DI-engine definitions, so as to use utilities in DI-engine. + It changes ``step``, ``reset`` and ``info`` method of ``gym.Env``, while others are straightly delivered. + + Arguments: + - env (BaseDriveEnv): The environment to be wrapped. + - cfg (Dict): Config dict. + """ + config = dict() + + def __init__(self, env: BaseDriveEnv, cfg: Dict = None, **kwargs) -> None: + if cfg is None: + self._cfg = self.__class__.default_config() + elif 'cfg_type' not in cfg: + self._cfg = self.__class__.default_config() + self._cfg = deep_merge_dicts(self._cfg, cfg) + else: + self._cfg = cfg + self.env = env + if not hasattr(self.env, 'reward_space'): + self.reward_space = gym.spaces.Box(low=-float('inf'), high=float('inf'), shape=(1, )) + if 'show_bird_view' in self._cfg and self._cfg['show_bird_view'] is True: + self.show_bird_view = True + else: + self.show_bird_view = False + self.action_space = self.env.action_space + self.env = env + + def reset(self, *args, **kwargs) -> Any: + """ + Overview: + Wrapper of ``reset`` method in env. The observations are converted to ``np.ndarray`` and final reward + are recorded. + Returns: + - Any: Observations from environment + """ + obs = self.env.reset(*args, **kwargs) + obs = to_ndarray(obs, dtype=np.float32) + if isinstance(obs, np.ndarray) and len(obs.shape) == 3: + obs = obs.transpose((2, 0, 1)) + elif isinstance(obs, dict): + vehicle_state = obs['vehicle_state'] + birdview = obs['birdview'].transpose((2, 0, 1)) + obs = {'vehicle_state': vehicle_state, 'birdview': birdview} + self._eval_episode_return = 0.0 + self._arrive_dest = False + return obs + + def step(self, action: Any = None) -> BaseEnvTimestep: + """ + Overview: + Wrapper of ``step`` method in env. This aims to convert the returns of ``gym.Env`` step method into + that of ``ding.envs.BaseEnv``, from ``(obs, reward, done, info)`` tuple to a ``BaseEnvTimestep`` + namedtuple defined in DI-engine. It will also convert actions, observations and reward into + ``np.ndarray``, and check legality if action contains control signal. + Arguments: + - action (Any, optional): Actions sent to env. Defaults to None. + Returns: + - BaseEnvTimestep: DI-engine format of env step returns. + """ + action = to_ndarray(action) + obs, rew, done, info = self.env.step(action) + if self.show_bird_view: + draw_multi_channels_top_down_observation(obs, show_time=0.5) + self._eval_episode_return += rew + obs = to_ndarray(obs, dtype=np.float32) + if isinstance(obs, np.ndarray) and len(obs.shape) == 3: + obs = obs.transpose((2, 0, 1)) + elif isinstance(obs, dict): + vehicle_state = obs['vehicle_state'] + birdview = obs['birdview'].transpose((2, 0, 1)) + obs = {'vehicle_state': vehicle_state, 'birdview': birdview} + rew = to_ndarray([rew], dtype=np.float32) + if done: + info['eval_episode_return'] = self._eval_episode_return + return BaseEnvTimestep(obs, rew, done, info) + + @property + def observation_space(self): + return gym.spaces.Box(0, 1, shape=(5, 84, 84), dtype=np.float32) + + def seed(self, seed: int, dynamic_seed: bool = True) -> None: + self._seed = seed + self._dynamic_seed = dynamic_seed + np.random.seed(self._seed) + + def enable_save_replay(self, replay_path: Optional[str] = None) -> None: + if replay_path is None: + replay_path = './video' + self._replay_path = replay_path + self.env = gym.wrappers.Monitor(self.env, self._replay_path, video_callable=lambda episode_id: True, force=True) + + @classmethod + def default_config(cls: type) -> EasyDict: + cfg = EasyDict(cls.config) + cfg.cfg_type = cls.__name__ + 'Config' + return copy.deepcopy(cfg) + + def __repr__(self) -> str: + return repr(self.env) + + def render(self): + self.env.render() + + def clone(self, caller: str): + cfg = copy.deepcopy(self._cfg) + return DriveEnvWrapper(self.env.clone(caller), cfg) diff --git a/dizoo/metadrive/metadrive_env.gif b/dizoo/metadrive/metadrive_env.gif new file mode 100644 index 0000000000..5e78923051 Binary files /dev/null and b/dizoo/metadrive/metadrive_env.gif differ diff --git a/dizoo/minigrid/__init__.py b/dizoo/minigrid/__init__.py index f2832cb85c..b05d8f7ceb 100644 --- a/dizoo/minigrid/__init__.py +++ b/dizoo/minigrid/__init__.py @@ -1,4 +1,4 @@ -from gym.envs.registration import register +from gymnasium.envs.registration import register register(id='MiniGrid-AKTDT-7x7-1-v0', entry_point='dizoo.minigrid.envs:AppleKeyToDoorTreasure_7x7_1') @@ -10,4 +10,6 @@ register(id='MiniGrid-AKTDT-19x19-v0', entry_point='dizoo.minigrid.envs:AppleKeyToDoorTreasure_19x19') -register(id='MiniGrid-AKTDT-19x19-3-v0', entry_point='dizoo.minigrid.envs:AppleKeyToDoorTreasure_19x19_3') \ No newline at end of file +register(id='MiniGrid-AKTDT-19x19-3-v0', entry_point='dizoo.minigrid.envs:AppleKeyToDoorTreasure_19x19_3') + +register(id='MiniGrid-NoisyTV-v0', entry_point='dizoo.minigrid.envs:NoisyTVEnv') diff --git a/dizoo/minigrid/config/minigrid_dreamer_config.py b/dizoo/minigrid/config/minigrid_dreamer_config.py new file mode 100644 index 0000000000..68afa2757c --- /dev/null +++ b/dizoo/minigrid/config/minigrid_dreamer_config.py @@ -0,0 +1,96 @@ +from easydict import EasyDict + +from ding.entry import serial_pipeline_dreamer + +cuda = False +collector_env_num = 8 +evaluator_env_num = 5 +minigrid_dreamer_config = dict( + exp_name='minigrid_dreamer_empty', + env=dict( + collector_env_num=collector_env_num, + evaluator_env_num=evaluator_env_num, + n_evaluator_episode=evaluator_env_num, + # typical MiniGrid env id: + # {'MiniGrid-Empty-8x8-v0', 'MiniGrid-FourRooms-v0', 'MiniGrid-DoorKey-8x8-v0','MiniGrid-DoorKey-16x16-v0'}, + # please refer to https://github.com/Farama-Foundation/MiniGrid for details. + env_id='MiniGrid-Empty-8x8-v0', + # env_id='MiniGrid-AKTDT-7x7-1-v0', + max_step=100, + stop_value=20, # run fixed env_steps + # stop_value=0.96, + flat_obs=True, + full_obs=True, + onehot_obs=True, + move_bonus=True, + ), + policy=dict( + cuda=cuda, + # it is better to put random_collect_size in policy.other + random_collect_size=2500, + model=dict( + action_shape=7, + # encoder_hidden_size_list=[256, 128, 64, 64], + # critic_head_hidden_size=64, + # actor_head_hidden_size=64, + actor_dist='onehot', + ), + learn=dict( + lambda_=0.95, + learning_rate=3e-5, + batch_size=16, + batch_length=64, + imag_sample=True, + discount=0.997, + reward_EMA=True, + ), + collect=dict( + n_sample=1, + unroll_len=1, + action_size=7, # has to be specified + collect_dyn_sample=True, + ), + eval=dict(evaluator=dict(eval_freq=5000, )), + other=dict( + # environment buffer + replay_buffer=dict(replay_buffer_size=500000, periodic_thruput_seconds=60), + ), + ), + world_model=dict( + pretrain=100, + train_freq=2, + cuda=cuda, + model=dict( + state_size=1344, + obs_type='vector', + action_size=7, + action_type='discrete', + encoder_hidden_size_list=[256, 128, 64, 64], + reward_size=1, + batch_size=16, + ), + ), +) + +minigrid_dreamer_config = EasyDict(minigrid_dreamer_config) + +minigrid_create_config = dict( + env=dict( + type='minigrid', + import_names=['dizoo.minigrid.envs.minigrid_env'], + ), + env_manager=dict(type='base'), + policy=dict( + type='dreamer', + import_names=['ding.policy.mbpolicy.dreamer'], + ), + replay_buffer=dict(type='sequence', ), + world_model=dict( + type='dreamer', + import_names=['ding.world_model.dreamer'], + ), +) +minigrid_create_config = EasyDict(minigrid_create_config) + +if __name__ == '__main__': + serial_pipeline_dreamer((minigrid_dreamer_config, minigrid_create_config), seed=0, max_env_step=500000) diff --git a/dizoo/minigrid/config/minigrid_icm_offppo_config.py b/dizoo/minigrid/config/minigrid_icm_offppo_config.py index f56d439ed8..92465fd1fb 100644 --- a/dizoo/minigrid/config/minigrid_icm_offppo_config.py +++ b/dizoo/minigrid/config/minigrid_icm_offppo_config.py @@ -1,38 +1,56 @@ from easydict import EasyDict -minigrid_ppo_icm_config = dict( - exp_name='minigrid_icm_offppo_seed0', +minigrid_icm_offppo_config = dict( + exp_name='minigrid_fourroom_icm_offppo_seed0', env=dict( collector_env_num=8, evaluator_env_num=5, n_evaluator_episode=5, - # minigrid env id: 'MiniGrid-Empty-8x8-v0', 'MiniGrid-FourRooms-v0','MiniGrid-DoorKey-16x16-v0' - env_id='MiniGrid-DoorKey-8x8-v0', - max_step=300, + # minigrid env id: 'MiniGrid-Empty-8x8-v0', 'MiniGrid-FourRooms-v0','MiniGrid-DoorKey-16x16-v0','MiniGrid-AKTDT-7x7-1-v0' + env_id='MiniGrid-FourRooms-v0', + max_step=100, stop_value=0.96, ), reward_model=dict( intrinsic_reward_type='add', - learning_rate=0.001, - obs_shape=2739, + # intrinsic_reward_weight means the relative weight of RND intrinsic_reward. + # Specifically for sparse reward env MiniGrid, in this env, + # if reach goal, the agent get reward ~1, otherwise 0, + # We could set the intrinsic_reward_weight approximately equal to the inverse of max_episode_steps. + # Please refer to rnd_reward_model for details. + intrinsic_reward_weight=0.001, + learning_rate=3e-4, + obs_shape=2835, batch_size=320, - update_per_collect=10, + update_per_collect=50, + clear_buffer_per_iters=int(1e3), + obs_norm=True, + obs_norm_clamp_max=5, + obs_norm_clamp_min=-5, + extrinsic_reward_norm=True, + extrinsic_reward_norm_max=1, ), policy=dict( cuda=True, + recompute_adv=True, + action_space='discrete', model=dict( - obs_shape=2739, + obs_shape=2835, action_shape=7, + action_space='discrete', encoder_hidden_size_list=[256, 128, 64, 64], + critic_head_hidden_size=64, + actor_head_hidden_size=64, ), learn=dict( - update_per_collect=10, + update_per_collect=1, batch_size=320, - learning_rate=0.0003, + learning_rate=3e-4, value_weight=0.5, entropy_weight=0.001, clip_ratio=0.2, - adv_norm=False, + adv_norm=True, + value_norm=True, ), collect=dict( n_sample=3200, @@ -40,11 +58,12 @@ discount_factor=0.99, gae_lambda=0.95, ), + eval=dict(evaluator=dict(eval_freq=200, )), ), ) -minigrid_ppo_icm_config = EasyDict(minigrid_ppo_icm_config) -main_config = minigrid_ppo_icm_config -minigrid_ppo_icm_create_config = dict( +minigrid_icm_offppo_config = EasyDict(minigrid_icm_offppo_config) +main_config = minigrid_icm_offppo_config +minigrid_icm_offppo_create_config = dict( env=dict( type='minigrid', import_names=['dizoo.minigrid.envs.minigrid_env'], @@ -53,10 +72,10 @@ policy=dict(type='ppo_offpolicy'), reward_model=dict(type='icm'), ) -minigrid_ppo_icm_create_config = EasyDict(minigrid_ppo_icm_create_config) -create_config = minigrid_ppo_icm_create_config +minigrid_icm_offppo_create_config = EasyDict(minigrid_icm_offppo_create_config) +create_config = minigrid_icm_offppo_create_config if __name__ == "__main__": # or you can enter `ding -m serial -c minigrid_icm_config.py -s 0` from ding.entry import serial_pipeline_reward_model_offpolicy - serial_pipeline_reward_model_offpolicy([main_config, create_config], seed=0) \ No newline at end of file + serial_pipeline_reward_model_offpolicy([main_config, create_config], seed=0, max_env_step=int(10e6)) diff --git a/dizoo/minigrid/config/minigrid_icm_onppo_config.py b/dizoo/minigrid/config/minigrid_icm_onppo_config.py new file mode 100644 index 0000000000..9e56b30cbe --- /dev/null +++ b/dizoo/minigrid/config/minigrid_icm_onppo_config.py @@ -0,0 +1,82 @@ +from easydict import EasyDict + +collector_env_num = 8 +evaluator_env_num = 5 +minigrid_icm_onppo_config = dict( + exp_name='minigrid_AKTDT-7x7_icm_onppo_seed0', + env=dict( + collector_env_num=collector_env_num, + evaluator_env_num=evaluator_env_num, + n_evaluator_episode=evaluator_env_num, + # minigrid env id: 'MiniGrid-Empty-8x8-v0', 'MiniGrid-FourRooms-v0','MiniGrid-DoorKey-16x16-v0','MiniGrid-AKTDT-7x7-1-v0' + env_id='MiniGrid-NoisyTV-v0', + max_step=100, + stop_value=12, # run fixed env_steps for MiniGrid-AKTDT-7x7-1-v0 + # stop_value=0.96, + ), + reward_model=dict( + intrinsic_reward_type='add', + # intrinsic_reward_weight means the relative weight of ICM intrinsic_reward. + # Specifically for sparse reward env MiniGrid, in this env, + # if reach goal, the agent get reward ~1, otherwise 0, + # We could set the intrinsic_reward_weight approximately equal to the inverse of max_episode_steps. + # Please refer to rnd_reward_model for details. + intrinsic_reward_weight=0.003, # 1/300 + learning_rate=3e-4, + obs_shape=2835, # 2715 in MiniGrid-AKTDT-7x7-1-v0 env + batch_size=320, + update_per_collect=50, + clear_buffer_per_iters=int(1e3), + extrinsic_reward_norm=True, + extrinsic_reward_norm_max=1, + ), + policy=dict( + cuda=True, + recompute_adv=True, + action_space='discrete', + model=dict( + obs_shape=2835, # 2715 in MiniGrid-AKTDT-7x7-1-v0 env + action_shape=7, + action_space='discrete', + encoder_hidden_size_list=[256, 128, 64, 64], + critic_head_hidden_size=64, + actor_head_hidden_size=64, + ), + learn=dict( + epoch_per_collect=10, + update_per_collect=1, + batch_size=320, + learning_rate=3e-4, + value_weight=0.5, + entropy_weight=0.001, + clip_ratio=0.2, + adv_norm=True, + value_norm=True, + ), + collect=dict( + n_sample=3200, + unroll_len=1, + discount_factor=0.99, + gae_lambda=0.95, + ), + eval=dict(evaluator=dict(eval_freq=1000, )), + ), +) +minigrid_icm_onppo_config = EasyDict(minigrid_icm_onppo_config) +main_config = minigrid_icm_onppo_config +minigrid_icm_onppo_create_config = dict( + env=dict( + type='minigrid', + import_names=['dizoo.minigrid.envs.minigrid_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='ppo'), + reward_model=dict(type='icm'), +) +minigrid_icm_onppo_create_config = EasyDict(minigrid_icm_onppo_create_config) +create_config = minigrid_icm_onppo_create_config + +if __name__ == "__main__": + # or you can enter `ding -m serial -c minigrid_icm_onppo_config.py -s 0` + from ding.entry import serial_pipeline_reward_model_onpolicy + serial_pipeline_reward_model_onpolicy([main_config, create_config], seed=0, max_env_step=int(10e6)) diff --git a/dizoo/minigrid/config/minigrid_ngu_config.py b/dizoo/minigrid/config/minigrid_ngu_config.py index 33960edcb0..c1aa47e1eb 100644 --- a/dizoo/minigrid/config/minigrid_ngu_config.py +++ b/dizoo/minigrid/config/minigrid_ngu_config.py @@ -9,7 +9,9 @@ collector_env_num=collector_env_num, evaluator_env_num=evaluator_env_num, n_evaluator_episode=evaluator_env_num, - # MiniGrid env id: 'MiniGrid-Empty-8x8-v0', 'MiniGrid-FourRooms-v0' + # typical MiniGrid env id: + # {'MiniGrid-Empty-8x8-v0', 'MiniGrid-FourRooms-v0', 'MiniGrid-DoorKey-8x8-v0','MiniGrid-DoorKey-16x16-v0'}, + # please refer to https://github.com/Farama-Foundation/MiniGrid for details. env_id='MiniGrid-DoorKey-16x16-v0', obs_plus_prev_action_reward=True, # use specific env wrapper for ngu policy max_step=300, @@ -18,7 +20,7 @@ rnd_reward_model=dict( intrinsic_reward_type='add', learning_rate=5e-4, - obs_shape=2739, + obs_shape=2835, action_shape=7, batch_size=320, # transitions update_per_collect=10, # 32*100/320=10 diff --git a/dizoo/minigrid/config/minigrid_offppo_config.py b/dizoo/minigrid/config/minigrid_offppo_config.py index b99001f64d..9ae43f5a84 100644 --- a/dizoo/minigrid/config/minigrid_offppo_config.py +++ b/dizoo/minigrid/config/minigrid_offppo_config.py @@ -5,7 +5,9 @@ env=dict( collector_env_num=8, evaluator_env_num=5, - # minigrid env id: 'MiniGrid-FourRooms-v0', 'MiniGrid-DoorKey-8x8-v0','MiniGrid-DoorKey-16x16-v0' + # typical MiniGrid env id: + # {'MiniGrid-Empty-8x8-v0', 'MiniGrid-FourRooms-v0', 'MiniGrid-DoorKey-8x8-v0','MiniGrid-DoorKey-16x16-v0'}, + # please refer to https://github.com/Farama-Foundation/MiniGrid for details. env_id='MiniGrid-Empty-8x8-v0', n_evaluator_episode=5, max_step=300, @@ -14,7 +16,7 @@ policy=dict( cuda=True, model=dict( - obs_shape=2739, + obs_shape=2835, action_shape=7, encoder_hidden_size_list=[256, 128, 64, 64], ), diff --git a/dizoo/minigrid/config/minigrid_onppo_config.py b/dizoo/minigrid/config/minigrid_onppo_config.py index 6fdd135827..de4f81cd8e 100644 --- a/dizoo/minigrid/config/minigrid_onppo_config.py +++ b/dizoo/minigrid/config/minigrid_onppo_config.py @@ -7,7 +7,9 @@ collector_env_num=8, evaluator_env_num=5, n_evaluator_episode=5, - # minigrid env id: 'MiniGrid-FourRooms-v0', 'MiniGrid-DoorKey-8x8-v0','MiniGrid-DoorKey-16x16-v0' + # typical MiniGrid env id: + # {'MiniGrid-Empty-8x8-v0', 'MiniGrid-FourRooms-v0', 'MiniGrid-DoorKey-8x8-v0','MiniGrid-DoorKey-16x16-v0'}, + # please refer to https://github.com/Farama-Foundation/MiniGrid for details. env_id='MiniGrid-Empty-8x8-v0', max_step=300, stop_value=0.96, @@ -17,7 +19,7 @@ recompute_adv=True, action_space='discrete', model=dict( - obs_shape=2739, + obs_shape=2835, action_shape=7, action_space='discrete', encoder_hidden_size_list=[256, 128, 64, 64], diff --git a/dizoo/minigrid/config/minigrid_onppo_stdim_config.py b/dizoo/minigrid/config/minigrid_onppo_stdim_config.py index c62908cf7b..f4e025984c 100644 --- a/dizoo/minigrid/config/minigrid_onppo_stdim_config.py +++ b/dizoo/minigrid/config/minigrid_onppo_stdim_config.py @@ -8,9 +8,9 @@ collector_env_num=collector_env_num, evaluator_env_num=evaluator_env_num, n_evaluator_episode=evaluator_env_num, - # MiniGrid env_id choices={'MiniGrid-Empty-8x8-v0', 'MiniGrid-FourRooms-v0',\ - # 'MiniGrid-DoorKey-8x8-v0', 'MiniGrid-DoorKey-16x16-v0', ...}, please refer to\ - # https://github.com/Farama-Foundation/gym-minigrid for details. + # typical MiniGrid env id: + # {'MiniGrid-Empty-8x8-v0', 'MiniGrid-FourRooms-v0', 'MiniGrid-DoorKey-8x8-v0','MiniGrid-DoorKey-16x16-v0'}, + # please refer to https://github.com/Farama-Foundation/MiniGrid for details. env_id='MiniGrid-Empty-8x8-v0', max_step=300, stop_value=0.96, @@ -20,7 +20,7 @@ recompute_adv=True, action_space='discrete', model=dict( - obs_shape=2739, + obs_shape=2835, action_shape=7, action_space='discrete', encoder_hidden_size_list=[256, 128, 64, 64], diff --git a/dizoo/minigrid/config/minigrid_r2d2_config.py b/dizoo/minigrid/config/minigrid_r2d2_config.py index 89eff09115..b650a503df 100644 --- a/dizoo/minigrid/config/minigrid_r2d2_config.py +++ b/dizoo/minigrid/config/minigrid_r2d2_config.py @@ -7,7 +7,9 @@ env=dict( collector_env_num=collector_env_num, evaluator_env_num=evaluator_env_num, - # minigrid env id: 'MiniGrid-FourRooms-v0', 'MiniGrid-DoorKey-8x8-v0','MiniGrid-DoorKey-16x16-v0' + # typical MiniGrid env id: + # {'MiniGrid-Empty-8x8-v0', 'MiniGrid-FourRooms-v0', 'MiniGrid-DoorKey-8x8-v0','MiniGrid-DoorKey-16x16-v0'}, + # please refer to https://github.com/Farama-Foundation/MiniGrid for details. env_id='MiniGrid-DoorKey-16x16-v0', n_evaluator_episode=5, max_step=300, @@ -19,7 +21,7 @@ priority=True, priority_IS_weight=True, model=dict( - obs_shape=2739, + obs_shape=2835, action_shape=7, encoder_hidden_size_list=[128, 128, 512], ), diff --git a/dizoo/minigrid/config/minigrid_rnd_onppo_config.py b/dizoo/minigrid/config/minigrid_rnd_onppo_config.py index ca13db078d..a39d3e4d8f 100644 --- a/dizoo/minigrid/config/minigrid_rnd_onppo_config.py +++ b/dizoo/minigrid/config/minigrid_rnd_onppo_config.py @@ -1,46 +1,47 @@ from easydict import EasyDict collector_env_num = 8 -evaluator_env_num = 8 +evaluator_env_num = 5 minigrid_ppo_rnd_config = dict( - exp_name='minigrid_empty8_rnd_onppo_seed0', + exp_name='minigrid_doorkey8x8_rnd_onppo_seed0', env=dict( collector_env_num=collector_env_num, evaluator_env_num=evaluator_env_num, n_evaluator_episode=evaluator_env_num, - # MiniGrid env id: 'MiniGrid-Empty-8x8-v0', 'MiniGrid-FourRooms-v0','MiniGrid-DoorKey-16x16-v0' - env_id='MiniGrid-Empty-8x8-v0', - max_step=300, - stop_value=0.96, + # typical MiniGrid env id: + # {'MiniGrid-Empty-8x8-v0', 'MiniGrid-FourRooms-v0', 'MiniGrid-DoorKey-8x8-v0','MiniGrid-DoorKey-16x16-v0'}, + # please refer to https://github.com/Farama-Foundation/MiniGrid for details. + env_id='MiniGrid-DoorKey-8x8-v0', + # env_id='MiniGrid-AKTDT-7x7-1-v0', + max_step=100, + stop_value=20, # run fixed env_steps + # stop_value=0.96, ), reward_model=dict( intrinsic_reward_type='add', - intrinsic_reward_weight=None, - # means the relative weight of RND intrinsic_reward. - # If intrinsic_reward_weight=None, we will automatically set it based on - # the absolute value of the difference between max and min extrinsic reward in the sampled mini-batch - # please refer to rnd_reward_model for details. + # intrinsic_reward_weight means the relative weight of RND intrinsic_reward. # Specifically for sparse reward env MiniGrid, in this env, - # if reach goal, the agent get reward ~1, otherwise 0. + # if reach goal, the agent get reward ~1, otherwise 0, # We could set the intrinsic_reward_weight approximately equal to the inverse of max_episode_steps. - intrinsic_reward_rescale=0.001, - # means the rescale value of RND intrinsic_reward only used when intrinsic_reward_weight is None - # please refer to rnd_reward_model for details. - learning_rate=5e-4, - obs_shape=2739, + # Please refer to rnd_reward_model for details. + intrinsic_reward_weight=0.003, # 1/300 + learning_rate=3e-4, + obs_shape=2835, batch_size=320, - update_per_collect=10, - clear_buffer_per_iters=10, - obs_norm=True, + update_per_collect=50, + clear_buffer_per_iters=int(1e3), + obs_norm=False, obs_norm_clamp_max=5, obs_norm_clamp_min=-5, + extrinsic_reward_norm=True, + extrinsic_reward_norm_max=1, ), policy=dict( recompute_adv=True, cuda=True, action_space='discrete', model=dict( - obs_shape=2739, + obs_shape=2835, action_shape=7, action_space='discrete', encoder_hidden_size_list=[256, 128, 64, 64], @@ -65,6 +66,7 @@ discount_factor=0.99, gae_lambda=0.95, ), + eval=dict(evaluator=dict(eval_freq=1000, )), ), ) minigrid_ppo_rnd_config = EasyDict(minigrid_ppo_rnd_config) @@ -82,6 +84,6 @@ create_config = minigrid_ppo_rnd_create_config if __name__ == "__main__": - # or you can enter `ding -m serial -c minigrid_rnd_config.py -s 0` + # or you can enter `ding -m serial -c minigrid_rnd_onppo_config.py -s 0` from ding.entry import serial_pipeline_reward_model_onpolicy - serial_pipeline_reward_model_onpolicy([main_config, create_config], seed=0) + serial_pipeline_reward_model_onpolicy([main_config, create_config], seed=0, max_env_step=int(10e6)) diff --git a/dizoo/minigrid/entry/minigrid_onppo_main.py b/dizoo/minigrid/entry/minigrid_onppo_main.py new file mode 100644 index 0000000000..eeb97b2fe7 --- /dev/null +++ b/dizoo/minigrid/entry/minigrid_onppo_main.py @@ -0,0 +1,87 @@ +import gym +from ding.worker import BaseLearner, SampleSerialCollector, InteractionSerialEvaluator +from ding.model import VAC +from ding.policy import PPOPolicy +from ding.envs import DingEnvWrapper, EvalEpisodeReturnWrapper, BaseEnvManager +from ding.config import compile_config +from ding.utils import set_pkg_seed + +from dizoo.minigrid.config.minigrid_onppo_config import minigrid_ppo_config +from minigrid.wrappers import FlatObsWrapper +import numpy as np +from tensorboardX import SummaryWriter +import os +import gymnasium + + +class MinigridWrapper(gym.Wrapper): + + def __init__(self, env): + super().__init__(env) + self._observation_space = gym.spaces.Box(low=float("-inf"), high=float("inf"), shape=(8, ), dtype=np.float32) + self._action_space = gym.spaces.Discrete(9) + self._action_space.seed(0) # default seed + self.reward_range = (float('-inf'), float('inf')) + self.max_steps = minigrid_ppo_config.env.max_step + + def step(self, action): + obs, reward, done, _, info = self.env.step(action) + self.cur_step += 1 + if self.cur_step > self.max_steps: + done = True + return obs, reward, done, info + + def reset(self): + self.cur_step = 0 + return self.env.reset()[0] + + +def wrapped_minigrid_env(): + return DingEnvWrapper( + gymnasium.make(minigrid_ppo_config.env.env_id), + cfg={ + 'env_wrapper': [ + lambda env: FlatObsWrapper(env), + lambda env: MinigridWrapper(env), + lambda env: EvalEpisodeReturnWrapper(env), + ] + } + ) + + +def main(cfg, seed=0, max_env_step=int(1e10), max_train_iter=int(1e10)): + cfg = compile_config( + cfg, BaseEnvManager, PPOPolicy, BaseLearner, SampleSerialCollector, InteractionSerialEvaluator, save_cfg=True + ) + collector_env_num, evaluator_env_num = cfg.env.collector_env_num, cfg.env.evaluator_env_num + collector_env = BaseEnvManager(env_fn=[wrapped_minigrid_env for _ in range(collector_env_num)], cfg=cfg.env.manager) + evaluator_env = BaseEnvManager(env_fn=[wrapped_minigrid_env for _ in range(evaluator_env_num)], cfg=cfg.env.manager) + + collector_env.seed(seed) + evaluator_env.seed(seed, dynamic_seed=False) + set_pkg_seed(seed, use_cuda=cfg.policy.cuda) + + model = VAC(**cfg.policy.model) + policy = PPOPolicy(cfg.policy, model=model) + tb_logger = SummaryWriter(os.path.join('./{}/log/'.format(cfg.exp_name), 'serial')) + learner = BaseLearner(cfg.policy.learn.learner, policy.learn_mode, tb_logger, exp_name=cfg.exp_name) + collector = SampleSerialCollector( + cfg.policy.collect.collector, collector_env, policy.collect_mode, tb_logger, exp_name=cfg.exp_name + ) + evaluator = InteractionSerialEvaluator( + cfg.policy.eval.evaluator, evaluator_env, policy.eval_mode, tb_logger, exp_name=cfg.exp_name + ) + + while True: + if evaluator.should_eval(learner.train_iter): + stop, reward = evaluator.eval(learner.save_checkpoint, learner.train_iter, collector.envstep) + if stop: + break + new_data = collector.collect(train_iter=learner.train_iter) + learner.train(new_data, collector.envstep) + if collector.envstep >= max_env_step or learner.train_iter >= max_train_iter: + break + + +if __name__ == '__main__': + main(minigrid_ppo_config) diff --git a/dizoo/minigrid/envs/__init__.py b/dizoo/minigrid/envs/__init__.py index 8db2e59889..46eb0846e1 100644 --- a/dizoo/minigrid/envs/__init__.py +++ b/dizoo/minigrid/envs/__init__.py @@ -1,2 +1,3 @@ from .minigrid_env import MiniGridEnv -from dizoo.minigrid.envs.app_key_to_door_treasure import AppleKeyToDoorTreasure, AppleKeyToDoorTreasure_13x13, AppleKeyToDoorTreasure_19x19, AppleKeyToDoorTreasure_13x13_1, AppleKeyToDoorTreasure_19x19_3, AppleKeyToDoorTreasure_7x7_1 \ No newline at end of file +from dizoo.minigrid.envs.app_key_to_door_treasure import AppleKeyToDoorTreasure, AppleKeyToDoorTreasure_13x13, AppleKeyToDoorTreasure_19x19, AppleKeyToDoorTreasure_13x13_1, AppleKeyToDoorTreasure_19x19_3, AppleKeyToDoorTreasure_7x7_1 +from dizoo.minigrid.envs.noisy_tv import NoisyTVEnv diff --git a/dizoo/minigrid/envs/app_key_to_door_treasure.py b/dizoo/minigrid/envs/app_key_to_door_treasure.py index 5817ee7460..3c17db3ae6 100644 --- a/dizoo/minigrid/envs/app_key_to_door_treasure.py +++ b/dizoo/minigrid/envs/app_key_to_door_treasure.py @@ -1,8 +1,9 @@ #!/usr/bin/env python # -*- coding: utf-8 -*- -from gym_minigrid.minigrid import * -from gym_minigrid.minigrid import WorldObj +from minigrid.minigrid_env import * +from minigrid.utils.rendering import * +from minigrid.core.world_object import WorldObj class Ball(WorldObj): @@ -27,7 +28,8 @@ def __init__(self, agent_pos=None, goal_pos=None, grid_size=19, apple=2): self._agent_default_pos = agent_pos self._goal_default_pos = goal_pos self.apple = apple - super().__init__(grid_size=grid_size, max_steps=100) + mission_space = MissionSpace(mission_func=lambda: "Reach the goal") + super().__init__(mission_space=mission_space, grid_size=grid_size, max_steps=100) def _gen_grid( self, width, height @@ -58,7 +60,7 @@ def _gen_grid( if i + 1 < 2: if j + 1 < 2: self.grid.vert_wall(xR, yT, room_h) - #pos = (xR, self._rand_int(yT + 1, yB)) + # pos = (xR, self._rand_int(yT + 1, yB)) else: self.grid.vert_wall(xR, yT, room_h) pos = (xR, self._rand_int(yT + 1, yB)) @@ -183,8 +185,8 @@ def step(self, action): done = True obs = self.gen_obs() - - return obs, reward, done, {} + # return is (observation, reward, terminated, truncated, info) + return obs, reward, done, done, {} class AppleKeyToDoorTreasure_13x13(AppleKeyToDoorTreasure): diff --git a/dizoo/minigrid/envs/minigrid_env.py b/dizoo/minigrid/envs/minigrid_env.py index 1689fd61db..923967394e 100644 --- a/dizoo/minigrid/envs/minigrid_env.py +++ b/dizoo/minigrid/envs/minigrid_env.py @@ -4,12 +4,13 @@ import copy import os import time -import gym +import gymnasium as gym + import numpy as np from matplotlib import animation import matplotlib.pyplot as plt -from gym_minigrid.wrappers import FlatObsWrapper, RGBImgPartialObsWrapper, ImgObsWrapper, ViewSizeWrapper -from gym_minigrid.window import Window +from minigrid.wrappers import FullyObsWrapper +from .minigrid_wrapper import ViewSizeWrapper, MoveBonus, OneHotObsWrapper, FlatObsWrapper from ding.envs import ObsPlusPrevActRewWrapper from ding.envs import BaseEnv, BaseEnvTimestep @@ -35,19 +36,31 @@ def __init__(self, cfg: dict) -> None: self._init_flag = False self._env_id = cfg.env_id self._flat_obs = cfg.flat_obs + self._full_obs = cfg.full_obs + self._onehot_obs = cfg.onehot_obs + self._move_bonus = cfg.move_bonus self._save_replay = False - # self._max_step = MINIGRID_INFO_DICT[self._env_id].max_step self._max_step = cfg.max_step def reset(self) -> np.ndarray: if not self._init_flag: - self._env = gym.make(self._env_id) + if self._save_replay: + self._env = gym.make(self._env_id, render_mode="rgb_array") # using the Gymnasium make method + else: + self._env = gym.make(self._env_id) + if self._env_id in ['MiniGrid-AKTDT-13x13-v0' or 'MiniGrid-AKTDT-13x13-1-v0']: # customize the agent field of view size, note this must be an odd number # This also related to the observation space, see gym_minigrid.wrappers for more details self._env = ViewSizeWrapper(self._env, agent_view_size=5) if self._env_id == 'MiniGrid-AKTDT-7x7-1-v0': self._env = ViewSizeWrapper(self._env, agent_view_size=3) + if self._full_obs: + self._env = FullyObsWrapper(self._env) + if self._onehot_obs: + self._env = OneHotObsWrapper(self._env) + if self._move_bonus: + self._env = MoveBonus(self._env) if self._flat_obs: self._env = FlatObsWrapper(self._env) # self._env = RGBImgPartialObsWrapper(self._env) @@ -55,23 +68,29 @@ def reset(self) -> np.ndarray: if hasattr(self._cfg, 'obs_plus_prev_action_reward') and self._cfg.obs_plus_prev_action_reward: self._env = ObsPlusPrevActRewWrapper(self._env) self._init_flag = True - self._observation_space = self._env.observation_space - # to be compatiable with subprocess env manager - if isinstance(self._observation_space, gym.spaces.Dict): - self._observation_space['obs'].dtype = np.dtype('float32') + if self._flat_obs: + self._observation_space = gym.spaces.Box(0, 1, shape=self._env.observation_space.shape, dtype=np.float32) else: - self._observation_space.dtype = np.dtype('float32') + self._observation_space = self._env.observation_space + # to be compatiable with subprocess env manager + if isinstance(self._observation_space, gym.spaces.Dict): + self._observation_space['obs'].dtype = np.dtype('float32') + else: + self._observation_space.dtype = np.dtype('float32') self._action_space = self._env.action_space self._reward_space = gym.spaces.Box( low=self._env.reward_range[0], high=self._env.reward_range[1], shape=(1, ), dtype=np.float32 ) + + self._eval_episode_return = 0 if hasattr(self, '_seed') and hasattr(self, '_dynamic_seed') and self._dynamic_seed: np_seed = 100 * np.random.randint(1, 1000) - self._env.seed(self._seed + np_seed) + self._seed = self._seed + np_seed + obs, _ = self._env.reset(seed=self._seed) # using the reset method of Gymnasium env elif hasattr(self, '_seed'): - self._env.seed(self._seed) - self._final_eval_reward = 0 - obs = self._env.reset() + obs, _ = self._env.reset(seed=self._seed) + else: + obs, _ = self._env.reset() obs = to_ndarray(obs) self._current_step = 0 if self._save_replay: @@ -94,15 +113,16 @@ def step(self, action: np.ndarray) -> BaseEnvTimestep: if action.shape == (1, ): action = action.squeeze() # 0-dim array if self._save_replay: - self._frames.append(self._env.render(mode='rgb_array')) - obs, rew, done, info = self._env.step(action) + self._frames.append(self._env.render()) + # using the step method of Gymnasium env, return is (observation, reward, terminated, truncated, info) + obs, rew, done, _, info = self._env.step(action) rew = float(rew) - self._final_eval_reward += rew + self._eval_episode_return += rew self._current_step += 1 if self._current_step >= self._max_step: done = True if done: - info['final_eval_reward'] = self._final_eval_reward + info['eval_episode_return'] = self._eval_episode_return info['current_step'] = self._current_step info['max_step'] = self._max_step if self._save_replay: diff --git a/dizoo/minigrid/envs/minigrid_wrapper.py b/dizoo/minigrid/envs/minigrid_wrapper.py new file mode 100644 index 0000000000..72683d159d --- /dev/null +++ b/dizoo/minigrid/envs/minigrid_wrapper.py @@ -0,0 +1,196 @@ +import gymnasium as gym +from gymnasium import spaces +from gymnasium.core import ObservationWrapper, Wrapper +import numpy as np +import operator +from functools import reduce +from minigrid.core.constants import COLOR_TO_IDX, OBJECT_TO_IDX, STATE_TO_IDX + + +class ViewSizeWrapper(ObservationWrapper): + """ + Wrapper to customize the agent field of view size. + This cannot be used with fully observable wrappers. + """ + + def __init__(self, env, agent_view_size=7): + super().__init__(env) + + assert agent_view_size % 2 == 1 + assert agent_view_size >= 3 + + self.agent_view_size = agent_view_size + + # Compute observation space with specified view size + new_image_space = gym.spaces.Box(low=0, high=255, shape=(agent_view_size, agent_view_size, 3), dtype="uint8") + + # Override the environment's observation spaceexit + self.observation_space = spaces.Dict({**self.observation_space.spaces, "image": new_image_space}) + + def observation(self, obs): + env = self.unwrapped + grid, vis_mask = env.gen_obs_grid(self.agent_view_size) + + # Encode the partially observable view into a numpy array + # print('grid:' + grid) + # print('vis_mask:' + vis_mask) + image = grid.encode(vis_mask) + return {**obs, "image": image} + + +class MoveBonus(Wrapper): + """ + Adds an movement bonus based on which positions + are visited on the grid. + + Example: + >>> import gymnasium as gym + >>> from minigrid.wrappers import PositionBonus + >>> env = gym.make("MiniGrid-Empty-5x5-v0") + >>> _, _ = env.reset(seed=0) + >>> _, reward, _, _, _ = env.step(1) + >>> print(reward) + 0 + >>> _, reward, _, _, _ = env.step(1) + >>> print(reward) + 0 + >>> env_bonus = MoveBonus(env) + >>> obs, _ = env_bonus.reset(seed=0) + >>> obs, reward, terminated, truncated, info = env_bonus.step(1) + >>> print(reward) + 1.0 + >>> obs, reward, terminated, truncated, info = env_bonus.step(1) + >>> print(reward) + 0.7071067811865475 + """ + + def __init__(self, env): + """A wrapper that adds an exploration bonus to less visited positions. + + Args: + env: The environment to apply the wrapper + """ + super().__init__(env) + self.goal_pos = (self.width - 2, self.height - 2) + self.scale = np.sqrt(self.width ** 2 + self.height ** 2) + + def step(self, action): + """Steps through the environment with `action`.""" + + cur_dis = np.sqrt( + (self.goal_pos[0] - self.env.agent_pos[0]) ** 2 + (self.goal_pos[1] - self.env.agent_pos[1]) ** 2 + ) + obs, reward, terminated, truncated, info = self.env.step(action) + tmp_dis = np.sqrt( + (self.goal_pos[0] - self.env.agent_pos[0]) ** 2 + (self.goal_pos[1] - self.env.agent_pos[1]) ** 2 + ) + + move_bonus = (cur_dis - tmp_dis) / self.scale + reward += move_bonus + + return obs, reward, terminated, truncated, info + + +class OneHotObsWrapper(ObservationWrapper): + """ + Wrapper to get a one-hot encoding of a partially observable + agent view as observation. + + Example: + >>> import gymnasium as gym + >>> from minigrid.wrappers import OneHotPartialObsWrapper + >>> env = gym.make("MiniGrid-Empty-5x5-v0") + >>> obs, _ = env.reset() + >>> obs["image"][0, :, :] + array([[2, 5, 0], + [2, 5, 0], + [2, 5, 0], + [2, 5, 0], + [2, 5, 0], + [2, 5, 0], + [2, 5, 0]], dtype=uint8) + >>> env = OneHotPartialObsWrapper(env) + >>> obs, _ = env.reset() + >>> obs["image"][0, :, :] + array([[0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0], + [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0], + [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0], + [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0], + [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0], + [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0], + [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0]], + dtype=uint8) + """ + + def __init__(self, env): + """A wrapper that makes the image observation a one-hot encoding of a partially observable agent view. + + Args: + env: The environment to apply the wrapper + """ + super().__init__(env) + + obs_shape = env.observation_space["image"].shape + + # Number of bits per cell + num_bits = len(OBJECT_TO_IDX) + len(COLOR_TO_IDX) + len(STATE_TO_IDX) + 1 + + new_image_space = spaces.Box(low=0, high=1, shape=(obs_shape[0], obs_shape[1], num_bits), dtype="float32") + self.observation_space = spaces.Dict({**self.observation_space.spaces, "image": new_image_space}) + + def observation(self, obs): + img = obs["image"] + out = np.zeros(self.observation_space.spaces["image"].shape, dtype="float32") + + for i in range(img.shape[0]): + for j in range(img.shape[1]): + type = img[i, j, 0] + color = img[i, j, 1] + state = img[i, j, 2] + + out[i, j, type] = 1 + out[i, j, len(OBJECT_TO_IDX) + color] = 1 + out[i, j, len(OBJECT_TO_IDX) + len(COLOR_TO_IDX) + state] = 1 + + return {**obs, "image": out} + + +class FlatObsWrapper(ObservationWrapper): + """ + Encode mission strings using a one-hot scheme, + and combine these with observed images into one flat array. + + This wrapper is not applicable to BabyAI environments, given that these have their own language component. + + Example: + >>> import gymnasium as gym + >>> import matplotlib.pyplot as plt + >>> from minigrid.wrappers import FlatObsWrapper + >>> env = gym.make("MiniGrid-LavaCrossingS11N5-v0") + >>> env_obs = FlatObsWrapper(env) + >>> obs, _ = env_obs.reset() + >>> obs.shape + (2835,) + """ + + def __init__(self, env): + super().__init__(env) + + imgSpace = env.observation_space.spaces["image"] + imgSize = reduce(operator.mul, imgSpace.shape, 1) + + self.observation_space = spaces.Box( + low=0, + high=255, + shape=(imgSize, ), + dtype="float32", + ) + + self.cachedStr: str = None + + def observation(self, obs): + img = obs["image"] + + img = img.flatten() + + return img diff --git a/dizoo/minigrid/envs/noisy_tv.py b/dizoo/minigrid/envs/noisy_tv.py new file mode 100644 index 0000000000..9c5c78e613 --- /dev/null +++ b/dizoo/minigrid/envs/noisy_tv.py @@ -0,0 +1,216 @@ +from minigrid.core.grid import Grid +from minigrid.core.mission import MissionSpace +from minigrid.minigrid_env import * +from minigrid.utils.rendering import * +from minigrid.core.world_object import WorldObj +import random + + +class NoisyTVEnv(MiniGridEnv): + """ + ### Description + + Classic four room reinforcement learning environment with random noise. The agent must + navigate in a maze composed of four rooms interconnected by 4 gaps in the + walls. To obtain a reward, the agent must reach the green goal square. Both + the agent and the goal square are randomly placed in any of the four rooms. + + ### Mission Space + + "reach the goal" + + ### Action Space + + | Num | Name | Action | + |-----|--------------|--------------| + | 0 | left | Turn left | + | 1 | right | Turn right | + | 2 | forward | Move forward | + | 3 | pickup | Unused | + | 4 | drop | Unused | + | 5 | toggle | Unused | + | 6 | done | Unused | + + ### Observation Encoding + + - Each tile is encoded as a 3 dimensional tuple: + `(OBJECT_IDX, COLOR_IDX, STATE)` + - `OBJECT_TO_IDX` and `COLOR_TO_IDX` mapping can be found in + [minigrid/minigrid.py](minigrid/minigrid.py) + - `STATE` refers to the door state with 0=open, 1=closed and 2=locked + + ### Rewards + + A reward of '1' is given for success, and '0' for failure. + Noisy reward are given upon reaching a noisy tile. Noise obeys Gaussian distribution. + + ### Termination + + The episode ends if any one of the following conditions is met: + + 1. The agent reaches the goal. + 2. Timeout (see `max_steps`). + + ### Registered Configurations + + - `MiniGrid-NoisyTV-v0` + + """ + + def __init__(self, agent_pos=None, goal_pos=None, noisy_tile_num=4, **kwargs): + self._agent_default_pos = agent_pos + self._goal_default_pos = goal_pos + self.size = 19 + self._noisy_tile_num = noisy_tile_num + self._noisy_tile_pos = [] + for i in range(self._noisy_tile_num): + pos2 = (self._rand_int(1, self.size - 1), self._rand_int(1, self.size - 1)) + while pos2 in self._noisy_tile_pos: + pos2 = (self._rand_int(1, self.size - 1), self._rand_int(1, self.size - 1)) + self._noisy_tile_pos.append(pos2) + mission_space = MissionSpace(mission_func=lambda: "reach the goal") + + super().__init__(mission_space=mission_space, width=self.size, height=self.size, max_steps=100, **kwargs) + + def _reward_noise(self): + """ + Compute the reward to be given upon reach a noisy tile + """ + return self._rand_float(0.05, 0.1) + + def _add_noise(self, obs): + """ + Add noise to obs['image'] + """ + image = obs['image'].astype(float) + for pos in self._noisy_tile_pos: + if self.in_view(pos[0], pos[1]): # if noisy tile is in the view of agent, the view of agent is 7*7. + relative_pos = self.relative_coords(pos[0], pos[1]) + image[relative_pos][1] += 0.5 + obs['image'] = image + return obs + + def _gen_grid(self, width, height): + # Create the grid + self.grid = Grid(width, height) + + # Generate the surrounding walls + self.grid.horz_wall(0, 0) + self.grid.horz_wall(0, height - 1) + self.grid.vert_wall(0, 0) + self.grid.vert_wall(width - 1, 0) + + room_w = width // 2 + room_h = height // 2 + + # For each row of rooms + for j in range(0, 2): + + # For each column + for i in range(0, 2): + xL = i * room_w + yT = j * room_h + xR = xL + room_w + yB = yT + room_h + + # Bottom wall and door + if i + 1 < 2: + self.grid.vert_wall(xR, yT, room_h) + pos = (xR, self._rand_int(yT + 1, yB)) + self.grid.set(*pos, None) + + # Bottom wall and door + if j + 1 < 2: + self.grid.horz_wall(xL, yB, room_w) + pos = (self._rand_int(xL + 1, xR), yB) + self.grid.set(*pos, None) + + # Randomize the player start position and orientation + if self._agent_default_pos is not None: + self.agent_pos = self._agent_default_pos + self.grid.set(*self._agent_default_pos, None) + # assuming random start direction + self.agent_dir = self._rand_int(0, 4) + else: + self.place_agent() + + if self._goal_default_pos is not None: + goal = Goal() + self.put_obj(goal, *self._goal_default_pos) + goal.init_pos, goal.cur_pos = self._goal_default_pos + else: + self.place_obj(Goal()) + + def step(self, action): + self.step_count += 1 + + reward = 0 + terminated = False + truncated = False + + # Get the position in front of the agent + fwd_pos = self.front_pos + + # Get the contents of the cell in front of the agent + fwd_cell = self.grid.get(*fwd_pos) + + # Rotate left + if action == self.actions.left: + self.agent_dir -= 1 + if self.agent_dir < 0: + self.agent_dir += 4 + + # Rotate right + elif action == self.actions.right: + self.agent_dir = (self.agent_dir + 1) % 4 + + # Move forward + elif action == self.actions.forward: + if fwd_cell is None or fwd_cell.can_overlap(): + self.agent_pos = tuple(fwd_pos) + if fwd_cell is not None and fwd_cell.type == "goal": + terminated = True + reward = self._reward() + if fwd_cell is not None and fwd_cell.type == "lava": + terminated = True + # if agent reach noisy tile, return noisy reward. + if self.agent_pos in self._noisy_tile_pos: + reward = self._reward_noise() + + # Pick up an object + elif action == self.actions.pickup: + if fwd_cell and fwd_cell.can_pickup(): + if self.carrying is None: + self.carrying = fwd_cell + self.carrying.cur_pos = np.array([-1, -1]) + self.grid.set(fwd_pos[0], fwd_pos[1], None) + + # Drop an object + elif action == self.actions.drop: + if not fwd_cell and self.carrying: + self.grid.set(fwd_pos[0], fwd_pos[1], self.carrying) + self.carrying.cur_pos = fwd_pos + self.carrying = None + + # Toggle/activate an object + elif action == self.actions.toggle: + if fwd_cell: + fwd_cell.toggle(self, fwd_pos) + + # Done action (not used by default) + elif action == self.actions.done: + pass + + else: + raise ValueError(f"Unknown action: {action}") + + if self.step_count >= self.max_steps: + truncated = True + + if self.render_mode == "human": + self.render() + + obs = self.gen_obs() + obs = self._add_noise(obs) + + return obs, reward, terminated, truncated, {} diff --git a/dizoo/minigrid/utils/eval.py b/dizoo/minigrid/utils/eval.py index 8f61703426..e3c6acb9fb 100644 --- a/dizoo/minigrid/utils/eval.py +++ b/dizoo/minigrid/utils/eval.py @@ -44,7 +44,7 @@ def eval( env.enable_save_replay(replay_path=replay_path) obs = env.reset() obs = {0: obs} - eval_reward = 0. + episode_return = 0. beta_index = {i: 0 for i in range(1)} beta_index = to_tensor(beta_index, dtype=torch.int64) @@ -73,7 +73,7 @@ def eval( timestep = timesteps[0] # print(timestep.info) - eval_reward += timestep.reward + episode_return += timestep.reward obs = timestep.obs obs = {0: obs} @@ -81,7 +81,7 @@ def eval( if timestep.done: print(timestep.info) break - print('Eval is over! The performance of your RL policy is {}'.format(eval_reward)) + print('Eval is over! The performance of your RL policy is {}'.format(episode_return)) if __name__ == "__main__": diff --git a/dizoo/mujoco/config/ant_ddpg_config.py b/dizoo/mujoco/config/ant_ddpg_config.py index 7973922571..c8698ebcc3 100644 --- a/dizoo/mujoco/config/ant_ddpg_config.py +++ b/dizoo/mujoco/config/ant_ddpg_config.py @@ -9,7 +9,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=6000, manager=dict(shared_memory=False, ), diff --git a/dizoo/mujoco/config/ant_gail_sac_config.py b/dizoo/mujoco/config/ant_gail_sac_config.py index c45186e47a..b7e7cd7d06 100644 --- a/dizoo/mujoco/config/ant_gail_sac_config.py +++ b/dizoo/mujoco/config/ant_gail_sac_config.py @@ -10,7 +10,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=6000, ), diff --git a/dizoo/mujoco/config/ant_onppo_config.py b/dizoo/mujoco/config/ant_onppo_config.py index b6eea20a65..036d651391 100644 --- a/dizoo/mujoco/config/ant_onppo_config.py +++ b/dizoo/mujoco/config/ant_onppo_config.py @@ -1,32 +1,38 @@ from easydict import EasyDict +import torch.nn as nn ant_ppo_config = dict( exp_name="ant_onppo_seed0", env=dict( env_id='Ant-v3', - norm_obs=dict(use_norm=False, ), - norm_reward=dict(use_norm=False, ), collector_env_num=10, evaluator_env_num=10, - use_act_scale=True, n_evaluator_episode=10, stop_value=6000, - manager=dict(shared_memory=False, ) ), policy=dict( cuda=True, recompute_adv=True, action_space='continuous', model=dict( + encoder_hidden_size_list=[128, 128], action_space='continuous', obs_shape=111, action_shape=8, + share_encoder=False, + actor_head_layer_num=0, + critic_head_layer_num=2, + critic_head_hidden_size=256, + actor_head_hidden_size=128, + activation=nn.Tanh(), + bound_type='tanh', ), learn=dict( epoch_per_collect=10, update_per_collect=1, - batch_size=320, + batch_size=128, learning_rate=3e-4, + lr_scheduler=dict(epoch_num=1500, min_lr_lambda=0), value_weight=0.5, entropy_weight=0.001, clip_ratio=0.2, @@ -40,7 +46,7 @@ grad_clip_value=0.5, ), collect=dict( - n_sample=3200, + n_sample=2048, unroll_len=1, discount_factor=0.99, gae_lambda=0.95, diff --git a/dizoo/mujoco/config/ant_ppo_config.py b/dizoo/mujoco/config/ant_ppo_config.py index c0e5203e49..25103592bb 100644 --- a/dizoo/mujoco/config/ant_ppo_config.py +++ b/dizoo/mujoco/config/ant_ppo_config.py @@ -9,7 +9,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=8, evaluator_env_num=10, - use_act_scale=True, n_evaluator_episode=10, stop_value=6000, ), diff --git a/dizoo/mujoco/config/ant_sac_config.py b/dizoo/mujoco/config/ant_sac_config.py index 88fa9e2b00..fd2881e91f 100644 --- a/dizoo/mujoco/config/ant_sac_config.py +++ b/dizoo/mujoco/config/ant_sac_config.py @@ -8,7 +8,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=6000, manager=dict(shared_memory=False, reset_inplace=True), diff --git a/dizoo/mujoco/config/ant_td3_config.py b/dizoo/mujoco/config/ant_td3_config.py index 6f3b48b321..ebcb3654ce 100644 --- a/dizoo/mujoco/config/ant_td3_config.py +++ b/dizoo/mujoco/config/ant_td3_config.py @@ -8,7 +8,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=6000, manager=dict(shared_memory=False, reset_inplace=True), diff --git a/dizoo/mujoco/config/ant_trex_onppo_config.py b/dizoo/mujoco/config/ant_trex_onppo_config.py index 99d508fed9..f3d5e96b75 100644 --- a/dizoo/mujoco/config/ant_trex_onppo_config.py +++ b/dizoo/mujoco/config/ant_trex_onppo_config.py @@ -9,7 +9,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=8, evaluator_env_num=10, - use_act_scale=True, n_evaluator_episode=10, stop_value=6000, ), diff --git a/dizoo/mujoco/config/ant_trex_sac_config.py b/dizoo/mujoco/config/ant_trex_sac_config.py index 2a1f01091f..6c0ef73097 100644 --- a/dizoo/mujoco/config/ant_trex_sac_config.py +++ b/dizoo/mujoco/config/ant_trex_sac_config.py @@ -9,7 +9,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=6000, ), diff --git a/dizoo/mujoco/config/halfcheetah_bco_config.py b/dizoo/mujoco/config/halfcheetah_bco_config.py index 5bdab4dec9..9253af06eb 100644 --- a/dizoo/mujoco/config/halfcheetah_bco_config.py +++ b/dizoo/mujoco/config/halfcheetah_bco_config.py @@ -8,7 +8,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=12000, ), diff --git a/dizoo/mujoco/config/halfcheetah_bdq_config.py b/dizoo/mujoco/config/halfcheetah_bdq_config.py new file mode 100644 index 0000000000..25fb65ba35 --- /dev/null +++ b/dizoo/mujoco/config/halfcheetah_bdq_config.py @@ -0,0 +1,71 @@ +from easydict import EasyDict + +halfcheetah_bdq_config = dict( + exp_name='halfcheetah_bdq_seed0', + env=dict( + env_id='HalfCheetah-v3', + norm_reward=dict(use_norm=False, ), + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=12000, + action_bins_per_branch=2, + ), + policy=dict( + cuda=False, + priority=False, + discount_factor=0.99, + nstep=1, + model=dict( + obs_shape=17, + num_branches=6, + action_bins_per_branch=2, # mean the action shape is 6, 2 discrete actions for each action dimension + encoder_hidden_size_list=[256, 256, 128], + ), + learn=dict( + batch_size=512, + learning_rate=3e-4, + ignore_done=True, + target_update_freq=500, + update_per_collect=20, + ), + collect=dict( + n_sample=256, + unroll_len=1, + ), + eval=dict(evaluator=dict(eval_freq=1000, )), + other=dict( + # Epsilon greedy with decay. + eps=dict( + # Decay type. Support ['exp', 'linear']. + type='exp', + start=1, + end=0.05, + decay=int(1e5), + ), + replay_buffer=dict(replay_buffer_size=int(1e6), ) + ), + ), +) +halfcheetah_bdq_config = EasyDict(halfcheetah_bdq_config) +main_config = halfcheetah_bdq_config + +halfcheetah_bdq_create_config = dict( + env=dict( + type='mujoco', + import_names=['dizoo.mujoco.envs.mujoco_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='bdq', ), +) +halfcheetah_bdq_create_config = EasyDict(halfcheetah_bdq_create_config) +create_config = halfcheetah_bdq_create_config + +if __name__ == "__main__": + # or you can enter `ding -m serial_onpolicy -c halfcheetah_onbdq_config.py -s 0` + from ding.entry import serial_pipeline + serial_pipeline( + (main_config, create_config), + seed=0, + max_env_step=10000000, + ) diff --git a/dizoo/mujoco/config/halfcheetah_d4pg_config.py b/dizoo/mujoco/config/halfcheetah_d4pg_config.py index 8cc428dbec..154bc27a46 100644 --- a/dizoo/mujoco/config/halfcheetah_d4pg_config.py +++ b/dizoo/mujoco/config/halfcheetah_d4pg_config.py @@ -8,7 +8,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=4, evaluator_env_num=4, - use_act_scale=True, n_evaluator_episode=8, stop_value=20000, ), diff --git a/dizoo/mujoco/config/halfcheetah_ddpg_config.py b/dizoo/mujoco/config/halfcheetah_ddpg_config.py index c162003e8e..c9031d36f4 100644 --- a/dizoo/mujoco/config/halfcheetah_ddpg_config.py +++ b/dizoo/mujoco/config/halfcheetah_ddpg_config.py @@ -8,7 +8,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=11000, ), @@ -63,4 +62,4 @@ if __name__ == "__main__": # or you can enter `ding -m serial -c halfcheetah_ddpg_config.py -s 0` from ding.entry import serial_pipeline - serial_pipeline((main_config, create_config), seed=0) \ No newline at end of file + serial_pipeline((main_config, create_config), seed=0) diff --git a/dizoo/mujoco/config/halfcheetah_gail_sac_config.py b/dizoo/mujoco/config/halfcheetah_gail_sac_config.py index aacb168f6f..bf64cd8c64 100644 --- a/dizoo/mujoco/config/halfcheetah_gail_sac_config.py +++ b/dizoo/mujoco/config/halfcheetah_gail_sac_config.py @@ -10,7 +10,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=12000, ), diff --git a/dizoo/mujoco/config/halfcheetah_gcl_sac_config.py b/dizoo/mujoco/config/halfcheetah_gcl_sac_config.py index db40e16151..367b7bcf03 100644 --- a/dizoo/mujoco/config/halfcheetah_gcl_sac_config.py +++ b/dizoo/mujoco/config/halfcheetah_gcl_sac_config.py @@ -8,7 +8,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=12000, ), diff --git a/dizoo/mujoco/config/halfcheetah_onppo_config.py b/dizoo/mujoco/config/halfcheetah_onppo_config.py index 4af1ad1ac7..b7646bbede 100644 --- a/dizoo/mujoco/config/halfcheetah_onppo_config.py +++ b/dizoo/mujoco/config/halfcheetah_onppo_config.py @@ -1,17 +1,15 @@ from easydict import EasyDict +import torch.nn as nn -collector_env_num = 1 -evaluator_env_num = 1 +collector_env_num = 8 +evaluator_env_num = 8 halfcheetah_ppo_config = dict( exp_name='halfcheetah_onppo_seed0', env=dict( env_id='HalfCheetah-v3', - norm_obs=dict(use_norm=False, ), - norm_reward=dict(use_norm=False, ), collector_env_num=collector_env_num, evaluator_env_num=evaluator_env_num, - use_act_scale=True, - n_evaluator_episode=1, + n_evaluator_episode=8, stop_value=12000, ), policy=dict( @@ -19,15 +17,24 @@ recompute_adv=True, action_space='continuous', model=dict( + encoder_hidden_size_list=[128, 128], action_space='continuous', + share_encoder=False, + actor_head_layer_num=0, + critic_head_layer_num=2, + critic_head_hidden_size=256, + actor_head_hidden_size=128, obs_shape=17, action_shape=6, + activation=nn.Tanh(), + bound_type='tanh', ), learn=dict( epoch_per_collect=10, update_per_collect=1, - batch_size=320, + batch_size=128, learning_rate=3e-4, + lr_scheduler=dict(epoch_num=1500, min_lr_lambda=0), value_weight=0.5, entropy_weight=0.001, clip_ratio=0.2, @@ -43,12 +50,12 @@ ), collect=dict( collector_env_num=collector_env_num, - n_sample=3200, + n_sample=2048, unroll_len=1, discount_factor=0.99, gae_lambda=0.95, ), - eval=dict(evaluator=dict(eval_freq=500, )), + eval=dict(evaluator=dict(eval_freq=5000, )), ), ) halfcheetah_ppo_config = EasyDict(halfcheetah_ppo_config) @@ -69,4 +76,4 @@ if __name__ == "__main__": # or you can enter `ding -m serial_onpolicy -c halfcheetah_onppo_config.py -s 0` from ding.entry import serial_pipeline_onpolicy - serial_pipeline_onpolicy((main_config, create_config), seed=0) \ No newline at end of file + serial_pipeline_onpolicy((main_config, create_config), seed=0) diff --git a/dizoo/mujoco/config/halfcheetah_sac_config.py b/dizoo/mujoco/config/halfcheetah_sac_config.py index b2bcd7f3ac..67ace8134b 100644 --- a/dizoo/mujoco/config/halfcheetah_sac_config.py +++ b/dizoo/mujoco/config/halfcheetah_sac_config.py @@ -8,7 +8,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=12000, ), diff --git a/dizoo/mujoco/config/halfcheetah_sqil_sac_config.py b/dizoo/mujoco/config/halfcheetah_sqil_sac_config.py index d01ede1891..ea6cb51b53 100644 --- a/dizoo/mujoco/config/halfcheetah_sqil_sac_config.py +++ b/dizoo/mujoco/config/halfcheetah_sqil_sac_config.py @@ -8,7 +8,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=12000, ), diff --git a/dizoo/mujoco/config/halfcheetah_td3_config.py b/dizoo/mujoco/config/halfcheetah_td3_config.py index 1c726a573b..47eb4ce5f1 100644 --- a/dizoo/mujoco/config/halfcheetah_td3_config.py +++ b/dizoo/mujoco/config/halfcheetah_td3_config.py @@ -1,13 +1,13 @@ from easydict import EasyDict halfcheetah_td3_config = dict( + exp_name='halfcheetah_td3_seed0', env=dict( - env_id='halfcheetah_td3_seed0', + env_id='HalfCheetah-v3', norm_obs=dict(use_norm=False, ), norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=11000, ), @@ -68,4 +68,4 @@ if __name__ == "__main__": # or you can enter `ding -m serial -c halfcheetah_td3_config.py -s 0 --env-step 1e7` from ding.entry import serial_pipeline - serial_pipeline((main_config, create_config), seed=0) \ No newline at end of file + serial_pipeline((main_config, create_config), seed=0) diff --git a/dizoo/mujoco/config/halfcheetah_trex_onppo_config.py b/dizoo/mujoco/config/halfcheetah_trex_onppo_config.py index 4bad899238..6d635c212d 100644 --- a/dizoo/mujoco/config/halfcheetah_trex_onppo_config.py +++ b/dizoo/mujoco/config/halfcheetah_trex_onppo_config.py @@ -8,7 +8,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=8, evaluator_env_num=10, - use_act_scale=True, n_evaluator_episode=10, stop_value=3000, ), diff --git a/dizoo/mujoco/config/halfcheetah_trex_sac_config.py b/dizoo/mujoco/config/halfcheetah_trex_sac_config.py index ec6e781c67..5f123682a0 100644 --- a/dizoo/mujoco/config/halfcheetah_trex_sac_config.py +++ b/dizoo/mujoco/config/halfcheetah_trex_sac_config.py @@ -8,7 +8,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=12000, ), diff --git a/dizoo/mujoco/config/hopper_bco_config.py b/dizoo/mujoco/config/hopper_bco_config.py index 4abf5fc78c..668e258e69 100644 --- a/dizoo/mujoco/config/hopper_bco_config.py +++ b/dizoo/mujoco/config/hopper_bco_config.py @@ -8,7 +8,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=6000, ), diff --git a/dizoo/mujoco/config/hopper_bdq_config.py b/dizoo/mujoco/config/hopper_bdq_config.py new file mode 100644 index 0000000000..34dbe21664 --- /dev/null +++ b/dizoo/mujoco/config/hopper_bdq_config.py @@ -0,0 +1,75 @@ +from easydict import EasyDict + +hopper_bdq_config = dict( + exp_name='hopper_bdq_seed0', + env=dict( + env_id='Hopper-v3', + norm_reward=dict(use_norm=False, ), + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=int(1e6), + action_bins_per_branch=4, + ), + policy=dict( + cuda=False, + priority=False, + discount_factor=0.99, + nstep=3, + model=dict( + obs_shape=11, + num_branches=3, + action_bins_per_branch=4, # mean the action shape is 3, 4 discrete actions for each action dimension + encoder_hidden_size_list=[256, 256, 128], + ), + learn=dict( + ignore_done=False, + batch_size=512, + learning_rate=3e-4, + # Frequency of target network update. + target_update_freq=500, + update_per_collect=20, + ), + collect=dict( + # You can use either "n_sample" or "n_episode" in collector.collect. + # Get "n_sample" samples per collect. + n_sample=256, + # Cut trajectories into pieces with length "unroll_len". + unroll_len=1, + ), + eval=dict(evaluator=dict(eval_freq=1000, )), + other=dict( + # Epsilon greedy with decay. + eps=dict( + # Decay type. Support ['exp', 'linear']. + type='exp', + start=1, + end=0.05, + decay=int(1e5), + ), + replay_buffer=dict(replay_buffer_size=int(1e6), ) + ), + ), +) +hopper_bdq_config = EasyDict(hopper_bdq_config) +main_config = hopper_bdq_config + +hopper_bdq_create_config = dict( + env=dict( + type='mujoco', + import_names=['dizoo.mujoco.envs.mujoco_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='bdq', ), +) +hopper_bdq_create_config = EasyDict(hopper_bdq_create_config) +create_config = hopper_bdq_create_config + +if __name__ == "__main__": + # or you can enter `ding -m serial_onpolicy -c hopper_bdq_config.py -s 0` + from ding.entry import serial_pipeline + serial_pipeline( + [main_config, create_config], + seed=0, + max_env_step=10000000, + ) diff --git a/dizoo/mujoco/config/hopper_cql_config.py b/dizoo/mujoco/config/hopper_cql_config.py index 4d86e2a57b..7713d23381 100644 --- a/dizoo/mujoco/config/hopper_cql_config.py +++ b/dizoo/mujoco/config/hopper_cql_config.py @@ -8,7 +8,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=6000, ), diff --git a/dizoo/mujoco/config/hopper_d4pg_config.py b/dizoo/mujoco/config/hopper_d4pg_config.py index eadbe4cbcd..e533ac684e 100644 --- a/dizoo/mujoco/config/hopper_d4pg_config.py +++ b/dizoo/mujoco/config/hopper_d4pg_config.py @@ -8,7 +8,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=4, evaluator_env_num=4, - use_act_scale=True, n_evaluator_episode=8, stop_value=5000, ), diff --git a/dizoo/mujoco/config/hopper_ddpg_config.py b/dizoo/mujoco/config/hopper_ddpg_config.py index 7c909ea2a0..c0a1a524e8 100644 --- a/dizoo/mujoco/config/hopper_ddpg_config.py +++ b/dizoo/mujoco/config/hopper_ddpg_config.py @@ -8,7 +8,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=6000, ), diff --git a/dizoo/mujoco/config/hopper_gail_sac_config.py b/dizoo/mujoco/config/hopper_gail_sac_config.py index b46a6d5c1d..26ef8b3816 100644 --- a/dizoo/mujoco/config/hopper_gail_sac_config.py +++ b/dizoo/mujoco/config/hopper_gail_sac_config.py @@ -10,7 +10,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=6000, ), diff --git a/dizoo/mujoco/config/hopper_gcl_config.py b/dizoo/mujoco/config/hopper_gcl_config.py index 95996ee597..214f44dbf7 100644 --- a/dizoo/mujoco/config/hopper_gcl_config.py +++ b/dizoo/mujoco/config/hopper_gcl_config.py @@ -8,7 +8,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=4, evaluator_env_num=10, - use_act_scale=True, n_evaluator_episode=10, stop_value=3000, ), diff --git a/dizoo/mujoco/config/hopper_onppo_config.py b/dizoo/mujoco/config/hopper_onppo_config.py index b60da05525..6d1c57e70e 100644 --- a/dizoo/mujoco/config/hopper_onppo_config.py +++ b/dizoo/mujoco/config/hopper_onppo_config.py @@ -1,14 +1,12 @@ from easydict import EasyDict +import torch.nn as nn hopper_onppo_config = dict( exp_name='hopper_onppo_seed0', env=dict( env_id='Hopper-v3', - norm_obs=dict(use_norm=False, ), - norm_reward=dict(use_norm=False, ), collector_env_num=8, evaluator_env_num=10, - use_act_scale=True, n_evaluator_episode=10, stop_value=4000, ), @@ -17,15 +15,24 @@ recompute_adv=True, action_space='continuous', model=dict( + encoder_hidden_size_list=[128, 128], obs_shape=11, action_shape=3, action_space='continuous', + share_encoder=False, + actor_head_layer_num=0, + critic_head_layer_num=2, + critic_head_hidden_size=256, + actor_head_hidden_size=128, + activation=nn.Tanh(), + bound_type='tanh', ), learn=dict( epoch_per_collect=10, update_per_collect=1, - batch_size=320, + batch_size=128, learning_rate=3e-4, + lr_scheduler=dict(epoch_num=1500, min_lr_lambda=0), value_weight=0.5, entropy_weight=0.001, clip_ratio=0.2, @@ -40,12 +47,12 @@ grad_clip_value=0.5, ), collect=dict( - n_sample=3200, + n_sample=2048, unroll_len=1, discount_factor=0.99, gae_lambda=0.95, ), - eval=dict(evaluator=dict(eval_freq=500, )), + eval=dict(evaluator=dict(eval_freq=5000, )), ), ) hopper_onppo_config = EasyDict(hopper_onppo_config) diff --git a/dizoo/mujoco/config/hopper_sac_config.py b/dizoo/mujoco/config/hopper_sac_config.py index 42769fb974..9835aff0f4 100644 --- a/dizoo/mujoco/config/hopper_sac_config.py +++ b/dizoo/mujoco/config/hopper_sac_config.py @@ -8,7 +8,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=6000, ), diff --git a/dizoo/mujoco/config/hopper_sac_data_generation_config.py b/dizoo/mujoco/config/hopper_sac_data_generation_config.py index 831459983d..9b32dd50c5 100644 --- a/dizoo/mujoco/config/hopper_sac_data_generation_config.py +++ b/dizoo/mujoco/config/hopper_sac_data_generation_config.py @@ -8,7 +8,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=10, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=6000, ), diff --git a/dizoo/mujoco/config/hopper_sqil_sac_config.py b/dizoo/mujoco/config/hopper_sqil_sac_config.py index 19a0920feb..172cb44ea0 100644 --- a/dizoo/mujoco/config/hopper_sqil_sac_config.py +++ b/dizoo/mujoco/config/hopper_sqil_sac_config.py @@ -10,7 +10,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=6000, ), diff --git a/dizoo/mujoco/config/hopper_td3_bc_config.py b/dizoo/mujoco/config/hopper_td3_bc_config.py index 19bd9f4aae..00a04075a8 100644 --- a/dizoo/mujoco/config/hopper_td3_bc_config.py +++ b/dizoo/mujoco/config/hopper_td3_bc_config.py @@ -5,13 +5,12 @@ env=dict( env_id='Hopper-v3', norm_obs=dict( - use_norm=True, + use_norm=True, offline_stats=dict(use_offline_stats=True, ), ), norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=6000, ), diff --git a/dizoo/mujoco/config/hopper_td3_config.py b/dizoo/mujoco/config/hopper_td3_config.py index dd4ed01f1a..5f72930ea0 100644 --- a/dizoo/mujoco/config/hopper_td3_config.py +++ b/dizoo/mujoco/config/hopper_td3_config.py @@ -8,7 +8,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=6000, ), diff --git a/dizoo/mujoco/config/hopper_td3_data_generation_config.py b/dizoo/mujoco/config/hopper_td3_data_generation_config.py index 847b464adf..97330419d4 100644 --- a/dizoo/mujoco/config/hopper_td3_data_generation_config.py +++ b/dizoo/mujoco/config/hopper_td3_data_generation_config.py @@ -8,7 +8,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=11000, ), diff --git a/dizoo/mujoco/config/hopper_trex_onppo_config.py b/dizoo/mujoco/config/hopper_trex_onppo_config.py index d02a149b3a..e69451fe3c 100644 --- a/dizoo/mujoco/config/hopper_trex_onppo_config.py +++ b/dizoo/mujoco/config/hopper_trex_onppo_config.py @@ -8,7 +8,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=8, evaluator_env_num=10, - use_act_scale=True, n_evaluator_episode=10, stop_value=3000, ), diff --git a/dizoo/mujoco/config/hopper_trex_sac_config.py b/dizoo/mujoco/config/hopper_trex_sac_config.py index 8f05c2b385..5c4aa6f2c1 100644 --- a/dizoo/mujoco/config/hopper_trex_sac_config.py +++ b/dizoo/mujoco/config/hopper_trex_sac_config.py @@ -8,7 +8,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=6000, ), diff --git a/dizoo/mujoco/config/mbrl/halfcheetah_mbsac_mbpo_config.py b/dizoo/mujoco/config/mbrl/halfcheetah_mbsac_mbpo_config.py index 773896dc16..1ee4ac165b 100644 --- a/dizoo/mujoco/config/mbrl/halfcheetah_mbsac_mbpo_config.py +++ b/dizoo/mujoco/config/mbrl/halfcheetah_mbsac_mbpo_config.py @@ -18,7 +18,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=100000, ), diff --git a/dizoo/mujoco/config/mbrl/halfcheetah_sac_mbpo_config.py b/dizoo/mujoco/config/mbrl/halfcheetah_sac_mbpo_config.py index b1ac3d9079..7c22eb0aa1 100644 --- a/dizoo/mujoco/config/mbrl/halfcheetah_sac_mbpo_config.py +++ b/dizoo/mujoco/config/mbrl/halfcheetah_sac_mbpo_config.py @@ -18,7 +18,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=100000, ), diff --git a/dizoo/mujoco/config/mbrl/halfcheetah_stevesac_mbpo_config.py b/dizoo/mujoco/config/mbrl/halfcheetah_stevesac_mbpo_config.py index 6bfe32cb88..c692252cfc 100644 --- a/dizoo/mujoco/config/mbrl/halfcheetah_stevesac_mbpo_config.py +++ b/dizoo/mujoco/config/mbrl/halfcheetah_stevesac_mbpo_config.py @@ -18,7 +18,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=100000, ), diff --git a/dizoo/mujoco/config/mbrl/hopper_mbsac_mbpo_config.py b/dizoo/mujoco/config/mbrl/hopper_mbsac_mbpo_config.py index 7e3abf7942..7a37fe91e0 100644 --- a/dizoo/mujoco/config/mbrl/hopper_mbsac_mbpo_config.py +++ b/dizoo/mujoco/config/mbrl/hopper_mbsac_mbpo_config.py @@ -18,7 +18,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=100000, ), diff --git a/dizoo/mujoco/config/mbrl/hopper_sac_mbpo_config.py b/dizoo/mujoco/config/mbrl/hopper_sac_mbpo_config.py index a673b55b39..5f22ffed84 100644 --- a/dizoo/mujoco/config/mbrl/hopper_sac_mbpo_config.py +++ b/dizoo/mujoco/config/mbrl/hopper_sac_mbpo_config.py @@ -18,7 +18,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=100000, ), diff --git a/dizoo/mujoco/config/mbrl/hopper_stevesac_mbpo_config.py b/dizoo/mujoco/config/mbrl/hopper_stevesac_mbpo_config.py index 7514eed11e..d22a0e42f1 100644 --- a/dizoo/mujoco/config/mbrl/hopper_stevesac_mbpo_config.py +++ b/dizoo/mujoco/config/mbrl/hopper_stevesac_mbpo_config.py @@ -18,7 +18,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=100000, ), diff --git a/dizoo/mujoco/config/mbrl/walker2d_mbsac_mbpo_config.py b/dizoo/mujoco/config/mbrl/walker2d_mbsac_mbpo_config.py index ab51b8cd83..4e8a348c26 100644 --- a/dizoo/mujoco/config/mbrl/walker2d_mbsac_mbpo_config.py +++ b/dizoo/mujoco/config/mbrl/walker2d_mbsac_mbpo_config.py @@ -18,7 +18,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=4, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=100000, ), diff --git a/dizoo/mujoco/config/mbrl/walker2d_sac_mbpo_config.py b/dizoo/mujoco/config/mbrl/walker2d_sac_mbpo_config.py index badd5f398d..654451c266 100644 --- a/dizoo/mujoco/config/mbrl/walker2d_sac_mbpo_config.py +++ b/dizoo/mujoco/config/mbrl/walker2d_sac_mbpo_config.py @@ -18,7 +18,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=4, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=100000, ), diff --git a/dizoo/mujoco/config/mbrl/walker2d_stevesac_mbpo_config.py b/dizoo/mujoco/config/mbrl/walker2d_stevesac_mbpo_config.py index ae575586e3..0b7502478f 100644 --- a/dizoo/mujoco/config/mbrl/walker2d_stevesac_mbpo_config.py +++ b/dizoo/mujoco/config/mbrl/walker2d_stevesac_mbpo_config.py @@ -18,7 +18,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=4, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=100000, ), diff --git a/dizoo/mujoco/config/walker2d_d4pg_config.py b/dizoo/mujoco/config/walker2d_d4pg_config.py index 6a55726f03..31c6ff7d94 100644 --- a/dizoo/mujoco/config/walker2d_d4pg_config.py +++ b/dizoo/mujoco/config/walker2d_d4pg_config.py @@ -8,7 +8,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=4, evaluator_env_num=4, - use_act_scale=True, n_evaluator_episode=8, stop_value=7000, ), diff --git a/dizoo/mujoco/config/walker2d_ddpg_config.py b/dizoo/mujoco/config/walker2d_ddpg_config.py index 5b0c362abf..efe9bf9a39 100644 --- a/dizoo/mujoco/config/walker2d_ddpg_config.py +++ b/dizoo/mujoco/config/walker2d_ddpg_config.py @@ -8,7 +8,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=6000, ), diff --git a/dizoo/mujoco/config/walker2d_gail_ddpg_config.py b/dizoo/mujoco/config/walker2d_gail_ddpg_config.py index aee01432e5..779f65f63b 100644 --- a/dizoo/mujoco/config/walker2d_gail_ddpg_config.py +++ b/dizoo/mujoco/config/walker2d_gail_ddpg_config.py @@ -8,7 +8,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=6000, ), diff --git a/dizoo/mujoco/config/walker2d_gail_sac_config.py b/dizoo/mujoco/config/walker2d_gail_sac_config.py index 28fd16a0a9..7bd2de9022 100644 --- a/dizoo/mujoco/config/walker2d_gail_sac_config.py +++ b/dizoo/mujoco/config/walker2d_gail_sac_config.py @@ -10,7 +10,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=6000, ), diff --git a/dizoo/mujoco/config/walker2d_gcl_config.py b/dizoo/mujoco/config/walker2d_gcl_config.py index 84b1048a62..1b0b56fa32 100644 --- a/dizoo/mujoco/config/walker2d_gcl_config.py +++ b/dizoo/mujoco/config/walker2d_gcl_config.py @@ -8,7 +8,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=8, evaluator_env_num=10, - use_act_scale=True, n_evaluator_episode=10, stop_value=3000, ), diff --git a/dizoo/mujoco/config/walker2d_onppo_config.py b/dizoo/mujoco/config/walker2d_onppo_config.py index 2166e28998..58f5e9054f 100644 --- a/dizoo/mujoco/config/walker2d_onppo_config.py +++ b/dizoo/mujoco/config/walker2d_onppo_config.py @@ -1,17 +1,15 @@ from easydict import EasyDict +import torch.nn as nn -collector_env_num = 1 -evaluator_env_num = 1 +collector_env_num = 8 +evaluator_env_num = 8 walker2d_onppo_config = dict( exp_name='walker2d_onppo_seed0', env=dict( env_id='Walker2d-v3', - norm_obs=dict(use_norm=False, ), - norm_reward=dict(use_norm=False, ), collector_env_num=collector_env_num, evaluator_env_num=evaluator_env_num, - use_act_scale=True, - n_evaluator_episode=10, + n_evaluator_episode=8, stop_value=6000, ), policy=dict( @@ -19,15 +17,24 @@ recompute_adv=True, action_space='continuous', model=dict( + encoder_hidden_size_list=[128, 128], action_space='continuous', + share_encoder=False, + actor_head_layer_num=0, + critic_head_layer_num=2, + critic_head_hidden_size=256, + actor_head_hidden_size=128, obs_shape=17, action_shape=6, + activation=nn.Tanh(), + bound_type='tanh', ), learn=dict( epoch_per_collect=10, update_per_collect=1, - batch_size=320, + batch_size=128, learning_rate=3e-4, + lr_scheduler=dict(epoch_num=1500, min_lr_lambda=0), value_weight=0.5, entropy_weight=0.001, clip_ratio=0.2, @@ -44,12 +51,12 @@ ), collect=dict( collector_env_num=collector_env_num, - n_sample=3200, + n_sample=2048, unroll_len=1, discount_factor=0.99, gae_lambda=0.95, ), - eval=dict(evaluator=dict(eval_freq=500, )), + eval=dict(evaluator=dict(eval_freq=5000, )), ), ) walker2d_onppo_config = EasyDict(walker2d_onppo_config) diff --git a/dizoo/mujoco/config/walker2d_sac_config.py b/dizoo/mujoco/config/walker2d_sac_config.py index cb583ee3c2..f5a5a3127e 100644 --- a/dizoo/mujoco/config/walker2d_sac_config.py +++ b/dizoo/mujoco/config/walker2d_sac_config.py @@ -8,7 +8,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=6000, ), diff --git a/dizoo/mujoco/config/walker2d_sqil_sac_config.py b/dizoo/mujoco/config/walker2d_sqil_sac_config.py index 95185bf95e..59967f9f33 100644 --- a/dizoo/mujoco/config/walker2d_sqil_sac_config.py +++ b/dizoo/mujoco/config/walker2d_sqil_sac_config.py @@ -10,7 +10,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=6000, ), diff --git a/dizoo/mujoco/config/walker2d_td3_config.py b/dizoo/mujoco/config/walker2d_td3_config.py index 18a90cfd85..1c1bffcedf 100644 --- a/dizoo/mujoco/config/walker2d_td3_config.py +++ b/dizoo/mujoco/config/walker2d_td3_config.py @@ -8,7 +8,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=6000, ), diff --git a/dizoo/mujoco/config/walker2d_trex_onppo_config.py b/dizoo/mujoco/config/walker2d_trex_onppo_config.py index b293b993e0..c53c1efb4b 100644 --- a/dizoo/mujoco/config/walker2d_trex_onppo_config.py +++ b/dizoo/mujoco/config/walker2d_trex_onppo_config.py @@ -8,7 +8,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=8, evaluator_env_num=10, - use_act_scale=True, n_evaluator_episode=10, stop_value=3000, ), diff --git a/dizoo/mujoco/config/walker2d_trex_sac_config.py b/dizoo/mujoco/config/walker2d_trex_sac_config.py index aaf3c87a32..fdd1cab65e 100644 --- a/dizoo/mujoco/config/walker2d_trex_sac_config.py +++ b/dizoo/mujoco/config/walker2d_trex_sac_config.py @@ -8,7 +8,6 @@ norm_reward=dict(use_norm=False, ), collector_env_num=1, evaluator_env_num=8, - use_act_scale=True, n_evaluator_episode=8, stop_value=6000, ), diff --git a/dizoo/mujoco/entry/mujoco_d4pg_main.py b/dizoo/mujoco/entry/mujoco_d4pg_main.py index 7469d930c7..5bb4b72e35 100644 --- a/dizoo/mujoco/entry/mujoco_d4pg_main.py +++ b/dizoo/mujoco/entry/mujoco_d4pg_main.py @@ -9,7 +9,6 @@ from ding.policy import D4PGPolicy from ding.model.template.qac_dist import QACDIST from ding.utils import set_pkg_seed -from dizoo.classic_control.pendulum.envs import PendulumEnv from dizoo.mujoco.envs.mujoco_env import MujocoEnv from dizoo.classic_control.pendulum.config.pendulum_ppo_config import pendulum_ppo_config from dizoo.mujoco.config.hopper_d4pg_config import hopper_d4pg_config @@ -24,6 +23,7 @@ def main(cfg, seed=0, max_iterations=int(1e10)): SampleSerialCollector, InteractionSerialEvaluator, AdvancedReplayBuffer, + MujocoEnv, save_cfg=True ) collector_env_num, evaluator_env_num = cfg.env.collector_env_num, cfg.env.evaluator_env_num diff --git a/dizoo/mujoco/entry/mujoco_ddpg_eval.py b/dizoo/mujoco/entry/mujoco_ddpg_eval.py index cacd34b726..c3092e24e0 100644 --- a/dizoo/mujoco/entry/mujoco_ddpg_eval.py +++ b/dizoo/mujoco/entry/mujoco_ddpg_eval.py @@ -7,16 +7,17 @@ from ding.config import compile_config from ding.worker import BaseLearner, SampleSerialCollector, InteractionSerialEvaluator, AdvancedReplayBuffer -from ding.envs import BaseEnvManager, DingEnvWrapper +from ding.envs import BaseEnvManager from ding.envs import get_vec_env_setting, create_env_manager from ding.policy import DDPGPolicy -from ding.model import QAC +from ding.model import ContinuousQAC from ding.utils import set_pkg_seed from ding.rl_utils import get_epsilon_greedy_fn -from dizoo.mujoco.config.ant_ddpg_config import ant_ddpg_config, ant_ddpg_create_config +from dizoo.mujoco.envs.mujoco_env import MujocoEnv +from dizoo.mujoco.config.ant_ddpg_config import ant_ddpg_config -def main(rl_cfg, seed=0): - main_cfg, create_cfg =rl_cfg + +def main(main_cfg, seed=0): cfg = compile_config( main_cfg, BaseEnvManager, @@ -25,13 +26,10 @@ def main(rl_cfg, seed=0): SampleSerialCollector, InteractionSerialEvaluator, AdvancedReplayBuffer, - create_cfg=create_cfg, + MujocoEnv, save_cfg=True ) - create_cfg.policy.type = create_cfg.policy.type + '_command' - env_fn = None - cfg = compile_config(cfg, seed=seed, env=env_fn, auto=True, create_cfg=create_cfg, save_cfg=True) # Create main components: env, policy env_fn, collector_env_cfg, evaluator_env_cfg = get_vec_env_setting(cfg.env) evaluator_env = create_env_manager(cfg.env.manager, [partial(env_fn, cfg=c) for c in evaluator_env_cfg]) @@ -43,7 +41,7 @@ def main(rl_cfg, seed=0): set_pkg_seed(seed, use_cuda=cfg.policy.cuda) # Set up RL Policy - model = QAC(**cfg.policy.model) + model = ContinuousQAC(**cfg.policy.model) policy = DDPGPolicy(cfg.policy, model=model) policy.eval_mode.load_state_dict(torch.load(cfg.policy.load_path, map_location='cpu')) @@ -56,4 +54,4 @@ def main(rl_cfg, seed=0): if __name__ == "__main__": - main(rl_cfg=(ant_ddpg_config, ant_ddpg_create_config),seed=0) + main(ant_ddpg_config, seed=0) diff --git a/dizoo/mujoco/entry/mujoco_ddpg_main.py b/dizoo/mujoco/entry/mujoco_ddpg_main.py index e157ef9eb3..e8313ec4ed 100644 --- a/dizoo/mujoco/entry/mujoco_ddpg_main.py +++ b/dizoo/mujoco/entry/mujoco_ddpg_main.py @@ -7,11 +7,9 @@ from ding.worker import BaseLearner, SampleSerialCollector, InteractionSerialEvaluator, AdvancedReplayBuffer from ding.envs import BaseEnvManager, DingEnvWrapper from ding.policy import DDPGPolicy -from ding.model import QAC +from ding.model import ContinuousQAC from ding.utils import set_pkg_seed -from dizoo.classic_control.pendulum.envs import PendulumEnv from dizoo.mujoco.envs.mujoco_env import MujocoEnv -from dizoo.classic_control.pendulum.config.pendulum_ppo_config import pendulum_ppo_config from dizoo.mujoco.config.hopper_ddpg_config import hopper_ddpg_config @@ -24,6 +22,7 @@ def main(cfg, seed=0, max_iterations=int(1e10)): SampleSerialCollector, InteractionSerialEvaluator, AdvancedReplayBuffer, + MujocoEnv, save_cfg=True ) collector_env_num, evaluator_env_num = cfg.env.collector_env_num, cfg.env.evaluator_env_num @@ -38,7 +37,7 @@ def main(cfg, seed=0, max_iterations=int(1e10)): evaluator_env.seed(seed, dynamic_seed=False) set_pkg_seed(seed, use_cuda=cfg.policy.cuda) - model = QAC(**cfg.policy.model) + model = ContinuousQAC(**cfg.policy.model) policy = DDPGPolicy(cfg.policy, model=model) tb_logger = SummaryWriter(os.path.join('./log/', 'serial')) learner = BaseLearner(cfg.policy.learn.learner, policy.learn_mode, tb_logger) diff --git a/dizoo/mujoco/entry/mujoco_trex_main.py b/dizoo/mujoco/entry/mujoco_trex_main.py deleted file mode 100644 index 877b696972..0000000000 --- a/dizoo/mujoco/entry/mujoco_trex_main.py +++ /dev/null @@ -1,23 +0,0 @@ -import argparse -import torch -from ding.entry import trex_collecting_data -from dizoo.mujoco.config.halfcheetah_trex_onppo_config import main_config, create_config -from ding.entry import serial_pipeline_preference_based_irl_onpolicy - -# Note serial_pipeline_preference_based_irl_onpolicy is for on policy ppo whereas serial_pipeline_preference_based_irl -# is for sac. Note before run this file, please add the correpsonding path in the config, all path expect exp_name -# should be abs path. - -parser = argparse.ArgumentParser() -parser.add_argument( - '--cfg', - type=str, - default='please enter abs path for halfcheetah_trex_onppo_default_config.py or halfcheetah_trex_sac_default_config.py' -) -parser.add_argument('--seed', type=int, default=0) -parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu') -args = parser.parse_args() - -trex_collecting_data(args) -# if run sac, please import the relevant config and use serial_pipeline_trex -serial_pipeline_preference_based_irl_onpolicy([main_config, create_config]) diff --git a/dizoo/mujoco/envs/__init__.py b/dizoo/mujoco/envs/__init__.py index fac36665b9..0ed00b309a 100644 --- a/dizoo/mujoco/envs/__init__.py +++ b/dizoo/mujoco/envs/__init__.py @@ -1 +1,2 @@ from .mujoco_env import MujocoEnv +from .mujoco_disc_env import MujocoDiscEnv diff --git a/dizoo/mujoco/envs/mujoco_disc_env.py b/dizoo/mujoco/envs/mujoco_disc_env.py new file mode 100644 index 0000000000..442b5b2535 --- /dev/null +++ b/dizoo/mujoco/envs/mujoco_disc_env.py @@ -0,0 +1,166 @@ +import copy +import os +from itertools import product +from typing import Union, List, Optional + +import gym +import numpy as np +from easydict import EasyDict + +from ding.envs import BaseEnv, BaseEnvTimestep +from ding.envs.common import save_frames_as_gif +from ding.torch_utils import to_ndarray +from ding.utils import ENV_REGISTRY +from .mujoco_wrappers import wrap_mujoco + + +@ENV_REGISTRY.register('mujoco-disc') +class MujocoDiscEnv(BaseEnv): + """ + Overview: + The modified Mujoco environment with manually discretized action space. For each dimension, equally dividing the + original continuous action into ``each_dim_disc_size`` bins and using their Cartesian product to obtain + handcrafted discrete actions. + """ + + @classmethod + def default_config(cls: type) -> EasyDict: + cfg = EasyDict(copy.deepcopy(cls.config)) + cfg.cfg_type = cls.__name__ + 'Dict' + return cfg + + config = dict( + action_clip=False, + delay_reward_step=0, + replay_path=None, + save_replay_gif=False, + replay_path_gif=None, + ) + + def __init__(self, cfg: dict) -> None: + self._cfg = cfg + self._action_clip = cfg.action_clip + self._delay_reward_step = cfg.delay_reward_step + self._init_flag = False + self._replay_path = None + self._replay_path_gif = cfg.replay_path_gif + self._save_replay_gif = cfg.save_replay_gif + + def reset(self) -> np.ndarray: + if not self._init_flag: + self._env = self._make_env() + self._env.observation_space.dtype = np.float32 # To unify the format of envs in DI-engine + self._observation_space = self._env.observation_space + self._raw_action_space = self._env.action_space + self._reward_space = gym.spaces.Box( + low=self._env.reward_range[0], high=self._env.reward_range[1], shape=(1, ), dtype=np.float32 + ) + self._init_flag = True + if hasattr(self, '_seed') and hasattr(self, '_dynamic_seed') and self._dynamic_seed: + np_seed = 100 * np.random.randint(1, 1000) + self._env.seed(self._seed + np_seed) + elif hasattr(self, '_seed'): + self._env.seed(self._seed) + if self._replay_path is not None: + self._env = gym.wrappers.RecordVideo( + self._env, + video_folder=self._replay_path, + episode_trigger=lambda episode_id: True, + name_prefix='rl-video-{}'.format(id(self)) + ) + if self._save_replay_gif: + self._frames = [] + obs = self._env.reset() + obs = to_ndarray(obs).astype('float32') + + # disc_to_cont: transform discrete action index to original continuous action + self.m = self._raw_action_space.shape[0] + self.n = self._cfg.each_dim_disc_size + self.K = self.n ** self.m + self.disc_to_cont = list(product(*[list(range(self.n)) for _ in range(self.m)])) + self._eval_episode_return = 0. + # the modified discrete action space + self._action_space = gym.spaces.Discrete(self.K) + + return obs + + def close(self) -> None: + if self._init_flag: + self._env.close() + self._init_flag = False + + def seed(self, seed: int, dynamic_seed: bool = True) -> None: + self._seed = seed + self._dynamic_seed = dynamic_seed + np.random.seed(self._seed) + + def step(self, action: Union[np.ndarray, list]) -> BaseEnvTimestep: + # disc_to_cont: transform discrete action index to original continuous action + action = [-1 + 2 / self.n * k for k in self.disc_to_cont[int(action)]] + action = to_ndarray(action) + + if self._save_replay_gif: + self._frames.append(self._env.render(mode='rgb_array')) + if self._action_clip: + action = np.clip(action, -1, 1) + obs, rew, done, info = self._env.step(action) + self._eval_episode_return += rew + + if done: + if self._save_replay_gif: + path = os.path.join( + self._replay_path_gif, '{}_episode_{}.gif'.format(self._cfg.env_id, self._save_replay_count) + ) + save_frames_as_gif(self._frames, path) + self._save_replay_count += 1 + info['eval_episode_return'] = self._eval_episode_return + + obs = to_ndarray(obs).astype(np.float32) + rew = to_ndarray([rew]).astype(np.float32) + return BaseEnvTimestep(obs, rew, done, info) + + def _make_env(self): + return wrap_mujoco( + self._cfg.env_id, + norm_obs=self._cfg.get('norm_obs', None), + norm_reward=self._cfg.get('norm_reward', None), + delay_reward_step=self._delay_reward_step + ) + + def enable_save_replay(self, replay_path: Optional[str] = None) -> None: + if replay_path is None: + replay_path = './video' + self._replay_path = replay_path + self._save_replay = True + self._save_replay_count = 0 + + def random_action(self) -> np.ndarray: + return self.action_space.sample() + + def __repr__(self) -> str: + return "DI-engine modified Mujoco Env({}) with manually discretized action space".format(self._cfg.env_id) + + @staticmethod + def create_collector_env_cfg(cfg: dict) -> List[dict]: + collector_cfg = copy.deepcopy(cfg) + collector_env_num = collector_cfg.pop('collector_env_num', 1) + return [collector_cfg for _ in range(collector_env_num)] + + @staticmethod + def create_evaluator_env_cfg(cfg: dict) -> List[dict]: + evaluator_cfg = copy.deepcopy(cfg) + evaluator_env_num = evaluator_cfg.pop('evaluator_env_num', 1) + evaluator_cfg.norm_reward.use_norm = False + return [evaluator_cfg for _ in range(evaluator_env_num)] + + @property + def observation_space(self) -> gym.spaces.Space: + return self._observation_space + + @property + def action_space(self) -> gym.spaces.Space: + return self._action_space + + @property + def reward_space(self) -> gym.spaces.Space: + return self._reward_space diff --git a/dizoo/mujoco/envs/mujoco_env.py b/dizoo/mujoco/envs/mujoco_env.py index 2ecf82d555..700142583b 100644 --- a/dizoo/mujoco/envs/mujoco_env.py +++ b/dizoo/mujoco/envs/mujoco_env.py @@ -1,13 +1,15 @@ -from typing import Any, Union, List, Optional import copy -import numpy as np -from easydict import EasyDict +import os +from typing import Union, List, Optional + import gym +import numpy as np import torch +from easydict import EasyDict from ding.envs import BaseEnv, BaseEnvTimestep -from ding.envs.common.common_function import affine_transform -from ding.torch_utils import to_ndarray, to_list +from ding.envs.common import save_frames_as_gif +from ding.torch_utils import to_ndarray from ding.utils import ENV_REGISTRY from .mujoco_wrappers import wrap_mujoco @@ -22,16 +24,41 @@ def default_config(cls: type) -> EasyDict: return cfg config = dict( - use_act_scale=False, + action_clip=False, delay_reward_step=0, + replay_path=None, + save_replay_gif=False, + replay_path_gif=None, + action_bins_per_branch=None, ) def __init__(self, cfg: dict) -> None: self._cfg = cfg - self._use_act_scale = cfg.use_act_scale + self._action_clip = cfg.action_clip self._delay_reward_step = cfg.delay_reward_step self._init_flag = False self._replay_path = None + self._replay_path_gif = cfg.replay_path_gif + self._save_replay_gif = cfg.save_replay_gif + self._action_bins_per_branch = cfg.action_bins_per_branch + + def map_action(self, action: Union[np.ndarray, list]) -> Union[np.ndarray, list]: + """ + Overview: + Map the discretized action index to the action in the original action space. + Arguments: + - action (:obj:`np.ndarray or list`): The discretized action index. \ + The value ranges is {0, 1, ..., self._action_bins_per_branch - 1}. + Returns: + - outputs (:obj:`list`): The action in the original action space. \ + The value ranges is [-1, 1]. + Examples: + >>> inputs = [2, 0, 4] + >>> self._action_bins_per_branch = 5 + >>> outputs = map_action(inputs) + >>> assert isinstance(outputs, list) and outputs == [0.0, -1.0, 1.0] + """ + return [2 * x / (self._action_bins_per_branch - 1) - 1 for x in action] def reset(self) -> np.ndarray: if not self._init_flag: @@ -58,6 +85,8 @@ def reset(self) -> np.ndarray: self._env.seed(self._seed) obs = self._env.reset() obs = to_ndarray(obs).astype('float32') + self._eval_episode_return = 0. + return obs def close(self) -> None: @@ -71,11 +100,24 @@ def seed(self, seed: int, dynamic_seed: bool = True) -> None: np.random.seed(self._seed) def step(self, action: Union[np.ndarray, list]) -> BaseEnvTimestep: + if self._action_bins_per_branch: + action = self.map_action(action) action = to_ndarray(action) - if self._use_act_scale: - action_range = {'min': self.action_space.low[0], 'max': self.action_space.high[0], 'dtype': np.float32} - action = affine_transform(action, min_val=action_range['min'], max_val=action_range['max']) + if self._save_replay_gif: + self._frames.append(self._env.render(mode='rgb_array')) + if self._action_clip: + action = np.clip(action, -1, 1) obs, rew, done, info = self._env.step(action) + self._eval_episode_return += rew + if done: + if self._save_replay_gif: + path = os.path.join( + self._replay_path_gif, '{}_episode_{}.gif'.format(self._cfg.env_id, self._save_replay_count) + ) + save_frames_as_gif(self._frames, path) + self._save_replay_count += 1 + info['eval_episode_return'] = self._eval_episode_return + obs = to_ndarray(obs).astype(np.float32) rew = to_ndarray([rew]).astype(np.float32) return BaseEnvTimestep(obs, rew, done, info) @@ -109,7 +151,6 @@ def create_collector_env_cfg(cfg: dict) -> List[dict]: def create_evaluator_env_cfg(cfg: dict) -> List[dict]: evaluator_cfg = copy.deepcopy(cfg) evaluator_env_num = evaluator_cfg.pop('evaluator_env_num', 1) - evaluator_cfg.norm_reward.use_norm = False return [evaluator_cfg for _ in range(evaluator_env_num)] @property diff --git a/dizoo/mujoco/envs/mujoco_wrappers.py b/dizoo/mujoco/envs/mujoco_wrappers.py index 6f7635ced1..377172f2f8 100644 --- a/dizoo/mujoco/envs/mujoco_wrappers.py +++ b/dizoo/mujoco/envs/mujoco_wrappers.py @@ -2,7 +2,7 @@ import gym import numpy as np -from ding.envs import ObsNormWrapper, RewardNormWrapper, DelayRewardWrapper, FinalEvalRewardEnv +from ding.envs import ObsNormWrapper, RewardNormWrapper, DelayRewardWrapper, EvalEpisodeReturnWrapper def wrap_mujoco( @@ -25,7 +25,7 @@ def wrap_mujoco( # import customized gym environment from . import mujoco_gym_env env = gym.make(env_id) - env = FinalEvalRewardEnv(env) + env = EvalEpisodeReturnWrapper(env) if norm_obs is not None and norm_obs.use_norm: env = ObsNormWrapper(env) if norm_reward is not None and norm_reward.use_norm: diff --git a/dizoo/mujoco/envs/test/test_mujoco_disc_env.py b/dizoo/mujoco/envs/test/test_mujoco_disc_env.py new file mode 100644 index 0000000000..e8a39cc9c2 --- /dev/null +++ b/dizoo/mujoco/envs/test/test_mujoco_disc_env.py @@ -0,0 +1,44 @@ +import pytest +import numpy as np +from easydict import EasyDict + +from ding.utils import set_pkg_seed +from dizoo.mujoco.envs import MujocoDiscEnv + + +@pytest.mark.envtest +def test_mujoco_env_eval_episode_return(): + set_pkg_seed(1234, use_cuda=False) + each_dim_disc_size = 2 + env = MujocoDiscEnv( + EasyDict( + { + 'env_id': 'Ant-v3', + 'action_clip': False, + 'each_dim_disc_size': each_dim_disc_size, + 'delay_reward_step': 4, + 'save_replay_gif': False, + 'replay_path_gif': None + } + ) + ) + env.seed(1234) + env.reset() + action_dim = env._raw_action_space.shape + eval_episode_return = np.array([0.], dtype=np.float32) + while True: + action = np.random.randint(0, each_dim_disc_size ** action_dim[0], 1) + timestep = env.step(action) + eval_episode_return += timestep.reward + # print("{}(dtype: {})".format(timestep.reward, timestep.reward.dtype)) + if timestep.done: + print( + "{}({}), {}({})".format( + timestep.info['eval_episode_return'], type(timestep.info['eval_episode_return']), + eval_episode_return, type(eval_episode_return) + ) + ) + # timestep.reward and the cumulative reward in wrapper EvalEpisodeReturn are not the same. + assert abs(timestep.info['eval_episode_return'].item() - eval_episode_return.item()) / \ + abs(timestep.info['eval_episode_return'].item()) < 1e-5 + break diff --git a/dizoo/mujoco/envs/test/test_mujoco_env.py b/dizoo/mujoco/envs/test/test_mujoco_env.py index 45bc92317b..34bd311850 100644 --- a/dizoo/mujoco/envs/test/test_mujoco_env.py +++ b/dizoo/mujoco/envs/test/test_mujoco_env.py @@ -11,7 +11,17 @@ @pytest.mark.parametrize('delay_reward_step', [1, 10]) def test_mujoco_env_delay_reward(delay_reward_step): set_pkg_seed(1234, use_cuda=False) - env = MujocoEnv(EasyDict({'env_id': 'Ant-v3', 'use_act_scale': False, 'delay_reward_step': delay_reward_step})) + env = MujocoEnv( + EasyDict( + { + 'env_id': 'Ant-v3', + 'action_clip': False, + 'delay_reward_step': delay_reward_step, + 'save_replay_gif': False, + 'replay_path_gif': None + } + ) + ) env.seed(1234) env.reset() action_dim = env.action_space.shape @@ -28,26 +38,36 @@ def test_mujoco_env_delay_reward(delay_reward_step): @pytest.mark.envtest -def test_mujoco_env_final_eval_reward(): +def test_mujoco_env_eval_episode_return(): set_pkg_seed(1234, use_cuda=False) - env = MujocoEnv(EasyDict({'env_id': 'Ant-v3', 'use_act_scale': False, 'delay_reward_step': 4})) + env = MujocoEnv( + EasyDict( + { + 'env_id': 'Ant-v3', + 'action_clip': False, + 'delay_reward_step': 4, + 'save_replay_gif': False, + 'replay_path_gif': None + } + ) + ) env.seed(1234) env.reset() action_dim = env.action_space.shape - final_eval_reward = np.array([0.], dtype=np.float32) + eval_episode_return = np.array([0.], dtype=np.float32) while True: action = np.random.random(size=action_dim) timestep = env.step(action) - final_eval_reward += timestep.reward + eval_episode_return += timestep.reward # print("{}(dtype: {})".format(timestep.reward, timestep.reward.dtype)) if timestep.done: print( "{}({}), {}({})".format( - timestep.info['final_eval_reward'], type(timestep.info['final_eval_reward']), final_eval_reward, - type(final_eval_reward) + timestep.info['eval_episode_return'], type(timestep.info['eval_episode_return']), + eval_episode_return, type(eval_episode_return) ) ) - # timestep.reward and the cumulative reward in wrapper FinalEvalReward are not the same. - assert abs(timestep.info['final_eval_reward'].item() - final_eval_reward.item()) / \ - abs(timestep.info['final_eval_reward'].item()) < 1e-5 + # timestep.reward and the cumulative reward in wrapper EvalEpisodeReturn are not the same. + assert abs(timestep.info['eval_episode_return'].item() - eval_episode_return.item()) / \ + abs(timestep.info['eval_episode_return'].item()) < 1e-5 break diff --git a/dizoo/mujoco/example/mujoco_sac.py b/dizoo/mujoco/example/mujoco_sac.py index 0349d8d161..f2833187f0 100644 --- a/dizoo/mujoco/example/mujoco_sac.py +++ b/dizoo/mujoco/example/mujoco_sac.py @@ -1,5 +1,5 @@ from ditk import logging -from ding.model import QAC +from ding.model import ContinuousQAC from ding.policy import SACPolicy from ding.envs import DingEnvWrapper, SubprocessEnvManagerV2 from ding.data import DequeBuffer @@ -26,7 +26,7 @@ def main(): set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) - model = QAC(**cfg.policy.model) + model = ContinuousQAC(**cfg.policy.model) buffer_ = DequeBuffer(size=cfg.policy.other.replay_buffer.replay_buffer_size) policy = SACPolicy(cfg.policy, model=model) @@ -36,8 +36,8 @@ def main(): ) task.use(data_pusher(cfg, buffer_)) task.use(OffPolicyLearner(cfg, policy.learn_mode, buffer_)) - task.use(CkptSaver(cfg, policy, train_freq=500)) - task.use(termination_checker(max_train_iter=int(3e6))) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=500)) + task.use(termination_checker(max_env_step=int(3e6))) task.run() diff --git a/dizoo/multiagent_mujoco/config/ant_maddpg_config.py b/dizoo/multiagent_mujoco/config/ant_maddpg_config.py new file mode 100644 index 0000000000..ed6744e818 --- /dev/null +++ b/dizoo/multiagent_mujoco/config/ant_maddpg_config.py @@ -0,0 +1,63 @@ +from easydict import EasyDict + +ant_ddpg_default_config = dict( + exp_name='multi_mujoco_ant_2x4_ddpg', + env=dict( + scenario='Ant-v2', + agent_conf="2x4d", + agent_obsk=2, + add_agent_id=False, + episode_limit=1000, + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=6000, + ), + policy=dict( + cuda=True, + random_collect_size=0, + multi_agent=True, + model=dict( + agent_obs_shape=54, + global_obs_shape=111, + action_shape=4, + action_space='regression', + actor_head_hidden_size=256, + critic_head_hidden_size=256, + ), + learn=dict( + update_per_collect=10, + batch_size=256, + learning_rate_actor=1e-3, + learning_rate_critic=1e-3, + target_theta=0.005, + discount_factor=0.99, + ), + collect=dict( + n_sample=400, + noise_sigma=0.1, + ), + eval=dict(evaluator=dict(eval_freq=500, )), + other=dict(replay_buffer=dict(replay_buffer_size=100000, ), ), + ), +) + +ant_ddpg_default_config = EasyDict(ant_ddpg_default_config) +main_config = ant_ddpg_default_config + +ant_ddpg_default_create_config = dict( + env=dict( + type='mujoco_multi', + import_names=['dizoo.multiagent_mujoco.envs.multi_mujoco_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='ddpg'), + replay_buffer=dict(type='naive', ), +) +ant_ddpg_default_create_config = EasyDict(ant_ddpg_default_create_config) +create_config = ant_ddpg_default_create_config + +if __name__ == '__main__': + # or you can enter `ding -m serial -c ant_maddpg_config.py -s 0` + from ding.entry.serial_entry import serial_pipeline + serial_pipeline((main_config, create_config), seed=0, max_env_step=int(1e7)) diff --git a/dizoo/multiagent_mujoco/config/ant_mappo_config.py b/dizoo/multiagent_mujoco/config/ant_mappo_config.py index c4559f5752..d11c31be8d 100644 --- a/dizoo/multiagent_mujoco/config/ant_mappo_config.py +++ b/dizoo/multiagent_mujoco/config/ant_mappo_config.py @@ -41,10 +41,8 @@ ), # used in state_num of hidden_state learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, - epoch_per_collect=10, - batch_size=3200, + epoch_per_collect=3, + batch_size=800, learning_rate=5e-4, # ============================================================== # The following configs is algorithm-specific @@ -52,19 +50,18 @@ # (float) The loss weight of value network, policy network weight is set to 1 value_weight=0.5, # (float) The loss weight of entropy regularization, policy network weight is set to 1 - entropy_weight=0.0, + entropy_weight=0.001, # (float) PPO clip ratio, defaults to 0.2 - clip_ratio=0.5, + clip_ratio=0.2, # (bool) Whether to use advantage norm in a whole training batch - adv_norm=False, + adv_norm=True, value_norm=True, ppo_param_init=True, grad_clip_type='clip_norm', - grad_clip_value=10, - ignore_done=False, + grad_clip_value=5, ), collect=dict(env_num=collector_env_num, n_sample=3200), - eval=dict(env_num=evaluator_env_num, evaluator=dict(eval_freq=200, )), + eval=dict(env_num=evaluator_env_num, evaluator=dict(eval_freq=1000, )), ), ) main_config = EasyDict(main_config) @@ -78,8 +75,6 @@ ) create_config = EasyDict(create_config) - if __name__ == '__main__': - from ding.entry import serial_pipeline_onpolicy - serial_pipeline_onpolicy((main_config, create_config), seed=0) + serial_pipeline_onpolicy((main_config, create_config), seed=0, max_env_step=int(1e7)) diff --git a/dizoo/multiagent_mujoco/config/ant_masac_config.py b/dizoo/multiagent_mujoco/config/ant_masac_config.py index ccb879fda7..9316b095c0 100644 --- a/dizoo/multiagent_mujoco/config/ant_masac_config.py +++ b/dizoo/multiagent_mujoco/config/ant_masac_config.py @@ -1,7 +1,7 @@ from easydict import EasyDict ant_sac_default_config = dict( - exp_name='multi_mujoco_ant_2x4', + exp_name='multi_mujoco_ant_2x4_sac', env=dict( scenario='Ant-v2', agent_conf="2x4d", @@ -21,7 +21,6 @@ agent_obs_shape=54, global_obs_shape=111, action_shape=4, - twin_critic=True, action_space='reparameterization', actor_head_hidden_size=256, critic_head_hidden_size=256, @@ -32,20 +31,11 @@ learning_rate_q=1e-3, learning_rate_policy=1e-3, learning_rate_alpha=3e-4, - ignore_done=False, target_theta=0.005, discount_factor=0.99, - alpha=0.2, - reparameterization=True, - auto_alpha=True, - log_space=True, ), - collect=dict( - n_sample=400, - unroll_len=1, - ), - command=dict(), - eval=dict(evaluator=dict(eval_freq=100, )), + collect=dict(n_sample=400, ), + eval=dict(evaluator=dict(eval_freq=500, )), other=dict(replay_buffer=dict(replay_buffer_size=100000, ), ), ), ) @@ -59,10 +49,7 @@ import_names=['dizoo.multiagent_mujoco.envs.multi_mujoco_env'], ), env_manager=dict(type='subprocess'), - policy=dict( - type='sac', - import_names=['ding.policy.sac'], - ), + policy=dict(type='sac'), replay_buffer=dict(type='naive', ), ) ant_sac_default_create_config = EasyDict(ant_sac_default_create_config) @@ -71,4 +58,4 @@ if __name__ == '__main__': # or you can enter `ding -m serial -c ant_masac_config.py -s 0` from ding.entry.serial_entry import serial_pipeline - serial_pipeline((main_config, create_config), seed=0) + serial_pipeline((main_config, create_config), seed=0, max_env_step=int(1e7)) diff --git a/dizoo/multiagent_mujoco/config/ant_matd3_config.py b/dizoo/multiagent_mujoco/config/ant_matd3_config.py new file mode 100644 index 0000000000..4575f40de5 --- /dev/null +++ b/dizoo/multiagent_mujoco/config/ant_matd3_config.py @@ -0,0 +1,66 @@ +from easydict import EasyDict + +ant_td3_default_config = dict( + exp_name='multi_mujoco_ant_2x4_td3', + env=dict( + scenario='Ant-v2', + agent_conf="2x4d", + agent_obsk=2, + add_agent_id=False, + episode_limit=1000, + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=6000, + ), + policy=dict( + cuda=True, + random_collect_size=0, + multi_agent=True, + model=dict( + agent_obs_shape=54, + global_obs_shape=111, + action_shape=4, + action_space='regression', + actor_head_hidden_size=256, + critic_head_hidden_size=256, + twin_critic=True, + ), + learn=dict( + update_per_collect=10, + batch_size=256, + learning_rate_actor=1e-3, + learning_rate_critic=1e-3, + target_theta=0.005, + discount_factor=0.99, + actor_update_freq=2, + noise=True, + ), + collect=dict( + n_sample=400, + noise_sigma=0.1, + ), + eval=dict(evaluator=dict(eval_freq=500, )), + other=dict(replay_buffer=dict(replay_buffer_size=100000, ), ), + ), +) + +ant_td3_default_config = EasyDict(ant_td3_default_config) +main_config = ant_td3_default_config + +ant_td3_default_create_config = dict( + env=dict( + type='mujoco_multi', + import_names=['dizoo.multiagent_mujoco.envs.multi_mujoco_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='td3'), + replay_buffer=dict(type='naive', ), +) +ant_td3_default_create_config = EasyDict(ant_td3_default_create_config) +create_config = ant_td3_default_create_config + +if __name__ == '__main__': + # or you can enter `ding -m serial -c ant_matd3_config.py -s 0` + from ding.entry.serial_entry import serial_pipeline + serial_pipeline((main_config, create_config), seed=0, max_env_step=int(1e7)) diff --git a/dizoo/multiagent_mujoco/config/halfcheetah_happo_config.py b/dizoo/multiagent_mujoco/config/halfcheetah_happo_config.py new file mode 100644 index 0000000000..c849551d94 --- /dev/null +++ b/dizoo/multiagent_mujoco/config/halfcheetah_happo_config.py @@ -0,0 +1,83 @@ +from easydict import EasyDict + +collector_env_num = 8 +evaluator_env_num = 8 +n_agent = 2 + +main_config = dict( + exp_name='HAPPO_result/debug/multi_mujoco_halfcheetah_2x3_happo', + env=dict( + scenario='HalfCheetah-v2', + agent_conf="2x3", + agent_obsk=2, + add_agent_id=False, + episode_limit=1000, + collector_env_num=collector_env_num, + evaluator_env_num=evaluator_env_num, + n_evaluator_episode=8, + stop_value=6000, + ), + policy=dict( + cuda=True, + multi_agent=True, + agent_num=n_agent, + action_space='continuous', + model=dict( + action_space='continuous', + agent_num=n_agent, + agent_obs_shape=8, + global_obs_shape=17, + action_shape=3, + use_lstm=False, + ), + learn=dict( + epoch_per_collect=5, + # batch_size=3200, + batch_size=800, + learning_rate=5e-4, + critic_learning_rate=5e-4, + # ============================================================== + # The following configs is algorithm-specific + # ============================================================== + # (float) The loss weight of value network, policy network weight is set to 1 + value_weight=0.5, + # (float) The loss weight of entropy regularization, policy network weight is set to 1 + # entropy_weight=0.001, + entropy_weight=0.001, + # (float) PPO clip ratio, defaults to 0.2 + clip_ratio=0.2, + # (bool) Whether to use advantage norm in a whole training batch + adv_norm=True, + value_norm=True, + ppo_param_init=True, + grad_clip_type='clip_norm', + grad_clip_value=3, + ignore_done=True, + # ignore_done=False, + ), + collect=dict( + n_sample=3200, + unroll_len=1, + env_num=collector_env_num, + ), + eval=dict( + env_num=evaluator_env_num, + evaluator=dict(eval_freq=1000, ), + ), + other=dict(), + ), +) +main_config = EasyDict(main_config) +create_config = dict( + env=dict( + type='mujoco_multi', + import_names=['dizoo.multiagent_mujoco.envs.multi_mujoco_env'], + ), + env_manager=dict(type='base'), + policy=dict(type='happo'), +) +create_config = EasyDict(create_config) + +if __name__ == '__main__': + from ding.entry import serial_pipeline_onpolicy + serial_pipeline_onpolicy((main_config, create_config), seed=0, max_env_step=int(1e7)) diff --git a/dizoo/multiagent_mujoco/config/halfcheetah_mappo_config.py b/dizoo/multiagent_mujoco/config/halfcheetah_mappo_config.py new file mode 100644 index 0000000000..b6db3feea7 --- /dev/null +++ b/dizoo/multiagent_mujoco/config/halfcheetah_mappo_config.py @@ -0,0 +1,80 @@ +from easydict import EasyDict + +collector_env_num = 8 +evaluator_env_num = 8 + +main_config = dict( + exp_name='HAPPO_result/multi_mujoco_halfcheetah_2x3_mappo', + env=dict( + scenario='HalfCheetah-v2', + agent_conf="2x3", + agent_obsk=2, + add_agent_id=False, + episode_limit=1000, + collector_env_num=collector_env_num, + evaluator_env_num=evaluator_env_num, + n_evaluator_episode=8, + stop_value=6000, + ), + policy=dict( + cuda=True, + multi_agent=True, + action_space='continuous', + model=dict( + # (int) agent_num: The number of the agent. + # For SMAC 3s5z, agent_num=8; for 2c_vs_64zg, agent_num=2. + agent_num=2, + # (int) obs_shape: The shapeension of observation of each agent. + # For 3s5z, obs_shape=150; for 2c_vs_64zg, agent_num=404. + # (int) global_obs_shape: The shapeension of global observation. + # For 3s5z, obs_shape=216; for 2c_vs_64zg, agent_num=342. + agent_obs_shape=8, + #global_obs_shape=216, + global_obs_shape=17, + # (int) action_shape: The number of action which each agent can take. + # action_shape= the number of common action (6) + the number of enemies. + # For 3s5z, obs_shape=14 (6+8); for 2c_vs_64zg, agent_num=70 (6+64). + action_shape=3, + # (List[int]) The size of hidden layer + # hidden_size_list=[64], + action_space='continuous' + ), + # used in state_num of hidden_state + learn=dict( + epoch_per_collect=5, + batch_size=800, + learning_rate=5e-4, + # ============================================================== + # The following configs is algorithm-specific + # ============================================================== + # (float) The loss weight of value network, policy network weight is set to 1 + value_weight=0.5, + # (float) The loss weight of entropy regularization, policy network weight is set to 1 + entropy_weight=0.001, + # (float) PPO clip ratio, defaults to 0.2 + clip_ratio=0.2, + # (bool) Whether to use advantage norm in a whole training batch + adv_norm=True, + value_norm=True, + ppo_param_init=True, + grad_clip_type='clip_norm', + grad_clip_value=5, + ), + collect=dict(env_num=collector_env_num, n_sample=3200), + eval=dict(env_num=evaluator_env_num, evaluator=dict(eval_freq=1000, )), + ), +) +main_config = EasyDict(main_config) +create_config = dict( + env=dict( + type='mujoco_multi', + import_names=['dizoo.multiagent_mujoco.envs.multi_mujoco_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='ppo'), +) +create_config = EasyDict(create_config) + +if __name__ == '__main__': + from ding.entry import serial_pipeline_onpolicy + serial_pipeline_onpolicy((main_config, create_config), seed=0, max_env_step=int(1e7)) diff --git a/dizoo/multiagent_mujoco/config/walker2d_happo_config.py b/dizoo/multiagent_mujoco/config/walker2d_happo_config.py new file mode 100644 index 0000000000..a947a25589 --- /dev/null +++ b/dizoo/multiagent_mujoco/config/walker2d_happo_config.py @@ -0,0 +1,91 @@ +from easydict import EasyDict +import os +collector_env_num = 8 +evaluator_env_num = 8 +n_agent = 2 + +main_config = dict( + exp_name='HAPPO_result/debug/multi_mujoco_walker_2x3_happo', + env=dict( + scenario='Walker2d-v2', + agent_conf="2x3", + agent_obsk=2, + add_agent_id=False, + episode_limit=1000, + collector_env_num=collector_env_num, + evaluator_env_num=evaluator_env_num, + n_evaluator_episode=8, + stop_value=6000, + ), + policy=dict( + cuda=True, + multi_agent=True, + agent_num=n_agent, + action_space='continuous', + model=dict( + action_space='continuous', + agent_num=n_agent, + agent_obs_shape=8, + global_obs_shape=17, + action_shape=3, + use_lstm=False, + ), + learn=dict( + epoch_per_collect=5, + # batch_size=3200, + # batch_size=800, + batch_size=320, + # batch_size=100, + learning_rate=5e-4, + critic_learning_rate=5e-3, + # learning_rate=3e-3, + # ============================================================== + # The following configs is algorithm-specific + # ============================================================== + # (float) The loss weight of value network, policy network weight is set to 1 + # value_weight=0.5, + value_weight=1, + # (float) The loss weight of entropy regularization, policy network weight is set to 1 + # entropy_weight=0.001, + entropy_weight=0.003, + # entropy_weight=0.005, + # entropy_weight=0.01, + # (float) PPO clip ratio, defaults to 0.2 + clip_ratio=0.2, + # (bool) Whether to use advantage norm in a whole training batch + adv_norm=True, + value_norm=True, + ppo_param_init=True, + grad_clip_type='clip_norm', + # grad_clip_value=5, + grad_clip_value=10, + # ignore_done=True, + ignore_done=False, + ), + collect=dict( + n_sample=3200, + # n_sample=4000, + unroll_len=1, + env_num=collector_env_num, + ), + eval=dict( + env_num=evaluator_env_num, + evaluator=dict(eval_freq=1000, ), + ), + other=dict(), + ), +) +main_config = EasyDict(main_config) +create_config = dict( + env=dict( + type='mujoco_multi', + import_names=['dizoo.multiagent_mujoco.envs.multi_mujoco_env'], + ), + env_manager=dict(type='base'), + policy=dict(type='happo'), +) +create_config = EasyDict(create_config) + +if __name__ == '__main__': + from ding.entry import serial_pipeline_onpolicy + serial_pipeline_onpolicy((main_config, create_config), seed=0, max_env_step=int(1e7)) diff --git a/dizoo/multiagent_mujoco/envs/multi_mujoco_env.py b/dizoo/multiagent_mujoco/envs/multi_mujoco_env.py index 79adfc14fd..8ed538909c 100644 --- a/dizoo/multiagent_mujoco/envs/multi_mujoco_env.py +++ b/dizoo/multiagent_mujoco/envs/multi_mujoco_env.py @@ -27,7 +27,7 @@ def reset(self) -> np.ndarray: self._env = MujocoMulti(env_args=self._cfg) self._init_flag = True obs = self._env.reset() - self._final_eval_reward = 0. + self._eval_episode_return = 0. # TODO: # self.env_info for scenario='Ant-v2', agent_conf="2x4d", @@ -78,10 +78,10 @@ def seed(self, seed: int, dynamic_seed: bool = True) -> None: def step(self, action: Union[np.ndarray, list]) -> BaseEnvTimestep: action = to_ndarray(action) obs, rew, done, info = self._env.step(action) - self._final_eval_reward += rew + self._eval_episode_return += rew rew = to_ndarray([rew]) # wrapped to be transfered to a array with shape (1,) if done: - info['final_eval_reward'] = self._final_eval_reward + info['eval_episode_return'] = self._eval_episode_return return BaseEnvTimestep(obs, rew, done, info) def random_action(self) -> np.ndarray: diff --git a/dizoo/overcooked/envs/overcooked_env.py b/dizoo/overcooked/envs/overcooked_env.py index 0455749290..2de769f0b3 100644 --- a/dizoo/overcooked/envs/overcooked_env.py +++ b/dizoo/overcooked/envs/overcooked_env.py @@ -93,9 +93,9 @@ def step(self, action): next_state, reward, done, env_info = self.base_env.step(joint_action) reward = np.array([float(reward)]) - self._final_eval_reward += reward + self._eval_episode_return += reward if self._shape_reward: - self._final_eval_reward += sum(env_info['shaped_r_by_agent']) + self._eval_episode_return += sum(env_info['shaped_r_by_agent']) reward += sum(env_info['shaped_r_by_agent']) ob_p0, ob_p1 = self.featurize_fn(self.mdp, next_state) @@ -110,7 +110,7 @@ def step(self, action): both_agents_ob = np.stack(both_agents_ob) env_info["policy_agent_idx"] = self.agent_idx - env_info["final_eval_reward"] = self._final_eval_reward + env_info["eval_episode_return"] = self._eval_episode_return env_info["other_agent_env_idx"] = 1 - self.agent_idx action_mask = self.get_action_mask() @@ -130,7 +130,7 @@ def obs_preprocess(self, obs): def reset(self): self.base_env.reset() - self._final_eval_reward = 0 + self._eval_episode_return = 0 self.mdp = self.base_env.mdp # random init agent index self.agent_idx = np.random.choice([0, 1]) @@ -248,9 +248,9 @@ def step(self, action): next_state, reward, done, env_info = self.base_env.step(joint_action) reward = np.array([float(reward)]) - self._final_eval_reward += reward + self._eval_episode_return += reward if self._shape_reward: - self._final_eval_reward += sum(env_info['shaped_r_by_agent']) + self._eval_episode_return += sum(env_info['shaped_r_by_agent']) reward += sum(env_info['shaped_r_by_agent']) ob_p0, ob_p1 = self.featurize_fn(self.mdp, next_state) ob_p0, ob_p1 = self.obs_preprocess(ob_p0), self.obs_preprocess(ob_p1) @@ -264,7 +264,7 @@ def step(self, action): both_agents_ob = np.stack(both_agents_ob) env_info["policy_agent_idx"] = self.agent_idx - env_info["final_eval_reward"] = self._final_eval_reward + env_info["eval_episode_return"] = self._eval_episode_return env_info["other_agent_env_idx"] = 1 - self.agent_idx action_mask = self.get_action_mask() @@ -280,7 +280,7 @@ def obs_preprocess(self, obs): def reset(self): self.base_env.reset() - self._final_eval_reward = 0 + self._eval_episode_return = 0 self.mdp = self.base_env.mdp # random init agent index self.agent_idx = np.random.choice([0, 1]) diff --git a/dizoo/overcooked/envs/test_overcooked_env.py b/dizoo/overcooked/envs/test_overcooked_env.py index 2275034a20..8e238a8ced 100644 --- a/dizoo/overcooked/envs/test_overcooked_env.py +++ b/dizoo/overcooked/envs/test_overcooked_env.py @@ -26,7 +26,7 @@ def test_overcook(self, action_mask): else: assert obs.shape == env.observation_space.shape assert timestep.done - sum_rew += timestep.info['final_eval_reward'][0] + sum_rew += timestep.info['eval_episode_return'][0] print("sum reward is:", sum_rew) @pytest.mark.parametrize("concat_obs", [True, False]) @@ -40,5 +40,5 @@ def test_overcook_game(self, concat_obs): obs = timestep.obs assert obs.shape == env.observation_space.shape assert timestep.done - print("agent 0 sum reward is:", timestep.info[0]['final_eval_reward']) - print("agent 1 sum reward is:", timestep.info[1]['final_eval_reward']) + print("agent 0 sum reward is:", timestep.info[0]['eval_episode_return']) + print("agent 1 sum reward is:", timestep.info[1]['eval_episode_return']) diff --git a/dizoo/petting_zoo/config/__init__.py b/dizoo/petting_zoo/config/__init__.py index ba30e7d2d1..9838dcaa00 100644 --- a/dizoo/petting_zoo/config/__init__.py +++ b/dizoo/petting_zoo/config/__init__.py @@ -6,3 +6,4 @@ from .ptz_simple_spread_qtran_config import ptz_simple_spread_qtran_config, ptz_simple_spread_qtran_create_config from .ptz_simple_spread_vdn_config import ptz_simple_spread_vdn_config, ptz_simple_spread_vdn_create_config from .ptz_simple_spread_wqmix_config import ptz_simple_spread_wqmix_config, ptz_simple_spread_wqmix_create_config +from .ptz_simple_spread_madqn_config import ptz_simple_spread_madqn_config, ptz_simple_spread_madqn_create_config # noqa diff --git a/dizoo/petting_zoo/config/ptz_pistonball_qmix_config.py b/dizoo/petting_zoo/config/ptz_pistonball_qmix_config.py new file mode 100644 index 0000000000..b6bfc0818e --- /dev/null +++ b/dizoo/petting_zoo/config/ptz_pistonball_qmix_config.py @@ -0,0 +1,77 @@ +from easydict import EasyDict + +n_pistons = 20 +collector_env_num = 8 +evaluator_env_num = 8 +max_env_step = 3e6 + +main_config = dict( + exp_name=f'data_pistonball/ptz_pistonball_n{n_pistons}_qmix_seed0', + env=dict( + env_family='butterfly', + env_id='pistonball_v6', + n_pistons=n_pistons, + max_cycles=125, + agent_obs_only=False, + continuous_actions=False, + collector_env_num=collector_env_num, + evaluator_env_num=evaluator_env_num, + n_evaluator_episode=evaluator_env_num, + stop_value=1e6, + manager=dict(shared_memory=False, ), + ), + policy=dict( + cuda=True, + model=dict( + agent_num=n_pistons, + obs_shape=(3, 457, 120), # RGB image observation shape for each piston agent + global_obs_shape=(3, 560, 880), # Global state shape + action_shape=3, # Discrete actions (0, 1, 2) + hidden_size_list=[32, 64, 128, 256], + mixer=True, + ), + learn=dict( + update_per_collect=20, + batch_size=32, + learning_rate=0.0001, + clip_value=5, + target_update_theta=0.001, + discount_factor=0.99, + double_q=True, + ), + collect=dict( + n_sample=16, + unroll_len=5, + env_num=collector_env_num, + ), + eval=dict(env_num=evaluator_env_num), + other=dict( + eps=dict( + type='exp', + start=1.0, + end=0.05, + decay=100000, + ), + replay_buffer=dict(replay_buffer_size=5000, ), + ), + ), +) +main_config = EasyDict(main_config) + +create_config = dict( + env=dict( + import_names=['dizoo.petting_zoo.envs.petting_zoo_pistonball_env'], + type='petting_zoo_pistonball', + ), + env_manager=dict(type='subprocess'), + policy=dict(type='qmix'), +) +create_config = EasyDict(create_config) + +ptz_pistonball_qmix_config = main_config +ptz_pistonball_qmix_create_config = create_config + +if __name__ == '__main__': + # or you can enter `ding -m serial -c ptz_pistonball_qmix_config.py -s 0` + from ding.entry import serial_pipeline + serial_pipeline((main_config, create_config), seed=0, max_env_step=max_env_step) diff --git a/dizoo/petting_zoo/config/ptz_simple_spread_happo_config.py b/dizoo/petting_zoo/config/ptz_simple_spread_happo_config.py new file mode 100644 index 0000000000..d1ff088326 --- /dev/null +++ b/dizoo/petting_zoo/config/ptz_simple_spread_happo_config.py @@ -0,0 +1,88 @@ +from easydict import EasyDict + +n_agent = 3 +n_landmark = n_agent +collector_env_num = 8 +evaluator_env_num = 8 +main_config = dict( + exp_name='ptz_simple_spread_happo_seed0', + env=dict( + env_family='mpe', + env_id='simple_spread_v2', + n_agent=n_agent, + n_landmark=n_landmark, + max_cycles=25, + agent_obs_only=False, + agent_specific_global_state=True, + continuous_actions=False, + collector_env_num=collector_env_num, + evaluator_env_num=evaluator_env_num, + n_evaluator_episode=evaluator_env_num, + stop_value=0, + ), + policy=dict( + cuda=True, + multi_agent=True, + agent_num=n_agent, + action_space='discrete', + model=dict( + action_space='discrete', + agent_num=n_agent, + agent_obs_shape=2 + 2 + n_landmark * 2 + (n_agent - 1) * 2 + (n_agent - 1) * 2, + global_obs_shape=2 + 2 + n_landmark * 2 + (n_agent - 1) * 2 + (n_agent - 1) * 2 + n_agent * (2 + 2) + + n_landmark * 2 + n_agent * (n_agent - 1) * 2, + action_shape=5, + use_lstm=False, + ), + learn=dict( + multi_gpu=False, + epoch_per_collect=5, + batch_size=3200, + learning_rate=5e-4, + critic_learning_rate=5e-4, + # ============================================================== + # The following configs is algorithm-specific + # ============================================================== + # (float) The loss weight of value network, policy network weight is set to 1 + value_weight=0.5, + # (float) The loss weight of entropy regularization, policy network weight is set to 1 + entropy_weight=0.01, + # (float) PPO clip ratio, defaults to 0.2 + clip_ratio=0.2, + # (bool) Whether to use advantage norm in a whole training batch + adv_norm=False, + value_norm=True, + ppo_param_init=True, + grad_clip_type='clip_norm', + grad_clip_value=10, + ignore_done=False, + ), + collect=dict( + n_sample=3200, + unroll_len=1, + env_num=collector_env_num, + ), + eval=dict( + env_num=evaluator_env_num, + evaluator=dict(eval_freq=50, ), + ), + other=dict(), + ), +) +main_config = EasyDict(main_config) +create_config = dict( + env=dict( + import_names=['dizoo.petting_zoo.envs.petting_zoo_simple_spread_env'], + type='petting_zoo', + ), + env_manager=dict(type='base'), + policy=dict(type='happo'), +) +create_config = EasyDict(create_config) +ptz_simple_spread_happo_config = main_config +ptz_simple_spread_happo_create_config = create_config + +if __name__ == '__main__': + # or you can enter `ding -m serial_onpolicy -c ptz_simple_spread_happo_config.py -s 0` + from ding.entry import serial_pipeline_onpolicy + serial_pipeline_onpolicy((main_config, create_config), seed=0) diff --git a/dizoo/petting_zoo/config/ptz_simple_spread_maddpg_config.py b/dizoo/petting_zoo/config/ptz_simple_spread_maddpg_config.py new file mode 100644 index 0000000000..822d28d966 --- /dev/null +++ b/dizoo/petting_zoo/config/ptz_simple_spread_maddpg_config.py @@ -0,0 +1,81 @@ +from easydict import EasyDict + +n_agent = 3 +n_landmark = n_agent +collector_env_num = 8 +evaluator_env_num = 8 +main_config = dict( + exp_name='ptz_simple_spread_maddpg_seed0', + env=dict( + env_family='mpe', + env_id='simple_spread_v2', + n_agent=n_agent, + n_landmark=n_landmark, + max_cycles=25, + agent_obs_only=False, + agent_specific_global_state=True, + continuous_actions=True, # ddpg only support continuous action space + act_scale=True, # necessary for continuous action space + collector_env_num=collector_env_num, + evaluator_env_num=evaluator_env_num, + n_evaluator_episode=evaluator_env_num, + stop_value=0, + ), + policy=dict( + cuda=True, + multi_agent=True, + random_collect_size=5000, + model=dict( + agent_obs_shape=2 + 2 + n_landmark * 2 + (n_agent - 1) * 2 + (n_agent - 1) * 2, + global_obs_shape=2 + 2 + n_landmark * 2 + (n_agent - 1) * 2 + (n_agent - 1) * 2 + n_agent * (2 + 2) + + n_landmark * 2 + n_agent * (n_agent - 1) * 2, + action_shape=5, + action_space='regression', + twin_critic=False, + ), + learn=dict( + update_per_collect=50, + batch_size=320, + # learning_rates + learning_rate_q=5e-4, + learning_rate_policy=5e-4, + target_theta=0.005, + discount_factor=0.99, + ), + collect=dict( + n_sample=1600, + env_num=collector_env_num, + ), + eval=dict( + env_num=evaluator_env_num, + evaluator=dict(eval_freq=500, ), + ), + other=dict( + eps=dict( + type='linear', + start=1, + end=0.05, + decay=100000, + ), + replay_buffer=dict(replay_buffer_size=int(1e6), ) + ), + ), +) + +main_config = EasyDict(main_config) +create_config = dict( + env=dict( + import_names=['dizoo.petting_zoo.envs.petting_zoo_simple_spread_env'], + type='petting_zoo', + ), + env_manager=dict(type='subprocess'), + policy=dict(type='ddpg'), +) +create_config = EasyDict(create_config) +ptz_simple_spread_maddpg_config = main_config +ptz_simple_spread_maddpg_create_config = create_config + +if __name__ == '__main__': + # or you can enter `ding -m serial_entry -c ptz_simple_spread_maddpg_config.py -s 0` + from ding.entry import serial_pipeline + serial_pipeline((main_config, create_config), seed=0, max_env_step=int(1e6)) diff --git a/dizoo/petting_zoo/config/ptz_simple_spread_madqn_config.py b/dizoo/petting_zoo/config/ptz_simple_spread_madqn_config.py new file mode 100644 index 0000000000..8ddb636abf --- /dev/null +++ b/dizoo/petting_zoo/config/ptz_simple_spread_madqn_config.py @@ -0,0 +1,83 @@ +from easydict import EasyDict + +n_agent = 3 +n_landmark = n_agent +collector_env_num = 8 +evaluator_env_num = 8 +main_config = dict( + exp_name='ptz_simple_spread_madqn_seed0', + env=dict( + env_family='mpe', + env_id='simple_spread_v2', + n_agent=n_agent, + n_landmark=n_landmark, + max_cycles=25, + agent_obs_only=False, + agent_specific_global_state=True, + continuous_actions=False, + collector_env_num=collector_env_num, + evaluator_env_num=evaluator_env_num, + n_evaluator_episode=evaluator_env_num, + stop_value=0, + ), + policy=dict( + cuda=True, + nstep=3, + model=dict( + obs_shape=2 + 2 + n_landmark * 2 + (n_agent - 1) * 2 + (n_agent - 1) * 2, + global_obs_shape=2 + 2 + n_landmark * 2 + (n_agent - 1) * 2 + (n_agent - 1) * 2 + n_agent * (2 + 2) + + n_landmark * 2 + n_agent * (n_agent - 1) * 2, + agent_num=n_agent, + action_shape=5, + global_cooperation=True, + hidden_size_list=[256, 256], + ), + learn=dict( + update_per_collect=20, + batch_size=32, + learning_rate=0.0005, + clip_value=5, + target_update_theta=0.008, + discount_factor=0.95, + ), + collect=dict( + collector=dict(get_train_sample=True, ), + n_episode=32, + unroll_len=10, + env_num=collector_env_num, + ), + command=dict(), + eval=dict( + env_num=evaluator_env_num, + evaluator=dict(eval_freq=1000, ), + ), + other=dict( + eps=dict( + type='linear', + start=1, + end=0.05, + decay=10000, + ), + replay_buffer=dict(replay_buffer_size=15000, ), + ), + ), +) + +main_config = EasyDict(main_config) +create_config = dict( + env=dict( + import_names=['dizoo.petting_zoo.envs.petting_zoo_simple_spread_env'], + type='petting_zoo', + ), + env_manager=dict(type='subprocess'), + policy=dict(type='madqn'), + collector=dict(type='episode'), +) +create_config = EasyDict(create_config) +ptz_simple_spread_madqn_config = main_config +ptz_simple_spread_madqn_create_config = create_config + +if __name__ == '__main__': + # or you can enter `ding -m serial_entry -c ptz_simple_spread_masac_config.py -s 0` + from ding.entry import serial_pipeline + serial_pipeline((main_config, create_config), seed=0) diff --git a/dizoo/petting_zoo/config/ptz_simple_spread_masac_config.py b/dizoo/petting_zoo/config/ptz_simple_spread_masac_config.py index ce000975c1..806d7608a1 100644 --- a/dizoo/petting_zoo/config/ptz_simple_spread_masac_config.py +++ b/dizoo/petting_zoo/config/ptz_simple_spread_masac_config.py @@ -24,25 +24,17 @@ cuda=True, on_policy=False, multi_agent=True, - # priority=True, - # priority_IS_weight=False, - random_collect_size=0, + random_collect_size=5000, model=dict( agent_obs_shape=2 + 2 + n_landmark * 2 + (n_agent - 1) * 2 + (n_agent - 1) * 2, global_obs_shape=2 + 2 + n_landmark * 2 + (n_agent - 1) * 2 + (n_agent - 1) * 2 + n_agent * (2 + 2) + n_landmark * 2 + n_agent * (n_agent - 1) * 2, action_shape=5, - # SAC concerned twin_critic=True, - actor_head_hidden_size=256, - critic_head_hidden_size=256, ), learn=dict( update_per_collect=50, batch_size=320, - # ============================================================== - # The following configs is algorithm-specific - # ============================================================== # learning_rates learning_rate_q=5e-4, learning_rate_policy=5e-4, @@ -51,16 +43,12 @@ discount_factor=0.99, alpha=0.2, auto_alpha=True, - log_space=True, - ignore_down=False, target_entropy=-2, ), collect=dict( n_sample=1600, - unroll_len=1, env_num=collector_env_num, ), - command=dict(), eval=dict( env_num=evaluator_env_num, evaluator=dict(eval_freq=50, ), @@ -84,7 +72,7 @@ type='petting_zoo', ), env_manager=dict(type='subprocess'), - policy=dict(type='sac_discrete'), + policy=dict(type='discrete_sac'), ) create_config = EasyDict(create_config) ptz_simple_spread_masac_config = main_config diff --git a/dizoo/petting_zoo/entry/ptz_simple_spread_eval.py b/dizoo/petting_zoo/entry/ptz_simple_spread_eval.py new file mode 100644 index 0000000000..e7f298ad39 --- /dev/null +++ b/dizoo/petting_zoo/entry/ptz_simple_spread_eval.py @@ -0,0 +1,12 @@ +from dizoo.petting_zoo.config.ptz_simple_spread_mappo_config import main_config, create_config +from ding.entry import eval + + +def main(): + ckpt_path = './ckpt_best.pth.tar' + replay_path = './replay_videos' + eval((main_config, create_config), seed=0, load_path=ckpt_path, replay_path=replay_path) + + +if __name__ == "__main__": + main() diff --git a/dizoo/petting_zoo/envs/petting_zoo_pistonball_env.py b/dizoo/petting_zoo/envs/petting_zoo_pistonball_env.py new file mode 100644 index 0000000000..d21ed049db --- /dev/null +++ b/dizoo/petting_zoo/envs/petting_zoo_pistonball_env.py @@ -0,0 +1,239 @@ +import copy +from functools import reduce +from typing import Dict, List, Optional + +import gymnasium as gym +import numpy as np +from ding.envs import BaseEnv, BaseEnvTimestep +from ding.envs.common.common_function import affine_transform +from ding.torch_utils import to_ndarray +from ding.utils import ENV_REGISTRY +from dizoo.petting_zoo.envs.petting_zoo_simple_spread_env import PTZRecordVideo +from pettingzoo.butterfly import pistonball_v6 + + +@ENV_REGISTRY.register('petting_zoo_pistonball') +class PettingZooPistonballEnv(BaseEnv): + """ + DI-engine PettingZoo environment adapter for the Pistonball environment. + This class integrates the `pistonball_v6` environment into the DI-engine + framework, supporting both continuous and discrete actions. + """ + + def __init__(self, cfg: dict) -> None: + self._cfg = cfg + self._init_flag = False + self._replay_path = None + self._num_pistons = self._cfg.get('n_pistons', 20) + self._continuous_actions = self._cfg.get('continuous_actions', False) + self._max_cycles = self._cfg.get('max_cycles', 125) + self._act_scale = self._cfg.get('act_scale', False) + self._agent_specific_global_state = self._cfg.get('agent_specific_global_state', False) + if self._act_scale: + assert self._continuous_actions, 'Action scaling only applies to continuous action spaces.' + self._channel_first = self._cfg.get('channel_first', True) + self.normalize_reward = self._cfg.normalize_reward + + def reset(self) -> np.ndarray: + """ + Resets the environment and returns the initial observations. + """ + if not self._init_flag: + # Initialize the pistonball environment + parallel_env = pistonball_v6.parallel_env + self._env = parallel_env( + n_pistons=self._num_pistons, continuous=self._continuous_actions, max_cycles=self._max_cycles + ) + self._env.reset() + self._agents = self._env.agents + + # Define action and observation spaces + self._action_space = gym.spaces.Dict({agent: self._env.action_space(agent) for agent in self._agents}) + single_agent_obs_space = self._env.observation_space(self._agents[0]) + single_agent_action_space = self._env.action_space(self._agents[0]) + + if isinstance(single_agent_action_space, gym.spaces.Box): + self._action_dim = single_agent_action_space.shape + elif isinstance(single_agent_action_space, gym.spaces.Discrete): + self._action_dim = (single_agent_action_space.n, ) + else: + raise Exception('Only support `Box` or `Discrete` obs space for single agent.') + + if isinstance(single_agent_obs_space, gym.spaces.Box): + self._obs_shape = single_agent_obs_space.shape + else: + raise ValueError("Only support `Box` observation space for each agent.") + + self._observation_space = gym.spaces.Box( + low=0, high=255, shape=(self._num_pistons, *self._obs_shape), dtype=np.uint8 + ) + + self._reward_space = gym.spaces.Dict( + { + agent: gym.spaces.Box(low=float('-inf'), high=float('inf'), shape=(1, ), dtype=np.float32) + for agent in self._agents + } + ) + + if self._replay_path is not None: + self._env.render_mode = 'rgb_array' + self._env = PTZRecordVideo( + self._env, self._replay_path, name_prefix=f'rl-video-{id(self)}', disable_logger=True + ) + self._init_flag = True + + if hasattr(self, '_seed'): + obs = self._env.reset(seed=self._seed) + else: + obs = self._env.reset() + + self._eval_episode_return = 0.0 + self._step_count = 0 + obs_n = self._process_obs(obs) + return obs_n + + def close(self) -> None: + """ + Closes the environment. + """ + if self._init_flag: + self._env.close() + self._init_flag = False + + def render(self) -> None: + """ + Renders the environment. + """ + self._env.render() + + def seed(self, seed: int, dynamic_seed: bool = True) -> None: + """ + Sets the seed for the environment. + """ + self._seed = seed + self._dynamic_seed = dynamic_seed + np.random.seed(self._seed) + + def step(self, action: np.ndarray) -> BaseEnvTimestep: + """ + Steps through the environment using the provided action. + """ + self._step_count += 1 + assert isinstance(action, np.ndarray), type(action) + action = self._process_action(action) + if self._act_scale: + for agent in self._agents: + action[agent] = affine_transform( + action[agent], min_val=self.action_space[agent].low, max_val=self.action_space[agent].high + ) + + obs, rew, done, trunc, info = self._env.step(action) + obs_n = self._process_obs(obs) + rew_n = np.array([sum([rew[agent] for agent in self._agents])]) + rew_n = rew_n.astype(np.float32) + + if self.normalize_reward: + # TODO: more elegant scale factor + rew_n = rew_n / (self._num_pistons * 50) + + self._eval_episode_return += rew_n.item() + + done_n = reduce(lambda x, y: x and y, done.values()) or self._step_count >= self._max_cycles + if done_n: + info['eval_episode_return'] = self._eval_episode_return + + return BaseEnvTimestep(obs_n, rew_n, done_n, info) + + def enable_save_replay(self, replay_path: Optional[str] = None) -> None: + """ + Enables video recording during the episode. + """ + if replay_path is None: + replay_path = './video' + self._replay_path = replay_path + + def _process_obs(self, obs: Dict[str, np.ndarray]) -> np.ndarray: + """ + Processes the observations into the required format. + """ + # Process agent observations, transpose if channel_first is True + obs = np.array( + [np.transpose(obs[agent], (2, 0, 1)) if self._channel_first else obs[agent] for agent in self._agents], + dtype=np.uint8 + ) + + # Return only agent observations if configured to do so + if self._cfg.get('agent_obs_only', False): + return obs + + # Initialize return dictionary + ret = {'agent_state': (obs / 255.0).astype(np.float32)} + + # Obtain global state, transpose if channel_first is True + global_state = self._env.state() + if self._channel_first: + global_state = global_state.transpose(2, 0, 1) + ret['global_state'] = (global_state / 255.0).astype(np.float32) + + # Handle agent-specific global states by repeating the global state for each agent + if self._agent_specific_global_state: + ret['global_state'] = np.tile(np.expand_dims(ret['global_state'], axis=0), (self._num_pistons, 1, 1, 1)) + + # Set action mask for each agent + ret['action_mask'] = np.ones((self._num_pistons, *self._action_dim), dtype=np.float32) + + return ret + + def _process_action(self, action: np.ndarray) -> Dict[str, np.ndarray]: + """ + Processes the action array into a dictionary format for each agent. + """ + dict_action = {} + for i, agent in enumerate(self._agents): + dict_action[agent] = action[i] + return dict_action + + def random_action(self) -> np.ndarray: + """ + Generates a random action for each agent. + """ + random_action = self.action_space.sample() + for k in random_action: + if isinstance(random_action[k], np.ndarray): + pass + elif isinstance(random_action[k], int): + random_action[k] = to_ndarray([random_action[k]], dtype=np.int64) + return random_action + + @property + def agents(self) -> List[str]: + return self._agents + + @property + def observation_space(self) -> gym.spaces.Space: + return self._observation_space + + @property + def action_space(self) -> gym.spaces.Space: + return self._action_space + + @property + def reward_space(self) -> gym.spaces.Space: + return self._reward_space + + @staticmethod + def create_collector_env_cfg(cfg: dict) -> List[dict]: + collector_env_num = cfg.pop('collector_env_num') + cfg = copy.deepcopy(cfg) + cfg.normalize_reward = True + return [cfg for _ in range(collector_env_num)] + + @staticmethod + def create_evaluator_env_cfg(cfg: dict) -> List[dict]: + evaluator_env_num = cfg.pop('evaluator_env_num') + cfg = copy.deepcopy(cfg) + cfg.normalize_reward = False + return [cfg for _ in range(evaluator_env_num)] + + def __repr__(self) -> str: + return "DI-engine PettingZoo Pistonball Env" diff --git a/dizoo/petting_zoo/envs/petting_zoo_simple_spread_env.py b/dizoo/petting_zoo/envs/petting_zoo_simple_spread_env.py index f530cddd83..4f5916f33b 100644 --- a/dizoo/petting_zoo/envs/petting_zoo_simple_spread_env.py +++ b/dizoo/petting_zoo/envs/petting_zoo_simple_spread_env.py @@ -1,5 +1,5 @@ from typing import Any, List, Union, Optional, Dict -import gym +import gymnasium as gym import numpy as np import pettingzoo from functools import reduce @@ -8,6 +8,52 @@ from ding.torch_utils import to_ndarray, to_list from ding.envs.common.common_function import affine_transform from ding.utils import ENV_REGISTRY, import_module +from pettingzoo.utils.conversions import parallel_wrapper_fn +from pettingzoo.mpe._mpe_utils.simple_env import SimpleEnv, make_env +from pettingzoo.mpe.simple_spread.simple_spread import Scenario + + +# Custom wrapper for recording videos in PettingZoo environments +class PTZRecordVideo(gym.wrappers.RecordVideo): + + def step(self, action): + """Steps through the environment using action, recording observations if :attr:`self.recording`.""" + # gymnasium==0.27.1 + observations, rewards, terminateds, truncateds, infos = self.env.step(action) + + # Because pettingzoo returns a dict of terminated and truncated, we need to check if any of the values are True + if not (self.terminated is True or self.truncated is True): # the first location for modifications + # increment steps and episodes + self.step_id += 1 + if not self.is_vector_env: + if terminateds or truncateds: + self.episode_id += 1 + self.terminated = terminateds + self.truncated = truncateds + elif terminateds[0] or truncateds[0]: + self.episode_id += 1 + self.terminated = terminateds[0] + self.truncated = truncateds[0] + + # Capture the video frame if recording + if self.recording: + assert self.video_recorder is not None + self.video_recorder.capture_frame() + self.recorded_frames += 1 + if self.video_length > 0: + if self.recorded_frames > self.video_length: + self.close_video_recorder() + else: + if not self.is_vector_env: + if terminateds is True or truncateds is True: # the second location for modifications + self.close_video_recorder() + elif terminateds[0] or truncateds[0]: + self.close_video_recorder() + + elif self._video_enabled(): + self.start_video_recorder() + + return observations, rewards, terminateds, truncateds, infos @ENV_REGISTRY.register('petting_zoo') @@ -32,10 +78,10 @@ def __init__(self, cfg: dict) -> None: def reset(self) -> np.ndarray: if not self._init_flag: - # In order to align with the simple spread in Multiagent Particle Env (MPE), - # instead of adopting the pettingzoo interface directly, + # In order to align with the simple spread in Multiagent Particle Env (MPE), + # instead of adopting the pettingzoo interface directly, # we have redefined the way rewards are calculated - + # import_module(['pettingzoo.{}.{}'.format(self._env_family, self._env_id)]) # self._env = pettingzoo.__dict__[self._env_family].__dict__[self._env_id].parallel_env( # N=self._cfg.n_agent, continuous_actions=self._continuous_actions, max_cycles=self._max_cycles @@ -48,28 +94,17 @@ def reset(self) -> np.ndarray: self._env = parallel_env( N=self._cfg.n_agent, continuous_actions=self._continuous_actions, max_cycles=self._max_cycles ) - # dynamic seed reduces training speed greatly - # if hasattr(self, '_seed') and hasattr(self, '_dynamic_seed') and self._dynamic_seed: - # np_seed = 100 * np.random.randint(1, 1000) - # self._env.seed(self._seed + np_seed) - if hasattr(self, '_seed'): - self._env.seed(self._seed) - if self._replay_path is not None: - self._env = gym.wrappers.Monitor( - self._env, self._replay_path, video_callable=lambda episode_id: True, force=True - ) - obs = self._env.reset() - if not self._init_flag: + self._env.reset() self._agents = self._env.agents self._action_space = gym.spaces.Dict({agent: self._env.action_space(agent) for agent in self._agents}) - single_agent_obs_space = self._env.action_space(self._agents[0]) - if isinstance(single_agent_obs_space, gym.spaces.Box): - self._action_dim = single_agent_obs_space.shape - elif isinstance(single_agent_obs_space, gym.spaces.Discrete): - self._action_dim = (single_agent_obs_space.n, ) + single_agent_action_space = self._env.action_space(self._agents[0]) + if isinstance(single_agent_action_space, gym.spaces.Box): + self._action_dim = single_agent_action_space.shape + elif isinstance(single_agent_action_space, gym.spaces.Discrete): + self._action_dim = (single_agent_action_space.n, ) else: - raise Exception('Only support `Box` or `Discrte` obs space for single agent.') + raise Exception('Only support `Box` or `Discrete` obs space for single agent.') # only for env 'simple_spread_v2', n_agent = 5 # now only for the case that each agent in the team have the same obs structure and corresponding shape. @@ -140,9 +175,18 @@ def reset(self) -> np.ndarray: for agent in self._agents } ) + if self._replay_path is not None: + self._env.render_mode = 'rgb_array' + self._env = PTZRecordVideo( + self._env, self._replay_path, name_prefix=f'rl-video-{id(self)}', disable_logger=True + ) self._init_flag = True - # self._final_eval_reward = {agent: 0. for agent in self._agents} - self._final_eval_reward = 0. + if hasattr(self, '_seed'): + obs = self._env.reset(seed=self._seed) + else: + obs = self._env.reset() + # self._eval_episode_return = {agent: 0. for agent in self._agents} + self._eval_episode_return = 0. self._step_count = 0 obs_n = self._process_obs(obs) return obs_n @@ -173,16 +217,17 @@ def step(self, action: np.ndarray) -> BaseEnvTimestep: action[agent], min_val=self.action_space[agent].low, max_val=self.action_space[agent].high ) - obs, rew, done, info = self._env.step(action) + obs, rew, done, trunc, info = self._env.step(action) obs_n = self._process_obs(obs) rew_n = np.array([sum([rew[agent] for agent in self._agents])]) + rew_n = rew_n.astype(np.float32) # collide_sum = 0 # for i in range(self._num_agents): # collide_sum += info['n'][i][1] # collide_penalty = self._cfg.get('collide_penal', self._num_agent) # rew_n += collide_sum * (1.0 - collide_penalty) # rew_n = rew_n / (self._cfg.get('max_cycles', 25) * self._num_agent) - self._final_eval_reward += rew_n.item() + self._eval_episode_return += rew_n.item() # occupied_landmarks = info['n'][0][3] # if self._step_count >= self._max_step or occupied_landmarks >= self._n_agent \ @@ -193,9 +238,9 @@ def step(self, action: np.ndarray) -> BaseEnvTimestep: done_n = reduce(lambda x, y: x and y, done.values()) or self._step_count >= self._max_cycles # for agent in self._agents: - # self._final_eval_reward[agent] += rew[agent] + # self._eval_episode_return[agent] += rew[agent] if done_n: # or reduce(lambda x, y: x and y, done.values()) - info['final_eval_reward'] = self._final_eval_reward + info['eval_episode_return'] = self._eval_episode_return # for agent in rew: # rew[agent] = to_ndarray([rew[agent]]) return BaseEnvTimestep(obs_n, rew_n, done_n, info) @@ -267,7 +312,7 @@ def _process_obs(self, obs: 'torch.Tensor') -> np.ndarray: # noqa 1 ) # action_mask: All actions are of use(either 1 for discrete or 5 for continuous). Thus all 1. - ret['action_mask'] = np.ones((self._num_agents, *self._action_dim)) + ret['action_mask'] = np.ones((self._num_agents, *self._action_dim)).astype(np.float32) return ret def _process_action(self, action: 'torch.Tensor') -> Dict[str, np.ndarray]: # noqa @@ -308,18 +353,14 @@ def reward_space(self) -> gym.spaces.Space: return self._reward_space -from pettingzoo.utils.conversions import parallel_wrapper_fn -from pettingzoo.mpe._mpe_utils.simple_env import SimpleEnv, make_env -from pettingzoo.mpe.scenarios.simple_spread import Scenario - - class simple_spread_raw_env(SimpleEnv): def __init__(self, N=3, local_ratio=0.5, max_cycles=25, continuous_actions=False): assert 0. <= local_ratio <= 1., "local_ratio is a proportion. Must be between 0 and 1." scenario = Scenario() world = scenario.make_world(N) - super().__init__(scenario, world, max_cycles, continuous_actions, local_ratio) + super().__init__(scenario, world, max_cycles, continuous_actions=continuous_actions, local_ratio=local_ratio) + self.render_mode = 'rgb_array' self.metadata['name'] = "simple_spread_v2" def _execute_world_step(self): @@ -355,3 +396,17 @@ def _execute_world_step(self): reward = agent_reward self.rewards[agent.name] = reward + + def render(self): + if self.render_mode is None: + gym.logger.warn("You are calling render method without specifying any render mode.") + return + import pygame + + self.enable_render(self.render_mode) + + self.draw() + observation = np.array(pygame.surfarray.pixels3d(self.screen)) + if self.render_mode == "human": + pygame.display.flip() + return (np.transpose(observation, axes=(1, 0, 2)) if self.render_mode == "rgb_array" else None) diff --git a/dizoo/petting_zoo/envs/test_petting_zoo_pistonball_env.py b/dizoo/petting_zoo/envs/test_petting_zoo_pistonball_env.py new file mode 100644 index 0000000000..1e9a88292d --- /dev/null +++ b/dizoo/petting_zoo/envs/test_petting_zoo_pistonball_env.py @@ -0,0 +1,106 @@ +from easydict import EasyDict +import pytest +import numpy as np +from dizoo.petting_zoo.envs.petting_zoo_pistonball_env import PettingZooPistonballEnv + + +@pytest.mark.envtest +class TestPettingZooPistonballEnv: + + def test_agent_obs_only(self): + n_pistons = 20 + env = PettingZooPistonballEnv( + EasyDict( + dict( + n_pistons=n_pistons, + max_cycles=125, + agent_obs_only=True, + continuous_actions=True, + act_scale=False, + ) + ) + ) + env.seed(123) + assert env._seed == 123 + obs = env.reset() + assert obs.shape == (n_pistons, 3, 457, 120) + for i in range(10): + random_action = env.random_action() + random_action = np.array([random_action[agent] for agent in random_action]) + timestep = env.step(random_action) + # print(timestep) + assert isinstance(timestep.obs, np.ndarray), timestep.obs + assert timestep.obs.shape == (n_pistons, 3, 457, 120) + assert isinstance(timestep.done, bool), timestep.done + assert isinstance(timestep.reward, np.ndarray), timestep.reward + assert timestep.reward.dtype == np.float32 + print(env.observation_space, env.action_space, env.reward_space) + env.close() + + def test_dict_obs(self): + n_pistons = 20 + env = PettingZooPistonballEnv( + EasyDict( + dict( + n_pistons=n_pistons, + max_cycles=125, + agent_obs_only=False, + agent_specific_global_state=False, + continuous_actions=True, + act_scale=False, + ) + ) + ) + env.seed(123) + assert env._seed == 123 + obs = env.reset() + for k, v in obs.items(): + print(k, v.shape) + for i in range(10): + random_action = env.random_action() + random_action = np.array([random_action[agent] for agent in random_action]) + timestep = env.step(random_action) + # print(timestep) + assert isinstance(timestep.obs['agent_state'], np.ndarray), timestep.obs['agent_state'] + assert isinstance(timestep.obs['global_state'], np.ndarray), timestep.obs['global_state'] + assert timestep.obs['agent_state'].shape == (n_pistons, 3, 457, 120) + assert timestep.obs['global_state'].shape == (3, 560, 880) + assert isinstance(timestep.done, bool), timestep.done + assert isinstance(timestep.reward, np.ndarray), timestep.reward + print(env.observation_space, env.action_space, env.reward_space) + env.close() + + def test_agent_specific_global_state(self): + n_pistons = 20 + env = PettingZooPistonballEnv( + EasyDict( + dict( + n_pistons=n_pistons, + max_cycles=125, + agent_obs_only=False, + continuous_actions=True, + agent_specific_global_state=True, + act_scale=False, + ) + ) + ) + env.seed(123) + assert env._seed == 123 + obs = env.reset() + for k, v in obs.items(): + print(k, v.shape) + for i in range(10): + random_action = env.random_action() + random_action = np.array([random_action[agent] for agent in random_action]) + timestep = env.step(random_action) + # print(timestep) + assert isinstance(timestep.obs['agent_state'], np.ndarray), timestep.obs['agent_state'] + assert isinstance(timestep.obs['global_state'], np.ndarray), timestep.obs['global_state'] + assert timestep.obs['agent_state'].shape == (n_pistons, 3, 457, 120) + assert timestep.obs['global_state'].shape == (n_pistons, 3, 560, 880) + assert timestep.obs['global_state'].shape == (n_pistons, 3, 560, 880) + + assert isinstance(timestep.done, bool), timestep.done + assert isinstance(timestep.reward, np.ndarray), timestep.reward + print(env.observation_space, env.action_space, env.reward_space) + env.close() diff --git a/dizoo/petting_zoo/envs/test_petting_zoo_simple_spread_env.py b/dizoo/petting_zoo/envs/test_petting_zoo_simple_spread_env.py index 9adc71e7f9..22117cf85f 100644 --- a/dizoo/petting_zoo/envs/test_petting_zoo_simple_spread_env.py +++ b/dizoo/petting_zoo/envs/test_petting_zoo_simple_spread_env.py @@ -1,10 +1,10 @@ +from easydict import EasyDict +import pytest import numpy as np import pettingzoo from ding.utils import import_module -from easydict import EasyDict -import pytest -from dizoo.petting_zoo.envs.petting_zoo_env import PettingZooEnv +from dizoo.petting_zoo.envs.petting_zoo_simple_spread_env import PettingZooEnv @pytest.mark.envtest @@ -39,6 +39,7 @@ def test_agent_obs_only(self): assert timestep.obs.shape == (n_agent, 2 + 2 + (n_agent - 1) * 2 + n_agent * 2 + (n_agent - 1) * 2) assert isinstance(timestep.done, bool), timestep.done assert isinstance(timestep.reward, np.ndarray), timestep.reward + assert timestep.reward.dtype == np.float32 print(env.observation_space, env.action_space, env.reward_space) env.close() @@ -80,6 +81,7 @@ def test_dict_obs(self): assert timestep.obs['agent_alone_padding_state'].shape == ( n_agent, 2 + 2 + n_landmark * 2 + (n_agent - 1) * 2 + (n_agent - 1) * 2 ) + assert timestep.obs['action_mask'].dtype == np.float32 assert isinstance(timestep.done, bool), timestep.done assert isinstance(timestep.reward, np.ndarray), timestep.reward print(env.observation_space, env.action_space, env.reward_space) diff --git a/dizoo/pomdp/envs/atari_env.py b/dizoo/pomdp/envs/atari_env.py index 13b4cc733b..d73ec6008d 100644 --- a/dizoo/pomdp/envs/atari_env.py +++ b/dizoo/pomdp/envs/atari_env.py @@ -55,7 +55,7 @@ def reset(self) -> Sequence: self._env.seed(self._seed) obs = self._env.reset() obs = to_ndarray(obs) - self._final_eval_reward = 0. + self._eval_episode_return = 0. return obs def close(self) -> None: @@ -72,11 +72,11 @@ def step(self, action: np.ndarray) -> BaseEnvTimestep: assert isinstance(action, np.ndarray), type(action) action = action.item() obs, rew, done, info = self._env.step(action) - self._final_eval_reward += rew + self._eval_episode_return += rew obs = to_ndarray(obs) rew = to_ndarray([rew]) # wrapped to be transfered to a array with shape (1,) if done: - info['final_eval_reward'] = self._final_eval_reward + info['eval_episode_return'] = self._eval_episode_return return BaseEnvTimestep(obs, rew, done, info) def _make_env(self, only_info=False): diff --git a/dizoo/pomdp/envs/test_atari_env.py b/dizoo/pomdp/envs/test_atari_env.py index 67ea102bbd..d18a98fac1 100644 --- a/dizoo/pomdp/envs/test_atari_env.py +++ b/dizoo/pomdp/envs/test_atari_env.py @@ -30,6 +30,6 @@ def test_env(): assert timestep.reward.shape == (1, ) # assert isinstance(timestep, tuple) if timestep.done: - assert 'final_eval_reward' in timestep.info, timestep.info + assert 'eval_episode_return' in timestep.info, timestep.info break pong_env.close() diff --git a/dizoo/procgen/README.md b/dizoo/procgen/README.md index 9631f75122..4582510a05 100644 --- a/dizoo/procgen/README.md +++ b/dizoo/procgen/README.md @@ -10,11 +10,11 @@ Procedural generation controls the number of platform sections, their correspond ## Train Coinrun with DI-engine -DI-engine can achive 10 return on average within 2M episodes by DQN. The tuned example can be found in `dizoo/procgen/entry/coinrun_dqn_config.py`. The training episode reward is as follows. +DI-engine can achive 10 return on average within 2M episodes by DQN. The tuned example can be found in `dizoo/procgen/entry/coinrun_dqn_config.py`. The training episode return is as follows. ![tb](./coinrun_dqn.svg) -DI-engine can achive 10 return on average within 2M episodes by PPO. The tuned example can be found in `dizoo/procgen/entry/coinrun_ppo_config.py`. The training episode reward is as follows. +DI-engine can achive 10 return on average within 2M episodes by PPO. The tuned example can be found in `dizoo/procgen/entry/coinrun_ppo_config.py`. The training episode return is as follows. ![tb](./coinrun_ppo.svg) @@ -28,6 +28,6 @@ Procedural generation controls the level layout by generating mazes using Kruska ## Train Maze with DI-engine -DI-engine can achive 10 return on average within 7M episodes by DQN. The tuned example can be found in `dizoo/procgen/entry/maze_dqn_config.py`. The training episode reward is as follows. +DI-engine can achive 10 return on average within 7M episodes by DQN. The tuned example can be found in `dizoo/procgen/entry/maze_dqn_config.py`. The training episode return is as follows. ![tb](./maze_dqn.svg) diff --git a/dizoo/procgen/__init__.py b/dizoo/procgen/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/dizoo/procgen/entry/__init__.py b/dizoo/procgen/config/__init__.py similarity index 100% rename from dizoo/procgen/entry/__init__.py rename to dizoo/procgen/config/__init__.py diff --git a/dizoo/procgen/entry/bigfish_plr_config.py b/dizoo/procgen/config/bigfish_plr_config.py similarity index 91% rename from dizoo/procgen/entry/bigfish_plr_config.py rename to dizoo/procgen/config/bigfish_plr_config.py index 4feb12d877..d39afd7097 100644 --- a/dizoo/procgen/entry/bigfish_plr_config.py +++ b/dizoo/procgen/config/bigfish_plr_config.py @@ -1,6 +1,6 @@ from easydict import EasyDict -bigfish_plr_default_config = dict( +bigfish_plr_config = dict( exp_name='bigfish_plr_seed1', env=dict( is_train=True, @@ -42,8 +42,8 @@ temperature=0.1, ), ) -bigfish_plr_default_config = EasyDict(bigfish_plr_default_config) -main_config = bigfish_plr_default_config +bigfish_plr_config = EasyDict(bigfish_plr_config) +main_config = bigfish_plr_config bigfish_plr_create_config = dict( env=dict( diff --git a/dizoo/procgen/entry/bigfish_ppg_config.py b/dizoo/procgen/config/bigfish_ppg_config.py similarity index 91% rename from dizoo/procgen/entry/bigfish_ppg_config.py rename to dizoo/procgen/config/bigfish_ppg_config.py index e519787b0c..19ac4be698 100644 --- a/dizoo/procgen/entry/bigfish_ppg_config.py +++ b/dizoo/procgen/config/bigfish_ppg_config.py @@ -1,6 +1,6 @@ from easydict import EasyDict -bigfish_ppg_default_config = dict( +bigfish_ppg_config = dict( exp_name='bigfish_ppg_seed0', env=dict( is_train=True, @@ -37,8 +37,8 @@ other=dict(), ), ) -bigfish_ppg_default_config = EasyDict(bigfish_ppg_default_config) -main_config = bigfish_ppg_default_config +bigfish_ppg_config = EasyDict(bigfish_ppg_config) +main_config = bigfish_ppg_config bigfish_ppg_create_config = dict( env=dict( diff --git a/dizoo/procgen/entry/coinrun_dqn_config.py b/dizoo/procgen/config/coinrun_dqn_config.py similarity index 89% rename from dizoo/procgen/entry/coinrun_dqn_config.py rename to dizoo/procgen/config/coinrun_dqn_config.py index 1a7f7ee653..4e788672b6 100644 --- a/dizoo/procgen/entry/coinrun_dqn_config.py +++ b/dizoo/procgen/config/coinrun_dqn_config.py @@ -1,6 +1,6 @@ from easydict import EasyDict -coinrun_dqn_default_config = dict( +coinrun_dqn_config = dict( env=dict( env_id='coinrun', collector_env_num=4, @@ -36,8 +36,8 @@ ), ), ) -coinrun_dqn_default_config = EasyDict(coinrun_dqn_default_config) -main_config = coinrun_dqn_default_config +coinrun_dqn_config = EasyDict(coinrun_dqn_config) +main_config = coinrun_dqn_config coinrun_dqn_create_config = dict( env=dict( diff --git a/dizoo/procgen/entry/coinrun_ppg_config.py b/dizoo/procgen/config/coinrun_ppg_config.py similarity index 91% rename from dizoo/procgen/entry/coinrun_ppg_config.py rename to dizoo/procgen/config/coinrun_ppg_config.py index 785076a738..793c2128b0 100644 --- a/dizoo/procgen/entry/coinrun_ppg_config.py +++ b/dizoo/procgen/config/coinrun_ppg_config.py @@ -1,6 +1,6 @@ from easydict import EasyDict -coinrun_ppg_default_config = dict( +coinrun_ppg_config = dict( exp_name='coinrun_ppg_seed0', env=dict( is_train=True, @@ -37,8 +37,8 @@ other=dict(), ), ) -coinrun_ppg_default_config = EasyDict(coinrun_ppg_default_config) -main_config = coinrun_ppg_default_config +coinrun_ppg_config = EasyDict(coinrun_ppg_config) +main_config = coinrun_ppg_config coinrun_ppg_create_config = dict( env=dict( diff --git a/dizoo/procgen/entry/coinrun_ppo_config.py b/dizoo/procgen/config/coinrun_ppo_config.py similarity index 90% rename from dizoo/procgen/entry/coinrun_ppo_config.py rename to dizoo/procgen/config/coinrun_ppo_config.py index 160e8624fe..4a04396b42 100644 --- a/dizoo/procgen/entry/coinrun_ppo_config.py +++ b/dizoo/procgen/config/coinrun_ppo_config.py @@ -1,6 +1,6 @@ from easydict import EasyDict -coinrun_ppo_default_config = dict( +coinrun_ppo_config = dict( env=dict( is_train=True, env_id='coinrun', @@ -39,8 +39,8 @@ ), ), ) -coinrun_ppo_default_config = EasyDict(coinrun_ppo_default_config) -main_config = coinrun_ppo_default_config +coinrun_ppo_config = EasyDict(coinrun_ppo_config) +main_config = coinrun_ppo_config coinrun_ppo_create_config = dict( env=dict( diff --git a/dizoo/procgen/entry/maze_dqn_config.py b/dizoo/procgen/config/maze_dqn_config.py similarity index 90% rename from dizoo/procgen/entry/maze_dqn_config.py rename to dizoo/procgen/config/maze_dqn_config.py index 4dacb44f75..a4cf674ca1 100644 --- a/dizoo/procgen/entry/maze_dqn_config.py +++ b/dizoo/procgen/config/maze_dqn_config.py @@ -1,6 +1,6 @@ from easydict import EasyDict -maze_dqn_default_config = dict( +maze_dqn_config = dict( env=dict( collector_env_num=4, env_id='maze', @@ -37,8 +37,8 @@ ), ), ) -maze_dqn_default_config = EasyDict(maze_dqn_default_config) -main_config = maze_dqn_default_config +maze_dqn_config = EasyDict(maze_dqn_config) +main_config = maze_dqn_config maze_dqn_create_config = dict( env=dict( diff --git a/dizoo/procgen/entry/maze_ppg_config.py b/dizoo/procgen/config/maze_ppg_config.py similarity index 92% rename from dizoo/procgen/entry/maze_ppg_config.py rename to dizoo/procgen/config/maze_ppg_config.py index d651be66bf..bb51dfd85d 100644 --- a/dizoo/procgen/entry/maze_ppg_config.py +++ b/dizoo/procgen/config/maze_ppg_config.py @@ -1,6 +1,6 @@ from easydict import EasyDict -maze_ppg_default_config = dict( +maze_ppg_config = dict( exp_name='maze_ppg_seed0', env=dict( is_train=True, @@ -40,8 +40,8 @@ other=dict(), ), ) -maze_ppg_default_config = EasyDict(maze_ppg_default_config) -main_config = maze_ppg_default_config +maze_ppg_config = EasyDict(maze_ppg_config) +main_config = maze_ppg_config maze_ppg_create_config = dict( env=dict( diff --git a/dizoo/procgen/entry/maze_ppo_config.py b/dizoo/procgen/config/maze_ppo_config.py similarity index 91% rename from dizoo/procgen/entry/maze_ppo_config.py rename to dizoo/procgen/config/maze_ppo_config.py index 1970fd36fc..d116305f76 100644 --- a/dizoo/procgen/entry/maze_ppo_config.py +++ b/dizoo/procgen/config/maze_ppo_config.py @@ -1,6 +1,6 @@ from easydict import EasyDict -maze_ppo_default_config = dict( +maze_ppo_config = dict( env=dict( # frame_stack=4, is_train=True, @@ -40,8 +40,8 @@ ), ), ) -maze_ppo_default_config = EasyDict(maze_ppo_default_config) -main_config = maze_ppo_default_config +maze_ppo_config = EasyDict(maze_ppo_config) +main_config = maze_ppo_config maze_ppo_create_config = dict( env=dict( diff --git a/dizoo/procgen/entry/coinrun_onppo_main.py b/dizoo/procgen/entry/coinrun_onppo_main.py new file mode 100644 index 0000000000..ca132b1fa7 --- /dev/null +++ b/dizoo/procgen/entry/coinrun_onppo_main.py @@ -0,0 +1,113 @@ +import os +from functools import partial + +import gym +import numpy as np +from easydict import EasyDict +from tensorboardX import SummaryWriter + +from ding.torch_utils import to_ndarray +from ding.worker import BaseLearner, SampleSerialCollector, InteractionSerialEvaluator +from ding.model import VAC +from ding.policy import PPOPolicy +from ding.envs import DingEnvWrapper, EvalEpisodeReturnWrapper, BaseEnvManager +from ding.config import compile_config +from ding.utils import set_pkg_seed +from dizoo.procgen.config.coinrun_ppo_config import coinrun_ppo_config + + +class CoinrunWrapper(gym.Wrapper): + + def __init__(self, env, cfg): + super().__init__(env) + cfg = EasyDict(cfg) + self._cfg = cfg + self._observation_space = gym.spaces.Box( + low=np.zeros(shape=(3, 64, 64)), high=np.ones(shape=(3, 64, 64)) * 255, shape=(3, 64, 64), dtype=np.float32 + ) + self._action_space = gym.spaces.Discrete(15) + self._reward_space = gym.spaces.Box(low=float("-inf"), high=float("inf"), shape=(1, ), dtype=np.float32) + + def _process_obs(self, obs): + obs = to_ndarray(obs) + obs = np.transpose(obs, (2, 0, 1)) + obs = obs.astype(np.float32) + return obs + + def step(self, action): + obs, reward, done, info = self.env.step(action) + return self._process_obs(obs), reward, bool(done), info + + def reset(self): + obs = self.env.reset() + return self._process_obs(obs) + + +def wrapped_procgen_env(cfg): + default_cfg = dict( + control_level=True, + start_level=0, + num_levels=0, + env_id='coinrun', + ) + default_cfg.update(cfg) + default_cfg = EasyDict(default_cfg) + + return DingEnvWrapper( + gym.make( + 'procgen:procgen-' + default_cfg.env_id + '-v0', + start_level=default_cfg.start_level, + num_levels=default_cfg.num_levels + ) if default_cfg.control_level else + gym.make('procgen:procgen-' + default_cfg.env_id + '-v0', start_level=0, num_levels=1), + cfg={ + 'env_wrapper': [ + lambda env: CoinrunWrapper(env, default_cfg), + lambda env: EvalEpisodeReturnWrapper(env), + ] + } + ) + + +def main(cfg, seed=0, max_env_step=int(1e10), max_train_iter=int(1e10)): + cfg = compile_config( + cfg, BaseEnvManager, PPOPolicy, BaseLearner, SampleSerialCollector, InteractionSerialEvaluator, save_cfg=True + ) + collector_env_num, evaluator_env_num = cfg.env.collector_env_num, cfg.env.evaluator_env_num + collector_env = BaseEnvManager( + env_fn=[partial(wrapped_procgen_env, cfg=coinrun_ppo_config.env) for _ in range(collector_env_num)], + cfg=cfg.env.manager + ) + evaluator_env = BaseEnvManager( + env_fn=[partial(wrapped_procgen_env, cfg=coinrun_ppo_config.env) for _ in range(evaluator_env_num)], + cfg=cfg.env.manager + ) + + collector_env.seed(seed) + evaluator_env.seed(seed, dynamic_seed=False) + set_pkg_seed(seed, use_cuda=cfg.policy.cuda) + + model = VAC(**cfg.policy.model) + policy = PPOPolicy(cfg.policy, model=model) + tb_logger = SummaryWriter(os.path.join('./{}/log/'.format(cfg.exp_name), 'serial')) + learner = BaseLearner(cfg.policy.learn.learner, policy.learn_mode, tb_logger, exp_name=cfg.exp_name) + collector = SampleSerialCollector( + cfg.policy.collect.collector, collector_env, policy.collect_mode, tb_logger, exp_name=cfg.exp_name + ) + evaluator = InteractionSerialEvaluator( + cfg.policy.eval.evaluator, evaluator_env, policy.eval_mode, tb_logger, exp_name=cfg.exp_name + ) + + while True: + if evaluator.should_eval(learner.train_iter): + stop, reward = evaluator.eval(learner.save_checkpoint, learner.train_iter, collector.envstep) + if stop: + break + new_data = collector.collect(train_iter=learner.train_iter) + learner.train(new_data, collector.envstep) + if collector.envstep >= max_env_step or learner.train_iter >= max_train_iter: + break + + +if __name__ == '__main__': + main(coinrun_ppo_config) diff --git a/dizoo/procgen/envs/procgen_env.py b/dizoo/procgen/envs/procgen_env.py index 165a82c0a8..4b194f1d82 100644 --- a/dizoo/procgen/envs/procgen_env.py +++ b/dizoo/procgen/envs/procgen_env.py @@ -59,7 +59,7 @@ def reset(self) -> np.ndarray: self._env = gym.make(self._env_name, start_level=self._start_level, num_levels=self._num_levels) else: self._env = gym.make(self._env_name, start_level=self._seed, num_levels=1) - self._final_eval_reward = 0 + self._eval_episode_return = 0 obs = self._env.reset() obs = to_ndarray(obs) obs = np.transpose(obs, (2, 0, 1)) @@ -80,9 +80,9 @@ def step(self, action: np.ndarray) -> BaseEnvTimestep: if action.shape == (1, ): action = action.squeeze() # 0-dim array obs, rew, done, info = self._env.step(action) - self._final_eval_reward += rew + self._eval_episode_return += rew if done: - info['final_eval_reward'] = self._final_eval_reward + info['eval_episode_return'] = self._eval_episode_return obs = to_ndarray(obs) obs = np.transpose(obs, (2, 0, 1)) obs = obs.astype(np.float32) diff --git a/dizoo/pybullet/envs/pybullet_env.py b/dizoo/pybullet/envs/pybullet_env.py index b74c93cb85..25def74a23 100644 --- a/dizoo/pybullet/envs/pybullet_env.py +++ b/dizoo/pybullet/envs/pybullet_env.py @@ -310,7 +310,7 @@ def reset(self) -> np.ndarray: self._env.seed(self._seed) obs = self._env.reset() obs = to_ndarray(obs).astype('float32') - self._final_eval_reward = 0. + self._eval_episode_return = 0. return obs def close(self) -> None: @@ -329,11 +329,11 @@ def step(self, action: Union[np.ndarray, list]) -> BaseEnvTimestep: action_range = self.info().act_space.value action = affine_transform(action, min_val=action_range['min'], max_val=action_range['max']) obs, rew, done, info = self._env.step(action) - self._final_eval_reward += rew + self._eval_episode_return += rew obs = to_ndarray(obs).astype('float32') rew = to_ndarray([rew]) # wrapped to be transfered to a array with shape (1,) if done: - info['final_eval_reward'] = self._final_eval_reward + info['eval_episode_return'] = self._eval_episode_return return BaseEnvTimestep(obs, rew, done, info) def info(self) -> BaseEnvInfo: diff --git a/dizoo/rocket/README.md b/dizoo/rocket/README.md new file mode 100644 index 0000000000..c9e49e47d2 --- /dev/null +++ b/dizoo/rocket/README.md @@ -0,0 +1,10 @@ +# Install + +```shell +pip install git+https://github.com/nighood/rocket-recycling@master#egg=rocket_recycling +``` + +# Chek Install +```shell +pytest -sv test_rocket_env.py +``` diff --git a/dizoo/rocket/__init__.py b/dizoo/rocket/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/dizoo/rocket/config/__init__.py b/dizoo/rocket/config/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/dizoo/rocket/entry/__init__.py b/dizoo/rocket/entry/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/dizoo/rocket/entry/rocket_hover_onppo_main_v2.py b/dizoo/rocket/entry/rocket_hover_onppo_main_v2.py new file mode 100644 index 0000000000..503312f47b --- /dev/null +++ b/dizoo/rocket/entry/rocket_hover_onppo_main_v2.py @@ -0,0 +1,95 @@ +import os +import gym +import numpy as np +from tensorboardX import SummaryWriter +import torch +from rocket_recycling.rocket import Rocket + +from ditk import logging +from ding.model import VAC +from ding.policy import PPOPolicy +from ding.envs import DingEnvWrapper, BaseEnvManagerV2, EvalEpisodeReturnWrapper +from ding.config import compile_config +from ding.framework import task +from ding.framework.context import OnlineRLContext +from ding.framework.middleware import multistep_trainer, StepCollector, interaction_evaluator, CkptSaver, \ + gae_estimator, termination_checker +from ding.utils import set_pkg_seed +from dizoo.rocket.config.rocket_hover_ppo_config import main_config, create_config + + +class RocketHoverWrapper(gym.Wrapper): + + def __init__(self, env): + super().__init__(env) + self._observation_space = gym.spaces.Box(low=float("-inf"), high=float("inf"), shape=(8, ), dtype=np.float32) + self._action_space = gym.spaces.Discrete(9) + self._action_space.seed(0) # default seed + self.reward_range = (float('-inf'), float('inf')) + + +def wrapped_rocket_env(task, max_steps): + return DingEnvWrapper( + Rocket(task=task, max_steps=max_steps), + cfg={'env_wrapper': [ + lambda env: RocketHoverWrapper(env), + lambda env: EvalEpisodeReturnWrapper(env), + ]} + ) + + +def main(): + logging.getLogger().setLevel(logging.INFO) + main_config.policy.cuda = True + print('torch.cuda.is_available(): ', torch.cuda.is_available()) + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + num_seed = 3 + for seed_i in range(num_seed): + main_config.exp_name = f'task_rocket_hovering_onppo_seed{seed_i}' + tb_logger = SummaryWriter(os.path.join('./{}/log/'.format(cfg.exp_name), 'seed' + str(seed_i))) + with task.start(async_mode=False, ctx=OnlineRLContext()): + collector_env = BaseEnvManagerV2( + env_fn=[ + lambda: wrapped_rocket_env(cfg.env.task, cfg.env.max_steps) + for _ in range(cfg.env.collector_env_num) + ], + cfg=cfg.env.manager + ) + evaluator_env = BaseEnvManagerV2( + env_fn=[ + lambda: wrapped_rocket_env(cfg.env.task, cfg.env.max_steps) + for _ in range(cfg.env.evaluator_env_num) + ], + cfg=cfg.env.manager + ) + + # evaluator_env.enable_save_replay() + + set_pkg_seed(seed_i, use_cuda=cfg.policy.cuda) + + model = VAC(**cfg.policy.model) + policy = PPOPolicy(cfg.policy, model=model) + + def _add_scalar(ctx): + if ctx.eval_value != -np.inf: + tb_logger.add_scalar('evaluator_step/reward', ctx.eval_value, global_step=ctx.env_step) + collector_rewards = [ctx.trajectories[i]['reward'] for i in range(len(ctx.trajectories))] + collector_mean_reward = sum(collector_rewards) / len(ctx.trajectories) + collector_max_reward = max(collector_rewards) + collector_min_reward = min(collector_rewards) + tb_logger.add_scalar('collecter_step/mean_reward', collector_mean_reward, global_step=ctx.env_step) + tb_logger.add_scalar('collecter_step/max_reward', collector_max_reward, global_step=ctx.env_step) + tb_logger.add_scalar('collecter_step/min_reward', collector_min_reward, global_step=ctx.env_step) + + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) + task.use(StepCollector(cfg, policy.collect_mode, collector_env)) + task.use(_add_scalar) + task.use(gae_estimator(cfg, policy.collect_mode)) + task.use(multistep_trainer(cfg, policy.learn_mode)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) + task.use(termination_checker(max_env_step=int(10e7))) + task.run() + + +if __name__ == "__main__": + main() diff --git a/dizoo/rocket/entry/rocket_hover_ppo_main.py b/dizoo/rocket/entry/rocket_hover_ppo_main.py index 0a7904f2b4..13f5714483 100644 --- a/dizoo/rocket/entry/rocket_hover_ppo_main.py +++ b/dizoo/rocket/entry/rocket_hover_ppo_main.py @@ -30,12 +30,10 @@ def main(): tb_logger = SummaryWriter(os.path.join('./{}/log/'.format(cfg.exp_name), 'seed' + str(seed_i))) with task.start(async_mode=False, ctx=OnlineRLContext()): collector_env = BaseEnvManagerV2( - env_fn=[lambda: RocketEnv(cfg.env) for _ in range(cfg.env.collector_env_num)], - cfg=cfg.env.manager + env_fn=[lambda: RocketEnv(cfg.env) for _ in range(cfg.env.collector_env_num)], cfg=cfg.env.manager ) evaluator_env = BaseEnvManagerV2( - env_fn=[lambda: RocketEnv(cfg.env) for _ in range(cfg.env.evaluator_env_num)], - cfg=cfg.env.manager + env_fn=[lambda: RocketEnv(cfg.env) for _ in range(cfg.env.evaluator_env_num)], cfg=cfg.env.manager ) # evaluator_env.enable_save_replay() @@ -61,7 +59,7 @@ def _add_scalar(ctx): task.use(_add_scalar) task.use(gae_estimator(cfg, policy.collect_mode)) task.use(multistep_trainer(cfg, policy.learn_mode)) - task.use(CkptSaver(cfg, policy, train_freq=100)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) task.use(termination_checker(max_env_step=int(10e7))) task.run() diff --git a/dizoo/rocket/entry/rocket_landing_onppo_main_v2.py b/dizoo/rocket/entry/rocket_landing_onppo_main_v2.py new file mode 100644 index 0000000000..bd682ecf42 --- /dev/null +++ b/dizoo/rocket/entry/rocket_landing_onppo_main_v2.py @@ -0,0 +1,95 @@ +import os +import torch +import gym +import numpy as np +from tensorboardX import SummaryWriter +from rocket_recycling.rocket import Rocket + +from ditk import logging +from ding.model import VAC +from ding.policy import PPOPolicy +from ding.envs import DingEnvWrapper, BaseEnvManagerV2, EvalEpisodeReturnWrapper +from ding.config import compile_config +from ding.framework import task +from ding.framework.context import OnlineRLContext +from ding.framework.middleware import multistep_trainer, StepCollector, interaction_evaluator, CkptSaver, \ +gae_estimator, termination_checker +from ding.utils import set_pkg_seed +from dizoo.rocket.config.rocket_landing_ppo_config import main_config, create_config + + +class RocketLandingWrapper(gym.Wrapper): + + def __init__(self, env): + super().__init__(env) + self._observation_space = gym.spaces.Box(low=float("-inf"), high=float("inf"), shape=(8, ), dtype=np.float32) + self._action_space = gym.spaces.Discrete(9) + self._action_space.seed(0) # default seed + self.reward_range = (float('-inf'), float('inf')) + + +def wrapped_rocket_env(task, max_steps): + return DingEnvWrapper( + Rocket(task=task, max_steps=max_steps), + cfg={'env_wrapper': [ + lambda env: RocketLandingWrapper(env), + lambda env: EvalEpisodeReturnWrapper(env), + ]} + ) + + +def main(): + logging.getLogger().setLevel(logging.INFO) + main_config.exp_name = 'rocket_landing_ppo_nseed' + main_config.policy.cuda = True + print('torch.cuda.is_available(): ', torch.cuda.is_available()) + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + num_seed = 4 + for seed_i in range(num_seed): + tb_logger = SummaryWriter(os.path.join('./{}/log/'.format(cfg.exp_name), 'seed' + str(seed_i))) + with task.start(async_mode=False, ctx=OnlineRLContext()): + collector_env = BaseEnvManagerV2( + env_fn=[ + lambda: wrapped_rocket_env(cfg.env.task, cfg.env.max_steps) + for _ in range(cfg.env.collector_env_num) + ], + cfg=cfg.env.manager + ) + evaluator_env = BaseEnvManagerV2( + env_fn=[ + lambda: wrapped_rocket_env(cfg.env.task, cfg.env.max_steps) + for _ in range(cfg.env.evaluator_env_num) + ], + cfg=cfg.env.manager + ) + + # evaluator_env.enable_save_replay() + + set_pkg_seed(seed_i, use_cuda=cfg.policy.cuda) + + model = VAC(**cfg.policy.model) + policy = PPOPolicy(cfg.policy, model=model) + + def _add_scalar(ctx): + if ctx.eval_value != -np.inf: + tb_logger.add_scalar('evaluator_step/reward', ctx.eval_value, global_step=ctx.env_step) + collector_rewards = [ctx.trajectories[i]['reward'] for i in range(len(ctx.trajectories))] + collector_mean_reward = sum(collector_rewards) / len(ctx.trajectories) + collector_max_reward = max(collector_rewards) + collector_min_reward = min(collector_rewards) + tb_logger.add_scalar('collecter_step/mean_reward', collector_mean_reward, global_step=ctx.env_step) + tb_logger.add_scalar('collecter_step/max_reward', collector_max_reward, global_step=ctx.env_step) + tb_logger.add_scalar('collecter_step/min_reward', collector_min_reward, global_step=ctx.env_step) + + task.use(interaction_evaluator(cfg, policy.eval_mode, evaluator_env)) + task.use(StepCollector(cfg, policy.collect_mode, collector_env)) + task.use(gae_estimator(cfg, policy.collect_mode)) + task.use(multistep_trainer(cfg, policy.learn_mode)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) + # task.use(_add_scalar) + task.use(termination_checker(max_env_step=int(3e6))) + task.run() + + +if __name__ == "__main__": + main() diff --git a/dizoo/rocket/entry/rocket_landing_ppo_main.py b/dizoo/rocket/entry/rocket_landing_ppo_main.py index 12d99a648c..bf8ebb5162 100644 --- a/dizoo/rocket/entry/rocket_landing_ppo_main.py +++ b/dizoo/rocket/entry/rocket_landing_ppo_main.py @@ -27,15 +27,13 @@ def main(): cfg = compile_config(main_config, create_cfg=create_config, auto=True) num_seed = 4 for seed_i in range(num_seed): - tb_logger = SummaryWriter(os.path.join('./{}/log/'.format(cfg.exp_name), 'seed'+str(seed_i))) + tb_logger = SummaryWriter(os.path.join('./{}/log/'.format(cfg.exp_name), 'seed' + str(seed_i))) with task.start(async_mode=False, ctx=OnlineRLContext()): collector_env = BaseEnvManagerV2( - env_fn=[lambda: RocketEnv(cfg.env) for _ in range(cfg.env.collector_env_num)], - cfg=cfg.env.manager + env_fn=[lambda: RocketEnv(cfg.env) for _ in range(cfg.env.collector_env_num)], cfg=cfg.env.manager ) evaluator_env = BaseEnvManagerV2( - env_fn=[lambda: RocketEnv(cfg.env) for _ in range(cfg.env.evaluator_env_num)], - cfg=cfg.env.manager + env_fn=[lambda: RocketEnv(cfg.env) for _ in range(cfg.env.evaluator_env_num)], cfg=cfg.env.manager ) # evaluator_env.enable_save_replay() @@ -60,7 +58,7 @@ def _add_scalar(ctx): task.use(StepCollector(cfg, policy.collect_mode, collector_env)) task.use(gae_estimator(cfg, policy.collect_mode)) task.use(multistep_trainer(cfg, policy.learn_mode)) - task.use(CkptSaver(cfg, policy, train_freq=100)) + task.use(CkptSaver(policy, cfg.exp_name, train_freq=100)) task.use(_add_scalar) task.use(termination_checker(max_env_step=int(3e6))) task.run() diff --git a/dizoo/rocket/envs/__init__.py b/dizoo/rocket/envs/__init__.py new file mode 100644 index 0000000000..47b076e8af --- /dev/null +++ b/dizoo/rocket/envs/__init__.py @@ -0,0 +1 @@ +from .rocket_env import RocketEnv diff --git a/dizoo/rocket/envs/rocket_env.py b/dizoo/rocket/envs/rocket_env.py index 69533299d9..dd77fdafa7 100644 --- a/dizoo/rocket/envs/rocket_env.py +++ b/dizoo/rocket/envs/rocket_env.py @@ -1,31 +1,26 @@ from typing import Any, List, Union, Optional import time +import os +import imageio import gym import copy import numpy as np from easydict import EasyDict +from rocket_recycling.rocket import Rocket from ding.envs import BaseEnv, BaseEnvTimestep from ding.torch_utils import to_ndarray, to_list from ding.utils import ENV_REGISTRY from ding.envs import ObsPlusPrevActRewWrapper -from rocket_recycling.rocket import Rocket -import os -import imageio -@ENV_REGISTRY.register('rocket') +@ENV_REGISTRY.register('rocket', force_overwrite=True) class RocketEnv(BaseEnv): def __init__(self, cfg: dict = {}) -> None: self._cfg = cfg self._init_flag = False self._save_replay = False - self._observation_space = gym.spaces.Box( - low=float("-inf"), - high=float("inf"), - shape=(8, ), - dtype=np.float32 - ) + self._observation_space = gym.spaces.Box(low=float("-inf"), high=float("inf"), shape=(8, ), dtype=np.float32) self._action_space = gym.spaces.Discrete(9) self._action_space.seed(0) # default seed self._reward_space = gym.spaces.Box(low=float("-inf"), high=float("inf"), shape=(1, ), dtype=np.float32) @@ -41,7 +36,7 @@ def reset(self) -> np.ndarray: elif hasattr(self, '_seed'): self._env.seed(self._seed) self._action_space.seed(self._seed) - self._final_eval_reward = 0 + self._eval_episode_return = 0 obs = self._env.reset() obs = to_ndarray(obs) if self._save_replay: @@ -64,21 +59,19 @@ def step(self, action: Union[int, np.ndarray]) -> BaseEnvTimestep: obs, rew, done, info = self._env.step(action) self._env.render() - self._final_eval_reward += rew + self._eval_episode_return += rew if self._save_replay: self._frames.extend(self._env.render()) if done: - info['final_eval_reward'] = self._final_eval_reward + info['eval_episode_return'] = self._eval_episode_return if self._save_replay: - path = os.path.join( - self._replay_path, '{}_episode.gif'.format(self._save_replay_count) - ) + path = os.path.join(self._replay_path, '{}_episode.gif'.format(self._save_replay_count)) self.display_frames_as_gif(self._frames, path) self._save_replay_count += 1 obs = to_ndarray(obs) # wrapped to be transfered to a array with shape (1,) - rew = to_ndarray([rew]).astype(np.float32) + rew = to_ndarray([rew]).astype(np.float32) return BaseEnvTimestep(obs, rew, done, info) def enable_save_replay(self, replay_path: Optional[str] = None) -> None: @@ -95,6 +88,9 @@ def random_action(self) -> np.ndarray: random_action = to_ndarray([random_action], dtype=np.int64) return random_action + def clone(self, caller: str) -> 'RocketEnv': + return RocketEnv(copy.deepcopy(self._cfg)) + @property def observation_space(self) -> gym.spaces.Space: return self._observation_space diff --git a/dizoo/rocket/envs/test_rocket_env.py b/dizoo/rocket/envs/test_rocket_env.py index f41a2dbd71..a8bf030fe7 100644 --- a/dizoo/rocket/envs/test_rocket_env.py +++ b/dizoo/rocket/envs/test_rocket_env.py @@ -1,6 +1,6 @@ import pytest import numpy as np -from dizoo.rocket.envs.rocket_env import RocketEnv +from dizoo.rocket.envs import RocketEnv from easydict import EasyDict @@ -12,7 +12,7 @@ def test_hover(self): env.seed(314, dynamic_seed=False) assert env._seed == 314 obs = env.reset() - assert obs.shape == (8,) + assert obs.shape == (8, ) for _ in range(5): env.reset() np.random.seed(314) @@ -28,8 +28,8 @@ def test_hover(self): print('timestep', timestep, '\n') assert isinstance(timestep.obs, np.ndarray) assert isinstance(timestep.done, bool) - assert timestep.obs.shape == (8,) - assert timestep.reward.shape == (1,) + assert timestep.obs.shape == (8, ) + assert timestep.reward.shape == (1, ) assert timestep.reward >= env.reward_space.low assert timestep.reward <= env.reward_space.high print(env.observation_space, env.action_space, env.reward_space) diff --git a/dizoo/slime_volley/envs/slime_volley_env.py b/dizoo/slime_volley/envs/slime_volley_env.py index 7938ae0383..866cc8e5d7 100644 --- a/dizoo/slime_volley/envs/slime_volley_env.py +++ b/dizoo/slime_volley/envs/slime_volley_env.py @@ -49,7 +49,7 @@ def step(self, action: Union[np.ndarray, List[np.ndarray]]) -> BaseEnvTimestep: # So we have to put two actions into one tuple. obs1, rew, done, info = self._env.step((action1, action2)) obs1 = to_ndarray(obs1).astype(np.float32) - self._final_eval_reward += rew + self._eval_episode_return += rew # info ('ale.lives', 'ale.otherLives', 'otherObs', 'state', 'otherState') if self._agent_vs_agent: info = [ @@ -63,14 +63,14 @@ def step(self, action: Union[np.ndarray, List[np.ndarray]]) -> BaseEnvTimestep: } ] if done: - info[0]['final_eval_reward'] = self._final_eval_reward - info[1]['final_eval_reward'] = -self._final_eval_reward - info[0]['result'] = self.get_episode_result(self._final_eval_reward) - info[1]['result'] = self.get_episode_result(-self._final_eval_reward) + info[0]['eval_episode_return'] = self._eval_episode_return + info[1]['eval_episode_return'] = -self._eval_episode_return + info[0]['result'] = self.get_episode_result(self._eval_episode_return) + info[1]['result'] = self.get_episode_result(-self._eval_episode_return) else: if done: - info['final_eval_reward'] = self._final_eval_reward - info['result'] = self.get_episode_result(self._final_eval_reward) + info['eval_episode_return'] = self._eval_episode_return + info['result'] = self.get_episode_result(self._eval_episode_return) reward = to_ndarray([rew]).astype(np.float32) if self._agent_vs_agent: obs2 = info[1]['obs'] @@ -82,8 +82,8 @@ def step(self, action: Union[np.ndarray, List[np.ndarray]]) -> BaseEnvTimestep: else: return BaseEnvTimestep(obs1, reward, done, info) - def get_episode_result(self, final_eval_reward: float): - if final_eval_reward > 0: # due to using 5 games (lives) in this env, the final_eval_reward can't be zero. + def get_episode_result(self, eval_episode_return: float): + if eval_episode_return > 0: # due to using 5 games (lives) in this env, the eval_episode_return can't be zero. return "wins" else: return "losses" @@ -122,7 +122,7 @@ def reset(self): self._env.seed(self._seed + np_seed) elif hasattr(self, '_seed'): self._env.seed(self._seed) - self._final_eval_reward = 0 + self._eval_episode_return = 0 obs = self._env.reset() obs = to_ndarray(obs).astype(np.float32) if self._agent_vs_agent: diff --git a/dizoo/slime_volley/envs/test_slime_volley_env.py b/dizoo/slime_volley/envs/test_slime_volley_env.py index 5ae7c31396..88a089a7eb 100644 --- a/dizoo/slime_volley/envs/test_slime_volley_env.py +++ b/dizoo/slime_volley/envs/test_slime_volley_env.py @@ -10,7 +10,7 @@ class TestSlimeVolley: @pytest.mark.parametrize('agent_vs_agent', [True, False]) def test_slime_volley(self, agent_vs_agent): - total_rew = 0 + total_return = 0 env = SlimeVolleyEnv(EasyDict({'env_id': 'SlimeVolley-v0', 'agent_vs_agent': agent_vs_agent})) # env.enable_save_replay('replay_video') obs1 = env.reset() @@ -21,13 +21,13 @@ def test_slime_volley(self, agent_vs_agent): action = env.random_action() observations, rewards, done, infos = env.step(action) if agent_vs_agent: - total_rew += rewards[0] + total_return += rewards[0] else: - total_rew += rewards + total_return += rewards obs1, obs2 = observations[0], observations[1] assert obs1.shape == obs2.shape, (obs1.shape, obs2.shape) if agent_vs_agent: agent_lives, opponent_lives = infos[0]['ale.lives'], infos[1]['ale.lives'] if agent_vs_agent: assert agent_lives == 0 or opponent_lives == 0, (agent_lives, opponent_lives) - print("total reward is:", total_rew) + print("total return is:", total_return) diff --git a/dizoo/smac/config/smac_10m11m_mappo_config.py b/dizoo/smac/config/smac_10m11m_mappo_config.py index 04cceea103..0dd28240ee 100644 --- a/dizoo/smac/config/smac_10m11m_mappo_config.py +++ b/dizoo/smac/config/smac_10m11m_mappo_config.py @@ -51,8 +51,6 @@ ), # used in state_num of hidden_state learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, epoch_per_collect=5, batch_size=3200, learning_rate=5e-4, diff --git a/dizoo/smac/config/smac_10m11m_masac_config.py b/dizoo/smac/config/smac_10m11m_masac_config.py index af80d1db0d..b58f3f823f 100644 --- a/dizoo/smac/config/smac_10m11m_masac_config.py +++ b/dizoo/smac/config/smac_10m11m_masac_config.py @@ -26,7 +26,7 @@ ), policy=dict( cuda=True, - on_policy=False, + multi_agent=True, random_collect_size=0, model=dict( agent_obs_shape=132, @@ -52,7 +52,6 @@ collect=dict( env_num=collector_env_num, n_sample=1600, - unroll_len=1, ), command=dict(), eval=dict( @@ -80,7 +79,7 @@ import_names=['dizoo.smac.envs.smac_env'], ), env_manager=dict(type='base'), - policy=dict(type='sac_discrete', ), + policy=dict(type='discrete_sac', ), ) SMAC_10m11m_masac_default_create_config = EasyDict(SMAC_10m11m_masac_default_create_config) create_config = SMAC_10m11m_masac_default_create_config @@ -88,4 +87,4 @@ if __name__ == '__main__': from ding.entry import serial_pipeline - serial_pipeline((main_config, create_config), seed=0) + serial_pipeline((main_config, create_config), seed=0, max_env_step=int(1e7)) diff --git a/dizoo/smac/config/smac_25m_mappo_config.py b/dizoo/smac/config/smac_25m_mappo_config.py index ba83b3caa8..fd9e6638a8 100644 --- a/dizoo/smac/config/smac_25m_mappo_config.py +++ b/dizoo/smac/config/smac_25m_mappo_config.py @@ -51,8 +51,6 @@ ), # used in state_num of hidden_state learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, epoch_per_collect=5, batch_size=3200, learning_rate=5e-4, diff --git a/dizoo/smac/config/smac_25m_masac_config.py b/dizoo/smac/config/smac_25m_masac_config.py index 5b6e279a0b..67e5db1f68 100644 --- a/dizoo/smac/config/smac_25m_masac_config.py +++ b/dizoo/smac/config/smac_25m_masac_config.py @@ -26,7 +26,7 @@ ), policy=dict( cuda=True, - on_policy=False, + multi_agent=True, random_collect_size=0, model=dict( agent_obs_shape=306, @@ -80,7 +80,7 @@ import_names=['dizoo.smac.envs.smac_env'], ), env_manager=dict(type='base'), - policy=dict(type='sac_discrete', ), + policy=dict(type='discrete_sac', ), ) SMAC_25m_masac_default_create_config = EasyDict(SMAC_25m_masac_default_create_config) create_config = SMAC_25m_masac_default_create_config diff --git a/dizoo/smac/config/smac_27m30m_mappo_config.py b/dizoo/smac/config/smac_27m30m_mappo_config.py index db0fab9079..14caadd256 100644 --- a/dizoo/smac/config/smac_27m30m_mappo_config.py +++ b/dizoo/smac/config/smac_27m30m_mappo_config.py @@ -51,8 +51,6 @@ ), # used in state_num of hidden_state learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, epoch_per_collect=5, batch_size=3200, learning_rate=5e-4, diff --git a/dizoo/smac/config/smac_2c64zg_mappo_config.py b/dizoo/smac/config/smac_2c64zg_mappo_config.py index 185a867900..ac29489d9c 100644 --- a/dizoo/smac/config/smac_2c64zg_mappo_config.py +++ b/dizoo/smac/config/smac_2c64zg_mappo_config.py @@ -48,8 +48,6 @@ ), # used in state_num of hidden_state learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, epoch_per_collect=5, batch_size=3200, learning_rate=5e-4, diff --git a/dizoo/smac/config/smac_2c64zg_masac_config.py b/dizoo/smac/config/smac_2c64zg_masac_config.py index fc729c23a5..b71e9529df 100644 --- a/dizoo/smac/config/smac_2c64zg_masac_config.py +++ b/dizoo/smac/config/smac_2c64zg_masac_config.py @@ -26,7 +26,7 @@ ), policy=dict( cuda=True, - on_policy=False, + multi_agent=True, random_collect_size=0, model=dict( agent_obs_shape=404, @@ -80,7 +80,7 @@ import_names=['dizoo.smac.envs.smac_env'], ), env_manager=dict(type='base'), - policy=dict(type='sac_discrete', ), + policy=dict(type='discrete_sac', ), ) SMAC_2c64zg_masac_default_create_config = EasyDict(SMAC_2c64zg_masac_default_create_config) create_config = SMAC_2c64zg_masac_default_create_config diff --git a/dizoo/smac/config/smac_2c64zg_qmix_config.py b/dizoo/smac/config/smac_2c64zg_qmix_config.py index 544c2f6556..228d8b878a 100644 --- a/dizoo/smac/config/smac_2c64zg_qmix_config.py +++ b/dizoo/smac/config/smac_2c64zg_qmix_config.py @@ -33,7 +33,6 @@ dueling=False, ), learn=dict( - multi_gpu=False, update_per_collect=20, batch_size=32, learning_rate=0.0005, diff --git a/dizoo/smac/config/smac_2s3z_qmix_config.py b/dizoo/smac/config/smac_2s3z_qmix_config.py index 61dd746dcb..8796048da0 100644 --- a/dizoo/smac/config/smac_2s3z_qmix_config.py +++ b/dizoo/smac/config/smac_2s3z_qmix_config.py @@ -33,7 +33,6 @@ dueling=False, ), learn=dict( - multi_gpu=False, update_per_collect=20, batch_size=32, learning_rate=0.0005, diff --git a/dizoo/smac/config/smac_2s3z_qtran_config.py b/dizoo/smac/config/smac_2s3z_qtran_config.py index 8daf27239b..0c66056a1b 100644 --- a/dizoo/smac/config/smac_2s3z_qtran_config.py +++ b/dizoo/smac/config/smac_2s3z_qtran_config.py @@ -33,7 +33,6 @@ dueling=False, ), learn=dict( - multi_gpu=False, update_per_collect=20, batch_size=32, learning_rate=0.0005, diff --git a/dizoo/smac/config/smac_3m_masac_config.py b/dizoo/smac/config/smac_3m_masac_config.py index 54d63d3c51..9613a83cce 100644 --- a/dizoo/smac/config/smac_3m_masac_config.py +++ b/dizoo/smac/config/smac_3m_masac_config.py @@ -26,6 +26,7 @@ ), policy=dict( cuda=True, + multi_agent=True, random_collect_size=0, model=dict( agent_obs_shape=42, @@ -41,7 +42,6 @@ learning_rate_q=5e-4, learning_rate_policy=5e-4, learning_rate_alpha=5e-4, - ignore_done=False, target_theta=0.005, discount_factor=0.99, alpha=0.2, @@ -63,7 +63,7 @@ start=1, end=0.05, decay=100000, - ), # TODO(pu) + ), replay_buffer=dict(replay_buffer_size=1000000, ), ), ), @@ -78,7 +78,7 @@ import_names=['dizoo.smac.envs.smac_env'], ), env_manager=dict(type='base'), - policy=dict(type='sac_discrete', ), + policy=dict(type='discrete_sac', ), ) SMAC_3m_masac_default_create_config = EasyDict(SMAC_3m_masac_default_create_config) create_config = SMAC_3m_masac_default_create_config diff --git a/dizoo/smac/config/smac_3s5z_collaq_config.py b/dizoo/smac/config/smac_3s5z_collaq_config.py index 7e55299647..7f1bf124b4 100644 --- a/dizoo/smac/config/smac_3s5z_collaq_config.py +++ b/dizoo/smac/config/smac_3s5z_collaq_config.py @@ -39,7 +39,6 @@ dueling=False, ), learn=dict( - multi_gpu=False, update_per_collect=20, batch_size=32, learning_rate=0.0005, diff --git a/dizoo/smac/config/smac_3s5z_collaq_per_config.py b/dizoo/smac/config/smac_3s5z_collaq_per_config.py index f2ebd71fef..5530e53724 100644 --- a/dizoo/smac/config/smac_3s5z_collaq_per_config.py +++ b/dizoo/smac/config/smac_3s5z_collaq_per_config.py @@ -43,7 +43,6 @@ dueling=False, ), learn=dict( - multi_gpu=False, update_per_collect=20, batch_size=32, learning_rate=0.0005, diff --git a/dizoo/smac/config/smac_3s5z_madqn_config.py b/dizoo/smac/config/smac_3s5z_madqn_config.py new file mode 100644 index 0000000000..5e771baf09 --- /dev/null +++ b/dizoo/smac/config/smac_3s5z_madqn_config.py @@ -0,0 +1,84 @@ +from ding.entry import serial_pipeline +from easydict import EasyDict + +agent_num = 8 +collector_env_num = 4 +evaluator_env_num = 8 + +main_config = dict( + exp_name='smac_3s5z_madqn_seed0', + env=dict( + map_name='3s5z', + difficulty=7, + reward_only_positive=True, + mirror_opponent=False, + agent_num=agent_num, + collector_env_num=collector_env_num, + evaluator_env_num=evaluator_env_num, + stop_value=0.999, + n_evaluator_episode=32, + special_global_state=True, + manager=dict(shared_memory=False, ), + ), + policy=dict( + nstep=1, + model=dict( + agent_num=agent_num, + obs_shape=150, + global_obs_shape=295, + global_cooperation=True, + action_shape=14, + hidden_size_list=[256, 256], + ), + learn=dict( + update_per_collect=20, + batch_size=64, + learning_rate=0.0005, + clip_value=5, + target_update_theta=0.008, + discount_factor=0.95, + ), + collect=dict( + collector=dict(get_train_sample=True, ), + n_episode=32, + unroll_len=10, + env_num=collector_env_num, + ), + eval=dict(env_num=evaluator_env_num, evaluator=dict(eval_freq=1000, )), + other=dict( + eps=dict( + type='linear', + start=1, + end=0.05, + decay=10000, + ), + replay_buffer=dict(replay_buffer_size=15000, ), + ), + ), +) +main_config = EasyDict(main_config) +create_config = dict( + env=dict( + type='smac', + import_names=['dizoo.smac.envs.smac_env'], + ), + env_manager=dict(type='subprocess'), + policy=dict(type='madqn'), + collector=dict(type='episode'), +) +create_config = EasyDict(create_config) + + +def train(args): + config = [main_config, create_config] + serial_pipeline(config, seed=args.seed, max_env_step=1e7) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument('--seed', '-s', type=int, default=1) + args = parser.parse_args() + + train(args) diff --git a/dizoo/smac/config/smac_3s5z_mappo_config.py b/dizoo/smac/config/smac_3s5z_mappo_config.py index c76f4c1390..d57055b93d 100644 --- a/dizoo/smac/config/smac_3s5z_mappo_config.py +++ b/dizoo/smac/config/smac_3s5z_mappo_config.py @@ -50,8 +50,6 @@ ), # used in state_num of hidden_state learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, epoch_per_collect=5, batch_size=3200, learning_rate=5e-4, diff --git a/dizoo/smac/config/smac_3s5z_masac_config.py b/dizoo/smac/config/smac_3s5z_masac_config.py index b133a4f5a7..456a9016f8 100644 --- a/dizoo/smac/config/smac_3s5z_masac_config.py +++ b/dizoo/smac/config/smac_3s5z_masac_config.py @@ -26,6 +26,7 @@ ), policy=dict( cuda=True, + multi_agent=True, random_collect_size=0, model=dict( agent_obs_shape=150, @@ -41,7 +42,6 @@ learning_rate_q=5e-4, learning_rate_policy=5e-4, learning_rate_alpha=5e-5, - ignore_done=False, target_theta=0.005, discount_factor=0.99, alpha=0.2, @@ -78,7 +78,7 @@ import_names=['dizoo.smac.envs.smac_env'], ), env_manager=dict(type='base'), - policy=dict(type='sac_discrete', ), + policy=dict(type='discrete_sac', ), ) smac_3s5z_masac_default_create_config = EasyDict(smac_3s5z_masac_default_create_config) create_config = smac_3s5z_masac_default_create_config diff --git a/dizoo/smac/config/smac_3s5z_qmix_config.py b/dizoo/smac/config/smac_3s5z_qmix_config.py index 5c38d79a82..cd622ae5f6 100644 --- a/dizoo/smac/config/smac_3s5z_qmix_config.py +++ b/dizoo/smac/config/smac_3s5z_qmix_config.py @@ -33,7 +33,6 @@ dueling=False, ), learn=dict( - multi_gpu=False, update_per_collect=20, batch_size=64, learning_rate=0.0005, diff --git a/dizoo/smac/config/smac_3s5z_qtran_config.py b/dizoo/smac/config/smac_3s5z_qtran_config.py index a7e6cbd35d..1190977354 100644 --- a/dizoo/smac/config/smac_3s5z_qtran_config.py +++ b/dizoo/smac/config/smac_3s5z_qtran_config.py @@ -33,7 +33,6 @@ dueling=False, ), learn=dict( - multi_gpu=False, update_per_collect=20, batch_size=32, learning_rate=0.0005, diff --git a/dizoo/smac/config/smac_3s5z_wqmix_config.py b/dizoo/smac/config/smac_3s5z_wqmix_config.py index 4bf5ab9e4c..92552f930b 100644 --- a/dizoo/smac/config/smac_3s5z_wqmix_config.py +++ b/dizoo/smac/config/smac_3s5z_wqmix_config.py @@ -32,7 +32,6 @@ dueling=False, ), learn=dict( - multi_gpu=False, update_per_collect=20, batch_size=32, learning_rate=0.0005, diff --git a/dizoo/smac/config/smac_3s5zvs3s6z_madqn_config.py b/dizoo/smac/config/smac_3s5zvs3s6z_madqn_config.py new file mode 100644 index 0000000000..438025241f --- /dev/null +++ b/dizoo/smac/config/smac_3s5zvs3s6z_madqn_config.py @@ -0,0 +1,84 @@ +from ding.entry import serial_pipeline +from easydict import EasyDict + +agent_num = 8 +collector_env_num = 4 +evaluator_env_num = 8 + +main_config = dict( + exp_name='smac_3s5zvs3s6z_madqn_seed0', + env=dict( + map_name='3s5z_vs_3s6z', + difficulty=7, + reward_only_positive=True, + mirror_opponent=False, + agent_num=agent_num, + collector_env_num=collector_env_num, + evaluator_env_num=evaluator_env_num, + stop_value=0.999, + n_evaluator_episode=32, + special_global_state=True, + manager=dict(shared_memory=False, ), + ), + policy=dict( + nstep=3, + model=dict( + agent_num=agent_num, + obs_shape=159, + global_obs_shape=314, + global_cooperation=True, + action_shape=15, + hidden_size_list=[256, 256], + ), + learn=dict( + update_per_collect=40, + batch_size=32, + learning_rate=0.0005, + clip_value=5, + target_update_theta=0.008, + discount_factor=0.95, + ), + collect=dict( + collector=dict(get_train_sample=True, ), + n_episode=32, + unroll_len=10, + env_num=collector_env_num, + ), + eval=dict(env_num=evaluator_env_num, evaluator=dict(eval_freq=1000, )), + other=dict( + eps=dict( + type='linear', + start=1, + end=0.05, + decay=100000, + ), + replay_buffer=dict(replay_buffer_size=30000, ), + ), + ), +) +main_config = EasyDict(main_config) +create_config = dict( + env=dict( + type='smac', + import_names=['dizoo.smac.envs.smac_env'], + ), + env_manager=dict(type='base'), + policy=dict(type='madqn'), + collector=dict(type='episode'), +) +create_config = EasyDict(create_config) + + +def train(args): + config = [main_config, create_config] + serial_pipeline(config, seed=args.seed, max_env_step=1e7) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument('--seed', '-s', type=int, default=0) + args = parser.parse_args() + + train(args) diff --git a/dizoo/smac/config/smac_3s5zvs3s6z_mappo_config.py b/dizoo/smac/config/smac_3s5zvs3s6z_mappo_config.py index 20f42682f6..a80527eeba 100644 --- a/dizoo/smac/config/smac_3s5zvs3s6z_mappo_config.py +++ b/dizoo/smac/config/smac_3s5zvs3s6z_mappo_config.py @@ -47,8 +47,6 @@ ), # used in state_num of hidden_state learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, epoch_per_collect=5, batch_size=3200, learning_rate=5e-4, diff --git a/dizoo/smac/config/smac_3s5zvs3s6z_masac_config.py b/dizoo/smac/config/smac_3s5zvs3s6z_masac_config.py index a2dbdf87c5..e28c375c4d 100644 --- a/dizoo/smac/config/smac_3s5zvs3s6z_masac_config.py +++ b/dizoo/smac/config/smac_3s5zvs3s6z_masac_config.py @@ -19,7 +19,6 @@ stop_value=0.99, death_mask=True, special_global_state=special_global_state, - # save_replay_episodes = 1, manager=dict( shared_memory=False, reset_timeout=6000, @@ -27,7 +26,7 @@ ), policy=dict( cuda=True, - on_policy=False, + multi_agent=True, random_collect_size=0, model=dict( agent_obs_shape=159, @@ -81,7 +80,7 @@ import_names=['dizoo.smac.envs.smac_env'], ), env_manager=dict(type='subprocess'), - policy=dict(type='sac_discrete', ), + policy=dict(type='discrete_sac', ), ) smac_3s5zvs3s6z_masac_default_create_config = EasyDict(smac_3s5zvs3s6z_masac_default_create_config) create_config = smac_3s5zvs3s6z_masac_default_create_config @@ -89,4 +88,4 @@ if __name__ == '__main__': from ding.entry import serial_pipeline - serial_pipeline((main_config, create_config), seed=0) + serial_pipeline((main_config, create_config), seed=0, max_env_step=int(1e7)) diff --git a/dizoo/smac/config/smac_5m6m_collaq_config.py b/dizoo/smac/config/smac_5m6m_collaq_config.py index 9ed38348ea..5f775b9dca 100644 --- a/dizoo/smac/config/smac_5m6m_collaq_config.py +++ b/dizoo/smac/config/smac_5m6m_collaq_config.py @@ -39,7 +39,6 @@ dueling=False, ), learn=dict( - multi_gpu=False, update_per_collect=20, batch_size=32, learning_rate=0.0005, diff --git a/dizoo/smac/config/smac_5m6m_madqn_config.py b/dizoo/smac/config/smac_5m6m_madqn_config.py new file mode 100644 index 0000000000..d05bb23dcb --- /dev/null +++ b/dizoo/smac/config/smac_5m6m_madqn_config.py @@ -0,0 +1,98 @@ +from ding.entry import serial_pipeline +from easydict import EasyDict + +agent_num = 5 +collector_env_num = 16 +evaluator_env_num = 8 + +main_config = dict( + exp_name='smac_5m6m_madqn_seed0', + env=dict( + map_name='5m_vs_6m', + difficulty=7, + reward_only_positive=True, + mirror_opponent=False, + agent_num=agent_num, + collector_env_num=collector_env_num, + evaluator_env_num=evaluator_env_num, + shared_memory=False, + special_global_state=True, + stop_value=0.999, + n_evaluator_episode=32, + ), + policy=dict( + nstep=3, + model=dict( + agent_num=agent_num, + obs_shape=72, + global_obs_shape=152, + action_shape=12, + hidden_size_list=[256, 256], + ), + learn=dict( + update_per_collect=40, + batch_size=32, + learning_rate=0.0005, + clip_value=10, + target_update_theta=0.008, + discount_factor=0.95, + ), + collect=dict( + collector=dict(get_train_sample=True, ), + n_episode=32, + unroll_len=10, + env_num=collector_env_num, + ), + eval=dict(env_num=evaluator_env_num, evaluator=dict(eval_freq=1000, )), + other=dict( + eps=dict( + type='linear', + start=1, + end=0.05, + decay=50000, + ), + replay_buffer=dict(replay_buffer_size=50000, ), + ), + ), +) +main_config = EasyDict(main_config) +create_config = dict( + env=dict( + type='smac', + import_names=['dizoo.smac.envs.smac_env'], + ), + env_manager=dict(type='base'), + policy=dict(type='madqn'), + collector=dict(type='episode'), +) +create_config = EasyDict(create_config) + + +def train(args): + config = [main_config, create_config] + serial_pipeline(config, seed=args.seed) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument('--seed', '-s', type=int, default=0) + args = parser.parse_args() + + train(args) + + +def train(args): + config = [main_config, create_config] + serial_pipeline(config, seed=args.seed, max_env_step=1e7) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument('--seed', '-s', type=int, default=0) + args = parser.parse_args() + + train(args) diff --git a/dizoo/smac/config/smac_5m6m_mappo_config.py b/dizoo/smac/config/smac_5m6m_mappo_config.py index 395e4604e2..d4e2d968e3 100644 --- a/dizoo/smac/config/smac_5m6m_mappo_config.py +++ b/dizoo/smac/config/smac_5m6m_mappo_config.py @@ -49,8 +49,6 @@ ), # used in state_num of hidden_state learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, epoch_per_collect=10, batch_size=3200, learning_rate=5e-4, diff --git a/dizoo/smac/config/smac_5m6m_masac_config.py b/dizoo/smac/config/smac_5m6m_masac_config.py index 8dad040e33..d4b4bf92b8 100644 --- a/dizoo/smac/config/smac_5m6m_masac_config.py +++ b/dizoo/smac/config/smac_5m6m_masac_config.py @@ -26,6 +26,7 @@ ), policy=dict( cuda=True, + multi_agent=True, random_collect_size=0, model=dict( agent_obs_shape=72, @@ -78,7 +79,7 @@ import_names=['dizoo.smac.envs.smac_env'], ), env_manager=dict(type='base'), - policy=dict(type='sac_discrete', ), + policy=dict(type='discrete_sac', ), ) SMAC_5m6m_masac_default_create_config = EasyDict(SMAC_5m6m_masac_default_create_config) create_config = SMAC_5m6m_masac_default_create_config diff --git a/dizoo/smac/config/smac_5m6m_qmix_config.py b/dizoo/smac/config/smac_5m6m_qmix_config.py index 92b09355aa..b52a019ad8 100644 --- a/dizoo/smac/config/smac_5m6m_qmix_config.py +++ b/dizoo/smac/config/smac_5m6m_qmix_config.py @@ -33,7 +33,6 @@ dueling=False, ), learn=dict( - multi_gpu=False, update_per_collect=20, batch_size=32, learning_rate=0.0005, diff --git a/dizoo/smac/config/smac_5m6m_qtran_config.py b/dizoo/smac/config/smac_5m6m_qtran_config.py index ecf20a380f..960e8affb8 100644 --- a/dizoo/smac/config/smac_5m6m_qtran_config.py +++ b/dizoo/smac/config/smac_5m6m_qtran_config.py @@ -33,7 +33,6 @@ dueling=False, ), learn=dict( - multi_gpu=False, update_per_collect=20, batch_size=32, learning_rate=0.0005, diff --git a/dizoo/smac/config/smac_5m6m_wqmix_config.py b/dizoo/smac/config/smac_5m6m_wqmix_config.py index c03f753055..2e7586b5ae 100644 --- a/dizoo/smac/config/smac_5m6m_wqmix_config.py +++ b/dizoo/smac/config/smac_5m6m_wqmix_config.py @@ -32,7 +32,6 @@ dueling=False, ), learn=dict( - multi_gpu=False, update_per_collect=20, batch_size=32, learning_rate=0.0005, diff --git a/dizoo/smac/config/smac_8m9m_madqn_config.py b/dizoo/smac/config/smac_8m9m_madqn_config.py new file mode 100644 index 0000000000..672330df24 --- /dev/null +++ b/dizoo/smac/config/smac_8m9m_madqn_config.py @@ -0,0 +1,98 @@ +from ding.entry import serial_pipeline +from easydict import EasyDict + +agent_num = 8 +collector_env_num = 16 +evaluator_env_num = 8 + +main_config = dict( + exp_name='smac_8m9m_madqn_seed0', + env=dict( + map_name='8m_vs_9m', + difficulty=7, + reward_only_positive=True, + mirror_opponent=False, + agent_num=agent_num, + collector_env_num=collector_env_num, + evaluator_env_num=evaluator_env_num, + shared_memory=False, + special_global_state=True, + stop_value=0.999, + n_evaluator_episode=32, + ), + policy=dict( + nstep=3, + model=dict( + agent_num=agent_num, + obs_shape=108, + global_obs_shape=263, + action_shape=15, + hidden_size_list=[256, 256], + ), + learn=dict( + update_per_collect=40, + batch_size=32, + learning_rate=0.0005, + clip_value=10, + target_update_theta=0.008, + discount_factor=0.95, + ), + collect=dict( + collector=dict(get_train_sample=True, ), + n_episode=32, + unroll_len=20, + env_num=collector_env_num, + ), + eval=dict(env_num=evaluator_env_num, evaluator=dict(eval_freq=1000, )), + other=dict( + eps=dict( + type='linear', + start=1, + end=0.05, + decay=50000, + ), + replay_buffer=dict(replay_buffer_size=20000, ), + ), + ), +) +main_config = EasyDict(main_config) +create_config = dict( + env=dict( + type='smac', + import_names=['dizoo.smac.envs.smac_env'], + ), + env_manager=dict(type='base'), + policy=dict(type='madqn'), + collector=dict(type='episode'), +) +create_config = EasyDict(create_config) + + +def train(args): + config = [main_config, create_config] + serial_pipeline(config, seed=args.seed) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument('--seed', '-s', type=int, default=0) + args = parser.parse_args() + + train(args) + + +def train(args): + config = [main_config, create_config] + serial_pipeline(config, seed=args.seed, max_env_step=1e7) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument('--seed', '-s', type=int, default=0) + args = parser.parse_args() + + train(args) diff --git a/dizoo/smac/config/smac_8m9m_mappo_config.py b/dizoo/smac/config/smac_8m9m_mappo_config.py index c5120d5170..e46f059f4f 100644 --- a/dizoo/smac/config/smac_8m9m_mappo_config.py +++ b/dizoo/smac/config/smac_8m9m_mappo_config.py @@ -48,8 +48,6 @@ ), # used in state_num of hidden_state learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, epoch_per_collect=5, batch_size=3200, learning_rate=5e-4, diff --git a/dizoo/smac/config/smac_8m9m_masac_config.py b/dizoo/smac/config/smac_8m9m_masac_config.py index e394b8fc60..bbd7f3ec3b 100644 --- a/dizoo/smac/config/smac_8m9m_masac_config.py +++ b/dizoo/smac/config/smac_8m9m_masac_config.py @@ -1,5 +1,4 @@ from easydict import EasyDict -from ding.entry import serial_pipeline agent_num = 8 collector_env_num = 8 @@ -27,7 +26,7 @@ ), policy=dict( cuda=True, - on_policy=False, + multi_agent=True, random_collect_size=0, model=dict( agent_obs_shape=108, @@ -81,7 +80,7 @@ import_names=['dizoo.smac.envs.smac_env'], ), env_manager=dict(type='subprocess'), - policy=dict(type='sac_discrete', ), + policy=dict(type='discrete_sac', ), ) SMAC_8m9m_masac_default_create_config = EasyDict(SMAC_8m9m_masac_default_create_config) create_config = SMAC_8m9m_masac_default_create_config diff --git a/dizoo/smac/config/smac_MMM2_collaq_config.py b/dizoo/smac/config/smac_MMM2_collaq_config.py index 4ba21bef35..5f20178dff 100644 --- a/dizoo/smac/config/smac_MMM2_collaq_config.py +++ b/dizoo/smac/config/smac_MMM2_collaq_config.py @@ -43,7 +43,6 @@ dueling=False, ), learn=dict( - multi_gpu=False, update_per_collect=20, batch_size=32, learning_rate=0.0005, diff --git a/dizoo/smac/config/smac_MMM2_coma_config.py b/dizoo/smac/config/smac_MMM2_coma_config.py index 365861e715..d9305c59ac 100644 --- a/dizoo/smac/config/smac_MMM2_coma_config.py +++ b/dizoo/smac/config/smac_MMM2_coma_config.py @@ -32,7 +32,6 @@ actor_hidden_size_list=[64], ), learn=dict( - multi_gpu=False, update_per_collect=20, batch_size=32, learning_rate=0.0005, diff --git a/dizoo/smac/config/smac_MMM2_madqn_config.py b/dizoo/smac/config/smac_MMM2_madqn_config.py new file mode 100644 index 0000000000..fe8e96501c --- /dev/null +++ b/dizoo/smac/config/smac_MMM2_madqn_config.py @@ -0,0 +1,84 @@ +from ding.entry import serial_pipeline +from easydict import EasyDict + +agent_num = 10 +collector_env_num = 4 +evaluator_env_num = 8 + +main_config = dict( + exp_name='smac_MMM2_madqn_seed0', + env=dict( + map_name='MMM2', + difficulty=7, + reward_only_positive=True, + mirror_opponent=False, + agent_num=agent_num, + collector_env_num=collector_env_num, + evaluator_env_num=evaluator_env_num, + stop_value=0.999, + n_evaluator_episode=32, + special_global_state=True, + manager=dict(shared_memory=False, ), + ), + policy=dict( + nstep=1, + model=dict( + agent_num=agent_num, + obs_shape=204, + global_obs_shape=431, + global_cooperation=True, + action_shape=18, + hidden_size_list=[256, 256], + ), + learn=dict( + update_per_collect=40, + batch_size=64, + learning_rate=0.0005, + clip_value=5, + target_update_theta=0.008, + discount_factor=0.95, + ), + collect=dict( + collector=dict(get_train_sample=True, ), + n_episode=32, + unroll_len=20, + env_num=collector_env_num, + ), + eval=dict(env_num=evaluator_env_num, evaluator=dict(eval_freq=1000, )), + other=dict( + eps=dict( + type='linear', + start=1, + end=0.05, + decay=100000, + ), + replay_buffer=dict(replay_buffer_size=30000, ), + ), + ), +) +main_config = EasyDict(main_config) +create_config = dict( + env=dict( + type='smac', + import_names=['dizoo.smac.envs.smac_env'], + ), + env_manager=dict(type='base'), + policy=dict(type='madqn'), + collector=dict(type='episode'), +) +create_config = EasyDict(create_config) + + +def train(args): + config = [main_config, create_config] + serial_pipeline(config, seed=args.seed, max_env_step=1e7) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument('--seed', '-s', type=int, default=0) + args = parser.parse_args() + + train(args) diff --git a/dizoo/smac/config/smac_MMM2_mappo_config.py b/dizoo/smac/config/smac_MMM2_mappo_config.py index 2f318bc753..711ddb8f42 100644 --- a/dizoo/smac/config/smac_MMM2_mappo_config.py +++ b/dizoo/smac/config/smac_MMM2_mappo_config.py @@ -48,8 +48,6 @@ ), # used in state_num of hidden_state learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, epoch_per_collect=5, batch_size=1600, learning_rate=5e-4, diff --git a/dizoo/smac/config/smac_MMM2_masac_config.py b/dizoo/smac/config/smac_MMM2_masac_config.py index 8e99fc464e..3652343318 100644 --- a/dizoo/smac/config/smac_MMM2_masac_config.py +++ b/dizoo/smac/config/smac_MMM2_masac_config.py @@ -26,6 +26,7 @@ ), policy=dict( cuda=True, + multi_agent=True, random_collect_size=0, model=dict( agent_obs_shape=204, @@ -78,7 +79,7 @@ import_names=['dizoo.smac.envs.smac_env'], ), env_manager=dict(type='base'), - policy=dict(type='sac_discrete', ), + policy=dict(type='discrete_sac', ), ) SMAC_MMM2_masac_default_create_config = EasyDict(SMAC_MMM2_masac_default_create_config) create_config = SMAC_MMM2_masac_default_create_config diff --git a/dizoo/smac/config/smac_MMM2_qmix_config.py b/dizoo/smac/config/smac_MMM2_qmix_config.py index b59b30160a..9b2c21d926 100644 --- a/dizoo/smac/config/smac_MMM2_qmix_config.py +++ b/dizoo/smac/config/smac_MMM2_qmix_config.py @@ -33,7 +33,6 @@ dueling=False, ), learn=dict( - multi_gpu=False, update_per_collect=20, batch_size=32, learning_rate=0.0005, diff --git a/dizoo/smac/config/smac_MMM2_wqmix_config.py b/dizoo/smac/config/smac_MMM2_wqmix_config.py index 7db6a4be47..50c0acbad3 100644 --- a/dizoo/smac/config/smac_MMM2_wqmix_config.py +++ b/dizoo/smac/config/smac_MMM2_wqmix_config.py @@ -32,7 +32,6 @@ dueling=False, ), learn=dict( - multi_gpu=False, update_per_collect=20, batch_size=32, learning_rate=0.0005, diff --git a/dizoo/smac/config/smac_MMM_collaq_config.py b/dizoo/smac/config/smac_MMM_collaq_config.py index b002e060b0..c2ed633359 100644 --- a/dizoo/smac/config/smac_MMM_collaq_config.py +++ b/dizoo/smac/config/smac_MMM_collaq_config.py @@ -43,7 +43,6 @@ dueling=False, ), learn=dict( - multi_gpu=False, update_per_collect=20, batch_size=32, learning_rate=0.0005, diff --git a/dizoo/smac/config/smac_MMM_madqn_config.py b/dizoo/smac/config/smac_MMM_madqn_config.py new file mode 100644 index 0000000000..892f1f5217 --- /dev/null +++ b/dizoo/smac/config/smac_MMM_madqn_config.py @@ -0,0 +1,84 @@ +from ding.entry import serial_pipeline +from easydict import EasyDict + +agent_num = 10 +collector_env_num = 4 +evaluator_env_num = 8 + +main_config = dict( + exp_name='smac_MMM_madqn_seed0', + env=dict( + map_name='MMM', + difficulty=7, + reward_only_positive=True, + mirror_opponent=False, + agent_num=agent_num, + collector_env_num=collector_env_num, + evaluator_env_num=evaluator_env_num, + stop_value=0.999, + n_evaluator_episode=32, + special_global_state=True, + manager=dict(shared_memory=False, ), + ), + policy=dict( + nstep=1, + model=dict( + agent_num=agent_num, + obs_shape=186, + global_obs_shape=389, + global_cooperation=True, + action_shape=16, + hidden_size_list=[256, 256], + ), + learn=dict( + update_per_collect=20, + batch_size=64, + learning_rate=0.0005, + clip_value=5, + target_update_theta=0.008, + discount_factor=0.95, + ), + collect=dict( + collector=dict(get_train_sample=True, ), + n_episode=32, + unroll_len=10, + env_num=collector_env_num, + ), + eval=dict(env_num=evaluator_env_num, evaluator=dict(eval_freq=1000, )), + other=dict( + eps=dict( + type='linear', + start=1, + end=0.05, + decay=10000, + ), + replay_buffer=dict(replay_buffer_size=15000, ), + ), + ), +) +main_config = EasyDict(main_config) +create_config = dict( + env=dict( + type='smac', + import_names=['dizoo.smac.envs.smac_env'], + ), + env_manager=dict(type='base'), + policy=dict(type='madqn'), + collector=dict(type='episode'), +) +create_config = EasyDict(create_config) + + +def train(args): + config = [main_config, create_config] + serial_pipeline(config, seed=args.seed, max_env_step=1e7) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument('--seed', '-s', type=int, default=0) + args = parser.parse_args() + + train(args) diff --git a/dizoo/smac/config/smac_MMM_mappo_config.py b/dizoo/smac/config/smac_MMM_mappo_config.py index 22c16e1eef..8559052e81 100644 --- a/dizoo/smac/config/smac_MMM_mappo_config.py +++ b/dizoo/smac/config/smac_MMM_mappo_config.py @@ -49,8 +49,6 @@ ), # used in state_num of hidden_state learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, epoch_per_collect=5, batch_size=320, learning_rate=5e-4, diff --git a/dizoo/smac/config/smac_MMM_masac_config.py b/dizoo/smac/config/smac_MMM_masac_config.py index 2a1b0bd8e4..13e41dddff 100644 --- a/dizoo/smac/config/smac_MMM_masac_config.py +++ b/dizoo/smac/config/smac_MMM_masac_config.py @@ -26,6 +26,7 @@ ), policy=dict( cuda=True, + multi_agent=True, random_collect_size=0, model=dict( agent_obs_shape=186, @@ -78,7 +79,7 @@ import_names=['dizoo.smac.envs.smac_env'], ), env_manager=dict(type='base'), - policy=dict(type='sac_discrete', ), + policy=dict(type='discrete_sac', ), ) MMM_masac_default_create_config = EasyDict(MMM_masac_default_create_config) create_config = MMM_masac_default_create_config diff --git a/dizoo/smac/config/smac_MMM_qmix_config.py b/dizoo/smac/config/smac_MMM_qmix_config.py index b2532624a8..4ac50e3c34 100644 --- a/dizoo/smac/config/smac_MMM_qmix_config.py +++ b/dizoo/smac/config/smac_MMM_qmix_config.py @@ -33,7 +33,6 @@ dueling=False, ), learn=dict( - multi_gpu=False, update_per_collect=20, batch_size=32, learning_rate=0.0005, diff --git a/dizoo/smac/config/smac_MMM_qtran_config.py b/dizoo/smac/config/smac_MMM_qtran_config.py index 4c33dc21df..5b0a5bcd13 100644 --- a/dizoo/smac/config/smac_MMM_qtran_config.py +++ b/dizoo/smac/config/smac_MMM_qtran_config.py @@ -33,7 +33,6 @@ dueling=False, ), learn=dict( - multi_gpu=False, update_per_collect=20, batch_size=32, learning_rate=0.0005, diff --git a/dizoo/smac/config/smac_MMM_wqmix_config.py b/dizoo/smac/config/smac_MMM_wqmix_config.py index 9ef4574e28..a4fa3b4c36 100644 --- a/dizoo/smac/config/smac_MMM_wqmix_config.py +++ b/dizoo/smac/config/smac_MMM_wqmix_config.py @@ -32,7 +32,6 @@ dueling=False, ), learn=dict( - multi_gpu=False, update_per_collect=20, batch_size=32, learning_rate=0.0005, diff --git a/dizoo/smac/config/smac_corridor_mappo_config.py b/dizoo/smac/config/smac_corridor_mappo_config.py index 70aa71dcc4..e160c0c106 100644 --- a/dizoo/smac/config/smac_corridor_mappo_config.py +++ b/dizoo/smac/config/smac_corridor_mappo_config.py @@ -51,8 +51,6 @@ ), # used in state_num of hidden_state learn=dict( - # (bool) Whether to use multi gpu - multi_gpu=False, epoch_per_collect=5, batch_size=3200, learning_rate=5e-4, diff --git a/dizoo/smac/config/smac_corridor_masac_config.py b/dizoo/smac/config/smac_corridor_masac_config.py index 405c7638ce..4acf39215b 100644 --- a/dizoo/smac/config/smac_corridor_masac_config.py +++ b/dizoo/smac/config/smac_corridor_masac_config.py @@ -1,5 +1,4 @@ from easydict import EasyDict -from ding.entry import serial_pipeline agent_num = 6 collector_env_num = 8 @@ -27,7 +26,7 @@ ), policy=dict( cuda=True, - on_policy=False, + multi_agent=True, random_collect_size=0, model=dict( agent_obs_shape=192, @@ -81,7 +80,7 @@ import_names=['dizoo.smac.envs.smac_env'], ), env_manager=dict(type='subprocess'), - policy=dict(type='sac_discrete', ), + policy=dict(type='discrete_sac', ), ) smac_corridor_masac_default_create_config = EasyDict(smac_corridor_masac_default_create_config) create_config = smac_corridor_masac_default_create_config @@ -89,4 +88,4 @@ if __name__ == '__main__': from ding.entry import serial_pipeline - serial_pipeline((main_config, create_config), seed=0) + serial_pipeline((main_config, create_config), seed=0, max_env_step=int(1e7)) diff --git a/dizoo/smac/envs/fake_smac_env.py b/dizoo/smac/envs/fake_smac_env.py index 41cc53eec5..cc0199d57f 100644 --- a/dizoo/smac/envs/fake_smac_env.py +++ b/dizoo/smac/envs/fake_smac_env.py @@ -38,7 +38,7 @@ def step(self, action): done = self.step_count >= 314 info = {} if done: - info['final_eval_reward'] = 0.71 + info['eval_episode_return'] = 0.71 self.step_count += 1 return FakeSMACEnvTimestep(obs, reward, done, info) diff --git a/dizoo/smac/envs/smac_env.py b/dizoo/smac/envs/smac_env.py index 1eff5db124..f08d5e096f 100644 --- a/dizoo/smac/envs/smac_env.py +++ b/dizoo/smac/envs/smac_env.py @@ -552,10 +552,10 @@ def step(self, actions, force_return_two_player=False): self._final_eval_fake_reward += rewards new_infos["battle_lost"] = infos[OPPONENT_AGENT]["battle_won"] new_infos["draw"] = infos["draw"] - new_infos['final_eval_reward'] = infos['final_eval_reward'] + new_infos['eval_episode_return'] = infos['eval_episode_return'] if 'episode_info' in infos: new_infos['episode_info'] = infos['episode_info'] - new_infos['fake_final_eval_reward'] = infos['fake_final_eval_reward'] + new_infos['fake_eval_episode_return'] = infos['fake_eval_episode_return'] infos = new_infos if self.obs_alone: agent_state, agent_alone_state, agent_alone_padding_state = self.get_obs() @@ -611,8 +611,8 @@ def _collect_step_data(self, game_end_code, action): OPPONENT_AGENT: { "battle_won": False }, - 'final_eval_reward': 0., - 'fake_final_eval_reward': 0. + 'eval_episode_return': 0., + 'fake_eval_episode_return': 0. } if game_end_code is not None: @@ -625,7 +625,7 @@ def _collect_step_data(self, game_end_code, action): self.win_counted = True info[ORIGINAL_AGENT]["battle_won"] = True info[OPPONENT_AGENT]["battle_won"] = False - info['final_eval_reward'] = 1. + info['eval_episode_return'] = 1. elif game_end_code == -1 and not self.defeat_counted: self.defeat_counted = True info[ORIGINAL_AGENT]["battle_won"] = False @@ -639,7 +639,7 @@ def _collect_step_data(self, game_end_code, action): info[OPPONENT_AGENT]["episode_limit"] = True self.battles_game += 1 self.timeouts += 1 - # info['final_eval_reward'] = -0.5 + # info['eval_episode_return'] = -0.5 # if sum(u.health + u.shield for u in self.agents.values()) >= \ # sum(u.health + u.shield for u in self.enemies.values()): @@ -678,7 +678,7 @@ def _collect_step_data(self, game_end_code, action): # Test purpose # reward = {k: 0 * v + 100 for k, v in reward.items()} - info['fake_final_eval_reward'] = reward[ORIGINAL_AGENT] + info['fake_eval_episode_return'] = reward[ORIGINAL_AGENT] return reward, {ORIGINAL_AGENT: terminated, OPPONENT_AGENT: terminated, "__all__": terminated}, info def close(self): diff --git a/dizoo/smac/envs/test_smac_env.py b/dizoo/smac/envs/test_smac_env.py index f9d3ad683e..bf3b99ee37 100644 --- a/dizoo/smac/envs/test_smac_env.py +++ b/dizoo/smac/envs/test_smac_env.py @@ -103,8 +103,8 @@ def main(policy, map_name="3m", two_player=False): env.reset() print('reset over') terminated = False - episode_reward_me = 0 - episode_reward_op = 0 + episode_return_me = 0 + episode_return_op = 0 env_info = env.info() print('begin new episode') @@ -127,8 +127,8 @@ def main(policy, map_name="3m", two_player=False): assert reward.shape == (1, ) print('reward', reward) assert isinstance(terminated, bool) - episode_reward_me += reward["me"] if two_player else reward - episode_reward_op += reward["opponent"] if two_player else 0 + episode_return_me += reward["me"] if two_player else reward + episode_return_op += reward["opponent"] if two_player else 0 terminated = terminated["me"] if two_player else terminated if two_player: @@ -141,8 +141,8 @@ def main(policy, map_name="3m", two_player=False): draw += int(infos["draw"]) print( - "Total reward in episode {} = {} (me), {} (opponent). Me win {}, Draw {}, Opponent win {}, total {}." - "".format(e, episode_reward_me, episode_reward_op, me_win, draw, op_win, e + 1) + "Total return in episode {} = {} (me), {} (opponent). Me win {}, Draw {}, Opponent win {}, total {}." + "".format(e, episode_return_me, episode_return_op, me_win, draw, op_win, e + 1) ) env.close() diff --git a/dizoo/smac/utils/eval.py b/dizoo/smac/utils/eval.py index d8a13c8342..1e112e84a7 100644 --- a/dizoo/smac/utils/eval.py +++ b/dizoo/smac/utils/eval.py @@ -46,20 +46,20 @@ def eval( policy.load_state_dict(state_dict) obs = env.reset() - eval_reward = 0. + episode_return = 0. while True: policy_output = policy.forward({0: obs}) action = policy_output[0]['action'] print(action) timestep = env.step(action) - eval_reward += timestep.reward + episode_return += timestep.reward obs = timestep.obs if timestep.done: print(timestep.info) break env.save_replay(replay_dir='.', prefix=env._map_name) - print('Eval is over! The performance of your RL policy is {}'.format(eval_reward)) + print('Eval is over! The performance of your RL policy is {}'.format(episode_return)) if __name__ == "__main__": diff --git a/dizoo/sokoban/__init__.py b/dizoo/sokoban/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/dizoo/sokoban/envs/sokoban_env.py b/dizoo/sokoban/envs/sokoban_env.py index b4723ecafe..295259e702 100644 --- a/dizoo/sokoban/envs/sokoban_env.py +++ b/dizoo/sokoban/envs/sokoban_env.py @@ -50,17 +50,17 @@ def reset(self) -> np.ndarray: self._env.seed(self._seed) obs = self._env.reset() obs = to_ndarray(obs).astype('float32') - self._final_eval_reward = 0. + self._eval_episode_return = 0. return obs def step(self, action: np.array): action = to_ndarray(action) obs, rew, done, info = self._env.step(int(action)) - self._final_eval_reward += rew + self._eval_episode_return += rew obs = to_ndarray(obs).astype('float32') rew = to_ndarray([rew]) # wrapped to be transfered to a array with shape (1,) if done: - info['final_eval_reward'] = self._final_eval_reward + info['eval_episode_return'] = self._eval_episode_return return BaseEnvTimestep(obs, rew, done, info) def _make_env(self, only_info=False): diff --git a/dizoo/tabmwp/README.md b/dizoo/tabmwp/README.md new file mode 100644 index 0000000000..410aed8e6f --- /dev/null +++ b/dizoo/tabmwp/README.md @@ -0,0 +1,16 @@ +## TabMWP Env + +## Dataset + +The **TabMWP** dataset contains 38,431 tabular math word problems. Each question in **TabMWP** is aligned with a tabular context, which is presented as an image, semi-structured text, and a structured table. There are two types of questions: *free-text* and *multi-choice*, and each problem is annotated with gold solutions to reveal the multi-step reasoning process. + +The environment is described in the paper [Dynamic Prompt Learning via Policy Gradient for Semi-structured Mathematical Reasoning](https://arxiv.org/abs/2209.14610) by Pan Lu, Liang Qiu, Kai-Wei Chang, Ying Nian Wu, Song-Chun Zhu, Tanmay Rajpurohit, Peter Clark, Ashwin Kalyan, 2023. + +You can find more details in [Prompt PG](https://github.com/lupantech/PromptPG) + +## Benchmark + +- We collect the responses of GPT-3 using a reduced dataset with 80 training samples and 16 candidates. In this way, there is no need for users to interact with GPT-3 using the API-key of openai. +- You can directly reproduce the benchmark by running ``python dizoo/tabmwp/configs/tabmwp_pg_config.py`` + +![origin](./benchmark.png) diff --git a/dizoo/tabmwp/__init__.py b/dizoo/tabmwp/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/dizoo/tabmwp/benchmark.png b/dizoo/tabmwp/benchmark.png new file mode 100644 index 0000000000..9dfccc56c0 Binary files /dev/null and b/dizoo/tabmwp/benchmark.png differ diff --git a/dizoo/tabmwp/config/tabmwp_awr_config.py b/dizoo/tabmwp/config/tabmwp_awr_config.py new file mode 100644 index 0000000000..7e3f22865f --- /dev/null +++ b/dizoo/tabmwp/config/tabmwp_awr_config.py @@ -0,0 +1,66 @@ +from easydict import EasyDict + +tabmwp_prompt_awr_config = dict( + exp_name='tabmwp_prompt_awr_seed0', + env=dict( + collector_env_num=1, + evaluator_env_num=1, + n_evaluator_episode=1, + stop_value=1, + cand_number=16, + train_number=80, + engine='text-davinci-002', + temperature=0., + max_tokens=512, + top_p=1., + frequency_penalty=0., + presence_penalty=0., + option_inds=["A", "B", "C", "D", "E", "F"], + # The API-key of openai. You can get your key in this website: https://platform.openai.com/ + api_key='', + enable_replay=True, + prompt_format='TQ-A', + seed=0, + ), + policy=dict( + cuda=True, + shot_number=2, + model=dict( + model_name="bert-base-uncased", + add_linear=True, + freeze_encoder=True, + embedding_size=128, + ), + learn=dict( + batch_size=10, + # (bool) Whether to normalize advantage. Default to False. + learning_rate=0.001, + # (float) loss weight of the value network, the weight of policy network is set to 1 + entropy_weight=0.001, + weight_decay=5e-3, + grad_norm=0.5, + ), + collect=dict( + # (int) collect n_sample data, train model 1 times + n_sample=20, + discount_factor=0., + ), + eval=dict(evaluator=dict(eval_freq=500, )), + ), +) +main_config = EasyDict(tabmwp_prompt_awr_config) + +tabmwp_prompt_awr_config = dict( + env=dict( + type='tabmwp', + import_names=['dizoo.tabmwp.envs.tabmwp_env'], + ), + env_manager=dict(type='base'), + policy=dict(type='prompt_awr'), + replay_buffer=dict(type='naive'), +) +create_config = EasyDict(tabmwp_prompt_awr_config) + +if __name__ == '__main__': + from ding.entry import serial_pipeline_onpolicy + serial_pipeline_onpolicy((main_config, create_config), seed=0) diff --git a/dizoo/tabmwp/config/tabmwp_pg_config.py b/dizoo/tabmwp/config/tabmwp_pg_config.py new file mode 100644 index 0000000000..acda7bcdbd --- /dev/null +++ b/dizoo/tabmwp/config/tabmwp_pg_config.py @@ -0,0 +1,66 @@ +from easydict import EasyDict + +tabmwp_prompt_pg_config = dict( + exp_name='tabmwp_prompt_pg_seed0', + env=dict( + collector_env_num=1, + evaluator_env_num=1, + n_evaluator_episode=1, + stop_value=1, + cand_number=16, + train_number=80, + engine='text-davinci-002', + temperature=0., + max_tokens=512, + top_p=1., + frequency_penalty=0., + presence_penalty=0., + option_inds=["A", "B", "C", "D", "E", "F"], + # The API-key of openai. You can get your key in this website: https://platform.openai.com/ + api_key='', + enable_replay=True, + prompt_format='TQ-A', + seed=0, + ), + policy=dict( + cuda=True, + shot_number=2, + model=dict( + model_name="bert-base-uncased", + add_linear=True, + freeze_encoder=True, + embedding_size=128, + ), + learn=dict( + batch_size=10, + # (bool) Whether to normalize advantage. Default to False. + learning_rate=0.001, + # (float) loss weight of the value network, the weight of policy network is set to 1 + entropy_weight=0.001, + weight_decay=5e-3, + grad_norm=0.5, + ), + collect=dict( + # (int) collect n_sample data, train model 1 times + n_sample=20, + discount_factor=0., + ), + eval=dict(evaluator=dict(eval_freq=500, )), + ), +) +main_config = EasyDict(tabmwp_prompt_pg_config) + +tabmwp_prompt_pg_config = dict( + env=dict( + type='tabmwp', + import_names=['dizoo.tabmwp.envs.tabmwp_env'], + ), + env_manager=dict(type='base'), + policy=dict(type='prompt_pg'), + replay_buffer=dict(type='naive'), +) +create_config = EasyDict(tabmwp_prompt_pg_config) + +if __name__ == '__main__': + from ding.entry import serial_pipeline_onpolicy + serial_pipeline_onpolicy((main_config, create_config), seed=0) diff --git a/dizoo/tabmwp/envs/__init__.py b/dizoo/tabmwp/envs/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/dizoo/tabmwp/envs/tabmwp_env.py b/dizoo/tabmwp/envs/tabmwp_env.py new file mode 100644 index 0000000000..fe32e02b35 --- /dev/null +++ b/dizoo/tabmwp/envs/tabmwp_env.py @@ -0,0 +1,266 @@ +import os +from functools import lru_cache + +import gym +import openai +import numpy as np + +from ding.utils import ENV_REGISTRY +from ding.envs import BaseEnv, BaseEnvTimestep +from dizoo.tabmwp.envs.utils import create_example_from_pid, build_prompt, get_gpt3_output, calc_rwkv, calc_internlm,\ + extract_prediction, normalize_answer, load_data + + +@ENV_REGISTRY.register('tabmwp') +class TabMWP(BaseEnv): + model = None + tokenizer = None + + def __init__(self, cfg): + self.cfg = cfg + self.enable_replay = cfg.enable_replay + self._init_flag = False + self.problems, self.cand_pids, self.train_pids = None, None, None + self.problem_id = 0 + self.cand_examples = [] + openai.api_key = cfg.api_key + self.observation_space = gym.spaces.Dict() + self.action_space = gym.spaces.Discrete(self.cfg.cand_number * (self.cfg.cand_number - 1)) + self.reward_space = gym.spaces.Box(low=-1, high=1, shape=(1, ), dtype=np.float32) + self.correct_num = 0 + + # Initialize language model if needed. + assert self.cfg.engine in ['text-davinci-002', 'glm-10B', 'rwkv-7B', 'internlm-7B'] + + try: + if self.cfg.engine == 'glm-10B' and TabMWP.model is None: + from transformers import AutoTokenizer, AutoModelForSeq2SeqLM + TabMWP.tokenizer = AutoTokenizer.from_pretrained("THUDM/glm-10b", trust_remote_code=True) + model = AutoModelForSeq2SeqLM.from_pretrained("THUDM/glm-10b", trust_remote_code=True) + TabMWP.model = model.half() + elif self.cfg.engine == 'rwkv-7B' and TabMWP.model is None: + from transformers import AutoTokenizer, RwkvForCausalLM + TabMWP.tokenizer = AutoTokenizer.from_pretrained("sgugger/rwkv-7b-pile", trust_remote_code=True) + model = RwkvForCausalLM.from_pretrained("sgugger/rwkv-7b-pile") + TabMWP.model = model.half() + elif self.cfg.engine == 'internlm-7B' and TabMWP.model is None: + from transformers import AutoTokenizer, AutoModelForCausalLM + TabMWP.tokenizer = AutoTokenizer.from_pretrained("internlm/internlm-7b", trust_remote_code=True) + model = AutoModelForCausalLM.from_pretrained("internlm/internlm-7b", trust_remote_code=True) + TabMWP.model = model.eval() + except ImportError: + import sys + from ditk import logging + logging.warning("not found transformer, please install it using: pip install transformers") + sys.exit(1) + + @lru_cache(maxsize=10000) + def get_output(self, inp: str) -> str: + inputs = TabMWP.tokenizer(inp + " [MASK].", return_tensors="pt") + inputs = TabMWP.tokenizer.build_inputs_for_generation(inputs, max_gen_length=512) + inputs = {key: value.cuda() for key, value in inputs.items()} + outputs = TabMWP.model.generate( + **inputs, + max_length=512, + eos_token_id=TabMWP.tokenizer.eop_token_id, + pad_token_id=TabMWP.tokenizer.eos_token_id + ) + outputs = TabMWP.tokenizer.decode(outputs[0].tolist()) + + t0 = outputs.find('<|startofpiece|>') + 16 + t1 = outputs.find('<|endofpiece|>') + + return outputs[t0:t1] + + def seed(self, seed: int, dynamic_seed: bool = False) -> None: + self.cfg.seed = seed + + def reset(self) -> dict: + self.problems, self.cand_pids, self.train_pids = load_data(self.cfg) + if TabMWP.model is not None: + TabMWP.model = TabMWP.model.cuda() + if self.enable_replay: + self.cand_pids = [ + '32889', '8044', '16892', '5408', '4051', '37355', '17962', '25807', '30602', '5514', '19270', '23713', + '17209', '33379', '34987', '11177' + ] + if self.cfg.seed == 0: # train + self.train_pids = [ + '14229', '3409', '29980', '799', '5086', '21778', '36441', '34146', '69', '33433', '26979', '18135', + '13347', '17679', '38426', '3454', '10432', '31011', '12162', '13063', '7812', '29661', '24482', + '4970', '4405', '17405', '27781', '26724', '5993', '16442', '30148', '15895', '6855', '29903', + '18107', '29504', '11106', '32964', '29891', '32104', '15712', '24287', '4997', '32581', '21020', + '17247', '31455', '13245', '15850', '10011', '10313', '10158', '1817', '33479', '35842', '14198', + '26039', '3791', '4909', '37056', '7144', '8185', '2131', '4398', '38199', '29520', '37329', + '21388', '28659', '15044', '28510', '12903', '11794', '37095', '32229', '22918', '31680', '15024', + '24607', '26930' + ] + model_io_path = 'dizoo/tabmwp/data/model_in_out_train.txt' + if not os.path.exists(model_io_path): + os.system( + f'wget https://opendilab.net/download/DI-zoo/tabmwp/model_in_out_train.txt -O ' + + model_io_path + ' --no-check-certificate' + ) + else: + self.train_pids = [ + '21037', '22976', '2224', '14145', '27962', '26553', '22110', '16541', '26044', '19492', '31882', + '11991', '27594', '7637', '15394', '7666', '5177', '33761', '13703', '29105' + ] + model_io_path = 'dizoo/tabmwp/data/model_in_out_eval.txt' + os.system( + f'wget https://opendilab.net/download/DI-zoo/tabmwp/model_in_out_eval.txt -O ' + model_io_path + + ' --no-check-certificate' + ) + + self.cfg.cand_number = len(self.cand_pids) + self.cfg.train_number = len(self.train_pids) + + self.results_memory = [] + with open(model_io_path, encoding="ISO-8859-1") as f: + tmp = f.read().split('\n') + for tt in tmp: + if len(tt.strip()) == 0: + continue + self.results_memory.append(eval(tt)) + + self.cand_examples = [] + self.correct_num = 0 + for pid in self.cand_pids: + example = create_example_from_pid(pid, self.problems, self.cfg, test=True) + self.cand_examples.append(example) + + self._init_flag = True + self.problem_id = 0 + train_sample = create_example_from_pid(self.train_pids[self.problem_id], self.problems, self.cfg, test=True) + obs = {'train_sample': train_sample, 'candidate_samples': self.cand_examples} + return obs + + def search_answer(self, pid, pids): + for item in self.results_memory: + if item['pid'] != pid: + continue + if item['shot_pids'] == pids: + return item['output'] + + raise ValueError('item does not exists.') + + def parse_all_answers(self): + self.cand_pids = [ + '32889', '8044', '16892', '5408', '4051', '37355', '17962', '25807', '30602', '5514', '19270', '23713', + '17209', '33379', '34987', '11177', '30218', '26066', '24169', '28492' + ] + self.train_pids = [ + '14229', '3409', '29980', '799', '5086', '21778', '36441', '34146', '69', '33433', '26979', '18135', + '13347', '17679', '38426', '3454', '10432', '31011', '12162', '13063', '7812', '29661', '24482', '4970', + '4405', '17405', '27781', '26724', '5993', '16442', '30148', '15895', '6855', '29903', '18107', '29504', + '11106', '32964', '29891', '32104', '15712', '24287', '4997', '32581', '21020', '17247', '31455', '13245', + '15850', '10011', '10313', '10158', '1817', '33479', '35842', '14198', '26039', '3791', '4909', '37056', + '7144', '8185', '2131', '4398', '38199', '29520', '37329', '21388', '28659', '15044', '28510', '12903', + '11794', '37095', '32229', '22918', '31680', '15024', '24607', '26930' + ] + self.problem_id = 0 + self.cfg.train_number = len(self.train_pids) + n = len(self.cand_pids) + + with open('sampled_pid.txt', 'w') as f: + f.write(str(self.cand_pids) + '\n') + f.write(str(self.train_pids) + '\n') + + with open('model_in_out.txt', 'w') as f: + while self.problem_id < self.cfg.train_number: + for i in range(n): + for j in range(n): + if i == j: + continue + shot_pids = [self.cand_pids[i], self.cand_pids[j]] + pid = self.train_pids[self.problem_id] + + # generate the prompt input + prompt = build_prompt(self.problems, shot_pids, pid, self.cfg) + + # get the output from LM + # assert self._args.engine == 'text-davinci-002' + output = get_gpt3_output(prompt, self.cfg) + + output_txt = {'shot_pids': shot_pids, 'pid': pid, 'prompt': prompt, 'output': output} + f.write(str(output_txt) + '\n') + print(self.problem_id, i, j) + + self.problem_id += 1 + + def close(self) -> None: + self._init_flag = False + + def step(self, action: np.array) -> BaseEnvTimestep: + shot_pids = [self.cand_pids[cid] for cid in action] + pid = self.train_pids[self.problem_id] + + # generate the prompt input + prompt = build_prompt(self.problems, shot_pids, pid, self.cfg) + + # get the output from LM + if self.enable_replay: + output = self.search_answer(pid, shot_pids) + elif self.cfg.engine == 'text-davinci-002': + output = get_gpt3_output(prompt, self.cfg) + elif self.cfg.engine == 'rwkv-7B': + output = calc_rwkv(self.model, self.tokenizer, prompt) + elif self.cfg.engine == 'internlm-7B': + output = calc_internlm(self.model, self.tokenizer, prompt, self.cfg) + else: + output = self.get_output(prompt) + + # extract the prediction from the output + prediction = extract_prediction(output, self.problems[pid]['choices'], self.cfg.option_inds) + + # normalize the number in the text + prediction_norm = normalize_answer(prediction, self.problems[pid]['unit']) + + if prediction_norm.lower() == normalize_answer(self.problems[pid]['answer'], + self.problems[pid]['unit']).lower(): + reward = 1 + self.correct_num += 1 + else: + reward = -1 + + self.problem_id += 1 + if self.problem_id == self.cfg.train_number: + done = True + info = {'eval_episode_return': self.correct_num / self.cfg.train_number} + else: + done = False + info = {} + + train_sample = create_example_from_pid(pid, self.problems, self.cfg, test=True) + obs = {'train_sample': train_sample, 'candidate_samples': self.cand_examples} + + return BaseEnvTimestep(obs, reward, done, info) + + def __repr__(self) -> str: + return "DI-engine tabmwp Env" + + +if __name__ == '__main__': + from easydict import EasyDict + env_cfg = EasyDict( + dict( + cand_number=16, + train_number=20, + engine='text-davinci-002', + temperature=0., + max_tokens=512, + top_p=1., + frequency_penalty=0., + presence_penalty=0., + option_inds=["A", "B", "C", "D", "E", "F"], + api_key='xxx', + prompt_format='TQ-A', + enable_replay=True, + seed=0, + ) + ) + env = TabMWP(env_cfg) + env.seed(0) + env.reset() + env.parse_all_answers() + env.search_answer('22976', ['32889', '8044']) diff --git a/dizoo/tabmwp/envs/test_tabmwp_env.py b/dizoo/tabmwp/envs/test_tabmwp_env.py new file mode 100644 index 0000000000..ca9020d971 --- /dev/null +++ b/dizoo/tabmwp/envs/test_tabmwp_env.py @@ -0,0 +1,25 @@ +from easydict import EasyDict +import pytest +from dizoo.tabmwp.envs.tabmwp_env import TabMWP + + +@pytest.mark.envtest +class TestSokoban: + + def test_tabmwp(self): + config = dict( + cand_number=20, + train_number=100, + engine='text-davinci-002', + temperature=0., + max_tokens=512, + top_p=1., + frequency_penalty=0., + presence_penalty=0., + option_inds=["A", "B", "C", "D", "E", "F"], + api_key='', + ) + config = EasyDict(config) + env = TabMWP(config) + env.seed(0) + env.close() diff --git a/dizoo/tabmwp/envs/utils.py b/dizoo/tabmwp/envs/utils.py new file mode 100644 index 0000000000..c97c183935 --- /dev/null +++ b/dizoo/tabmwp/envs/utils.py @@ -0,0 +1,354 @@ +import json +import os +import random +import re +import time +from functools import lru_cache +import torch + +import numpy as np +import openai +try: + import transformers +except ImportError: + import sys + from ditk import logging + logging.warning("not found transformer, please install it using: pip install transformers") + sys.exit(1) + + +def sample_logits(out: torch.Tensor, temperature: float = 1.0, top_p: float = 0.8) -> int: + # Sample an action given the logits. + probs = torch.softmax(out, dim=-1).cpu().numpy() + sorted_probs = np.sort(probs)[::-1] + cumulative_probs = np.cumsum(sorted_probs) + cutoff = float(sorted_probs[np.argmax(cumulative_probs > top_p)]) + probs[probs < cutoff] = 0 + if temperature != 1.0: + probs = probs.pow(1.0 / temperature) + probs = probs / np.sum(probs) + out = np.random.choice(a=len(probs), p=probs) + return out + + +def calc_rwkv( + model: transformers.RwkvForCausalLM, + tokenizer: transformers.AutoTokenizer, + prompt: str, + max_len: int = 10 +) -> str: + # Use RWKV to generate sentence. + orig_len = len(prompt) + inputs = tokenizer(prompt, return_tensors="pt").to('cuda') + outputs = model(**inputs, labels=inputs["input_ids"]) + out, state = outputs.logits, outputs.state + # Recurrent generation. + with torch.no_grad(): + for i in range(max_len): + token = sample_logits(out[0, -1]) + tmp = tokenizer.decode([token]) + prompt = prompt + tmp + inputs = tokenizer(prompt, return_tensors="pt").to('cuda') + outputs = model(**inputs, labels=inputs["input_ids"]) + out, state = outputs.logits, outputs.state + return prompt[orig_len:] + + +def calc_internlm(model, tokenizer, prompt: str, args): + inputs = tokenizer(prompt, return_tensors="pt") + for k, v in inputs.items(): + inputs[k] = v.cuda() + gen_kwargs = { + "max_length": args.max_tokens, + "top_p": args.top_p, + "temperature": args.temperature, + "do_sample": True, + "repetition_penalty": args.frequency_penalty + } + output = model.generate(**inputs, **gen_kwargs) + output = tokenizer.decode(output) + return output + + +def load_data(args: dict) -> tuple: + # Load tabmwp dataset. + random.seed(args.seed) + data_root = 'dizoo/tabmwp/data' + + if not os.path.exists(data_root): + os.mkdir(data_root) + + if not os.path.exists(os.path.join(data_root, f'problems_train.json')): + os.system( + f'wget https://opendilab.net/download/DI-zoo/tabmwp/problems_train.json -O ' + + os.path.join(data_root, f'problems_train.json') + ' --no-check-certificate' + ) + problems = json.load(open(os.path.join(data_root, f'problems_train.json'))) + + pids = list(problems.keys()) + samples = random.sample(pids, args.train_number + args.cand_number) # random sample + train_pids = samples[:args.train_number] + cand_pids = samples[args.train_number:] + return problems, cand_pids, train_pids + + +def get_gpt3_output(prompt: str, args: dict) -> str: + return call_gpt3( + args.engine, prompt, args.temperature, args.max_tokens, args.top_p, args.frequency_penalty, + args.presence_penalty + ) + + +@lru_cache(maxsize=10000) +def call_gpt3( + engine: str, prompt: str, temperature: float, max_tokens: int, top_p: float, frequency_penalty: float, + presence_penalty: float +) -> str: + patience = 100 + while True: + try: + response = openai.Completion.create( + engine=engine, + prompt=prompt, + temperature=temperature, + max_tokens=max_tokens, + top_p=top_p, + frequency_penalty=frequency_penalty, + presence_penalty=presence_penalty, + stop=["\n"] + ) + output = response["choices"][0]["text"].strip() + break + except Exception: + patience -= 1 + if not patience: + print("!!! running out of patience waiting for OpenAI") + else: + time.sleep(0.1) + return output + + +def get_table_text(problem: dict) -> str: + table = problem['table'] + title = problem['table_title'] + if title and len(title) > 0: + table = f"[TITLE]: {title}\n{table}" + return table + + +def get_question_text(problem: dict, option_inds: list) -> str: + question = problem['question'] + + unit = problem['unit'] + if unit and len(unit) > 0: + question = f"{question} (Unit: {unit})" + + choices = problem['choices'] + if choices and len(choices) > 0: + choice_list = [] + for i, c in enumerate(choices): + choice_list.append("({}) {}".format(option_inds[i], c)) + options = " ".join(choice_list) + question = f"{question}\nOptions: {options}" + + return question + + +def get_answer(problem: dict) -> str: + return problem['answer'] + + +def get_solution_text(problem: dict) -> str: + # GPT-3 can generate the solution with more tokens + solution = problem['solution'].replace("\n", "\\n") + return solution + + +def create_one_example( + format: str, table: str, question: str, answer: str, solution: str, test_example: bool = True +) -> str: + # Using template to generate one prompt example. + input_format, output_format = format.split("-") # e.g., "TQ-A" + + elements = { + "Q": f"Question: {question}", + "T": f"Table: {table}", + "S": f"Solution: {solution}", + "A": f"Answer: The answer is {answer}.", + "AS": f"Answer: The answer is {answer}. BECAUSE: {solution}", + "SA": f"Answer: {solution} The answer is {answer}." + } + + # Input + input = "\n".join(elements[label] for label in input_format) + + # Output + if test_example: + output = "Answer:" + else: + output = elements[output_format] + + # Prompt text + text = input + "\n" + output + text = text.replace(" ", " ").strip() + + return text + + +def build_prompt(problems: list, shot_pids: list, test_pid: int, args: dict) -> str: + # Given ids, generate the complete prompt. That is, the input to LM. + examples = [] + pids = shot_pids + [test_pid] + + # n-shot training examples + for pid in pids: + problem = problems[pid] + table = get_table_text(problem) + question = get_question_text(problem, args.option_inds) + answer = get_answer(problem) + solution = get_solution_text(problems[pid]) + + if pid == test_pid: + assert pid not in shot_pids + example = create_one_example(args.prompt_format, table, question, answer, solution, test_example=True) + else: + example = create_one_example(args.prompt_format, table, question, answer, solution, test_example=False) + + examples.append(example) + + # create the prompt input + prompt_input = '\n\n'.join(examples) + + return prompt_input + + +def extract_prediction(output: str, options: list, option_inds: list) -> str: + idx = output.find('\n') + if idx > 0: + output = output[:idx] + idx = output.find('=') + if idx > 0: + output = output[idx + 1:].strip() + # $\\frac{16}{95}$ -> 16/95 + output = re.sub(r"\$?\\frac\{([\d\.\,\-]+)\}\{([\d\.\,]+)\}\$?", r"\1/\2", output) + + output = re.sub(r"(? 0: + pred = res[0].upper() # e.g., "B" + if pred in option_inds: + ind = option_inds.index(pred) # 1 + if ind >= len(options): + ind = random.choice(range(len(options))) + predition = options[ind] + return predition + + # find the most similar options + scores = [score_string_similarity(x, output) for x in options] + max_idx = int(np.argmax(scores)) # json does not recognize NumPy data types + predition = options[max_idx] + return predition + + else: + # free_text QA problems, numeric answer + patterns = [ + # r'^\([A-Za-z]\) ([\s\S]+)$', # "(A) XXXXX" + # r'[Th]he answer is \([A-Za-z]\) ([\s\S]+)$', # "The answer is (B) XXXXX." + r'[Th]he answer is ([\s\S]+)$', # "The answer is XXXXX.", + r'[Th]he table shows that ([\d\$\.\,\/\:]+) ', + r' = ([\d\$\.\,\/\:]+)', # "= $1.40" + r'(?<= be| is) ([\-\d\$\.\,\/\:]{0,}[\d]+)', # "will be $1.40" + r'(?<= are| was) ([\-\d\$\.\,\/\:]{0,}[\d]+)', # "are $1.40" + r'(?<= were) ([\-\d\$\.\,\/\:]{0,}[\d]+)', # "are $1.40" + r' ([\d\$\.\,\/\:]+ [AP]\.M\.)', # 7:25 P.M. + r'([\-\d\$\.\,\/\:]{0,}[\d]+)', # 14.5 + ] + + for p in patterns: + pattern = re.compile(p) + res = pattern.findall(output) + if len(res) > 0: + predition = res[-1].strip() + if predition.endswith(".") and ".M." not in predition: + predition = predition[:-1] + return predition + + return output + + +def normalize_answer(text: str, unit: str) -> str: + # ["1,000", "123", "3/4", "56.456", "$56.4", "-3", "-10.02", "-3/2"] + + text = re.sub("^[\$]", "", text) + text = re.sub("[\,\.\,\/]$", "", text) + result = re.match("^[-+]?[\d,./]+$", text) + + if result is not None: + # is number? + text = text.replace(",", "") + result = re.match("[-+]?\d+$", text) + try: + if result is not None: + number = int(text) + elif "/" in text: + nums = text.split("/") + number = round(float(nums[0]) / float(nums[1]), 3) + else: + number = round(float(text), 3) + number = str(number) + number = re.sub(r"\.[0]+$", "", number) + return number + except: + return text + else: + # is text + if unit: + text = text.replace(unit, "").strip() + return text + + +def score_string_similarity(str1: str, str2: str) -> float: + if str1 == str2: + return 2.0 + if " " in str1 or " " in str2: + str1_split = str1.split(" ") + str2_split = str2.split(" ") + overlap = list(set(str1_split) & set(str2_split)) + return len(overlap) / max(len(str1_split), len(str2_split)) + else: + if str1 == str2: + return 1.0 + else: + return 0.0 + + +def create_example_from_pid(pid: int, problems: list, args: dict, test: bool = False) -> str: + problem = problems[pid] + table = get_table_text(problem) + question = get_question_text(problem, args.option_inds) + answer = get_answer(problem) + solution = get_solution_text(problems[pid]) + + if test: + example = create_one_example(args.prompt_format, table, question, answer, solution, test_example=True) + else: + example = create_one_example(args.prompt_format, table, question, answer, solution, test_example=False) + + return example diff --git a/dizoo/tabmwp/tabmwp.jpeg b/dizoo/tabmwp/tabmwp.jpeg new file mode 100644 index 0000000000..e3e4efe8f4 Binary files /dev/null and b/dizoo/tabmwp/tabmwp.jpeg differ diff --git a/dizoo/taxi/Taxi-v3_episode_0.gif b/dizoo/taxi/Taxi-v3_episode_0.gif new file mode 100644 index 0000000000..33cd611e4a Binary files /dev/null and b/dizoo/taxi/Taxi-v3_episode_0.gif differ diff --git a/dizoo/taxi/__init__.py b/dizoo/taxi/__init__.py new file mode 100644 index 0000000000..83f128dd69 --- /dev/null +++ b/dizoo/taxi/__init__.py @@ -0,0 +1 @@ +from .envs import * diff --git a/dizoo/taxi/config/__init__.py b/dizoo/taxi/config/__init__.py new file mode 100644 index 0000000000..6629536be8 --- /dev/null +++ b/dizoo/taxi/config/__init__.py @@ -0,0 +1 @@ +from .taxi_dqn_config import main_config, create_config diff --git a/dizoo/taxi/config/taxi_dqn_config.py b/dizoo/taxi/config/taxi_dqn_config.py new file mode 100644 index 0000000000..312442b9f1 --- /dev/null +++ b/dizoo/taxi/config/taxi_dqn_config.py @@ -0,0 +1,48 @@ +from easydict import EasyDict + +taxi_dqn_config = dict( + exp_name='taxi_dqn_seed0', + env=dict( + collector_env_num=8, + evaluator_env_num=8, + n_evaluator_episode=8, + stop_value=20, + max_episode_steps=60, + env_id="Taxi-v3" + ), + policy=dict( + cuda=True, + model=dict(obs_shape=34, action_shape=6, encoder_hidden_size_list=[128, 128]), + random_collect_size=5000, + nstep=3, + discount_factor=0.99, + learn=dict( + update_per_collect=10, + batch_size=64, + learning_rate=0.0001, + learner=dict(hook=dict(log_show_after_iter=1000, )), + ), + collect=dict(n_sample=32), + eval=dict(evaluator=dict(eval_freq=1000, )), + other=dict( + eps=dict(type="linear", start=1, end=0.05, decay=3000000), + replay_buffer=dict(replay_buffer_size=100000, ), + ), + ) +) +taxi_dqn_config = EasyDict(taxi_dqn_config) +main_config = taxi_dqn_config + +taxi_dqn_create_config = dict( + env=dict(type="taxi", import_names=["dizoo.taxi.envs.taxi_env"]), + env_manager=dict(type='base'), + policy=dict(type='dqn'), + replay_buffer=dict(type='deque', import_names=['ding.data.buffer.deque_buffer_wrapper']), +) + +taxi_dqn_create_config = EasyDict(taxi_dqn_create_config) +create_config = taxi_dqn_create_config + +if __name__ == "__main__": + from ding.entry import serial_pipeline + serial_pipeline((main_config, create_config), max_env_step=3000000, seed=0) diff --git a/dizoo/taxi/entry/taxi_dqn_deploy.py b/dizoo/taxi/entry/taxi_dqn_deploy.py new file mode 100644 index 0000000000..d2faba273e --- /dev/null +++ b/dizoo/taxi/entry/taxi_dqn_deploy.py @@ -0,0 +1,35 @@ +import gym +import torch +from easydict import EasyDict + +from ding.config import compile_config +from ding.envs import DingEnvWrapper +from ding.model import DQN +from ding.policy import DQNPolicy, single_env_forward_wrapper +from dizoo.taxi.config.taxi_dqn_config import create_config, main_config +from dizoo.taxi.envs.taxi_env import TaxiEnv + + +def main(main_config: EasyDict, create_config: EasyDict, ckpt_path: str) -> None: + main_config.exp_name = f'taxi_dqn_seed0_deploy' + cfg = compile_config(main_config, create_cfg=create_config, auto=True) + env = TaxiEnv(cfg.env) + env.enable_save_replay(replay_path=f'./{main_config.exp_name}/video') + model = DQN(**cfg.policy.model) + state_dict = torch.load(ckpt_path, map_location='cpu') + model.load_state_dict(state_dict['model']) + policy = DQNPolicy(cfg.policy, model=model).eval_mode + forward_fn = single_env_forward_wrapper(policy.forward) + obs = env.reset() + returns = 0. + while True: + action = forward_fn(obs) + obs, rew, done, info = env.step(action) + returns += rew + if done: + break + print(f'Deploy is finished, final epsiode return is: {returns}') + + +if __name__ == "__main__": + main(main_config=main_config, create_config=create_config, ckpt_path=f'./taxi_dqn_seed0/ckpt/ckpt_best.pth.tar') diff --git a/dizoo/taxi/envs/__init__.py b/dizoo/taxi/envs/__init__.py new file mode 100644 index 0000000000..188628ae48 --- /dev/null +++ b/dizoo/taxi/envs/__init__.py @@ -0,0 +1 @@ +from .taxi_env import TaxiEnv diff --git a/dizoo/taxi/envs/taxi_env.py b/dizoo/taxi/envs/taxi_env.py new file mode 100644 index 0000000000..c47c8c49f1 --- /dev/null +++ b/dizoo/taxi/envs/taxi_env.py @@ -0,0 +1,167 @@ +from typing import List, Optional +import os + +from easydict import EasyDict +from gym.spaces import Space, Discrete +from gym.spaces.box import Box +import gym +import numpy as np +import imageio + +from ditk import logging +from ding.envs.env.base_env import BaseEnv, BaseEnvTimestep +from ding.torch_utils import to_ndarray +from ding.utils import ENV_REGISTRY + + +@ENV_REGISTRY.register('taxi', force_overwrite=True) +class TaxiEnv(BaseEnv): + + def __init__(self, cfg: EasyDict) -> None: + + self._cfg = cfg + assert self._cfg.env_id == "Taxi-v3", "Your environment name is not Taxi-v3!" + self._init_flag = False + self._replay_path = None + self._save_replay = False + self._frames = [] + + def reset(self) -> np.ndarray: + if not self._init_flag: + self._env = gym.make( + id=self._cfg.env_id, render_mode="single_rgb_array", max_episode_steps=self._cfg.max_episode_steps + ) + self._observation_space = self._env.observation_space + self._action_space = self._env.action_space + self._reward_space = Box( + low=self._env.reward_range[0], high=self._env.reward_range[1], shape=(1, ), dtype=np.float32 + ) + self._init_flag = True + self._eval_episode_return = 0 + if hasattr(self, '_seed') and hasattr(self, '_dynamic_seed') and self._dynamic_seed: + np_seed = 100 * np.random.randint(1, 1000) + self._env_seed = self._seed + np_seed + elif hasattr(self, '_seed'): + self._env_seed = self._seed + if hasattr(self, '_seed'): + obs = self._env.reset(seed=self._env_seed) + else: + obs = self._env.reset() + + if self._save_replay: + picture = self._env.render() + self._frames.append(picture) + self._eval_episode_return = 0. + obs = self._encode_taxi(obs).astype(np.float32) + return obs + + def close(self) -> None: + if self._init_flag: + self._env.close() + self._init_flag = False + + def seed(self, seed: int, dynamic_seed: bool = True) -> None: + self._seed = seed + self._dynamic_seed = dynamic_seed + np.random.seed(self._seed) + + def step(self, action: np.ndarray) -> BaseEnvTimestep: + assert isinstance(action, np.ndarray), type(action) + action = action.item() + obs, rew, done, info = self._env.step(action) + self._eval_episode_return += rew + obs = self._encode_taxi(obs) + rew = to_ndarray([rew]) # Transformed to an array with shape (1, ) + if self._save_replay: + picture = self._env.render() + self._frames.append(picture) + if done: + info['eval_episode_return'] = self._eval_episode_return + if self._save_replay: + assert self._replay_path is not None, "your should have a path" + path = os.path.join( + self._replay_path, '{}_episode_{}.gif'.format(self._cfg.env_id, self._save_replay_count) + ) + self.frames_to_gif(self._frames, path) + self._frames = [] + self._save_replay_count += 1 + rew = rew.astype(np.float32) + obs = obs.astype(np.float32) + return BaseEnvTimestep(obs, rew, done, info) + + def enable_save_replay(self, replay_path: Optional[str] = None) -> None: + if replay_path is None: + replay_path = './video' + if not os.path.exists(replay_path): + os.makedirs(replay_path) + self._replay_path = replay_path + self._save_replay = True + self._save_replay_count = 0 + + def random_action(self) -> np.ndarray: + random_action = self.action_space.sample() + if isinstance(random_action, np.ndarray): + pass + elif isinstance(random_action, int): + random_action = to_ndarray([random_action], dtype=np.int64) + elif isinstance(random_action, dict): + random_action = to_ndarray(random_action) + else: + raise TypeError( + '`random_action` should be either int/np.ndarray or dict of int/np.ndarray, but get {}: {}'.format( + type(random_action), random_action + ) + ) + return random_action + + #todo encode the state into a vector + def _encode_taxi(self, obs: np.ndarray) -> np.ndarray: + taxi_row, taxi_col, passenger_location, destination = self._env.unwrapped.decode(obs) + encoded_obs = np.zeros(34) + encoded_obs[5 * taxi_row + taxi_col] = 1 + encoded_obs[25 + passenger_location] = 1 + encoded_obs[30 + destination] = 1 + return to_ndarray(encoded_obs) + + @property + def observation_space(self) -> Space: + return self._observation_space + + @property + def action_space(self) -> Space: + return self._action_space + + @property + def reward_space(self) -> Space: + return self._reward_space + + def __repr__(self) -> str: + return "DI-engine Taxi-v3 Env" + + @staticmethod + def frames_to_gif(frames: List[imageio.core.util.Array], gif_path: str, duration: float = 0.1) -> None: + """ + Overview: + Convert a list of frames into a GIF. + Arguments: + - frames (:obj:`List[imageio.core.util.Array]`): A list of frames, each frame is an image. + - gif_path (:obj:`str`): The path to save the GIF file. + - duration (:obj:`float`): Duration between each frame in the GIF (seconds). + """ + # Save all frames as temporary image files + temp_image_files = [] + for i, frame in enumerate(frames): + temp_image_file = f"frame_{i}.png" # Temporary file name + imageio.imwrite(temp_image_file, frame) # Save the frame as a PNG file + temp_image_files.append(temp_image_file) + + # Use imageio to convert temporary image files to GIF + with imageio.get_writer(gif_path, mode='I', duration=duration) as writer: + for temp_image_file in temp_image_files: + image = imageio.imread(temp_image_file) + writer.append_data(image) + + # Clean up temporary image files + for temp_image_file in temp_image_files: + os.remove(temp_image_file) + logging.info(f"GIF saved as {gif_path}") diff --git a/dizoo/taxi/envs/test_taxi_env.py b/dizoo/taxi/envs/test_taxi_env.py new file mode 100644 index 0000000000..ad3802ac5e --- /dev/null +++ b/dizoo/taxi/envs/test_taxi_env.py @@ -0,0 +1,36 @@ +import numpy as np +import pytest +from easydict import EasyDict +from dizoo.taxi import TaxiEnv + + +@pytest.mark.envtest +class TestTaxiEnv: + + def test_naive(self): + env = TaxiEnv(EasyDict({"env_id": "Taxi-v3", "max_episode_steps": 300})) + env.seed(314, dynamic_seed=False) + assert env._seed == 314 + obs = env.reset() + assert obs.shape == (34, ) + for _ in range(5): + env.reset() + np.random.seed(314) + print('=' * 60) + for i in range(10): + # Both ``env.random_action()``, and utilizing ``np.random`` as well as action space, + # can generate legal random action. + if i < 5: + random_action = np.array([env.action_space.sample()]) + else: + random_action = env.random_action() + timestep = env.step(random_action) + print(f"Your timestep in wrapped mode is: {timestep}") + assert isinstance(timestep.obs, np.ndarray) + assert isinstance(timestep.done, bool) + assert timestep.obs.shape == (34, ) + assert timestep.reward.shape == (1, ) + assert timestep.reward >= env.reward_space.low + assert timestep.reward <= env.reward_space.high + print(env.observation_space, env.action_space, env.reward_space) + env.close() diff --git a/docker/Dockerfile.base b/docker/Dockerfile.base index df1824bced..b4aebf8441 100644 --- a/docker/Dockerfile.base +++ b/docker/Dockerfile.base @@ -1,9 +1,9 @@ -FROM pytorch/pytorch:1.7.1-cuda11.0-cudnn8-runtime as base +FROM pytorch/pytorch:1.12.1-cuda11.3-cudnn8-runtime as base WORKDIR /ding RUN apt update \ - && apt install libgl1-mesa-glx libglib2.0-0 libsm6 libxext6 libxrender-dev swig curl git vim gcc \g++ make wget locales dnsutils -y \ + && apt install libgl1-mesa-glx libglib2.0-0 libsm6 libxext6 libxrender-dev swig curl git vim gcc \g++ make wget locales dnsutils zip unzip cmake -y \ && apt clean \ && rm -rf /var/cache/apt/* \ && sed -i '/en_US.UTF-8/s/^# //g' /etc/locale.gen \ @@ -13,6 +13,11 @@ ENV LANG en_US.UTF-8 ENV LANGUAGE en_US:UTF-8 ENV LC_ALL en_US.UTF-8 +RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y \ + && . $HOME/.cargo/env + +ENV PATH="/root/.cargo/bin:${PATH}" + ADD setup.py setup.py ADD dizoo dizoo ADD ding ding @@ -33,7 +38,7 @@ RUN apt-get update && \ python3.8 python3-pip python3.8-dev RUN apt update \ - && apt install libgl1-mesa-glx libglib2.0-0 libsm6 libxext6 libxrender-dev swig curl git vim gcc \g++ make wget locales dnsutils -y \ + && apt install libgl1-mesa-glx libglib2.0-0 libsm6 libxext6 libxrender-dev swig curl git vim gcc \g++ make wget locales dnsutils zip unzip cmake -y \ && apt clean \ && rm -rf /var/cache/apt/* \ && sed -i '/en_US.UTF-8/s/^# //g' /etc/locale.gen \ @@ -43,6 +48,13 @@ ENV LANG en_US.UTF-8 ENV LANGUAGE en_US:UTF-8 ENV LC_ALL en_US.UTF-8 +# 安装 Rust +RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y \ + && . $HOME/.cargo/env + +# 添加 Rust 到 PATH +ENV PATH="/root/.cargo/bin:${PATH}" + ADD setup.py setup.py ADD dizoo dizoo ADD ding ding diff --git a/docker/Dockerfile.env b/docker/Dockerfile.env index 931f1fee27..dbf89c7f3e 100644 --- a/docker/Dockerfile.env +++ b/docker/Dockerfile.env @@ -35,8 +35,10 @@ RUN mkdir -p /root/.mujoco \ ENV LD_LIBRARY_PATH /root/.mujoco/mjpro210/bin:/root/.mujoco/mujoco210/bin:${LD_LIBRARY_PATH} Run python3 -m pip install --upgrade pip \ + && pip3 install "cython<3" \ && pip3 install --no-cache-dir numpy \ - && pip3 install --no-cache-dir -U "gym[mujoco,mujoco_py]>=0.25.2" --user \ + && pip3 install --no-cache-dir -U "gym[mujoco,mujoco_py]==0.25.1" --user \ + && pip install gymnasium[mujoco] \ && python -c "import mujoco_py" FROM opendilab/di-star:latest as smac @@ -59,7 +61,7 @@ ENV DEBIAN_FRONTEND=noninteractive WORKDIR /ding -RUN apt-get update && apt-get install git cmake build-essential libgl1-mesa-dev libsdl2-dev \ +RUN apt-get update && apt-get install git build-essential libgl1-mesa-dev libsdl2-dev \ libsdl2-image-dev libsdl2-ttf-dev libsdl2-gfx-dev libboost-all-dev \ libdirectfb-dev libst-dev mesa-utils xvfb x11vnc -y \ && apt clean \ @@ -90,21 +92,21 @@ WORKDIR /ding RUN mkdir tempfile \ && cd tempfile \ - && wget https://github.com/rlworkgroup/metaworld/archive/refs/heads/master.zip -O metaworld_master.zip \ - && apt-get install unzip \ - && unzip metaworld_master.zip \ - && python3 -m pip install --no-cache-dir ./metaworld-master/ \ + && python3 -m pip install --no-cache-dir git+https://github.com/Farama-Foundation/Metaworld.git@b2a4cbb98e20081412cb4cc7ae3d4afc456a732a \ && cd .. \ && rm -rf tempfile +RUN apt-get install xvfb ffmpeg -y \ + && rm -rf /opt/conda/bin/ffmpeg \ + && ln -s /usr/bin/ffmpeg /opt/conda/bin/ffmpeg + FROM opendilab/ding:nightly as cityflow WORKDIR /ding RUN apt update \ && apt install -y \ - build-essential \ - cmake + build-essential RUN mkdir -p /root/.cityflow \ && cd /root/.cityflow \ @@ -118,3 +120,28 @@ RUN mkdir -p /root/.smartcross \ && cd DI-smartcross \ && pip install -e . + +FROM opendilab/ding:nightly as evogym + +WORKDIR /ding + +RUN apt update \ + && apt install -y \ + build-essential libglew-dev libglu1-mesa-dev xorg-dev + +RUN mkdir -p /root/.evogym \ + && cd /root/.evogym \ + && git clone --recurse-submodules https://github.com/PaParaZz1/evogym.git \ + && cd evogym \ + && pip3 install -r requirements.txt + +RUN cd /root/.evogym/evogym && python3 setup.py install + +FROM opendilab/ding:nightly-mujoco as d4rl + +WORKDIR /ding + +RUN git clone https://github.com/PaParaZz1/D4RL.git + +RUN cd D4RL \ + && pip install -e . diff --git a/docker/Dockerfile.rpc b/docker/Dockerfile.rpc new file mode 100644 index 0000000000..b9e9496548 --- /dev/null +++ b/docker/Dockerfile.rpc @@ -0,0 +1,23 @@ +FROM snsao/pytorch:tensorpipe-fix as base + +WORKDIR /ding + +RUN apt update \ + && apt install libgl1-mesa-glx libglib2.0-0 libsm6 libxext6 libxrender-dev swig curl git vim gcc \g++ make wget locales dnsutils -y \ + && apt clean \ + && rm -rf /var/cache/apt/* \ + && sed -i '/en_US.UTF-8/s/^# //g' /etc/locale.gen \ + && locale-gen + +ENV LANG en_US.UTF-8 +ENV LANGUAGE en_US:UTF-8 +ENV LC_ALL en_US.UTF-8 + +ADD setup.py setup.py +ADD dizoo dizoo +ADD ding ding +ADD README.md README.md + +RUN python3 -m pip install --upgrade pip \ + && python3 -m pip install --ignore-installed 'PyYAML<6.0' \ + && python3 -m pip install --no-cache-dir .[fast,test] diff --git a/pytest.ini b/pytest.ini index d88de33008..efdeaba023 100644 --- a/pytest.ini +++ b/pytest.ini @@ -8,6 +8,7 @@ markers = algotest benchmark envpooltest + other tmp norecursedirs = ding/hpc_rl/tests diff --git a/setup.py b/setup.py old mode 100755 new mode 100644 index ae2962de32..f3d60222f1 --- a/setup.py +++ b/setup.py @@ -50,44 +50,41 @@ }, python_requires=">=3.7", install_requires=[ - 'gym>=0.25.1, <0.26.0', # pypy incompatible; some environmrnt only support gym==0.22.0 - 'torch>=1.1.0, <=1.12.1', # If encountering pytorch errors, you need to do something like https://github.com/opendilab/DI-engine/discussions/81 - 'numpy>=1.18.0', - 'pandas', + 'setuptools<=66.1.1', + 'yapf==0.29.0', + 'gym==0.25.1', # pypy incompatible; some environments only support gym==0.22.0 + 'gymnasium', + 'torch>=1.1.0', + 'numpy>=1.18.0,<2', + 'DI-treetensor>=0.4.0', + 'DI-toolkit>=0.1.0', + 'trueskill', 'tensorboardX>=2.2', - 'requests>=2.25.1', + 'wandb<=0.19.0', + 'matplotlib', + 'easydict>=1.9', 'pyyaml', - 'easydict==1.9', - 'protobuf', - 'yapf==0.29.0', - 'flask~=1.1.2', - 'tqdm', - 'lz4', - 'scipy', + 'enum_tools', 'cloudpickle', + 'hickle', 'tabulate', 'click>=7.0.0', - 'URLObject>=2.4.0', - 'urllib3>=1.26.5', - 'responses~=0.12.1', - 'readerwriterlock', - 'enum_tools', - 'trueskill', - 'h5py', - 'mpire>=2.3.5', - 'pynng', - 'redis', - 'pettingzoo==1.12.0', - 'DI-treetensor>=0.4.0', - 'DI-toolkit>=0.0.2', - 'hbutils>=0.5.0', - 'MarkupSafe==2.0.1', # compatibility + 'flask<=2.0.3', # interaction + 'werkzeug<=2.0.3', # interaction + 'requests', # interaction + 'responses', # interaction + 'URLObject', # interaction + 'pynng', # parallel + 'sniffio', # parallel + 'redis', # parallel + 'mpire>=2.3.5', # parallel + 'einops', + 'transformers', + 'datasets', ], extras_require={ 'test': [ - 'gym[box2d]>=0.25.0', - 'opencv-python', # pypy incompatible - 'coverage>=5', + 'coverage>=5,<=7.0.1', 'mock>=4.0.3', 'pytest~=7.0.1', # required by gym>=0.25.0 'pytest-cov~=3.0.0', @@ -95,23 +92,37 @@ 'pytest-xdist>=1.34.0', 'pytest-rerunfailures~=10.2', 'pytest-timeout~=2.0.2', + 'readerwriterlock', + 'pandas', + 'lz4', + 'h5py', + 'scipy', + 'scikit-learn', + 'pettingzoo<=1.22.3', + 'pygame', + 'opencv-python', # pypy incompatible + 'pyecharts', ], 'style': [ 'yapf==0.29.0', 'flake8<=3.9.2', + 'importlib-metadata<5.0.0', # compatibility ], 'fast': [ 'numpy-stl', 'numba>=0.53.0', ], + 'video': [ + 'moviepy', + 'imageio[ffmpeg]', + ], 'dist': [ - 'redis==3.5.3', 'redis-py-cluster==2.1.0', ], 'common_env': [ 'ale-py', # >=0.7.5', # atari 'autorom', - 'gym[all]>=0.25.0' + 'gym[all]==0.25.1', 'cmake>=3.18.4', 'opencv-python', # pypy incompatible ], @@ -126,7 +137,7 @@ 'bsuite', ], 'minigrid_env': [ - 'gym-minigrid', + 'minigrid>=2.0.0', ], # 'd4rl_env': [ # 'd4rl @ git+https://github.com/rail-berkeley/d4rl@master#egg=d4rl', @@ -180,6 +191,7 @@ 'Programming Language :: Python :: 3.7', 'Programming Language :: Python :: 3.8', 'Programming Language :: Python :: 3.9', + 'Programming Language :: Python :: 3.10', 'Topic :: Scientific/Engineering :: Artificial Intelligence', ], )