Added possiblity to record and replay delta actions during teleoperation rather than absolute actions

Co-authored-by: Adil Zouitine <adilzouitinegm@gmail.com>

lerobot/common/policies/hilserl/classifier/modeling_classifier.py

@@ -147,6 +147,8 @@ class Classifier(
 
     def predict_reward(self, x, threshold=0.6):
         if self.config.num_classes == 2:
-            return (self.forward(x).probabilities > threshold).float()
+            probs = self.forward(x).probabilities
+            logging.info(f"Predicted reward images: {probs}")
+            return (probs > threshold).float()
         else:
             return torch.argmax(self.forward(x).probabilities, dim=1)

lerobot/common/robot_devices/control_utils.py

@@ -221,6 +221,7 @@ def record_episode(
         events=events,
         policy=policy,
         fps=fps,
+        record_delta_actions=record_delta_actions,
         teleoperate=policy is None,
         single_task=single_task,
     )
@@ -262,8 +263,12 @@ def control_loop(
     while timestamp < control_time_s:
         start_loop_t = time.perf_counter()
 
+        current_joint_positions = robot.follower_arms["main"].read("Present_Position")
+
         if teleoperate:
             observation, action = robot.teleop_step(record_data=True)
+            if record_delta_actions:
+                action["action"] = action["action"] - current_joint_positions
         else:
             observation = robot.capture_observation()
 
@@ -280,6 +285,9 @@ def control_loop(
             frame = {**observation, **action, "task": single_task}
             dataset.add_frame(frame)
 
+            # if frame["next.done"]:
+            # break
+
         if display_cameras and not is_headless():
             image_keys = [key for key in observation if "image" in key]
             for key in image_keys: