pushing model

README.md

@@ -16,7 +16,7 @@ model-index:
       type: Breakout-v5
     metrics:
     - type: mean_reward
-      value: 765.90 +/- 176.61
+      value: 828.90 +/- 104.97
       name: mean_reward
       verified: false
 ---
@@ -46,7 +46,7 @@ curl -OL https://huggingface.co/cleanrl/Breakout-v5-cleanba_ppo_envpool_impala_a
 curl -OL https://huggingface.co/cleanrl/Breakout-v5-cleanba_ppo_envpool_impala_atari_wrapper-seed2/raw/main/pyproject.toml
 curl -OL https://huggingface.co/cleanrl/Breakout-v5-cleanba_ppo_envpool_impala_atari_wrapper-seed2/raw/main/poetry.lock
 poetry install --all-extras
-python cleanba_ppo_envpool_impala_atari_wrapper.py --distributed --learner-device-ids 1 2 3 --track --save-model --upload-model --hf-entity cleanrl --env-id Breakout-v5 --seed 2
+python cleanba_ppo_envpool_impala_atari_wrapper.py --distributed --learner-device-ids 1 2 3 --track --wandb-project-name cleanba --save-model --upload-model --hf-entity cleanrl --env-id Breakout-v5 --seed 2
 ```
 
 # Hyperparameters
@@ -59,6 +59,7 @@ python cleanba_ppo_envpool_impala_atari_wrapper.py --distributed --learner-devic
  'batch_size': 15360,
  'capture_video': False,
  'clip_coef': 0.1,
+ 'concurrency': True,
  'cuda': True,
  'distributed': True,
  'ent_coef': 0.01,
@@ -99,7 +100,7 @@ python cleanba_ppo_envpool_impala_atari_wrapper.py --distributed --learner-devic
  'upload_model': True,
  'vf_coef': 0.5,
  'wandb_entity': None,
- 'wandb_project_name': 'cleanRL',
+ 'wandb_project_name': 'cleanba',
  'world_size': 2}
 ```
     

cleanba_ppo_envpool_impala_atari_wrapper.cleanrl_model

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:16c99b5c9a9dbf5062cc4867e43cde5e143e36257efd4303fe33a6e33a3415d4
-size 4364153
+oid sha256:fe673afdf2b5c4167541a185784366369acb8c1f69e91d2317641066f2acacaa
+size 4364166

cleanba_ppo_envpool_impala_atari_wrapper.py

@@ -1,4 +1,3 @@
-# docs and experiment results can be found at https://docs.cleanrl.dev/rl-algorithms/ppo/#ppo_atari_envpool_async_jax_scan_impalanet_machadopy
 import argparse
 import os
 import random
@@ -26,7 +25,7 @@ import numpy as np
 import optax
 from flax.linen.initializers import constant, orthogonal
 from flax.training.train_state import TrainState
-from torch.utils.tensorboard import SummaryWriter
+from tensorboardX import SummaryWriter
 
 
 def parse_args():
@@ -47,7 +46,7 @@ def parse_args():
     parser.add_argument("--wandb-entity", type=str, default=None,
         help="the entity (team) of wandb's project")
     parser.add_argument("--capture-video", type=lambda x: bool(strtobool(x)), default=False, nargs="?", const=True,
-        help="weather to capture videos of the agent performances (check out `videos` folder)")
+        help="whether to capture videos of the agent performances (check out `videos` folder)")
     parser.add_argument("--save-model", type=lambda x: bool(strtobool(x)), default=False, nargs="?", const=True,
         help="whether to save model into the `runs/{run_name}` folder")
     parser.add_argument("--upload-model", type=lambda x: bool(strtobool(x)), default=False, nargs="?", const=True,
@@ -97,6 +96,8 @@ def parse_args():
         help="the device ids that learner workers will use")
     parser.add_argument("--distributed", type=lambda x: bool(strtobool(x)), default=False, nargs="?", const=True,
         help="whether to use `jax.distirbuted`")
