fix

vwxyzjn · Oct 17, 2023 · 5f3f716 · 5f3f716
1 parent 9c0959c
commit 5f3f716
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 0 deletions.
diff --git a/cleanrl/c51_jax.py b/cleanrl/c51_jax.py
@@ -29,6 +29,8 @@ class Args:
     """the wandb's project name"""
     wandb_entity: str = None
     """the entity (team) of wandb's project"""
+    capture_video: bool = False
+    """whether to capture videos of the agent performances (check out `videos` folder)"""
     save_model: bool = False
     """whether to save model into the `runs/{run_name}` folder"""
     upload_model: bool = False

diff --git a/cleanrl/dqn_jax.py b/cleanrl/dqn_jax.py
@@ -29,6 +29,8 @@ class Args:
     """the wandb's project name"""
     wandb_entity: str = None
     """the entity (team) of wandb's project"""
+    capture_video: bool = False
+    """whether to capture videos of the agent performances (check out `videos` folder)"""
     save_model: bool = False
     """whether to save model into the `runs/{run_name}` folder"""
     upload_model: bool = False

diff --git a/cleanrl/ppo_rnd_envpool.py b/cleanrl/ppo_rnd_envpool.py
@@ -1,4 +1,5 @@
 # docs and experiment results can be found at https://docs.cleanrl.dev/rl-algorithms/ppo-rnd/#ppo_rnd_envpoolpy
+from dataclasses import dataclass
 import os
 import random
 import time