add an option for deepcopy to NaiveParallelNetwork and Network.copy

Author: runjerry

alf/algorithms/actor_bayes_critic_algorithm.py

@@ -63,6 +63,8 @@
     ],
     default_value=())
 
+AbcRolloutInfo = namedtuple("AbcRolloutInfo", ["action"], default_value=())
+
 AbcLossInfo = namedtuple(
     'AbcLossInfo', ['actor', 'explore', 'critic', 'alpha', 'explore_alpha'],
     default_value=())
@@ -379,7 +381,7 @@ def rollout_step(self, inputs: TimeStep, state: AbcState):
             explore=state.explore,
             critic=state.critic)
         return AlgStep(
-            output=action, state=new_state, info=AbcInfo(action=action))
+            output=action, state=new_state, info=AbcRolloutInfo(action=action))
 
     def _consensus_q_for_actor_train(self, critics, explore, info=()):
         """Get q_value for _actor_train_step. 
@@ -474,8 +476,7 @@ def actor_loss_fn(dqda, action):
         return critics_state, actor_info
 
     def _compute_critic_train_info(self, inputs: TimeStep,
-                                   state: AbcCriticState,
-                                   rollout_info: AbcInfo, action):
+                                   state: AbcCriticState, action):
         target_critics_dist, target_critics_state = self._target_critic_network(
             (inputs.observation, action), state.target_critics)
 
@@ -498,7 +499,7 @@ def _alpha_train_step(self, log_pi):
         return sum(nest.flatten(alpha_loss))
 
     def train_step(self, inputs: TimeStep, state: AbcState,
-                   rollout_info: AbcInfo):
+                   rollout_info: AbcRolloutInfo):
 
         self._training_started = True
 
@@ -549,7 +550,7 @@ def train_step(self, inputs: TimeStep, state: AbcState,
 
         # compute train_info for critic_module, trained in calc_loss
         critic_state, critic_train_info = self._compute_critic_train_info(
-            inputs, state.critic, rollout_info, action)
+            inputs, state.critic, action)
 
         if self._deterministic_actor or self._fixed_alpha:
             alpha_loss = ()

alf/networks/network.py

@@ -143,18 +143,20 @@ def singleton(self, singleton_instance=True):
         self._singleton_instance = singleton_instance
         return self
 
-    def copy(self, **kwargs):
+    def copy(self, deepcopy=False, **kwargs):
         """Create a  copy of this network or return the current instance.
 
         If ``self._singleton_instance`` is True, calling ``copy()`` will return
         ``self``; otherwise it will re-create and return a new ``Network``
         instance using the original arguments used by the constructor.
 
-        **NOTE** When re-creating ``Network``, Network layer weights are *never*
-        copied. This method recreates the ``Network`` instance with the same
-        arguments it was initialized with (excepting any new kwargs).
+        **NOTE** When re-creating ``Network``, by default (``deepcopy=False``), 
+        Network layer weights are *never* copied, instead, this method recreates the 
+        ``Network`` instance with the same arguments it was initialized with 
+        (excepting any new kwargs).
 
         Args:
+            deepcopy (bool): If True, Network layer weights are also copied.
             **kwargs: Args to override when recreating this network.  Commonly
                 overridden args include 'name'.
 
@@ -171,7 +173,8 @@ def _copy(a):
                         return a.copy()
                     elif isinstance(a, torch.nn.Module):
                         b = copy.deepcopy(a)
-                        alf.layers.reset_parameters(b)
+                        if not deepcopy:
+                            alf.layers.reset_parameters(b)
                         return b
                     else:
                         return a
@@ -264,10 +267,9 @@ def make_parallel(self, n):
 class NaiveParallelNetwork(Network):
     """Naive implementation of parallel network."""
 
-    def __init__(self, network, n, name=None):
+    def __init__(self, network, n, deepcopy=False, name=None):
         """
-        A parallel network has ``n`` copies of network with the same structure but
-        different indepently initialized parameters.
+        A parallel network has ``n`` copies of network with the same structure.
 
         ``NaiveParallelNetwork`` created ``n`` independent networks with the same
         structure as ``network`` and evaluate them separately in loop during
@@ -277,6 +279,9 @@ def __init__(self, network, n, name=None):
             network (Network): the parallel network will have ``n`` copies of
                 ``network``.
             n (int): ``n`` copies of ``network``
+            deepcopy (bool): if True, the copies of network will have same
+                initialized parameters, otherwise, they will have different
+                independently initialized parameters.
             name(str): a string that will be used as the name of the created
                 NaiveParallelNetwork instance. If ``None``, ``naive_parallel_``
                 followed by the ``network.name`` will be used by default.
@@ -288,7 +293,8 @@ def __init__(self, network, n, name=None):
         super().__init__(
             network.input_tensor_spec, state_spec=state_spec, name=name)
         self._networks = nn.ModuleList(
-            [network.copy(name=self.name + '_%d' % i) for i in range(n)])
+            [network.copy(
+                deepcopy=deepcopy, name=self.name + '_%d' % i) for i in range(n)])
         self._n = n
 
     def forward(self, inputs, state=()):