Separate q network target update

araffin · Sep 16, 2022 · f0cc8ff · f0cc8ff
1 parent 9704f1d
commit f0cc8ff
Showing 1 changed file with 18 additions and 1 deletion.
diff --git a/cleanrl/td3_droq_continuous_action_jax.py b/cleanrl/td3_droq_continuous_action_jax.py
@@ -343,6 +343,20 @@ def actor_loss(params):
             )
         )
 
+        # qf1_state = qf1_state.replace(
+        #     target_params=optax.incremental_update(
+        #         qf1_state.params, qf1_state.target_params, args.tau
+        #     )
+        # )
+        # qf2_state = qf2_state.replace(
+        #     target_params=optax.incremental_update(
+        #         qf2_state.params, qf2_state.target_params, args.tau
+        #     )
+        # )
+        return actor_state, (qf1_state, qf2_state), actor_loss_value, key
+
+    @jax.jit
+    def update_q_target_networks(qf1_state, qf2_state):
         qf1_state = qf1_state.replace(
             target_params=optax.incremental_update(
                 qf1_state.params, qf1_state.target_params, args.tau
@@ -353,7 +367,7 @@ def actor_loss(params):
                 qf2_state.params, qf2_state.target_params, args.tau
             )
         )
-        return actor_state, (qf1_state, qf2_state), actor_loss_value, key
+        return qf1_state, qf2_state
 
     start_time = time.time()
     n_updates = 0
@@ -433,6 +447,9 @@ def actor_loss(params):
                     key,
                 )
 
+                # TODO: check if we need to update actor target too
+                qf1_state, qf2_state = update_q_target_networks(qf1_state, qf2_state)
+
                 if n_updates % args.policy_frequency == 0:
                     (
                         actor_state,