logging: Sequences Containing Done considers critic mask

Author: jakegrigsby

amago/agent.py

@@ -608,7 +608,9 @@ def masked_avg(x_, dim=0):
                 "Max TD Target": td_target[where_mask].max(),
                 "TD Target (test-time gamma)": masked_avg(td_target, -1),
                 "Mean Reward (in training sequences)": masked_avg(r),
-                "Sequences Containing Done": d[:, :, 0, 0, 0].any(1).sum(),
+                "Sequences Containing Done": (d * mask.all(2, keepdim=True))
+                .any((1, 2, 3))
+                .sum(),
                 "Min Reward (in training sequences)": r[where_mask].min(),
                 "Max Reward (in training sequences)": r[where_mask].max(),
             }

docs/tutorial/customization.rst

@@ -113,14 +113,14 @@ RLDataset
 
 Actor Output Head
 ~~~~~~~~~~~~~~~~~~~
-**Implement** :py:class:`~amago.nets.actor_critic.BaseActorHead`
+**Implement**: :py:class:`~amago.nets.actor_critic.BaseActorHead`
 
 **Configure**: ``Agent.actor_type : MyActor``
 
 |
 
 Critic Output Head
 ~~~~~~~~~~~~~~~~~~~
-**Implement** : :py:class:`~amago.nets.actor_critic.BaseCriticHead`
+**Implement**: :py:class:`~amago.nets.actor_critic.BaseCriticHead`
 
 **Configure**: ``Agent.critic_type : MyCritic``