fix(pu): fix total_episode_count bug in collector

opendilab · Mar 30, 2024 · 52ba3f9 · 52ba3f9
1 parent 5fc7c45
commit 52ba3f9
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/lzero/worker/muzero_collector.py b/lzero/worker/muzero_collector.py
@@ -570,7 +570,6 @@ def collect(self,
 
                 self._env_info[env_id]['time'] += self._timer.value + interaction_duration
                 if timestep.done:
-                    self._total_episode_count += 1
                     reward = timestep.info['eval_episode_return']
                     info = {
                         'reward': reward,
@@ -686,7 +685,7 @@ def collect(self,
             collected_episode = allreduce_data(collected_episode, 'sum')
             collected_duration = allreduce_data(collected_duration, 'sum')
         self._total_envstep_count += collected_step
-        # self._total_episode_count += collected_episode
+        self._total_episode_count += collected_episode
         self._total_duration += collected_duration
 
         # log