update docs for actor/critic output customization

Author: jakegrigsby

amago/agent.py

@@ -194,6 +194,8 @@ def V(state, critic, action_dist, k) -> float:
         use_target_actor: If True, use a target actor to sample actions used in TD targets.
             Defaults to True.
         use_multigamma: If True, train on multiple discount horizons (:py:class:`~amago.agent.Multigammas`) in parallel. Defaults to True.
+        actor_type: Actor MLP head for producing action distributions. Defaults to :py:class:`~amago.nets.actor_critic.Actor`.
+        critic_type: Critic MLP head for producing Q-values. Defaults to :py:class:`~amago.nets.actor_critic.NCritics`.
     """
 
     def __init__(
@@ -708,7 +710,7 @@ class MultiTaskAgent(Agent):
 
     The combination of points 2 and 3 stresses accurate advantage estimates and motivates a change
     in the default value of num_actions_for_value_in_critic_loss from 1 --> 3. Arguments otherwise
-    follow the information listed in amago.agent.Agent.
+    follow the information listed in :py:class:`~amago.agent.Agent`.
     """
 
     def __init__(

amago/nets/actor_critic.py

@@ -309,9 +309,7 @@ def forward(
         self, state: torch.Tensor, action: torch.Tensor, log_dict: Optional[dict] = None
     ) -> torch.Tensor:
         assert state.ndim == 3
-        B, L, _ = state.shape
         assert action.ndim == 5
-        K = action.shape[0]
         out = self.critic_network_forward(state=state, action=action, log_dict=log_dict)
         return out
 

docs/tutorial/customization.rst

@@ -107,6 +107,20 @@ RLDataset
 
 **Implement**: :py:class:`~amago.nets.policy_dists.PolicyOutput`
 
-**Substitute**: ``config = {"amago.nets.actor_critic.Actor.continuous_dist_type" : MyPolicyOutput, ...}; use_config(config); experiment = Experiment(...)``
+**Configure**: ``config = {"amago.nets.actor_critic.Actor.continuous_dist_type" : MyPolicyOutput, ...}; use_config(config); experiment = Experiment(...)``
 
-|
+|
+
+Actor Output Head
+~~~~~~~~~~~~~~~~~~~
+**Implement** :py:class:`~amago.nets.actor_critic.BaseActorHead`
+
+**Configure**: ``Agent.actor_type : MyActor``
+
+|
+
+Critic Output Head
+~~~~~~~~~~~~~~~~~~~
+**Implement** : :py:class:`~amago.nets.actor_critic.BaseCriticHead`
+
+**Configure**: ``Agent.critic_type : MyCritic``