diff --git a/acme/agents/jax/sac/learning.py b/acme/agents/jax/sac/learning.py index ab786d9512..6b77df1c14 100644 --- 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.