Add Go1 evaluation and keyboard control

robogauge/tasks/gauge/goals/__init__.py

@@ -1,3 +1,4 @@
 from robogauge.tasks.gauge.goals.base_goal import BaseGoal
 from robogauge.tasks.gauge.goals.joystick_goal import JoystickGoal
+from robogauge.tasks.gauge.goals.keyboard_goal import KeyboardGoal
 from robogauge.tasks.gauge.goals.velocity_goals import MaxVelocityGoal, DiagonalVelocityGoal, TargetPosVelocityGoal

robogauge/tasks/gauge/goals/keyboard_goal.py

@@ -0,0 +1,149 @@
+# -*- coding: utf-8 -*-
+"""Keyboard teleoperation goal using an input listener separate from MuJoCo."""
+
+import atexit
+import queue
+import threading
+
+from robogauge.tasks.gauge.goal_data import GoalData, VelocityGoal
+from robogauge.tasks.gauge.goals.base_goal import BaseGoal
+from robogauge.tasks.simulator.sim_data import SimData
+from robogauge.utils.helpers import class_to_dict
+from robogauge.utils.logger import logger
+
+
+class _KeyboardState:
+    """Track held keys without replacing MuJoCo viewer shortcuts."""
+
+    def __init__(self):
+        self.events = queue.Queue()
+        self.running = False
+        self.held = set()
+        self.lock = threading.Lock()
+        self.thread = None
+        self.listener = None
+
+    @staticmethod
+    def _name(key):
+        try:
+            if hasattr(key, 'char') and key.char:
+                return key.char.lower()
+        except Exception:
+            pass
+        return str(key).lower()
+
+    def _worker(self):
+        while self.running:
+            try:
+                event_type, key = self.events.get(timeout=0.05)
+            except queue.Empty:
+                continue
+            name = self._name(key)
+            with self.lock:
+                if event_type == 'press':
+                    self.held.add(name)
+                else:
+                    self.held.discard(name)
+
+    def start(self):
+        try:
+            from pynput import keyboard
+        except ImportError as exc:
+            raise ImportError(
+                "Keyboard control requires 'pynput'. Install it with "
+                "'python -m pip install pynput'."
+            ) from exc
+
+        self.running = True
+        self.listener = keyboard.Listener(
+            on_press=lambda key: self.events.put(('press', key)),
+            on_release=lambda key: self.events.put(('release', key)),
+        )
+        self.listener.start()
+        self.thread = threading.Thread(target=self._worker, daemon=True)
+        self.thread.start()
+
+    def snapshot(self):
+        with self.lock:
+            return set(self.held)
+
+    def clear(self):
+        with self.lock:
+            self.held.clear()
+
+    def stop(self):
+        self.running = False
+        if self.listener is not None:
+            self.listener.stop()
+            self.listener = None
+
+
+class KeyboardGoal(BaseGoal):
+    """Produce velocity commands while the configured teleoperation keys are held."""
+
+    name = 'keyboard'
+
+    def __init__(self, max_velocity, goal_type='velocity', **kwargs):
+        super().__init__()
+        if goal_type != 'velocity':
+            raise NotImplementedError("Only 'velocity' goal type is implemented for KeyboardGoal.")
+
+        kwargs.pop('enabled', None)
+        if kwargs:
+            logger.warning(f"Unused kwargs in KeyboardGoal: {kwargs}")
+
+        self.max_velocity = class_to_dict(max_velocity)
+        self.total = 1
+        self.current_command = (0.0, 0.0, 0.0)
+        self.keyboard = _KeyboardState()
+        self.keyboard.start()
+        atexit.register(self.keyboard.stop)
+        logger.info(
+            "Keyboard control: Up/Down=forward/back, Left/Right=yaw, "
+            "','=left strafe, '.'=right strafe, K=stop. "
+            "MuJoCo viewer shortcuts remain independent."
+        )
+
+    def is_reset(self, sim_data: SimData) -> bool:
+        return False
+
+    def pre_get_goal(self, sim_data: SimData) -> bool:
+        return False
+
+    def reset_goal(self):
+        self.keyboard.clear()
+        self.current_command = (0.0, 0.0, 0.0)
+
+    @staticmethod
+    def _axis_command(held, positive_key, negative_key, limits):
+        positive = positive_key in held
+        negative = negative_key in held
+        if positive == negative:
+            return 0.0
+        return float(limits[1] if positive else limits[0])
+
+    def command_from_keys(self, held):
+        if 'k' in held:
+            return 0.0, 0.0, 0.0
+        return (
+            self._axis_command(held, 'key.up', 'key.down', self.max_velocity['lin_vel_x']),
+            self._axis_command(held, ',', '.', self.max_velocity['lin_vel_y']),
+            self._axis_command(held, 'key.left', 'key.right', self.max_velocity['ang_vel_yaw']),
+        )
+
+    def get_goal(self, sim_data: SimData) -> GoalData:
+        command = self.command_from_keys(self.keyboard.snapshot())
+        if command != self.current_command:
+            self.current_command = command
+            logger.info(
+                "Keyboard command: vx=%.2f, vy=%.2f, yaw=%.2f",
+                *command,
+            )
+        return GoalData(
+            goal_type='velocity',
+            velocity_goal=VelocityGoal(
+                lin_vel_x=command[0],
+                lin_vel_y=command[1],
+                ang_vel_yaw=command[2],
+            ),
+        )