Add a filename argument.

isri-aist · Jan 3, 2025 · f22c196 · f22c196
1 parent 38104c7
commit f22c196
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/robo_manip_baselines/teleop/lib/TeleopBase.py b/robo_manip_baselines/teleop/lib/TeleopBase.py
@@ -404,17 +404,16 @@ def manage_status(self):
         if key == 27:  # escape key
             self.quit_flag = True
 
-    def save_data(self):
-        filename = (
-            "teleop_data/{}_{:%Y%m%d_%H%M%S}/env{:0>1}/{}_env{:0>1}_{:0>3}.npz".format(
+    def save_data(self, filename=None):
+        if filename is None:
+            filename = "teleop_data/{}_{:%Y%m%d_%H%M%S}/env{:0>1}/{}_env{:0>1}_{:0>3}.npz".format(
                 self.demo_name,
                 self.datetime_now,
                 self.data_manager.world_idx,
                 self.demo_name,
                 self.data_manager.world_idx,
                 self.data_manager.data_idx,
             )
-        )
         if self.args.compress_rgb:
             print("[TeleopBase] Compress rgb images")
             for camera_name in self.env.unwrapped.camera_names: