SAC jax

benchmark/sac.sh

@@ -4,4 +4,13 @@ OMP_NUM_THREADS=1 xvfb-run -a poetry run python -m cleanrl_utils.benchmark \
     --env-ids HalfCheetah-v2 Walker2d-v2 Hopper-v2 \
     --command "poetry run python cleanrl/sac_continuous_action.py --track --capture-video" \
     --num-seeds 3 \
-    --workers 3
+    --workers 3
+
+poetry install --with mujoco,pybullet,jax
+poetry run pip install --upgrade "jax[cuda]==0.3.17" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+poetry run python -c "import mujoco_py"
+poetry run python -m cleanrl_utils.benchmark \
+    --env-ids HalfCheetah-v2 Walker2d-v2 Hopper-v2 \
+    --command "poetry run python cleanrl/sac_continuous_action_jax.py --track" \
+    --num-seeds 3 \
+    --workers 1

cleanrl/ddpg_continuous_action_jax.py

@@ -108,8 +108,8 @@ def __call__(self, x):
         return x
 
 
-class TrainState(TrainState):
-    target_params: flax.core.FrozenDict
+class RLTrainState(TrainState):
+    target_params: flax.core.FrozenDict = None
 
 
 if __name__ == "__main__":
@@ -163,13 +163,13 @@ class TrainState(TrainState):
         action_bias=action_bias,
     )
     qf1 = QNetwork()
-    actor_state = TrainState.create(
+    actor_state = RLTrainState.create(
         apply_fn=actor.apply,
         params=actor.init(actor_key, obs),
         target_params=actor.init(actor_key, obs),
         tx=optax.adam(learning_rate=args.learning_rate),
     )
-    qf1_state = TrainState.create(
+    qf1_state = RLTrainState.create(
         apply_fn=qf1.apply,
         params=qf1.init(qf1_key, obs, envs.action_space.sample()),
         target_params=qf1.init(qf1_key, obs, envs.action_space.sample()),
@@ -180,8 +180,8 @@ class TrainState(TrainState):
 
     @jax.jit
     def update_critic(
-        actor_state: TrainState,
-        qf1_state: TrainState,
+        actor_state: RLTrainState,
+        qf1_state: RLTrainState,
         observations: np.ndarray,
         actions: np.ndarray,
         next_observations: np.ndarray,
@@ -202,8 +202,8 @@ def mse_loss(params):
 
     @jax.jit
     def update_actor(
-        actor_state: TrainState,
-        qf1_state: TrainState,
+        actor_state: RLTrainState,
+        qf1_state: RLTrainState,
         observations: np.ndarray,
     ):
         def actor_loss(params):

cleanrl/dqn_atari_jax.py

@@ -119,8 +119,8 @@ def __call__(self, x):
         return x
 
 
-class TrainState(TrainState):
-    target_params: flax.core.FrozenDict
+class RLTrainState(TrainState):
+    target_params: flax.core.FrozenDict = None
 
 
 def linear_schedule(start_e: float, end_e: float, duration: int, t: int):
@@ -163,7 +163,7 @@ def linear_schedule(start_e: float, end_e: float, duration: int, t: int):
 
     q_network = QNetwork(action_dim=envs.single_action_space.n)
 
-    q_state = TrainState.create(
+    q_state = RLTrainState.create(
         apply_fn=q_network.apply,
         params=q_network.init(q_key, obs),
         target_params=q_network.init(q_key, obs),

cleanrl/dqn_jax.py

@@ -92,8 +92,8 @@ def __call__(self, x: jnp.ndarray):
         return x
 
 
-class TrainState(TrainState):
-    target_params: flax.core.FrozenDict
+class RLTrainState(TrainState):
+    target_params: flax.core.FrozenDict = None
 
 
 def linear_schedule(start_e: float, end_e: float, duration: int, t: int):
@@ -136,7 +136,7 @@ def linear_schedule(start_e: float, end_e: float, duration: int, t: int):
 
     q_network = QNetwork(action_dim=envs.single_action_space.n)
 
-    q_state = TrainState.create(
+    q_state = RLTrainState.create(
         apply_fn=q_network.apply,
         params=q_network.init(q_key, obs),
         target_params=q_network.init(q_key, obs),

cleanrl/sac_continuous_action.py

@@ -49,8 +49,6 @@ def parse_args():
         help="target smoothing coefficient (default: 0.005)")
     parser.add_argument("--batch-size", type=int, default=256,
         help="the batch size of sample from the reply memory")
-    parser.add_argument("--exploration-noise", type=float, default=0.1,
-        help="the scale of exploration noise")
     parser.add_argument("--learning-starts", type=int, default=5e3,
         help="timestep to start learning")
     parser.add_argument("--policy-lr", type=float, default=3e-4,
@@ -61,8 +59,6 @@ def parse_args():
         help="the frequency of training policy (delayed)")
     parser.add_argument("--target-network-frequency", type=int, default=1, # Denis Yarats' implementation delays this by 2.
         help="the frequency of updates for the target nerworks")
-    parser.add_argument("--noise-clip", type=float, default=0.5,
-        help="noise clip parameter of the Target Policy Smoothing Regularization")
     parser.add_argument("--alpha", type=float, default=0.2,
             help="Entropy regularization coefficient.")
     parser.add_argument("--autotune", type=lambda x:bool(strtobool(x)), default=True, nargs="?", const=True,