Log rewards statistics in SAC agents

(similar to PPO agents)
google-deepmind · May 23, 2022 · 5e38355 · 5e38355
1 parent 71b2ab8
commit 5e38355
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/acme/agents/jax/sac/learning.py b/acme/agents/jax/sac/learning.py
@@ -222,6 +222,10 @@ def update_step(
               jax.tree_map(lambda x: jnp.std(x, axis=0),
                            transitions.next_observation)))
 
+      metrics['rewards_mean'] = jnp.mean(
+          jnp.abs(jnp.mean(transitions.reward, axis=0)))
+      metrics['rewards_std'] = jnp.std(transitions.reward, axis=0)
+
       return new_state, metrics
 
     # General learner book-keeping and loggers.