From bc1208353a659f0e91a87e39a7950da8ffc2d660 Mon Sep 17 00:00:00 2001
From: silencht <silencht@qq.com>
Date: Fri, 5 Sep 2025 21:29:37 +0800
Subject: [PATCH] [update] optim data record.

---
 teleop/teleop_hand_and_arm.py  | 10 ++++++----
 teleop/utils/episode_writer.py | 17 +++++++++--------
 2 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/teleop/teleop_hand_and_arm.py b/teleop/teleop_hand_and_arm.py
index b926891..f76a1fc 100644
--- a/teleop/teleop_hand_and_arm.py
+++ b/teleop/teleop_hand_and_arm.py
@@ -54,7 +54,7 @@ listen_keyboard_thread.start()
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument('--task_dir', type = str, default = './utils/data', help = 'path to save data')
+    parser.add_argument('--task_dir', type = str, default = './utils/data/', help = 'path to save data')
     parser.add_argument('--frequency', type = float, default = 60.0, help = 'save data\'s frequency')
 
     # basic control parameters
@@ -62,10 +62,12 @@ if __name__ == '__main__':
     parser.add_argument('--arm', type=str, choices=['G1_29', 'G1_23', 'H1_2', 'H1'], default='G1_29', help='Select arm controller')
     parser.add_argument('--ee', type=str, choices=['dex1', 'dex3', 'inspire1', 'brainco'], help='Select end effector controller')
     # mode flags
-    parser.add_argument('--record', action = 'store_true', help = 'Enable data recording')
     parser.add_argument('--motion', action = 'store_true', help = 'Enable motion control mode')
     parser.add_argument('--headless', action='store_true', help='Enable headless mode (no display)')
     parser.add_argument('--sim', action = 'store_true', help = 'Enable isaac simulation mode')
+    parser.add_argument('--record', action = 'store_true', help = 'Enable data recording')
+    parser.add_argument('--tast-name', type = str, default = 'pick cube', help = 'task name for recording')
+    parser.add_argument('--task-goal', type = str, default = 'e.g. pick the red cube on the table.', help = 'task goal for recording')
 
     args = parser.parse_args()
     logger_mp.info(f"args: {args}")
@@ -196,9 +198,9 @@ if __name__ == '__main__':
     
     # record + headless mode
     if args.record and args.headless:
-        recorder = EpisodeWriter(task_dir = args.task_dir, frequency = args.frequency, rerun_log = False)
+        recorder = EpisodeWriter(task_dir = args.task_dir + args.tast_name, task_goal = args.task_goal, frequency = args.frequency, rerun_log = False)
     elif args.record and not args.headless:
-        recorder = EpisodeWriter(task_dir = args.task_dir, frequency = args.frequency, rerun_log = True)
+        recorder = EpisodeWriter(task_dir = args.task_dir + args.tast_name, task_goal = args.task_goal, frequency = args.frequency, rerun_log = True)
         
     try:
         logger_mp.info("Please enter the start signal (enter 'r' to start the subsequent program)")
diff --git a/teleop/utils/episode_writer.py b/teleop/utils/episode_writer.py
index 2dbb07c..8955d10 100644
--- a/teleop/utils/episode_writer.py
+++ b/teleop/utils/episode_writer.py
@@ -11,12 +11,20 @@ import logging_mp
 logger_mp = logging_mp.get_logger(__name__)
 
 class EpisodeWriter():
-    def __init__(self, task_dir, frequency=30, image_size=[640, 480], rerun_log = True):
+    def __init__(self, task_dir, task_goal=None, frequency=30, image_size=[640, 480], rerun_log = True):
         """
         image_size: [width, height]
         """
         logger_mp.info("==> EpisodeWriter initializing...\n")
         self.task_dir = task_dir
+        self.text = {
+            "goal": "Pick up the red cup on the table.",
+            "desc": "task description",
+            "steps":"step1: do this; step2: do that; ...",
+        }
+        if task_goal is not None:
+            self.text['goal'] = task_goal
+
         self.frequency = frequency
         self.image_size = image_size
 
@@ -39,7 +47,6 @@ class EpisodeWriter():
             os.makedirs(self.task_dir)
             logger_mp.info(f"==> episode directory does not exist, now create one.\n")
         self.data_info()
-        self.text_desc()
 
         self.is_available = True  # Indicates whether the class is available for new operations
         # Initialize the queue and worker thread
@@ -73,12 +80,6 @@ class EpisodeWriter():
                 }, 
                 "sim_state": ""
             }
-    def text_desc(self):
-        self.text = {
-            "goal": "Pick up the red cup on the table.",
-            "desc": "Pick up the cup from the table and place it in another position. The operation should be smooth and the water in the cup should not spill out",
-            "steps":"step1: searching for cups. step2: go to the target location. step3: pick up the cup",
-        }
 
  
     def create_episode(self):