vwxyzjn commited on
Commit
e5d2cdc
1 Parent(s): dfc6a1b

pushing model

Browse files
.gitattributes CHANGED
@@ -32,3 +32,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
32
  *.zip filter=lfs diff=lfs merge=lfs -text
33
  *.zst filter=lfs diff=lfs merge=lfs -text
34
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
32
  *.zip filter=lfs diff=lfs merge=lfs -text
33
  *.zst filter=lfs diff=lfs merge=lfs -text
34
  *tfevents* filter=lfs diff=lfs merge=lfs -text
35
+ cleanba_ppo_envpool_impala_atari_wrapper_naturecnn.cleanrl_model filter=lfs diff=lfs merge=lfs -text
README.md ADDED
@@ -0,0 +1,100 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ tags:
3
+ - Surround-v5
4
+ - deep-reinforcement-learning
5
+ - reinforcement-learning
6
+ - custom-implementation
7
+ library_name: cleanrl
8
+ model-index:
9
+ - name: PPO
10
+ results:
11
+ - task:
12
+ type: reinforcement-learning
13
+ name: reinforcement-learning
14
+ dataset:
15
+ name: Surround-v5
16
+ type: Surround-v5
17
+ metrics:
18
+ - type: mean_reward
19
+ value: -4.20 +/- 2.96
20
+ name: mean_reward
21
+ verified: false
22
+ ---
23
+
24
+ # (CleanRL) **PPO** Agent Playing **Surround-v5**
25
+
26
+ This is a trained model of a PPO agent playing Surround-v5.
27
+ The model was trained by using [CleanRL](https://github.com/vwxyzjn/cleanrl) and the most up-to-date training code can be
28
+ found [here](https://github.com/vwxyzjn/cleanrl/blob/master/cleanrl/cleanba_ppo_envpool_impala_atari_wrapper_naturecnn.py).
29
+
30
+ ## Get Started
31
+
32
+ To use this model, please install the `cleanrl` package with the following command:
33
+
34
+ ```
35
+ pip install "cleanrl[jax,envpool,atari]"
36
+ python -m cleanrl_utils.enjoy --exp-name cleanba_ppo_envpool_impala_atari_wrapper_naturecnn --env-id Surround-v5
37
+ ```
38
+
39
+ Please refer to the [documentation](https://docs.cleanrl.dev/get-started/zoo/) for more detail.
40
+
41
+
42
+ ## Command to reproduce the training
43
+
44
+ ```bash
45
+ curl -OL https://huggingface.co/cleanrl/Surround-v5-cleanba_ppo_envpool_impala_atari_wrapper_naturecnn-seed2/raw/main/cleanba_ppo_envpool_impala_atari_wrapper_naturecnn.py
46
+ curl -OL https://huggingface.co/cleanrl/Surround-v5-cleanba_ppo_envpool_impala_atari_wrapper_naturecnn-seed2/raw/main/pyproject.toml
47
+ curl -OL https://huggingface.co/cleanrl/Surround-v5-cleanba_ppo_envpool_impala_atari_wrapper_naturecnn-seed2/raw/main/poetry.lock
48
+ poetry install --all-extras
49
+ python cleanba_ppo_envpool_impala_atari_wrapper_naturecnn.py --distributed --learner-device-ids 1 --track --wandb-project-name cleanba --save-model --upload-model --hf-entity cleanrl --env-id Surround-v5 --seed 2
50
+ ```
51
+
52
+ # Hyperparameters
53
+ ```python
54
+ {'actor_device_ids': [0],
55
+ 'actor_devices': ['gpu:0'],
56
+ 'anneal_lr': True,
57
+ 'async_batch_size': 20,
58
+ 'async_update': 3,
59
+ 'batch_size': 15360,
60
+ 'capture_video': False,
61
+ 'clip_coef': 0.1,
62
+ 'cuda': True,
63
+ 'distributed': True,
64
+ 'ent_coef': 0.01,
65
+ 'env_id': 'Surround-v5',
66
+ 'exp_name': 'cleanba_ppo_envpool_impala_atari_wrapper_naturecnn',
67
+ 'gae_lambda': 0.95,
68
+ 'gamma': 0.99,
69
+ 'global_learner_decices': ['gpu:1', 'gpu:3'],
70
+ 'hf_entity': 'cleanrl',
71
+ 'learner_device_ids': [1],
72
+ 'learner_devices': ['gpu:1'],
73
+ 'learning_rate': 0.00025,
74
+ 'local_batch_size': 7680,
75
+ 'local_minibatch_size': 1920,
76
+ 'local_num_envs': 60,
77
+ 'local_rank': 0,
78
+ 'max_grad_norm': 0.5,
79
+ 'minibatch_size': 3840,
80
+ 'norm_adv': True,
81
+ 'num_envs': 120,
82
+ 'num_minibatches': 4,
83
+ 'num_steps': 128,
84
+ 'num_updates': 3255,
85
+ 'profile': False,
86
+ 'save_model': True,
87
+ 'seed': 2,
88
+ 'target_kl': None,
89
+ 'test_actor_learner_throughput': False,
90
+ 'torch_deterministic': True,
91
+ 'total_timesteps': 50000000,
92
+ 'track': True,
93
+ 'update_epochs': 4,
94
+ 'upload_model': True,
95
+ 'vf_coef': 0.5,
96
+ 'wandb_entity': None,
97
+ 'wandb_project_name': 'cleanba',
98
+ 'world_size': 2}
99
+ ```
100
+
cleanba_ppo_envpool_impala_atari_wrapper_naturecnn.cleanrl_model ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:d14b0aad9ccf2e8a898249532775bef07f1ba42d41884594206a236b85343352
3
+ size 6750043
cleanba_ppo_envpool_impala_atari_wrapper_naturecnn.py ADDED
@@ -0,0 +1,816 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import argparse
2
+ import os
3
+ import random
4
+ import time
5
+ import uuid
6
+ from collections import deque
7
+ from distutils.util import strtobool
8
+ from functools import partial
9
+
10
+ os.environ[
11
+ "XLA_PYTHON_CLIENT_MEM_FRACTION"
12
+ ] = "0.6" # see https://github.com/google/jax/discussions/6332#discussioncomment-1279991
13
+ os.environ["XLA_FLAGS"] = "--xla_cpu_multi_thread_eigen=false " "intra_op_parallelism_threads=1"
14
+ import queue
15
+ import threading
16
+
17
+ import envpool
18
+ import flax
19
+ import flax.linen as nn
20
+ import gym
21
+ import jax
22
+ import jax.numpy as jnp
23
+ import numpy as np
24
+ import optax
25
+ from flax.linen.initializers import constant, orthogonal
26
+ from flax.training.train_state import TrainState
27
+ from tensorboardX import SummaryWriter
28
+
29
+
30
+ def parse_args():
31
+ # fmt: off
32
+ parser = argparse.ArgumentParser()
33
+ parser.add_argument("--exp-name", type=str, default=os.path.basename(__file__).rstrip(".py"),
34
+ help="the name of this experiment")
35
+ parser.add_argument("--seed", type=int, default=1,
36
+ help="seed of the experiment")
37
+ parser.add_argument("--torch-deterministic", type=lambda x: bool(strtobool(x)), default=True, nargs="?", const=True,
38
+ help="if toggled, `torch.backends.cudnn.deterministic=False`")
39
+ parser.add_argument("--cuda", type=lambda x: bool(strtobool(x)), default=True, nargs="?", const=True,
40
+ help="if toggled, cuda will be enabled by default")
41
+ parser.add_argument("--track", type=lambda x: bool(strtobool(x)), default=False, nargs="?", const=True,
42
+ help="if toggled, this experiment will be tracked with Weights and Biases")
43
+ parser.add_argument("--wandb-project-name", type=str, default="cleanRL",
44
+ help="the wandb's project name")
45
+ parser.add_argument("--wandb-entity", type=str, default=None,
46
+ help="the entity (team) of wandb's project")
47
+ parser.add_argument("--capture-video", type=lambda x: bool(strtobool(x)), default=False, nargs="?", const=True,
48
+ help="weather to capture videos of the agent performances (check out `videos` folder)")
49
+ parser.add_argument("--save-model", type=lambda x: bool(strtobool(x)), default=False, nargs="?", const=True,
50
+ help="whether to save model into the `runs/{run_name}` folder")
51
+ parser.add_argument("--upload-model", type=lambda x: bool(strtobool(x)), default=False, nargs="?", const=True,
52
+ help="whether to upload the saved model to huggingface")
53
+ parser.add_argument("--hf-entity", type=str, default="",
54
+ help="the user or org name of the model repository from the Hugging Face Hub")
55
+
56
+ # Algorithm specific arguments
57
+ parser.add_argument("--env-id", type=str, default="Breakout-v5",
58
+ help="the id of the environment")
59
+ parser.add_argument("--total-timesteps", type=int, default=50000000,
60
+ help="total timesteps of the experiments")
61
+ parser.add_argument("--learning-rate", type=float, default=2.5e-4,
62
+ help="the learning rate of the optimizer")
63
+ parser.add_argument("--local-num-envs", type=int, default=60,
64
+ help="the number of parallel game environments")
65
+ parser.add_argument("--async-batch-size", type=int, default=20,
66
+ help="the envpool's batch size in the async mode")
67
+ parser.add_argument("--num-steps", type=int, default=128,
68
+ help="the number of steps to run in each environment per policy rollout")
69
+ parser.add_argument("--anneal-lr", type=lambda x: bool(strtobool(x)), default=True, nargs="?", const=True,
70
+ help="Toggle learning rate annealing for policy and value networks")
71
+ parser.add_argument("--gamma", type=float, default=0.99,
72
+ help="the discount factor gamma")
73
+ parser.add_argument("--gae-lambda", type=float, default=0.95,
74
+ help="the lambda for the general advantage estimation")
75
+ parser.add_argument("--num-minibatches", type=int, default=4,
76
+ help="the number of mini-batches")
77
+ parser.add_argument("--update-epochs", type=int, default=4,
78
+ help="the K epochs to update the policy")
79
+ parser.add_argument("--norm-adv", type=lambda x: bool(strtobool(x)), default=True, nargs="?", const=True,
80
+ help="Toggles advantages normalization")
81
+ parser.add_argument("--clip-coef", type=float, default=0.1,
82
+ help="the surrogate clipping coefficient")
83
+ parser.add_argument("--ent-coef", type=float, default=0.01,
84
+ help="coefficient of the entropy")
85
+ parser.add_argument("--vf-coef", type=float, default=0.5,
86
+ help="coefficient of the value function")
87
+ parser.add_argument("--max-grad-norm", type=float, default=0.5,
88
+ help="the maximum norm for the gradient clipping")
89
+ parser.add_argument("--target-kl", type=float, default=None,
90
+ help="the target KL divergence threshold")
91
+
92
+ parser.add_argument("--actor-device-ids", type=int, nargs="+", default=[0], # type is actually List[int]
93
+ help="the device ids that actor workers will use (currently only support 1 device)")
94
+ parser.add_argument("--learner-device-ids", type=int, nargs="+", default=[0], # type is actually List[int]
95
+ help="the device ids that learner workers will use")
96
+ parser.add_argument("--distributed", type=lambda x: bool(strtobool(x)), default=False, nargs="?", const=True,
97
+ help="whether to use `jax.distirbuted`")
98
+ parser.add_argument("--profile", type=lambda x: bool(strtobool(x)), default=False, nargs="?", const=True,
99
+ help="whether to call block_until_ready() for profiling")
100
+ parser.add_argument("--test-actor-learner-throughput", type=lambda x: bool(strtobool(x)), default=False, nargs="?", const=True,
101
+ help="whether to test actor-learner throughput by removing the actor-learner communication")
102
+ args = parser.parse_args()
103
+ args.local_batch_size = int(args.local_num_envs * args.num_steps)
104
+ args.local_minibatch_size = int(args.local_batch_size // args.num_minibatches)
105
+ args.num_updates = args.total_timesteps // args.local_batch_size
106
+ args.async_update = int(args.local_num_envs / args.async_batch_size)
107
+ assert len(args.actor_device_ids) == 1, "only 1 actor_device_ids is supported now"
108
+ # fmt: on
109
+ return args
110
+
111
+
112
+ ATARI_MAX_FRAMES = int(
113
+ 108000 / 4
114
+ ) # 108000 is the max number of frames in an Atari game, divided by 4 to account for frame skipping
115
+
116
+
117
+ def make_env(env_id, seed, num_envs, async_batch_size=1):
118
+ def thunk():
119
+ envs = envpool.make(
120
+ env_id,
121
+ env_type="gym",
122
+ num_envs=num_envs,
123
+ batch_size=async_batch_size,
124
+ episodic_life=True, # Espeholt et al., 2018, Tab. G.1
125
+ repeat_action_probability=0, # Hessel et al., 2022 (Muesli) Tab. 10
126
+ noop_max=30, # Espeholt et al., 2018, Tab. C.1 "Up to 30 no-ops at the beginning of each episode."
127
+ full_action_space=False, # Espeholt et al., 2018, Appendix G., "Following related work, experts use game-specific action sets."
128
+ max_episode_steps=ATARI_MAX_FRAMES, # Hessel et al. 2018 (Rainbow DQN), Table 3, Max frames per episode
129
+ reward_clip=True,
130
+ seed=seed,
131
+ )
132
+ envs.num_envs = num_envs
133
+ envs.single_action_space = envs.action_space
134
+ envs.single_observation_space = envs.observation_space
135
+ envs.is_vector_env = True
136
+ return envs
137
+
138
+ return thunk
139
+
140
+
141
+ class Network(nn.Module):
142
+ @nn.compact
143
+ def __call__(self, x):
144
+ x = jnp.transpose(x, (0, 2, 3, 1))
145
+ x = x / (255.0)
146
+ x = nn.Conv(
147
+ 32,
148
+ kernel_size=(8, 8),
149
+ strides=(4, 4),
150
+ padding="VALID",
151
+ kernel_init=orthogonal(np.sqrt(2)),
152
+ bias_init=constant(0.0),
153
+ )(x)
154
+ x = nn.relu(x)
155
+ x = nn.Conv(
156
+ 64,
157
+ kernel_size=(4, 4),
158
+ strides=(2, 2),
159
+ padding="VALID",
160
+ kernel_init=orthogonal(np.sqrt(2)),
161
+ bias_init=constant(0.0),
162
+ )(x)
163
+ x = nn.relu(x)
164
+ x = nn.Conv(
165
+ 64,
166
+ kernel_size=(3, 3),
167
+ strides=(1, 1),
168
+ padding="VALID",
169
+ kernel_init=orthogonal(np.sqrt(2)),
170
+ bias_init=constant(0.0),
171
+ )(x)
172
+ x = nn.relu(x)
173
+ x = x.reshape((x.shape[0], -1))
174
+ x = nn.Dense(512, kernel_init=orthogonal(np.sqrt(2)), bias_init=constant(0.0))(x)
175
+ x = nn.relu(x)
176
+ return x
177
+
178
+
179
+ class Critic(nn.Module):
180
+ @nn.compact
181
+ def __call__(self, x):
182
+ return nn.Dense(1, kernel_init=orthogonal(1), bias_init=constant(0.0))(x)
183
+
184
+
185
+ class Actor(nn.Module):
186
+ action_dim: int
187
+
188
+ @nn.compact
189
+ def __call__(self, x):
190
+ return nn.Dense(self.action_dim, kernel_init=orthogonal(0.01), bias_init=constant(0.0))(x)
191
+
192
+
193
+ @flax.struct.dataclass
194
+ class AgentParams:
195
+ network_params: flax.core.FrozenDict
196
+ actor_params: flax.core.FrozenDict
197
+ critic_params: flax.core.FrozenDict
198
+
199
+
200
+ @partial(jax.jit, static_argnums=(3))
201
+ def get_action_and_value(
202
+ params: TrainState,
203
+ next_obs: np.ndarray,
204
+ key: jax.random.PRNGKey,
205
+ action_dim: int,
206
+ ):
207
+ next_obs = jnp.array(next_obs)
208
+ hidden = Network().apply(params.network_params, next_obs)
209
+ logits = Actor(action_dim).apply(params.actor_params, hidden)
210
+ # sample action: Gumbel-softmax trick
211
+ # see https://stats.stackexchange.com/questions/359442/sampling-from-a-categorical-distribution
212
+ key, subkey = jax.random.split(key)
213
+ u = jax.random.uniform(subkey, shape=logits.shape)
214
+ action = jnp.argmax(logits - jnp.log(-jnp.log(u)), axis=1)
215
+ logprob = jax.nn.log_softmax(logits)[jnp.arange(action.shape[0]), action]
216
+ value = Critic().apply(params.critic_params, hidden)
217
+ return next_obs, action, logprob, value.squeeze(), key
218
+
219
+
220
+ def prepare_data(
221
+ obs: list,
222
+ dones: list,
223
+ values: list,
224
+ actions: list,
225
+ logprobs: list,
226
+ env_ids: list,
227
+ rewards: list,
228
+ ):
229
+ obs = jnp.asarray(obs)
230
+ dones = jnp.asarray(dones)
231
+ values = jnp.asarray(values)
232
+ actions = jnp.asarray(actions)
233
+ logprobs = jnp.asarray(logprobs)
234
+ env_ids = jnp.asarray(env_ids)
235
+ rewards = jnp.asarray(rewards)
236
+
237
+ # TODO: in an unlikely event, one of the envs might have not stepped at all, which may results in unexpected behavior
238
+ T, B = env_ids.shape
239
+ index_ranges = jnp.arange(T * B, dtype=jnp.int32)
240
+ next_index_ranges = jnp.zeros_like(index_ranges, dtype=jnp.int32)
241
+ last_env_ids = jnp.zeros(args.local_num_envs, dtype=jnp.int32) - 1
242
+
243
+ def f(carry, x):
244
+ last_env_ids, next_index_ranges = carry
245
+ env_id, index_range = x
246
+ next_index_ranges = next_index_ranges.at[last_env_ids[env_id]].set(
247
+ jnp.where(last_env_ids[env_id] != -1, index_range, next_index_ranges[last_env_ids[env_id]])
248
+ )
249
+ last_env_ids = last_env_ids.at[env_id].set(index_range)
250
+ return (last_env_ids, next_index_ranges), None
251
+
252
+ (last_env_ids, next_index_ranges), _ = jax.lax.scan(
253
+ f,
254
+ (last_env_ids, next_index_ranges),
255
+ (env_ids.reshape(-1), index_ranges),
256
+ )
257
+
258
+ # rewards is off by one time step
259
+ rewards = rewards.reshape(-1)[next_index_ranges].reshape((args.num_steps) * args.async_update, args.async_batch_size)
260
+ advantages, returns, _, final_env_ids = compute_gae(env_ids, rewards, values, dones)
261
+ # b_inds = jnp.nonzero(final_env_ids.reshape(-1), size=(args.num_steps) * args.async_update * args.async_batch_size)[0] # useful for debugging
262
+ b_obs = obs.reshape((-1,) + obs.shape[2:])
263
+ b_actions = actions.reshape(-1)
264
+ b_logprobs = logprobs.reshape(-1)
265
+ b_advantages = advantages.reshape(-1)
266
+ b_returns = returns.reshape(-1)
267
+ return b_obs, b_actions, b_logprobs, b_advantages, b_returns
268
+
269
+
270
+ def rollout(
271
+ key: jax.random.PRNGKey,
272
+ args,
273
+ rollout_queue,
274
+ params_queue: queue.Queue,
275
+ writer,
276
+ learner_devices,
277
+ ):
278
+ envs = make_env(args.env_id, args.seed, args.local_num_envs, args.async_batch_size)()
279
+ len_actor_device_ids = len(args.actor_device_ids)
280
+ global_step = 0
281
+ # TRY NOT TO MODIFY: start the game
282
+ start_time = time.time()
283
+
284
+ # put data in the last index
285
+ episode_returns = np.zeros((args.local_num_envs,), dtype=np.float32)
286
+ returned_episode_returns = np.zeros((args.local_num_envs,), dtype=np.float32)
287
+ episode_lengths = np.zeros((args.local_num_envs,), dtype=np.float32)
288
+ returned_episode_lengths = np.zeros((args.local_num_envs,), dtype=np.float32)
289
+ envs.async_reset()
290
+
291
+ params_queue_get_time = deque(maxlen=10)
292
+ rollout_time = deque(maxlen=10)
293
+ rollout_queue_put_time = deque(maxlen=10)
294
+ actor_policy_version = 0
295
+ for update in range(1, args.num_updates + 2):
296
+ # NOTE: This is a major difference from the sync version:
297
+ # at the end of the rollout phase, the sync version will have the next observation
298
+ # ready for the value bootstrap, but the async version will not have it.
299
+ # for this reason we do `num_steps + 1`` to get the extra states for value bootstrapping.
300
+ # but note that the extra states are not used for the loss computation in the next iteration,
301
+ # while the sync version will use the extra state for the loss computation.
302
+ update_time_start = time.time()
303
+ obs = []
304
+ dones = []
305
+ actions = []
306
+ logprobs = []
307
+ values = []
308
+ env_ids = []
309
+ rewards = []
310
+ truncations = []
311
+ terminations = []
312
+ env_recv_time = 0
313
+ inference_time = 0
314
+ storage_time = 0
315
+ env_send_time = 0
316
+
317
+ # NOTE: `update != 2` is actually IMPORTANT — it allows us to start running policy collection
318
+ # concurrently with the learning process. It also ensures the actor's policy version is only 1 step
319
+ # behind the learner's policy version
320
+ params_queue_get_time_start = time.time()
321
+ if update != 2:
322
+ params = params_queue.get()
323
+ actor_policy_version += 1
324
+ params_queue_get_time.append(time.time() - params_queue_get_time_start)
325
+ writer.add_scalar("stats/params_queue_get_time", np.mean(params_queue_get_time), global_step)
326
+ rollout_time_start = time.time()
327
+ for _ in range(
328
+ args.async_update, (args.num_steps + 1) * args.async_update
329
+ ): # num_steps + 1 to get the states for value bootstrapping.
330
+ env_recv_time_start = time.time()
331
+ next_obs, next_reward, next_done, info = envs.recv()
332
+ env_recv_time += time.time() - env_recv_time_start
333
+ global_step += len(next_done) * len_actor_device_ids * args.world_size
334
+ env_id = info["env_id"]
335
+
336
+ inference_time_start = time.time()
337
+ next_obs, action, logprob, value, key = get_action_and_value(params, next_obs, key, envs.single_action_space.n)
338
+ inference_time += time.time() - inference_time_start
339
+
340
+ env_send_time_start = time.time()
341
+ envs.send(np.array(action), env_id)
342
+ env_send_time += time.time() - env_send_time_start
343
+ storage_time_start = time.time()
344
+ obs.append(next_obs)
345
+ dones.append(next_done)
346
+ values.append(value)
347
+ actions.append(action)
348
+ logprobs.append(logprob)
349
+ env_ids.append(env_id)
350
+ rewards.append(next_reward)
351
+
352
+ # info["TimeLimit.truncated"] has a bug https://github.com/sail-sg/envpool/issues/239
353
+ # so we use our own truncated flag
354
+ truncated = info["elapsed_step"] >= envs.spec.config.max_episode_steps
355
+ truncations.append(truncated)
356
+ terminations.append(info["terminated"])
357
+ episode_returns[env_id] += info["reward"]
358
+ returned_episode_returns[env_id] = np.where(
359
+ info["terminated"] + truncated, episode_returns[env_id], returned_episode_returns[env_id]
360
+ )
361
+ episode_returns[env_id] *= (1 - info["terminated"]) * (1 - truncated)
362
+ episode_lengths[env_id] += 1
363
+ returned_episode_lengths[env_id] = np.where(
364
+ info["terminated"] + truncated, episode_lengths[env_id], returned_episode_lengths[env_id]
365
+ )
366
+ episode_lengths[env_id] *= (1 - info["terminated"]) * (1 - truncated)
367
+ storage_time += time.time() - storage_time_start
368
+ if args.profile:
369
+ action.block_until_ready()
370
+ rollout_time.append(time.time() - rollout_time_start)
371
+ writer.add_scalar("stats/rollout_time", np.mean(rollout_time), global_step)
372
+
373
+ avg_episodic_return = np.mean(returned_episode_returns)
374
+ writer.add_scalar("charts/avg_episodic_return", avg_episodic_return, global_step)
375
+ writer.add_scalar("charts/avg_episodic_length", np.mean(returned_episode_lengths), global_step)
376
+ print(f"global_step={global_step}, avg_episodic_return={avg_episodic_return}")
377
+ print("SPS:", int(global_step / (time.time() - start_time)))
378
+ writer.add_scalar("charts/SPS", int(global_step / (time.time() - start_time)), global_step)
379
+
380
+ writer.add_scalar("stats/truncations", np.sum(truncations), global_step)
381
+ writer.add_scalar("stats/terminations", np.sum(terminations), global_step)
382
+ writer.add_scalar("stats/env_recv_time", env_recv_time, global_step)
383
+ writer.add_scalar("stats/inference_time", inference_time, global_step)
384
+ writer.add_scalar("stats/storage_time", storage_time, global_step)
385
+ writer.add_scalar("stats/env_send_time", env_send_time, global_step)
386
+
387
+ payload = (
388
+ global_step,
389
+ actor_policy_version,
390
+ update,
391
+ obs,
392
+ dones,
393
+ values,
394
+ actions,
395
+ logprobs,
396
+ env_ids,
397
+ rewards,
398
+ )
399
+ if update == 1 or not args.test_actor_learner_throughput:
400
+ rollout_queue_put_time_start = time.time()
401
+ rollout_queue.put(payload)
402
+ rollout_queue_put_time.append(time.time() - rollout_queue_put_time_start)
403
+ writer.add_scalar("stats/rollout_queue_put_time", np.mean(rollout_queue_put_time), global_step)
404
+
405
+ writer.add_scalar(
406
+ "charts/SPS_update",
407
+ int(
408
+ args.local_num_envs
409
+ * args.num_steps
410
+ * len_actor_device_ids
411
+ * args.world_size
412
+ / (time.time() - update_time_start)
413
+ ),
414
+ global_step,
415
+ )
416
+
417
+
418
+ @partial(jax.jit, static_argnums=(3))
419
+ def get_action_and_value2(
420
+ params: flax.core.FrozenDict,
421
+ x: np.ndarray,
422
+ action: np.ndarray,
423
+ action_dim: int,
424
+ ):
425
+ hidden = Network().apply(params.network_params, x)
426
+ logits = Actor(action_dim).apply(params.actor_params, hidden)
427
+ logprob = jax.nn.log_softmax(logits)[jnp.arange(action.shape[0]), action]
428
+ logits = logits - jax.scipy.special.logsumexp(logits, axis=-1, keepdims=True)
429
+ logits = logits.clip(min=jnp.finfo(logits.dtype).min)
430
+ p_log_p = logits * jax.nn.softmax(logits)
431
+ entropy = -p_log_p.sum(-1)
432
+ value = Critic().apply(params.critic_params, hidden).squeeze()
433
+ return logprob, entropy, value
434
+
435
+
436
+ @jax.jit
437
+ def compute_gae(
438
+ env_ids: np.ndarray,
439
+ rewards: np.ndarray,
440
+ values: np.ndarray,
441
+ dones: np.ndarray,
442
+ ):
443
+ dones = jnp.asarray(dones)
444
+ values = jnp.asarray(values)
445
+ env_ids = jnp.asarray(env_ids)
446
+ rewards = jnp.asarray(rewards)
447
+
448
+ _, B = env_ids.shape
449
+ final_env_id_checked = jnp.zeros(args.local_num_envs, jnp.int32) - 1
450
+ final_env_ids = jnp.zeros(B, jnp.int32)
451
+ advantages = jnp.zeros(B)
452
+ lastgaelam = jnp.zeros(args.local_num_envs)
453
+ lastdones = jnp.zeros(args.local_num_envs) + 1
454
+ lastvalues = jnp.zeros(args.local_num_envs)
455
+
456
+ def compute_gae_once(carry, x):
457
+ lastvalues, lastdones, advantages, lastgaelam, final_env_ids, final_env_id_checked = carry
458
+ (
459
+ done,
460
+ value,
461
+ eid,
462
+ reward,
463
+ ) = x
464
+ nextnonterminal = 1.0 - lastdones[eid]
465
+ nextvalues = lastvalues[eid]
466
+ delta = jnp.where(final_env_id_checked[eid] == -1, 0, reward + args.gamma * nextvalues * nextnonterminal - value)
467
+ advantages = delta + args.gamma * args.gae_lambda * nextnonterminal * lastgaelam[eid]
468
+ final_env_ids = jnp.where(final_env_id_checked[eid] == 1, 1, 0)
469
+ final_env_id_checked = final_env_id_checked.at[eid].set(
470
+ jnp.where(final_env_id_checked[eid] == -1, 1, final_env_id_checked[eid])
471
+ )
472
+
473
+ # the last_ variables keeps track of the actual `num_steps`
474
+ lastgaelam = lastgaelam.at[eid].set(advantages)
475
+ lastdones = lastdones.at[eid].set(done)
476
+ lastvalues = lastvalues.at[eid].set(value)
477
+ return (lastvalues, lastdones, advantages, lastgaelam, final_env_ids, final_env_id_checked), (
478
+ advantages,
479
+ final_env_ids,
480
+ )
481
+
482
+ (_, _, _, _, final_env_ids, final_env_id_checked), (advantages, final_env_ids) = jax.lax.scan(
483
+ compute_gae_once,
484
+ (
485
+ lastvalues,
486
+ lastdones,
487
+ advantages,
488
+ lastgaelam,
489
+ final_env_ids,
490
+ final_env_id_checked,
491
+ ),
492
+ (
493
+ dones,
494
+ values,
495
+ env_ids,
496
+ rewards,
497
+ ),
498
+ reverse=True,
499
+ )
500
+ return advantages, advantages + values, final_env_id_checked, final_env_ids
501
+
502
+
503
+ def ppo_loss(params, x, a, logp, mb_advantages, mb_returns, action_dim):
504
+ newlogprob, entropy, newvalue = get_action_and_value2(params, x, a, action_dim)
505
+ logratio = newlogprob - logp
506
+ ratio = jnp.exp(logratio)
507
+ approx_kl = ((ratio - 1) - logratio).mean()
508
+
509
+ if args.norm_adv:
510
+ mb_advantages = (mb_advantages - mb_advantages.mean()) / (mb_advantages.std() + 1e-8)
511
+
512
+ # Policy loss
513
+ pg_loss1 = -mb_advantages * ratio
514
+ pg_loss2 = -mb_advantages * jnp.clip(ratio, 1 - args.clip_coef, 1 + args.clip_coef)
515
+ pg_loss = jnp.maximum(pg_loss1, pg_loss2).mean()
516
+
517
+ # Value loss
518
+ v_loss = 0.5 * ((newvalue - mb_returns) ** 2).mean()
519
+
520
+ entropy_loss = entropy.mean()
521
+ loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef
522
+ return loss, (pg_loss, v_loss, entropy_loss, jax.lax.stop_gradient(approx_kl))
523
+
524
+
525
+ @partial(jax.jit, static_argnums=(6))
526
+ def single_device_update(
527
+ agent_state: TrainState,
528
+ b_obs,
529
+ b_actions,
530
+ b_logprobs,
531
+ b_advantages,
532
+ b_returns,
533
+ action_dim,
534
+ key: jax.random.PRNGKey,
535
+ ):
536
+ ppo_loss_grad_fn = jax.value_and_grad(ppo_loss, has_aux=True)
537
+
538
+ def update_epoch(carry, _):
539
+ agent_state, key = carry
540
+ key, subkey = jax.random.split(key)
541
+
542
+ # taken from: https://github.com/google/brax/blob/main/brax/training/agents/ppo/train.py
543
+ def convert_data(x: jnp.ndarray):
544
+ x = jax.random.permutation(subkey, x)
545
+ x = jnp.reshape(x, (args.num_minibatches, -1) + x.shape[1:])
546
+ return x
547
+
548
+ def update_minibatch(agent_state, minibatch):
549
+ mb_obs, mb_actions, mb_logprobs, mb_advantages, mb_returns = minibatch
550
+ (loss, (pg_loss, v_loss, entropy_loss, approx_kl)), grads = ppo_loss_grad_fn(
551
+ agent_state.params,
552
+ mb_obs,
553
+ mb_actions,
554
+ mb_logprobs,
555
+ mb_advantages,
556
+ mb_returns,
557
+ action_dim,
558
+ )
559
+ grads = jax.lax.pmean(grads, axis_name="local_devices")
560
+ agent_state = agent_state.apply_gradients(grads=grads)
561
+ return agent_state, (loss, pg_loss, v_loss, entropy_loss, approx_kl, grads)
562
+
563
+ agent_state, (loss, pg_loss, v_loss, entropy_loss, approx_kl, grads) = jax.lax.scan(
564
+ update_minibatch,
565
+ agent_state,
566
+ (
567
+ convert_data(b_obs),
568
+ convert_data(b_actions),
569
+ convert_data(b_logprobs),
570
+ convert_data(b_advantages),
571
+ convert_data(b_returns),
572
+ ),
573
+ )
574
+ return (agent_state, key), (loss, pg_loss, v_loss, entropy_loss, approx_kl, grads)
575
+
576
+ (agent_state, key), (loss, pg_loss, v_loss, entropy_loss, approx_kl, _) = jax.lax.scan(
577
+ update_epoch, (agent_state, key), (), length=args.update_epochs
578
+ )
579
+ return agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, key
580
+
581
+
582
+ if __name__ == "__main__":
583
+ args = parse_args()
584
+ if args.distributed:
585
+ jax.distributed.initialize(
586
+ local_device_ids=range(len(args.learner_device_ids) + len(args.actor_device_ids)),
587
+ )
588
+ print(list(range(len(args.learner_device_ids) + len(args.actor_device_ids))))
589
+
590
+ args.world_size = jax.process_count()
591
+ args.local_rank = jax.process_index()
592
+ args.num_envs = args.local_num_envs * args.world_size
593
+ args.batch_size = args.local_batch_size * args.world_size
594
+ args.minibatch_size = args.local_minibatch_size * args.world_size
595
+ args.num_updates = args.total_timesteps // (args.local_batch_size * args.world_size)
596
+ args.async_update = int(args.local_num_envs / args.async_batch_size)
597
+ local_devices = jax.local_devices()
598
+ global_devices = jax.devices()
599
+ learner_devices = [local_devices[d_id] for d_id in args.learner_device_ids]
600
+ actor_devices = [local_devices[d_id] for d_id in args.actor_device_ids]
601
+ global_learner_decices = [
602
+ global_devices[d_id + process_index * len(local_devices)]
603
+ for process_index in range(args.world_size)
604
+ for d_id in args.learner_device_ids
605
+ ]
606
+ print("global_learner_decices", global_learner_decices)
607
+ args.global_learner_decices = [str(item) for item in global_learner_decices]
608
+ args.actor_devices = [str(item) for item in actor_devices]
609
+ args.learner_devices = [str(item) for item in learner_devices]
610
+
611
+ run_name = f"{args.env_id}__{args.exp_name}__{args.seed}__{uuid.uuid4()}"
612
+ if args.track and args.local_rank == 0:
613
+ import wandb
614
+
615
+ wandb.init(
616
+ project=args.wandb_project_name,
617
+ entity=args.wandb_entity,
618
+ sync_tensorboard=True,
619
+ config=vars(args),
620
+ name=run_name,
621
+ monitor_gym=True,
622
+ save_code=True,
623
+ )
624
+ writer = SummaryWriter(f"runs/{run_name}")
625
+ writer.add_text(
626
+ "hyperparameters",
627
+ "|param|value|\n|-|-|\n%s" % ("\n".join([f"|{key}|{value}|" for key, value in vars(args).items()])),
628
+ )
629
+
630
+ # TRY NOT TO MODIFY: seeding
631
+ random.seed(args.seed)
632
+ np.random.seed(args.seed)
633
+ key = jax.random.PRNGKey(args.seed)
634
+ key, network_key, actor_key, critic_key = jax.random.split(key, 4)
635
+
636
+ # env setup
637
+ envs = make_env(args.env_id, args.seed, args.local_num_envs, args.async_batch_size)()
638
+ assert isinstance(envs.single_action_space, gym.spaces.Discrete), "only discrete action space is supported"
639
+
640
+ def linear_schedule(count):
641
+ # anneal learning rate linearly after one training iteration which contains
642
+ # (args.num_minibatches * args.update_epochs) gradient updates
643
+ frac = 1.0 - (count // (args.num_minibatches * args.update_epochs)) / args.num_updates
644
+ return args.learning_rate * frac
645
+
646
+ network = Network()
647
+ actor = Actor(action_dim=envs.single_action_space.n)
648
+ critic = Critic()
649
+ network_params = network.init(network_key, np.array([envs.single_observation_space.sample()]))
650
+ agent_state = TrainState.create(
651
+ apply_fn=None,
652
+ params=AgentParams(
653
+ network_params,
654
+ actor.init(actor_key, network.apply(network_params, np.array([envs.single_observation_space.sample()]))),
655
+ critic.init(critic_key, network.apply(network_params, np.array([envs.single_observation_space.sample()]))),
656
+ ),
657
+ tx=optax.chain(
658
+ optax.clip_by_global_norm(args.max_grad_norm),
659
+ optax.inject_hyperparams(optax.adam)(
660
+ learning_rate=linear_schedule if args.anneal_lr else args.learning_rate, eps=1e-5
661
+ ),
662
+ ),
663
+ )
664
+ agent_state = flax.jax_utils.replicate(agent_state, devices=learner_devices)
665
+
666
+ multi_device_update = jax.pmap(
667
+ single_device_update,
668
+ axis_name="local_devices",
669
+ devices=global_learner_decices,
670
+ in_axes=(0, 0, 0, 0, 0, 0, None, None),
671
+ out_axes=(0, 0, 0, 0, 0, 0, None),
672
+ static_broadcasted_argnums=(6),
673
+ )
674
+
675
+ rollout_queue = queue.Queue(maxsize=1)
676
+ params_queues = []
677
+ for d_idx, d_id in enumerate(args.actor_device_ids):
678
+ params_queue = queue.Queue(maxsize=1)
679
+ params_queue.put(jax.device_put(flax.jax_utils.unreplicate(agent_state.params), local_devices[d_id]))
680
+ threading.Thread(
681
+ target=rollout,
682
+ args=(
683
+ jax.device_put(key, local_devices[d_id]),
684
+ args,
685
+ rollout_queue,
686
+ params_queue,
687
+ writer,
688
+ learner_devices,
689
+ ),
690
+ ).start()
691
+ params_queues.append(params_queue)
692
+
693
+ rollout_queue_get_time = deque(maxlen=10)
694
+ data_transfer_time = deque(maxlen=10)
695
+ learner_policy_version = 0
696
+ prepare_data = jax.jit(prepare_data, device=learner_devices[0])
697
+ while True:
698
+ learner_policy_version += 1
699
+ if learner_policy_version == 1 or not args.test_actor_learner_throughput:
700
+ rollout_queue_get_time_start = time.time()
701
+ (
702
+ global_step,
703
+ actor_policy_version,
704
+ update,
705
+ obs,
706
+ dones,
707
+ values,
708
+ actions,
709
+ logprobs,
710
+ env_ids,
711
+ rewards,
712
+ ) = rollout_queue.get()
713
+ rollout_queue_get_time.append(time.time() - rollout_queue_get_time_start)
714
+ writer.add_scalar("stats/rollout_queue_get_time", np.mean(rollout_queue_get_time), global_step)
715
+
716
+ data_transfer_time_start = time.time()
717
+ b_obs, b_actions, b_logprobs, b_advantages, b_returns = prepare_data(
718
+ obs,
719
+ dones,
720
+ values,
721
+ actions,
722
+ logprobs,
723
+ env_ids,
724
+ rewards,
725
+ )
726
+ b_obs = jnp.array_split(b_obs, len(learner_devices))
727
+ b_actions = jnp.array_split(b_actions, len(learner_devices))
728
+ b_logprobs = jnp.array_split(b_logprobs, len(learner_devices))
729
+ b_advantages = jnp.array_split(b_advantages, len(learner_devices))
730
+ b_returns = jnp.array_split(b_returns, len(learner_devices))
731
+ data_transfer_time.append(time.time() - data_transfer_time_start)
732
+ writer.add_scalar("stats/data_transfer_time", np.mean(data_transfer_time), global_step)
733
+
734
+ training_time_start = time.time()
735
+ (agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, key) = multi_device_update(
736
+ agent_state,
737
+ jax.device_put_sharded(b_obs, learner_devices),
738
+ jax.device_put_sharded(b_actions, learner_devices),
739
+ jax.device_put_sharded(b_logprobs, learner_devices),
740
+ jax.device_put_sharded(b_advantages, learner_devices),
741
+ jax.device_put_sharded(b_returns, learner_devices),
742
+ envs.single_action_space.n,
743
+ key,
744
+ )
745
+ if learner_policy_version == 1 or not args.test_actor_learner_throughput:
746
+ for d_idx, d_id in enumerate(args.actor_device_ids):
747
+ params_queues[d_idx].put(jax.device_put(flax.jax_utils.unreplicate(agent_state.params), local_devices[d_id]))
748
+ if args.profile:
749
+ v_loss[-1, -1, -1].block_until_ready()
750
+ writer.add_scalar("stats/training_time", time.time() - training_time_start, global_step)
751
+ writer.add_scalar("stats/rollout_queue_size", rollout_queue.qsize(), global_step)
752
+ writer.add_scalar("stats/params_queue_size", params_queue.qsize(), global_step)
753
+ print(
754
+ global_step,
755
+ f"actor_policy_version={actor_policy_version}, actor_update={update}, learner_policy_version={learner_policy_version}, training time: {time.time() - training_time_start}s",
756
+ )
757
+
758
+ # TRY NOT TO MODIFY: record rewards for plotting purposes
759
+ writer.add_scalar("charts/learning_rate", agent_state.opt_state[1].hyperparams["learning_rate"][0].item(), global_step)
760
+ writer.add_scalar("losses/value_loss", v_loss[-1, -1, -1].item(), global_step)
761
+ writer.add_scalar("losses/policy_loss", pg_loss[-1, -1, -1].item(), global_step)
762
+ writer.add_scalar("losses/entropy", entropy_loss[-1, -1, -1].item(), global_step)
763
+ writer.add_scalar("losses/approx_kl", approx_kl[-1, -1, -1].item(), global_step)
764
+ writer.add_scalar("losses/loss", loss[-1, -1, -1].item(), global_step)
765
+ if update >= args.num_updates:
766
+ break
767
+
768
+ if args.save_model and args.local_rank == 0:
769
+ if args.distributed:
770
+ jax.distributed.shutdown()
771
+ agent_state = flax.jax_utils.unreplicate(agent_state)
772
+ model_path = f"runs/{run_name}/{args.exp_name}.cleanrl_model"
773
+ with open(model_path, "wb") as f:
774
+ f.write(
775
+ flax.serialization.to_bytes(
776
+ [
777
+ vars(args),
778
+ [
779
+ agent_state.params.network_params,
780
+ agent_state.params.actor_params,
781
+ agent_state.params.critic_params,
782
+ ],
783
+ ]
784
+ )
785
+ )
786
+ print(f"model saved to {model_path}")
787
+ from cleanrl_utils.evals.ppo_envpool_jax_eval import evaluate
788
+
789
+ episodic_returns = evaluate(
790
+ model_path,
791
+ make_env,
792
+ args.env_id,
793
+ eval_episodes=10,
794
+ run_name=f"{run_name}-eval",
795
+ Model=(Network, Actor, Critic),
796
+ )
797
+ for idx, episodic_return in enumerate(episodic_returns):
798
+ writer.add_scalar("eval/episodic_return", episodic_return, idx)
799
+
800
+ if args.upload_model:
801
+ from cleanrl_utils.huggingface import push_to_hub
802
+
803
+ repo_name = f"{args.env_id}-{args.exp_name}-seed{args.seed}"
804
+ repo_id = f"{args.hf_entity}/{repo_name}" if args.hf_entity else repo_name
805
+ push_to_hub(
806
+ args,
807
+ episodic_returns,
808
+ repo_id,
809
+ "PPO",
810
+ f"runs/{run_name}",
811
+ f"videos/{run_name}-eval",
812
+ extra_dependencies=["jax", "envpool", "atari"],
813
+ )
814
+
815
+ envs.close()
816
+ writer.close()
events.out.tfevents.1677105271.ip-26-0-129-73 ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:1c4dd1849f967622eedad49595e0b0f71f241680d16b64cb76bfba4af1d592c6
3
+ size 4754162
poetry.lock ADDED
The diff for this file is too large to render. See raw diff
 
pyproject.toml ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [tool.poetry]
2
+ name = "cleanba"
3
+ version = "0.1.0"
4
+ description = ""
5
+ authors = ["Costa Huang <[email protected]>"]
6
+ readme = "README.md"
7
+ packages = [
8
+ { include = "cleanba" },
9
+ { include = "cleanrl_utils" },
10
+ ]
11
+
12
+ [tool.poetry.dependencies]
13
+ python = "^3.8"
14
+ tensorboard = "^2.12.0"
15
+ envpool = "^0.8.1"
16
+ jax = "0.3.25"
17
+ flax = "0.6.0"
18
+ optax = "0.1.3"
19
+ huggingface-hub = "^0.12.0"
20
+ jaxlib = "0.3.25"
21
+ wandb = "^0.13.10"
22
+ tensorboardx = "^2.5.1"
23
+ chex = "0.1.5"
24
+ gym = "0.23.1"
25
+ opencv-python = "^4.7.0.68"
26
+ moviepy = "^1.0.3"
27
+
28
+
29
+ [tool.poetry.group.dev.dependencies]
30
+ pre-commit = "^3.0.4"
31
+
32
+ [build-system]
33
+ requires = ["poetry-core"]
34
+ build-backend = "poetry.core.masonry.api"
replay.mp4 ADDED
Binary file (203 kB). View file
 
videos/Surround-v5__cleanba_ppo_envpool_impala_atari_wrapper_naturecnn__2__a8cc05aa-b6f7-49a3-8749-88a5281364a9-eval/0.mp4 ADDED
Binary file (203 kB). View file