Support actor model averaging for ActorCriticAlgorithm (#1820)

* Support actor model averaging for ActorCriticAlgorithm

model averaging is often helpful for evaluation.

* Address comments

Author: emailweixu

alf/algorithms/actor_critic_algorithm.py

@@ -14,6 +14,7 @@
 """Actor critic algorithm."""
 
 import torch
+from typing import Callable, List
 
 import alf
 from alf.algorithms.on_policy_algorithm import OnPolicyAlgorithm
@@ -50,6 +51,8 @@ def __init__(self,
                  config: TrainerConfig = None,
                  loss=None,
                  loss_class=ActorCriticLoss,
+                 actor_avg_fns: List[Callable] = [],
+                 predict_ema_id: int = -1,
                  optimizer=None,
                  checkpoint=None,
                  debug_summaries=False,
@@ -93,6 +96,30 @@ def __init__(self,
                 None, a default loss of class loss_class will be used.
             loss_class (type): the class of the loss. The signature of its
                 constructor: ``loss_class(debug_summaries)``
+            actor_avg_fns: a list of functions for doing model averaging for
+                the actor network. Each function will be responsible for
+                performing one model averaging. The function must take in the
+                current value of the `AveragedModel` parameter, the
+                current value of `model` parameter, and the number of models
+                already averaged. See alf.utils/model_averager.ema_avg_fn for an example.
+
+                Example:
+
+                .. code-block:: python
+
+                    actor_avg_fns=[
+                        partial(ema_avg_fn,
+                                ema_rate=1e-3,
+                                starting_average_after=300_000),
+                        partial(ema_avg_fn,
+                                ema_rate=1e-2,
+                                starting_average_after=300_000),
+                    ])
+
+            predict_ema_id: the index of the actor average model to be used
+                for prediction. -1 means the original actor network is used.
+                0 means the first averaged model in actor_avg_fns is used,
+                and so on
             optimizer (torch.optim.Optimizer): The optimizer for training
             checkpoint (None|str): a string in the format of "prefix@path",
                 where the "prefix" is the multi-step path to the contents in the
@@ -158,13 +185,31 @@ def _deployment_hook(state_dict, prefix: str, unused_loacl_metadata,
 
         self._register_load_state_dict_pre_hook(_deployment_hook)
 
+        self._actor_emas = torch.nn.ModuleList()
+        self._predict_ema_id = predict_ema_id
+        for actor_avg_fn in actor_avg_fns:
+            self._actor_emas.append(
+                alf.utils.model_averager.AveragedModel(self._actor_network,
+                                                       avg_fn=actor_avg_fn))
+
+    def after_update(self, root_inputs: TimeStep, info: ActorCriticInfo):
+        for actor_ema in self._actor_emas:
+            actor_ema.update_parameters(self._actor_network)
+
+    def _trainable_attributes_to_ignore(self):
+        return ['_actor_emas']
+
     def convert_train_state_to_predict_state(self, state):
         return state._replace(value=())
 
     def predict_step(self, inputs: TimeStep, state: ActorCriticState):
         """Predict for one step."""
-        action_dist, actor_state = self._actor_network(inputs.observation,
-                                                       state=state.actor)
+        if self._predict_ema_id == -1:
+            predict_model = self._actor_network
+        else:
+            predict_model = self._actor_emas[self._predict_ema_id]
+        action_dist, actor_state = predict_model(inputs.observation,
+                                                 state=state.actor)
 
         action = dist_utils.epsilon_greedy_sample(action_dist,
                                                   self._epsilon_greedy)

alf/utils/model_averager.py

@@ -12,18 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from numbers import Number
 import torch
+from typing import SupportsFloat
 from torch.optim.swa_utils import AveragedModel as _AveragedModel
-from typing import Union
 
 from alf.utils.schedulers import Scheduler
 
 
-def ema_avg_fn(averaged_model_parameter,
-               model_parameter,
-               num_averaged,
-               ema_rate: Union[Number, Scheduler],
+def ema_avg_fn(averaged_model_parameter: torch.Tensor,
+               model_parameter: torch.Tensor,
+               num_averaged: int,
+               ema_rate: SupportsFloat | Scheduler,
                starting_average_after=0,
                begin_with_simple_average=True):
     """Exponential moving average of model parameters.
@@ -45,7 +44,7 @@ def ema_avg_fn(averaged_model_parameter,
     """
     if num_averaged <= starting_average_after:
         return model_parameter
-    if not isinstance(ema_rate, Number):
+    if not isinstance(ema_rate, SupportsFloat):
         assert isinstance(
             ema_rate, Scheduler), ("ema_rate must be a number or a Scheduler")
         ema_rate = ema_rate()