Copybara import of the project:
by Jongwook Choi
Log rewards statistics in SAC agents
(similar to PPO agents)
COPYBARA_INTEGRATE_REVIEW=https://github.com/deepmind/acme/pull/232 from wookayin:sac-log-reward
@@ -222,6 +222,10 @@ class SACLearner(acme.Learner):
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.
