Save only if minedojo

Eclectic-Sheep · Jul 17, 2023 · f77389e · f77389e
1 parent 0671cc8
commit f77389e
Showing 1 changed file with 14 additions and 12 deletions.
diff --git a/sheeprl/algos/p2e_dv2/p2e_dv2.py b/sheeprl/algos/p2e_dv2/p2e_dv2.py
@@ -730,9 +730,10 @@ def main():
         )
 
     # Stats file info
-    stats_dir = os.path.join(log_dir, "stats")
-    os.makedirs(stats_dir, exist_ok=True)
-    stats_filename = os.path.join(stats_dir, "stats.jsonl")
+    if "minedojo" in args.env_id:
+        stats_dir = os.path.join(log_dir, "stats")
+        os.makedirs(stats_dir, exist_ok=True)
+        stats_filename = os.path.join(stats_dir, "stats.jsonl")
 
     # Get the first environment observation and start the optimization
     episode_steps = []
@@ -791,15 +792,16 @@ def main():
                     real_actions = np.array([real_act.cpu().argmax() for real_act in real_actions])
 
         # Save stats
-        with jsonlines.open(stats_filename, mode="a") as writer:
-            writer.write(
-                {
-                    "life_stats": infos["life_stats"],
-                    "location_stats": infos["location_stats"],
-                    "action": real_actions.tolist(),
-                    "biomeid": infos["biomeid"],
-                }
-            )
+        if "minedojo" in args.env_id:
+            with jsonlines.open(stats_filename, mode="a") as writer:
+                writer.write(
+                    {
+                        "life_stats": infos["life_stats"],
+                        "location_stats": infos["location_stats"],
+                        "action": real_actions.tolist(),
+                        "biomeid": infos["biomeid"],
+                    }
+                )
 
         step_data["is_first"] = copy.deepcopy(step_data["dones"])
         o, rewards, dones, truncated, infos = env.step(real_actions.reshape(env.action_space.shape))