on-policy throughput improvement

HorizonRobotics · Sep 15, 2023 · 4270417 · 4270417
1 parent a238283
commit 4270417
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/alf/algorithms/algorithm.py b/alf/algorithms/algorithm.py
@@ -1415,7 +1415,8 @@ def train_from_unroll(self, experience, train_info):
         loss_info, params = self.update_with_gradient(loss_info, valid_masks)
         self.after_update(experience.time_step, train_info)
         self.summarize_train(experience, train_info, loss_info, params)
-        return experience.step_type.shape.numel()
+        shape = alf.nest.get_nest_shape(experience)
+        return shape[0] * shape[1]
 
     @common.mark_replay
     def train_from_replay_buffer(self, update_global_counter=False):