MuJoCo Benchmark - DDPG, TD3, SAC (thu-ml#305)

Releasing Tianshou's SOTA benchmark of 9 out of 13 environments from the MuJoCo Gym task suite.

examples/mujoco/README.md

@@ -1,27 +1,135 @@
-# Mujoco Result
+# Tianshou's Mujoco Benchmark
 
+We benchmarked Tianshou algorithm implementations in 9 out of 13 environments from the MuJoCo Gym task suite<sup>[[1]](#footnote1)</sup>.
 
+For each supported algorithm and supported mujoco environments, we provide:
+- Default hyperparameters used for benchmark and scripts to reproduce the benchmark;
+- A comparison of performance (or code level details) with other open source implementations or classic papers;
+- Graphs and raw data that can be used for research purposes<sup>[[2]](#footnote2)</sup>;
+- Log details obtained during training<sup>[[2]](#footnote2)</sup>;
+- Pretrained agents<sup>[[2]](#footnote2)</sup>;
+- Some hints on how to tune the algorithm.
+
 
-## SAC (single run)
+Supported algorithms are listed below:
+- [Deep Deterministic Policy Gradient (DDPG)](https://arxiv.org/pdf/1509.02971.pdf), [commit id](https://github.com/thu-ml/tianshou/tree/v0.4.0)
+- [Twin Delayed DDPG (TD3)](https://arxiv.org/pdf/1802.09477.pdf), [commit id](https://github.com/thu-ml/tianshou/tree/v0.4.0)
+- [Soft Actor-Critic (SAC)](https://arxiv.org/pdf/1812.05905.pdf), [commit id](https://github.com/thu-ml/tianshou/tree/v0.4.0)
 
-The best reward computes from 100 episodes returns in the test phase.
+## Offpolicy algorithms
 
-SAC on Swimmer-v3 always stops at 47\~48.
+#### Usage
 
-| task           | 3M best reward    | parameters                                              | time cost (3M) |
-| -------------- | ----------------- | ------------------------------------------------------- | -------------- |
-| HalfCheetah-v3 | 10157.70 ± 171.70 | `python3 mujoco_sac.py --task HalfCheetah-v3`           | 2~3h           |
-| Walker2d-v3    | 5143.04 ± 15.57   | `python3 mujoco_sac.py --task Walker2d-v3`              | 2~3h           |
-| Hopper-v3      | 3604.19 ± 169.55  | `python3 mujoco_sac.py --task Hopper-v3`                | 2~3h           |
-| Humanoid-v3    | 6579.20 ± 1470.57 | `python3 mujoco_sac.py --task Humanoid-v3 --alpha 0.05` | 2~3h           |
-| Ant-v3         | 6281.65 ± 686.28  | `python3 mujoco_sac.py --task Ant-v3`                   | 2~3h           |
+Run
 
-![](results/sac/all.png)
+```bash
+$ python mujoco_sac.py --task Ant-v3
+```
 
-### Which parts are important?
+Logs is saved in `./log/` and can be monitored with tensorboard.
+
+```bash
+$ tensorboard --logdir log
+```
+
+You can also reproduce the benchmark (e.g. SAC in Ant-v3) with the example script we provide under `examples/mujoco/`:
+
+```bash
+$ ./run_experiments.sh Ant-v3
+```
+
+This will start 10 experiments with different seeds.
+
+#### Example benchmark
+
+<img src="./benchmark/Ant-v3/figure.png" width="500" height="450">
+
+Other graphs can be found under `/examples/mujuco/benchmark/`
+
+#### Hints
+
+In offpolicy algorithms(DDPG, TD3, SAC), the shared hyperparameters are almost the same<sup>[[8]](#footnote8)</sup>, and most hyperparameters are consistent with those used for benchmark in SpinningUp's implementations<sup>[[9]](#footnote9)</sup>.
+
+By comparison to both classic literature and open source implementations (e.g., SpinningUp)<sup>[[1]](#footnote1)</sup><sup>[[2]](#footnote2)</sup>, Tianshou's implementations of DDPG, TD3, and SAC are roughly at-parity with or better than the best reported results for these algorithms.
+
+### DDPG
+
+|      Environment       |     Tianshou      | [SpinningUp (PyTorch)](https://spinningup.openai.com/en/latest/spinningup/bench.html) | [TD3 paper (DDPG)](https://arxiv.org/abs/1802.09477) | [TD3 paper (OurDDPG)](https://arxiv.org/abs/1802.09477) |
+| :--------------------: | :---------------: | :----------------------------------------------------------: | :--------------------------------------------------: | :-----------------------------------------------------: |
+|          Ant           |     990.4±4.3     |                             ~840                             |                      **1005.3**                      |                          888.8                          |
+|      HalfCheetah       | **11718.7±465.6** |                            ~11000                            |                        3305.6                        |                         8577.3                          |
+|         Hopper         | **2197.0±971.6**  |                            ~1800                             |                      **2020.5**                      |                         1860.0                          |
+|        Walker2d        |   1400.6±905.0    |                            ~1950                             |                        1843.6                        |                       **3098.1**                        |
+|        Swimmer         |   **144.1±6.5**   |                             ~137                             |                          N                           |                            N                            |
+|        Humanoid        |  **177.3±77.6**   |                              N                               |                          N                           |                            N                            |
+|        Reacher         |   **-3.3±0.3**    |                              N                               |                        -6.51                         |                          -4.01                          |
+|    InvertedPendulum    |  **1000.0±0.0**   |                              N                               |                      **1000.0**                      |                       **1000.0**                        |
+| InvertedDoublePendulum |   8364.3±2778.9   |                              N                               |                      **9355.5**                      |                         8370.0                          |
+
+\* details<sup>[[5]](#footnote5)</sup><sup>[[6]](#footnote6)</sup><sup>[[7]](#footnote7)</sup>
+
+### TD3
+
+|      Environment       |     Tianshou      | [SpinningUp (Pytorch)](https://spinningup.openai.com/en/latest/spinningup/bench.html) |   [TD3 paper](https://arxiv.org/abs/1802.09477)    |
+| :--------------------: | :---------------: | :-------------------: | :--------------: |
+|          Ant           | **5116.4±799.9**  |         ~3800         |  4372.4±1000.3   |
+|      HalfCheetah       | **10201.2±772.8** |         ~9750         |   9637.0±859.1   |
+|         Hopper         |   3472.2±116.8    |         ~2860         | **3564.1±114.7** |
+|        Walker2d        |   3982.4±274.5    |         ~4000         | **4682.8±539.6** |
+|        Swimmer         |  **104.2±34.2**   |          ~78          |        N         |
+|        Humanoid        | **5189.5±178.5**  |           N           |        N         |
+|        Reacher         |   **-2.7±0.2**    |           N           |     -3.6±0.6     |
+|    InvertedPendulum    |  **1000.0±0.0**   |           N           |  **1000.0±0.0**  |
+| InvertedDoublePendulum |  **9349.2±14.3**  |           N           | **9337.5±15.0**  |
+
+\* details<sup>[[5]](#footnote5)</sup><sup>[[6]](#footnote6)</sup><sup>[[7]](#footnote7)</sup>
+
+### SAC
+
+|      Environment       |      Tianshou      | [SpinningUp (Pytorch)](https://spinningup.openai.com/en/latest/spinningup/bench.html) | [SAC paper](https://arxiv.org/abs/1801.01290) |
+| :--------------------: | :----------------: | :-------------------: | :---------: |
+|          Ant           |  **5850.2±475.7**  |         ~3980         |    ~3720    |
+|      HalfCheetah       | **12138.8±1049.3** |        ~11520         |   ~10400    |
+|         Hopper         |  **3542.2±51.5**   |         ~3150         |    ~3370    |
+|        Walker2d        |  **5007.0±251.5**  |         ~4250         |    ~3740    |
+|        Swimmer         |    **44.4±0.5**    |         ~41.7         |      N      |
+|        Humanoid        |  **5488.5±81.2**   |           N           |    ~5200    |
+|        Reacher         |    **-2.6±0.2**    |           N           |      N      |
+|    InvertedPendulum    |   **1000.0±0.0**   |           N           |      N      |
+| InvertedDoublePendulum |   **9359.5±0.4**   |           N           |      N      |
+
+\* details<sup>[[5]](#footnote5)</sup><sup>[[6]](#footnote6)</sup>
+
+#### Hints for SAC
 
 0. DO NOT share the same network with two critic networks.
-1. The sigma (of the Gaussian policy) MUST be conditioned on input.
+1. The sigma (of the Gaussian policy) should be conditioned on input.
 2. The network size should not be less than 256.
 3. The deterministic evaluation helps a lot :)
 
+## Onpolicy Algorithms
+
+TBD
+
+
+
+
+## Note
+
+<a name="footnote1">[1]</a>  Supported environments include HalfCheetah-v3, Hopper-v3, Swimmer-v3, Walker2d-v3, Ant-v3, Humanoid-v3, Reacher-v2, InvertedPendulum-v2 and InvertedDoublePendulum-v2. Pusher, Thrower, Striker and HumanoidStandup are not supported because they are not commonly seen in literatures.
+
+<a name="footnote2">[2]</a>  Pretrained agents, detailed graphs (single agent, single game) and log details can all be found [here](https://cloud.tsinghua.edu.cn/d/356e0f5d1e66426b9828/).
+
+<a name="footnote3">[3]</a>  We used the latest version of all mujoco environments in gym (0.17.3 with mujoco==2.0.2.13), but it's not often the case with other benchmarks. Please check for details yourself in the original paper. (Different version's outcomes are usually similar, though)
+
+<a name="footnote4">[4]</a>  We didn't compare offpolicy algorithms to OpenAI baselines [benchmark](https://github.com/openai/baselines/blob/master/benchmarks_mujoco1M.htm), because for now it seems that they haven't provided benchmark for offpolicy algorithms, but in [SpinningUp docs](https://spinningup.openai.com/en/latest/spinningup/bench.html) they stated that "SpinningUp implementations of DDPG, TD3, and SAC are roughly at-parity with the best-reported results for these algorithms", so we think lack of comparisons with OpenAI baselines is okay.
+
+<a name="footnote5">[5]</a>  ~ means the number is approximated from the graph because accurate numbers is not provided in the paper. N means graphs not provided.
+
+<a name="footnote6">[6]</a>  Reward metric: The meaning of the table value is the max average return over 10 trails (different seeds) ± a single standard deviation over trails. Each trial is averaged on another 10 test seeds. Only the first 1M steps data will be considered. The shaded region on the graph also represents a single standard deviation. It is the same as [TD3 evaluation method](https://github.com/sfujim/TD3/issues/34).
+
+<a name="footnote7">[7]</a>  In TD3 paper, shaded region represents only half of standard deviation.
+
+<a name="footnote8">[8]</a>  SAC's start-timesteps is set to 10000 by default while it is 25000 is DDPG/TD3. TD3's learning rate is set to 3e-4 while it is 1e-3 for DDPG/SAC. However, there is NO enough evidence to support our choice of such hyperparameters (we simply choose them because of SpinningUp) and you can try playing with those hyperparameters to see if you can improve performance. Do tell us if you can!
+
+<a name="footnote9">[9]</a>  We use batchsize of 256 in DDPG/TD3/SAC while SpinningUp use 100. Minor difference also lies with `start-timesteps`, data loop method `step_per_collect`, method to deal with/bootstrap truncated steps because of timelimit and unfinished/collecting episodes (contribute to performance improvement), etc.

examples/mujoco/mujoco_ddpg.py

@@ -0,0 +1,131 @@
+#!/usr/bin/env python3
+
+import os
+import gym
+import torch
+import datetime
+import argparse
+import numpy as np
+from torch.utils.tensorboard import SummaryWriter
+
+from tianshou.policy import DDPGPolicy
+from tianshou.utils import BasicLogger
+from tianshou.env import SubprocVectorEnv
+from tianshou.utils.net.common import Net
+from tianshou.exploration import GaussianNoise
+from tianshou.trainer import offpolicy_trainer
+from tianshou.utils.net.continuous import Actor, Critic
+from tianshou.data import Collector, ReplayBuffer, VectorReplayBuffer
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--task', type=str, default='Ant-v3')
+    parser.add_argument('--seed', type=int, default=0)
+    parser.add_argument('--buffer-size', type=int, default=1000000)
+    parser.add_argument('--hidden-sizes', type=int, nargs='*', default=[256, 256])
+    parser.add_argument('--actor-lr', type=float, default=1e-3)
+    parser.add_argument('--critic-lr', type=float, default=1e-3)
+    parser.add_argument('--gamma', type=float, default=0.99)
+    parser.add_argument('--tau', type=float, default=0.005)
+    parser.add_argument('--exploration-noise', type=float, default=0.1)
+    parser.add_argument("--start-timesteps", type=int, default=25000)
+    parser.add_argument('--epoch', type=int, default=200)
+    parser.add_argument('--step-per-epoch', type=int, default=5000)
+    parser.add_argument('--step-per-collect', type=int, default=1)
+    parser.add_argument('--update-per-step', type=int, default=1)
+    parser.add_argument('--n-step', type=int, default=1)
+    parser.add_argument('--batch-size', type=int, default=256)
+    parser.add_argument('--training-num', type=int, default=1)
+    parser.add_argument('--test-num', type=int, default=10)
+    parser.add_argument('--logdir', type=str, default='log')
+    parser.add_argument('--render', type=float, default=0.)
+    parser.add_argument(
+        '--device', type=str,
+        default='cuda' if torch.cuda.is_available() else 'cpu')
+    parser.add_argument('--resume-path', type=str, default=None)
+    return parser.parse_args()
+
+
+def test_ddpg(args=get_args()):
+    env = gym.make(args.task)
+    args.state_shape = env.observation_space.shape or env.observation_space.n
+    args.action_shape = env.action_space.shape or env.action_space.n
+    args.max_action = env.action_space.high[0]
+    args.exploration_noise = args.exploration_noise * args.max_action
+    print("Observations shape:", args.state_shape)
+    print("Actions shape:", args.action_shape)
+    print("Action range:", np.min(env.action_space.low),
+          np.max(env.action_space.high))
+    # train_envs = gym.make(args.task)
+    if args.training_num > 1:
+        train_envs = SubprocVectorEnv(
+            [lambda: gym.make(args.task) for _ in range(args.training_num)])
+    else:
+        train_envs = gym.make(args.task)
+    # test_envs = gym.make(args.task)
+    test_envs = SubprocVectorEnv(
+        [lambda: gym.make(args.task) for _ in range(args.test_num)])
+    # seed
+    np.random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    train_envs.seed(args.seed)
+    test_envs.seed(args.seed)
+    # model
+    net_a = Net(args.state_shape, hidden_sizes=args.hidden_sizes, device=args.device)
+    actor = Actor(
+        net_a, args.action_shape, max_action=args.max_action,
+        device=args.device).to(args.device)
+    actor_optim = torch.optim.Adam(actor.parameters(), lr=args.actor_lr)
+    net_c = Net(args.state_shape, args.action_shape,
+                hidden_sizes=args.hidden_sizes,
+                concat=True, device=args.device)
+    critic = Critic(net_c, device=args.device).to(args.device)
+    critic_optim = torch.optim.Adam(critic.parameters(), lr=args.critic_lr)
+    policy = DDPGPolicy(
+        actor, actor_optim, critic, critic_optim,
+        action_range=[env.action_space.low[0], env.action_space.high[0]],
+        tau=args.tau, gamma=args.gamma,
+        exploration_noise=GaussianNoise(sigma=args.exploration_noise),
+        estimation_step=args.n_step)
+    # load a previous policy
+    if args.resume_path:
+        policy.load_state_dict(torch.load(
+            args.resume_path, map_location=args.device
+        ))
+        print("Loaded agent from: ", args.resume_path)
+
+    # collector
+    if args.training_num > 1:
+        buffer = VectorReplayBuffer(args.buffer_size, len(train_envs))
+    else:
+        buffer = ReplayBuffer(args.buffer_size)
+    train_collector = Collector(policy, train_envs, buffer, exploration_noise=True)
+    test_collector = Collector(policy, test_envs)
+    train_collector.collect(n_step=args.start_timesteps, random=True)
+    # log
+    log_path = os.path.join(args.logdir, args.task, 'ddpg', 'seed_' + str(
+        args.seed) + '_' + datetime.datetime.now().strftime('%m%d-%H%M%S'))
+    writer = SummaryWriter(log_path)
+    writer.add_text("args", str(args))
+    logger = BasicLogger(writer)
+
+    def save_fn(policy):
+        torch.save(policy.state_dict(), os.path.join(log_path, 'policy.pth'))
+    # trainer
+    result = offpolicy_trainer(
+        policy, train_collector, test_collector, args.epoch,
+        args.step_per_epoch, args.step_per_collect, args.test_num,
+        args.batch_size, save_fn=save_fn, logger=logger,
+        update_per_step=args.update_per_step, test_in_train=False)
+
+    # Let's watch its performance!
+    policy.eval()
+    test_envs.seed(args.seed)
+    test_collector.reset()
+    result = test_collector.collect(n_episode=args.test_num, render=args.render)
+    print(f'Final reward: {result["rews"].mean()}, length: {result["lens"].mean()}')
+
+
+if __name__ == '__main__':
+    test_ddpg()