+    parser.add_argument("--concurrency", type=lambda x: bool(strtobool(x)), default=True, nargs="?", const=True,
+        help="whether to run the actor and learner concurrently")
     parser.add_argument("--profile", type=lambda x: bool(strtobool(x)), default=False, nargs="?", const=True,
         help="whether to call block_until_ready() for profiling")
     parser.add_argument("--test-actor-learner-throughput", type=lambda x: bool(strtobool(x)), default=False, nargs="?", const=True,
@@ -213,7 +214,7 @@ class AgentParams:
 
 @partial(jax.jit, static_argnums=(3))
 def get_action_and_value(
-    params: TrainState,
+    params: flax.core.FrozenDict,
     next_obs: np.ndarray,
     key: jax.random.PRNGKey,
     action_dim: int,
@@ -281,6 +282,20 @@ def prepare_data(
     return b_obs, b_actions, b_logprobs, b_advantages, b_returns
 
 
+@jax.jit
+def make_bulk_array(
+    obs: list,
+    values: list,
+    actions: list,
+    logprobs: list,
+):
+    obs = jnp.asarray(obs)
+    values = jnp.asarray(values)
+    actions = jnp.asarray(actions)
+    logprobs = jnp.asarray(logprobs)
+    return obs, values, actions, logprobs
+
+
 def rollout(
     key: jax.random.PRNGKey,
     args,
@@ -289,7 +304,7 @@ def rollout(
     writer,
     learner_devices,
 ):
-    envs = make_env(args.env_id, args.seed, args.local_num_envs, args.async_batch_size)()
+    envs = make_env(args.env_id, args.seed + jax.process_index(), args.local_num_envs, args.async_batch_size)()
     len_actor_device_ids = len(args.actor_device_ids)
     global_step = 0
     # TRY NOT TO MODIFY: start the game
@@ -332,9 +347,13 @@ def rollout(
         # concurrently with the learning process. It also ensures the actor's policy version is only 1 step
         # behind the learner's policy version
         params_queue_get_time_start = time.time()
-        if update != 2:
+        if not args.concurrency:
             params = params_queue.get()
             actor_policy_version += 1
+        else:
+            if update != 2:
+                params = params_queue.get()
+                actor_policy_version += 1
         params_queue_get_time.append(time.time() - params_queue_get_time_start)
         writer.add_scalar("stats/params_queue_get_time", np.mean(params_queue_get_time), global_step)
         rollout_time_start = time.time()
@@ -397,18 +416,29 @@ def rollout(
         writer.add_scalar("stats/inference_time", inference_time, global_step)
         writer.add_scalar("stats/storage_time", storage_time, global_step)
         writer.add_scalar("stats/env_send_time", env_send_time, global_step)
+        # `make_bulk_array` is actually important. It accumulates the data from the lists
+        # into single bulk arrays, which later makes transferring the data to the learner's
+        # device slightly faster. See https://wandb.ai/costa-huang/cleanRL/reports/data-transfer-optimization--VmlldzozNjU5MTg1
+        if args.learner_device_ids[0] != args.actor_device_ids[0]:
+            obs, values, actions, logprobs = make_bulk_array(
+                obs,
+                values,
+                actions,
+                logprobs,
+            )
 
         payload = (
             global_step,
             actor_policy_version,
             update,
             obs,
-            dones,
             values,
             actions,
             logprobs,
+            dones,
             env_ids,
             rewards,
+            np.mean(params_queue_get_time),
         )
         if update == 1 or not args.test_actor_learner_throughput:
             rollout_queue_put_time_start = time.time()
@@ -717,15 +747,21 @@ if __name__ == "__main__":
                 actor_policy_version,
                 update,
                 obs,
-                dones,
                 values,
                 actions,
                 logprobs,
+                dones,
                 env_ids,
                 rewards,
+                avg_params_queue_get_time,
             ) = rollout_queue.get()
             rollout_queue_get_time.append(time.time() - rollout_queue_get_time_start)
             writer.add_scalar("stats/rollout_queue_get_time", np.mean(rollout_queue_get_time), global_step)
+            writer.add_scalar(
+                "stats/rollout_params_queue_get_time_diff",
+                np.mean(rollout_queue_get_time) - avg_params_queue_get_time,
+                global_step,
+            )
 
         data_transfer_time_start = time.time()
         b_obs, b_actions, b_logprobs, b_advantages, b_returns = prepare_data(
@@ -780,11 +816,22 @@ if __name__ == "__main__":
             break
 
     if args.save_model and args.local_rank == 0:
+        if args.distributed:
+            jax.distributed.shutdown()
         agent_state = flax.jax_utils.unreplicate(agent_state)
         model_path = f"runs/{run_name}/{args.exp_name}.cleanrl_model"
         with open(model_path, "wb") as f:
             f.write(
-                flax.serialization.to_bytes([    vars(args),    [        agent_state.params.network_params,        agent_state.params.actor_params,        agent_state.params.critic_params,    ],])
+                flax.serialization.to_bytes(
+                    [
+                        vars(args),
+                        [
+                            agent_state.params.network_params,
+                            agent_state.params.actor_params,
+                            agent_state.params.critic_params,
+                        ],
+                    ]
+                )
             )
         print(f"model saved to {model_path}")
         from cleanrl_utils.evals.ppo_envpool_jax_eval import evaluate

events.out.tfevents.1676644135.ip-26-0-130-110.250950.0 → events.out.tfevents.1678205984.ip-26-0-135-192

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:111e3f8f0c43d1de2a9be38fd060876ae4e90b0616adcc536f88a18efc195cb8
-size 4754790
+oid sha256:9c446f81724a76b207165a02cc5f9be23e69fba7bb7745bb5adb6de55caab498
+size 5017749

pyproject.toml

@@ -1,178 +1,34 @@
 [tool.poetry]
-name = "cleanrl"
-version = "1.1.0"
-description = "High-quality single file implementation of Deep Reinforcement Learning algorithms with research-friendly features"
+name = "cleanba"
+version = "0.1.0"
+description = ""
 authors = ["Costa Huang <[email protected]>"]
+readme = "README.md"
 packages = [
-    { include = "cleanrl" },
+    { include = "cleanba" },
     { include = "cleanrl_utils" },
 ]
-keywords = ["reinforcement", "machine", "learning", "research"]
-license="MIT"
-readme = "README.md"
 
 [tool.poetry.dependencies]
-python = ">=3.7.1,<3.10"
-tensorboard = "^2.10.0"
-wandb = "^0.13.6"
+python = "^3.8"
+tensorboard = "^2.12.0"
+envpool = "^0.8.1"
+jax = "0.3.25"
+flax = "0.6.0"
+optax = "0.1.3"
+huggingface-hub = "^0.12.0"
+jaxlib = "0.3.25"
+wandb = "^0.13.10"
+tensorboardx = "^2.5.1"
+chex = "0.1.5"
 gym = "0.23.1"
-torch = ">=1.12.1"
-stable-baselines3 = "1.2.0"
-gymnasium = "^0.26.3"
+opencv-python = "^4.7.0.68"
 moviepy = "^1.0.3"
-pygame = "2.1.0"
-huggingface-hub = "^0.11.1"
 
-ale-py = {version = "0.7.4", optional = true}
-AutoROM = {extras = ["accept-rom-license"], version = "^0.4.2"}
-opencv-python = {version = "^4.6.0.66", optional = true}
-pybullet = {version = "3.1.8", optional = true}
-procgen = {version = "^0.10.7", optional = true}
-pytest = {version = "^7.1.3", optional = true}
-mujoco = {version = "^2.2", optional = true}
-imageio = {version = "^2.14.1", optional = true}
-free-mujoco-py = {version = "^2.1.6", optional = true}
-mkdocs-material = {version = "^8.4.3", optional = true}
-markdown-include = {version = "^0.7.0", optional = true}
-jax = {version = "^0.3.17", optional = true}
-jaxlib = {version = "^0.3.15", optional = true}
-flax = {version = "^0.6.0", optional = true}
-optuna = {version = "^3.0.1", optional = true}
-optuna-dashboard = {version = "^0.7.2", optional = true}
-rich = {version = "<12.0", optional = true}
-envpool = {version = "^0.8.1", optional = true}
-PettingZoo = {version = "1.18.1", optional = true}
-SuperSuit = {version = "3.4.0", optional = true}
-multi-agent-ale-py = {version = "0.1.11", optional = true}
-boto3 = {version = "^1.24.70", optional = true}
-awscli = {version = "^1.25.71", optional = true}
-shimmy = {version = "^0.1.0", optional = true}
-dm-control = {version = "^1.0.8", optional = true}
 
 [tool.poetry.group.dev.dependencies]
-pre-commit = "^2.20.0"
-
-[tool.poetry.group.atari]
-optional = true
-[tool.poetry.group.atari.dependencies]
-ale-py = "0.7.4"
-AutoROM = {extras = ["accept-rom-license"], version = "^0.4.2"}
-opencv-python = "^4.6.0.66"
-
-[tool.poetry.group.pybullet]
-optional = true
-[tool.poetry.group.pybullet.dependencies]
-pybullet = "3.1.8"
-
-[tool.poetry.group.procgen]
-optional = true
-[tool.poetry.group.procgen.dependencies]
-procgen = "^0.10.7"
-
-[tool.poetry.group.pytest]
-optional = true
-[tool.poetry.group.pytest.dependencies]
-pytest = "^7.1.3"
-
-[tool.poetry.group.mujoco]
-optional = true
-[tool.poetry.group.mujoco.dependencies]
-mujoco = "^2.2"
-imageio = "^2.14.1"
-
-[tool.poetry.group.mujoco_py]
-optional = true
-[tool.poetry.group.mujoco_py.dependencies]
-free-mujoco-py = "^2.1.6"
-
-[tool.poetry.group.docs]
-optional = true
-[tool.poetry.group.docs.dependencies]
-mkdocs-material = "^8.4.3"
-markdown-include = "^0.7.0"
-
-[tool.poetry.group.jax]
-optional = true
-[tool.poetry.group.jax.dependencies]
-jax = "^0.3.17"
-jaxlib = "^0.3.15"
-flax = "^0.6.0"
-
-[tool.poetry.group.optuna]
-optional = true
-[tool.poetry.group.optuna.dependencies]
-optuna = "^3.0.1"
-optuna-dashboard = "^0.7.2"
-rich = "<12.0"
-
-[tool.poetry.group.envpool]
-optional = true
-[tool.poetry.group.envpool.dependencies]
-envpool = "^0.8.1"
-
-[tool.poetry.group.pettingzoo]
-optional = true
-[tool.poetry.group.pettingzoo.dependencies]
-PettingZoo = "1.18.1"
-SuperSuit = "3.4.0"
-multi-agent-ale-py = "0.1.11"
-
-[tool.poetry.group.cloud]
-optional = true
-[tool.poetry.group.cloud.dependencies]
-boto3 = "^1.24.70"
-awscli = "^1.25.71"
-
-[tool.poetry.group.isaacgym]
-optional = true
-[tool.poetry.group.isaacgym.dependencies]
-isaacgymenvs = {git = "https://github.com/vwxyzjn/IsaacGymEnvs.git", rev = "poetry"}
-isaacgym = {path = "cleanrl/ppo_continuous_action_isaacgym/isaacgym", develop = true}
-
-[tool.poetry.group.dm_control]
-optional = true
-[tool.poetry.group.dm_control.dependencies]
-shimmy = "^0.1.0"
-dm-control = "^1.0.8"
-mujoco = "^2.2"
+pre-commit = "^3.0.4"
 
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
-
-[tool.poetry.extras]
-atari = ["ale-py", "AutoROM", "opencv-python"]
-pybullet = ["pybullet"]
-procgen = ["procgen"]
-plot = ["pandas", "seaborn"]
-pytest = ["pytest"]
-mujoco = ["mujoco", "imageio"]
-mujoco_py = ["free-mujoco-py"]
-jax = ["jax", "jaxlib", "flax"]
-docs = ["mkdocs-material", "markdown-include"]
-envpool = ["envpool"]
-optuna = ["optuna", "optuna-dashboard", "rich"]
-pettingzoo = ["PettingZoo", "SuperSuit", "multi-agent-ale-py"]
-cloud = ["boto3", "awscli"]
-dm_control = ["shimmy", "dm-control", "mujoco"]
-
-# dependencies for algorithm variant (useful when you want to run a specific algorithm)
-dqn = []
-dqn_atari = ["ale-py", "AutoROM", "opencv-python"]
-dqn_jax = ["jax", "jaxlib", "flax"]
-dqn_atari_jax = [
-    "ale-py", "AutoROM", "opencv-python", # atari
-    "jax", "jaxlib", "flax" # jax
-]
-c51 = []
-c51_atari = ["ale-py", "AutoROM", "opencv-python"]
-c51_jax = ["jax", "jaxlib", "flax"]
-c51_atari_jax = [
-    "ale-py", "AutoROM", "opencv-python", # atari
-    "jax", "jaxlib", "flax" # jax
-]
-ppo_atari_envpool_xla_jax_scan = [
-    "ale-py", "AutoROM", "opencv-python", # atari
-    "jax", "jaxlib", "flax", # jax
-    "envpool", # envpool
-]