Add Go1 evaluation and keyboard control

robogauge/tasks/__init__.py

@@ -8,6 +8,8 @@ from robogauge.tasks.robots import (
     Go2TerrainConfig,
     Go2LabTerrainConfig,
     Go2MoETerrainConfig,
+    Go1MoEConfig,
+    Go1MoETerrainConfig,
 )
 from robogauge.tasks.pipeline import BasePipeline
 from robogauge.tasks.gauge import BaseGaugeConfig
@@ -35,6 +37,15 @@ task_register.register('go2_moe.stairs_fd', BasePipeline, MujocoConfig, Go2Stair
 task_register.register('go2_moe.stairs_bd', BasePipeline, MujocoConfig, Go2StairsBackwardGaugeConfig, Go2MoETerrainConfig)
 task_register.register('go2_moe.obstacle', BasePipeline, MujocoConfig, Go2ObstacleGaugeConfig, Go2MoETerrainConfig)
 
+# Go1 MoE (the policy is evaluated in the Go1 MuJoCo dynamics model)
+task_register.register('go1_moe.flat', BasePipeline, MujocoConfig, Go2FlatGaugeConfig, Go1MoEConfig)
+task_register.register('go1_moe.slope_fd', BasePipeline, MujocoConfig, Go2SlopeForwardGaugeConfig, Go1MoETerrainConfig)
+task_register.register('go1_moe.slope_bd', BasePipeline, MujocoConfig, Go2SlopeBackwardGaugeConfig, Go1MoETerrainConfig)
+task_register.register('go1_moe.wave', BasePipeline, MujocoConfig, Go2WaveGaugeConfig, Go1MoETerrainConfig)
+task_register.register('go1_moe.stairs_fd', BasePipeline, MujocoConfig, Go2StairsForwardGaugeConfig, Go1MoETerrainConfig)
+task_register.register('go1_moe.stairs_bd', BasePipeline, MujocoConfig, Go2StairsBackwardGaugeConfig, Go1MoETerrainConfig)
+task_register.register('go1_moe.obstacle', BasePipeline, MujocoConfig, Go2ObstacleGaugeConfig, Go1MoETerrainConfig)
+
 # Go2 Lab
 task_register.register('go2_lab.flat', BasePipeline, MujocoConfig, Go2FlatGaugeConfig, Go2LabConfig)
 task_register.register('go2_lab.slope_fd', BasePipeline, MujocoConfig, Go2SlopeForwardGaugeConfig, Go2LabTerrainConfig)

robogauge/tasks/gauge/base_gauge.py

@@ -58,6 +58,9 @@ class BaseGauge:
             elif name == 'joystick':
                 self.goals.append(JoystickGoal(robot_cfg.commands, **kwargs))
                 log_str += f"  - Joystick Goal: {kwargs}\n"
+            elif name == 'keyboard':
+                self.goals.append(KeyboardGoal(robot_cfg.commands, **kwargs))
+                log_str += f"  - Keyboard Goal: {kwargs}\n"
             else:
                 raise NotImplementedError(f"Goal '{name}' is not implemented in BaseGauge.")
             self.info['goal'].append(name)

robogauge/tasks/gauge/base_gauge_config.py

@@ -55,6 +55,10 @@ class BaseGaugeConfig(Config):
             goal_type = 'velocity'  # 'velocity'
             dead_zone = 0.1  # joystick dead zone
 
+        class keyboard:  # goal controlled by a keyboard listener independent of MuJoCo
+            enabled = False
+            goal_type = 'velocity'
+
     class metrics:
         metric_dt = 0.1  # [s] frequency to compute metrics
         class dof_limits:

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],
+            ),
+        )

robogauge/tasks/pipeline/base_pipeline.py

@@ -17,7 +17,7 @@ from copy import deepcopy
 from robogauge.utils.logger import logger
 from robogauge.tasks.simulator import MujocoSimulator, MujocoConfig, SimData
 from robogauge.tasks.robots import (
-    BaseRobot, RobotConfig, Go2Config, Go2, Go2MoEConfig, Go2MoE
+    BaseRobot, RobotConfig, Go2Config, Go2, Go2MoEConfig, Go2MoE, Go1MoE
 )
 from robogauge.tasks.gauge import BaseGauge, BaseGaugeConfig
 from robogauge.tasks.gauge.goal_data import GoalData, VelocityGoal, PositionGoal

