From 5791d9a4f74d9a7aefc02aa94919ea84b4aac890 Mon Sep 17 00:00:00 2001
From: Ciyang <896681584@qq.com>
Date: Mon, 23 Mar 2026 14:13:08 +0800
Subject: [PATCH] Fix as number of images instead of eposide in rollouts

The directory rollouts/minedojo_task/image_from_agent/timestamp is usually saved as episode_x/image/step.png, while os.listdir retrieves the folder under episode_x. The  image_num is not the number of images but the episode.

By default, the rollout_image_num in config.yaml is set to 2000, which is consistent with the number of images used to train UNET in the paper. Therefore, pathlib should be used instead of os.listdir to recursively retrieve images.
---
 collect_rollouts.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/collect_rollouts.py b/collect_rollouts.py
index 822fb1a..aa3c4db 100644
--- a/collect_rollouts.py
+++ b/collect_rollouts.py
@@ -30,7 +30,6 @@ def collect_rollouts(config, envs, agent):
     agent_state = None
     reward = [0] * len(envs)
 
-    images = os.listdir(save_dir)
 
     image_num = 0
 
@@ -65,7 +64,7 @@ def collect_rollouts(config, envs, agent):
         done = np.stack(done)
         info = list(info)
 
-        images = os.listdir(save_dir)
+        images = list(pathlib.Path(save_dir).rglob("*.png"))
         image_num = len(images)
 
 
@@ -197,4 +196,4 @@ def recursive_update(base, update):
         arg_type = tools.args_type(value)
         parser.add_argument(f"--{key}", type=arg_type, default=arg_type(value))
 
-    main(parser.parse_args(remaining))
\ No newline at end of file
+    main(parser.parse_args(remaining))