Add movement scheduler and integrate into tracker

reachy_phone_home/movements.py

@@ -0,0 +1,196 @@
+"""Situation-based movement library for Reachy Mini."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+import random
+import time
+from typing import Callable, Dict, List, Optional
+
+from reachy_mini import ReachyMini
+from reachy_mini.motion.recorded_move import RecordedMove, RecordedMoves
+
+
+@dataclass(frozen=True)
+class Movement:
+    name: str
+    run: Callable[[], None]
+
+
+class SituationMovements:
+    """Expose recorded moves grouped by situation."""
+
+    EMOTIONS_DATASET = "pollen-robotics/reachy-mini-emotions-library"
+    DANCES_DATASET = "pollen-robotics/reachy-mini-dances-library"
+
+    def __init__(self, reachy: ReachyMini) -> None:
+        self.reachy = reachy
+        self._emotions: Optional[RecordedMoves] = None
+        self._dances: Optional[RecordedMoves] = None
+
+        self._groups: Dict[str, List[str]] = {
+            "ambient": [
+                "amazed1",
+                "attentive1",
+                "attentive2",
+                "cheerful1",
+                "curious1",
+                "helpful1",
+                "impatient1",
+            ],
+            "good_job": [
+                "dance1",
+                "grateful1",
+                "success1",
+                "proud1",
+            ],
+            "checking": [
+                "confused1",
+            ],
+            "bad": [
+                "disgusted1",
+                "displeased1",
+                "displeased2",
+            ],
+            "sad": [
+                "downcast1",
+                "lonely1",
+            ],
+            "response_yes": [
+                "enthusiastic1",
+                "inquiring1",
+                "inquiring2",
+                "welcoming2",
+                "understanding2",
+            ],
+            "response_no": [
+                "uncertain1",
+                "no1",
+            ],
+        }
+
+        self._dance_only = {"dance1"}
+
+    def _load_emotions(self) -> None:
+        if self._emotions is None:
+            self._emotions = RecordedMoves(self.EMOTIONS_DATASET)
+
+    def _load_dances(self) -> None:
+        if self._dances is None:
+            self._dances = RecordedMoves(self.DANCES_DATASET)
+
+    def _get_move(self, name: str) -> Optional[RecordedMove]:
+        if name in self._dance_only:
+            self._load_dances()
+            assert self._dances is not None
+            if name in self._dances.list_moves():
+                return self._dances.get(name)
+            return None
+
+        self._load_emotions()
+        assert self._emotions is not None
+        if name in self._emotions.list_moves():
+            return self._emotions.get(name)
+        return None
+
+    def _as_movements(self, names: List[str]) -> List[Movement]:
+        movements: List[Movement] = []
+        for name in names:
+            move = self._get_move(name)
+            if move is None:
+                continue
+            movements.append(
+                Movement(
+                    name=name,
+                    run=lambda m=move: self.reachy.play_move(
+                        m, initial_goto_duration=0.6, sound=False
+                    ),
+                )
+            )
+        return movements
+
+    def ambient(self) -> List[Movement]:
+        return self._as_movements(self._groups["ambient"])
+
+    def good_job(self) -> List[Movement]:
+        return self._as_movements(self._groups["good_job"])
+
+    def checking(self) -> List[Movement]:
+        return self._as_movements(self._groups["checking"])
+
+    def bad(self) -> List[Movement]:
+        return self._as_movements(self._groups["bad"])
+
+    def sad(self) -> List[Movement]:
+        return self._as_movements(self._groups["sad"])
+
+    def response_yes(self) -> List[Movement]:
+        return self._as_movements(self._groups["response_yes"])
+
+    def response_no(self) -> List[Movement]:
+        return self._as_movements(self._groups["response_no"])
+
+    def list_group(self, group: str) -> List[Movement]:
+        if group not in self._groups:
+            raise ValueError(f"Unknown group: {group}")
+        return self._as_movements(self._groups[group])
+
+
+class MovementScheduler:
+    """Trigger situation movements based on heartbeat and phone-use state."""
+
+    def __init__(
+        self,
+        movements: SituationMovements,
+        good_job_heartbeats: int = 3,
+        phone_use_bad_sec: float = 10.0,
+        phone_use_clear_sec: float = 3.0,
+        restore_head_duration: float = 0.6,
+    ) -> None:
+        self.movements = movements
+        self.good_job_heartbeats = int(good_job_heartbeats)
+        self.phone_use_bad_sec = float(phone_use_bad_sec)
+        self.phone_use_clear_sec = float(phone_use_clear_sec)
+        self.restore_head_duration = float(restore_head_duration)
+
+        self._good_job_streak = 0
+        self._phone_use_start: Optional[float] = None
+        self._phone_use_clear_start: Optional[float] = None
+
+        self._good_job_moves = movements.good_job()
+        self._bad_moves = movements.bad()
+
+    def _pick_and_run(self, pool: List[Movement]) -> None:
+        if not pool:
+            return
+        reachy = self.movements.reachy
+        prev_head = reachy.get_current_head_pose()
+        random.choice(pool).run()
+        reachy.goto_target(head=prev_head, duration=self.restore_head_duration)
+
+    def on_heartbeat(self, *, phone_tracked: bool, phone_use: bool, now: Optional[float] = None) -> None:
+        """Call on each heartbeat with current phone states."""
+        if now is None:
+            now = time.time()
+
+        if phone_tracked and not phone_use:
+            self._good_job_streak += 1
+            if self._good_job_streak >= self.good_job_heartbeats:
+                self._pick_and_run(self._good_job_moves)
+                self._good_job_streak = 0
+        else:
+            self._good_job_streak = 0
+
+        if phone_use:
+            self._phone_use_clear_start = None
+            if self._phone_use_start is None:
+                self._phone_use_start = now
+            elif now - self._phone_use_start >= self.phone_use_bad_sec:
+                self._pick_and_run(self._bad_moves)
+                self._phone_use_start = now
+        else:
+            if self._phone_use_clear_start is None:
+                self._phone_use_clear_start = now
+            elif now - self._phone_use_clear_start >= self.phone_use_clear_sec:
+                self._phone_use_start = None
+                self._phone_use_clear_start = None