temp

Author: juelg

python/rcsss/camera/hw.py

@@ -94,6 +94,10 @@ def start(self, warm_up: bool = True):
         self._thread.start()
 
     def record_video(self, path: Path, episode: int):
+        if self.recording_ongoing():
+            return
+        print("start recording")
+        self.clear_buffer()
         for camera in self.camera_names:
             self.writer[camera] = cv2.VideoWriter(
                 str(path / f"episode_{episode}_{camera}.mp4"),
@@ -104,19 +108,21 @@ def record_video(self, path: Path, episode: int):
             )
 
     def recording_ongoing(self) -> bool:
-        return len(self.writer) > 0
+        with self._buffer_lock:
+            return len(self.writer) > 0
 
     def stop_video(self):
         if len(self.writer) > 0:
-            for camera_key, writer in self.writer.items():
-                for i in range(self._next_ring_index):
-                    frameset = self._buffer[i]
-                    assert frameset is not None
-                    # rgb to bgr as expected by opencv
-                    writer.write(frameset.frames[camera_key].camera.color.data[:, :, ::-1])
-            for camera in self.camera_names:
-                self.writer[camera].release()
-            self.writer = {}
+            with self._buffer_lock:
+                for camera_key, writer in self.writer.items():
+                    for i in range(self._next_ring_index):
+                        frameset = self._buffer[i]
+                        assert frameset is not None
+                        # rgb to bgr as expected by opencv
+                        writer.write(frameset.frames[camera_key].camera.color.data[:, :, ::-1])
+                for camera in self.camera_names:
+                    self.writer[camera].release()
+                self.writer = {}
 
     def warm_up(self):
         for _ in range(self.config.warm_up_disposal_frames):
@@ -138,7 +144,8 @@ def polling_thread(self, warm_up: bool = True):
                     writer.write(frameset.frames[camera_key].camera.color.data[:, :, ::-1])
                 self._next_ring_index = (self._next_ring_index + 1) % self.config.max_buffer_frames
                 self._buffer_len = max(self._buffer_len + 1, self.config.max_buffer_frames)
-            self.rate(self.config.frame_rate)
+            # self.rate(self.config.frame_rate)
+            sleep(1 / self.config.frame_rate)
 
     def poll_frame_set(self) -> FrameSet:
         """Gather frames over all available cameras."""

python/rcsss/control/vive.py

@@ -41,7 +41,7 @@
 INCLUDE_ROTATION = True
 ROBOT_IP = "192.168.101.1"
 ROBOT_INSTANCE = RobotInstance.HARDWARE
-DEBUG = True
+DEBUG = False
 
 
 class Button(IntFlag):
@@ -223,6 +223,10 @@ def input_loop(env_rel, action_server: UDPViveActionServer, camera_set: RealSens
                 # record videos
                 video_path = env_rel.path / "videos"
                 video_path.mkdir(parents=True, exist_ok=True)
+                # camera_set.record_video(env_rel.path / "videos", env_rel.episode_count)
+
+                gif_path = env_rel.path / "gifs"
+                gif_path.mkdir(parents=True, exist_ok=True)
                 print(f'{env_rel.episode_count = }')
 
                 thread = threading.Thread(target=action_server.environment_step_loop)
@@ -252,7 +256,7 @@ def main():
                 "side": "243522070385",
             }
             camera_set = default_realsense(camera_dict)
-            camera_set = None
+            # camera_set = None
             env_rel = fr3_hw_env(
                 ip=ROBOT_IP,
                 camera_set = camera_set,
@@ -261,7 +265,7 @@ def main():
                 control_mode=ControlMode.CARTESIAN_TQuart,
                 # control_mode=ControlMode.JOINTS,
                 gripper_cfg=default_fr3_hw_gripper_cfg(),
-                max_relative_movement=(0.01, np.deg2rad(5)),
+                max_relative_movement=(0.5, np.deg2rad(90)),
                 # max_relative_movement=(0.5, np.deg2rad(90)),
                 # TODO: max should be always according to the last step
                 # max_relative_movement=np.deg2rad(20),
@@ -283,7 +287,7 @@ def main():
             env_rel.get_wrapper_attr("sim").open_gui()
 
         if not DEBUG:
-            env_rel = StorageWrapper(env_rel, path="/home/juelg/code/frankcsy/record_real_christmas", camera_set=camera_set)
+            env_rel = StorageWrapper(env_rel, path="/home/juelg/test_data", camera_set=camera_set)
             # ip_secondary = "192.168.102.1"
             # with Desk.fci(ip_secondary, user, pw):
             #     f = rcsss.hw.FR3(ip_secondary)
@@ -299,7 +303,7 @@ def main():
         with env_rel:
             with UDPViveActionServer(VIVE_HOST, VIVE_PORT, env_rel) as action_server:
                 if not DEBUG:
-                    input_loop(env_rel, action_server, None)
+                    input_loop(env_rel, action_server, camera_set)
 
                 else:
                     action_server.environment_step_loop()

python/rcsss/envs/wrappers.py

@@ -18,7 +18,7 @@
 class StorageWrapper(gym.Wrapper):
     # TODO: this should also record the instruction
     FILE = "episode_{}.npz"
-    GIF = "{}_episode_{}_{}.gif"
+    GIF = "gifs/{}_episode_{}_{}.gif"
     FOLDER = "experiment_{}"
     GIF_DURATION_S = 0.5
 
@@ -40,6 +40,7 @@ def __init__(
         self.record_numpy = record_numpy
         self.gif = gif
         self.camera_set = camera_set
+        print("camera_set", self.camera_set)
 
         # make folders
         self.path = Path(path) / self.FOLDER.format(self.timestamp)
@@ -69,9 +70,11 @@ def flush(self):
         if self.record_numpy:
             np.savez(self.path / self.FILE.format(self.episode_count), **self.data)
         if self.camera_set is not None and self.camera_set.recording_ongoing():
+            print("stopping video recording")
             self.camera_set.stop_video()
 
         if self.gif:
+            print("saving gifs")
             # for key in ["side", "wrist", "bird_eye", "openvla_view"]:
             for key in ["side", "right_side", "bird_eye", "left_side", "front"]:
                 if f"observation.frames.{key}.rgb" in self.data:
@@ -84,6 +87,7 @@ def flush(self):
                             continue
                         previous_timestamp = self.data["timestamp"][idx]
                         imgs.append(Image.fromarray(img))
+                    print(self.path / self.GIF.format(self.timestamp, self.episode_count, key))
                     imgs[0].save(
                         self.path / self.GIF.format(self.timestamp, self.episode_count, key),
                         save_all=True,
@@ -100,7 +104,8 @@ def step(self, action: dict) -> tuple[Any, SupportsFloat, bool, bool, dict[str,
         obs, reward, terminated, truncated, info = super().step(action)
         # write obs and action into data
         act_obs = {"action": action, "observation": obs}
-        flattened_act_obs = flatten(copy.deepcopy(act_obs), reducer="dot")
+        # act_obs = {"action": action}
+        flattened_act_obs = act_obs #flatten(act_obs, reducer="dot")
         # add timestamp
         flattened_act_obs["timestamp"] = datetime.now().timestamp()
         self.data["language_instruction"] = self.language_instruction
@@ -116,8 +121,10 @@ def reset(self, *, seed: int | None = None, options: dict[str, Any] | None = Non
         self.flush()
         self.step_count = 0
         re = super().reset(seed=seed, options=options)
-        if self.camera_set is not None:
-            self.camera_set.record_video(self.path, self.episode_count)
+        # if self.camera_set is not None:
+        #     print("starting video recording")
+        #     self.camera_set.clear_buffer()
+        #     self.camera_set.record_video(self.path / "videos", self.episode_count)
         return re
 
     def close(self):