robogauge/tasks/robots/__init__.py

@@ -5,3 +5,4 @@ from .go2.go2_lab_config import Go2LabConfig, Go2LabTerrainConfig
 from .go2.go2 import Go2
 from .go2.go2_moe_config import Go2MoEConfig, Go2MoETerrainConfig
 from .go2.go2_moe import Go2MoE
+from .go1 import Go1Config, Go1TerrainConfig, Go1, Go1MoEConfig, Go1MoETerrainConfig, Go1MoE

robogauge/tasks/robots/go1/__init__.py

@@ -0,0 +1,4 @@
+from .go1_config import Go1Config, Go1TerrainConfig
+from .go1 import Go1
+from .go1_moe_config import Go1MoEConfig, Go1MoETerrainConfig
+from .go1_moe import Go1MoE

robogauge/tasks/robots/go1/go1.py

@@ -0,0 +1,13 @@
+# -*- coding: utf-8 -*-
+"""Go1 policy adapter.
+
+The trained policy and the Go1 MuJoCo asset both use FR, FL, RR, RL joint
+groups. Isaac Gym exposed the training asset as FL, FR, RL, RR, but Go1Robot
+permuted its observations and actions during training.
+"""
+
+from robogauge.tasks.robots.go2.go2 import Go2
+
+
+class Go1(Go2):
+    pass

robogauge/tasks/robots/go1/go1_config.py

@@ -0,0 +1,38 @@
+# -*- coding: utf-8 -*-
+"""Go1 configuration for the RoboGauge MuJoCo evaluator."""
+
+from robogauge.tasks.robots.go2.go2_config import Go2Config
+
+
+class Go1Config(Go2Config):
+    robot_name = 'go1'
+    robot_class = 'Go1'
+
+    class assets(Go2Config.assets):
+        robot_xml = "{ROBOGAUGE_ROOT_DIR}/resources/robots/go1/go1.xml"
+        foot_geom_names = ['FR', 'FL', 'RR', 'RL']
+
+    class control(Go2Config.control):
+        model_path = "{ROBOGAUGE_ROOT_DIR}/resources/models/go1/policy.pt"
+        control_dt = 0.02
+        p_gains = [28.0] * 12
+        d_gains = [0.7] * 12
+
+        # Both this MuJoCo model and the trained policy use FR, FL, RR, RL.
+        # Go1Robot permuted Isaac Gym's FL, FR, RL, RR asset state into this
+        # order during training, so applying that permutation again here
+        # would swap every left/right leg pair.
+        default_dof_pos = [
+            -0.1, 0.8, -1.5,
+             0.1, 0.8, -1.5,
+            -0.1, 1.0, -1.5,
+             0.1, 1.0, -1.5,
+        ]
+        mj2model_dof_indices = list(range(12))
+
+
+class Go1TerrainConfig(Go1Config):
+    class commands(Go1Config.commands):
+        lin_vel_x = [-1.0, 1.0]
+        lin_vel_y = [-1.0, 1.0]
+        ang_vel_yaw = [-1.5, 1.5]

robogauge/tasks/robots/go1/go1_moe.py

@@ -0,0 +1,8 @@
+# -*- coding: utf-8 -*-
+"""MoE Go1 policy adapter."""
+
+from robogauge.tasks.robots.go2.go2_moe import Go2MoE
+
+
+class Go1MoE(Go2MoE):
+    pass

robogauge/tasks/robots/go1/go1_moe_config.py

@@ -0,0 +1,12 @@
+# -*- coding: utf-8 -*-
+"""MoE policy configuration for the trained Go1 checkpoint."""
+
+from robogauge.tasks.robots.go1.go1_config import Go1Config, Go1TerrainConfig
+
+
+class Go1MoEConfig(Go1Config):
+    robot_class = 'Go1MoE'
+
+
+class Go1MoETerrainConfig(Go1TerrainConfig, Go1MoEConfig):
+    robot_class = 'Go1MoE'