diff --git a/resources/models/go1/policy.pt b/resources/models/go1/policy.pt new file mode 100644 index 0000000..e99abf1 Binary files /dev/null and b/resources/models/go1/policy.pt differ diff --git a/resources/models/go1/policy_10k.pt b/resources/models/go1/policy_10k.pt new file mode 100644 index 0000000..a45529b Binary files /dev/null and b/resources/models/go1/policy_10k.pt differ diff --git a/resources/robots/go1/assets/calf.stl b/resources/robots/go1/assets/calf.stl new file mode 100644 index 0000000..3e8ac94 Binary files /dev/null and b/resources/robots/go1/assets/calf.stl differ diff --git a/resources/robots/go1/assets/hip.stl b/resources/robots/go1/assets/hip.stl new file mode 100644 index 0000000..ae590bf Binary files /dev/null and b/resources/robots/go1/assets/hip.stl differ diff --git a/resources/robots/go1/assets/thigh.stl b/resources/robots/go1/assets/thigh.stl new file mode 100644 index 0000000..353801b Binary files /dev/null and b/resources/robots/go1/assets/thigh.stl differ diff --git a/resources/robots/go1/assets/thigh_mirror.stl b/resources/robots/go1/assets/thigh_mirror.stl new file mode 100644 index 0000000..11ef770 Binary files /dev/null and b/resources/robots/go1/assets/thigh_mirror.stl differ diff --git a/resources/robots/go1/assets/trunk.stl b/resources/robots/go1/assets/trunk.stl new file mode 100644 index 0000000..2659eb1 Binary files /dev/null and b/resources/robots/go1/assets/trunk.stl differ diff --git a/resources/robots/go1/go1.xml b/resources/robots/go1/go1.xml new file mode 100644 index 0000000..ea60e88 --- /dev/null +++ b/resources/robots/go1/go1.xml @@ -0,0 +1,197 @@ + + + + diff --git a/robogauge/tasks/__init__.py b/robogauge/tasks/__init__.py index 2f70917..525db8b 100644 --- a/robogauge/tasks/__init__.py +++ b/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) diff --git a/robogauge/tasks/gauge/base_gauge.py b/robogauge/tasks/gauge/base_gauge.py index f938961..526439c 100644 --- a/robogauge/tasks/gauge/base_gauge.py +++ b/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) diff --git a/robogauge/tasks/gauge/base_gauge_config.py b/robogauge/tasks/gauge/base_gauge_config.py index 7f7ce8a..91edb9f 100644 --- a/robogauge/tasks/gauge/base_gauge_config.py +++ b/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: diff --git a/robogauge/tasks/gauge/goals/__init__.py b/robogauge/tasks/gauge/goals/__init__.py index df3b9df..6236c97 100644 --- a/robogauge/tasks/gauge/goals/__init__.py +++ b/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 diff --git a/robogauge/tasks/gauge/goals/keyboard_goal.py b/robogauge/tasks/gauge/goals/keyboard_goal.py new file mode 100644 index 0000000..89e27ce --- /dev/null +++ b/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], + ), + ) diff --git a/robogauge/tasks/pipeline/base_pipeline.py b/robogauge/tasks/pipeline/base_pipeline.py index 5d7c9ff..a53c767 100644 --- a/robogauge/tasks/pipeline/base_pipeline.py +++ b/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 diff --git a/robogauge/tasks/robots/__init__.py b/robogauge/tasks/robots/__init__.py index 3270665..a6c4dbf 100644 --- a/robogauge/tasks/robots/__init__.py +++ b/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 diff --git a/robogauge/tasks/robots/go1/__init__.py b/robogauge/tasks/robots/go1/__init__.py new file mode 100644 index 0000000..c232eec --- /dev/null +++ b/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 diff --git a/robogauge/tasks/robots/go1/go1.py b/robogauge/tasks/robots/go1/go1.py new file mode 100644 index 0000000..3daec05 --- /dev/null +++ b/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 diff --git a/robogauge/tasks/robots/go1/go1_config.py b/robogauge/tasks/robots/go1/go1_config.py new file mode 100644 index 0000000..5c99377 --- /dev/null +++ b/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] diff --git a/robogauge/tasks/robots/go1/go1_moe.py b/robogauge/tasks/robots/go1/go1_moe.py new file mode 100644 index 0000000..13c81cc --- /dev/null +++ b/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 diff --git a/robogauge/tasks/robots/go1/go1_moe_config.py b/robogauge/tasks/robots/go1/go1_moe_config.py new file mode 100644 index 0000000..0cc3865 --- /dev/null +++ b/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' diff --git a/setup.py b/setup.py index 300cf18..fa1273e 100644 --- a/setup.py +++ b/setup.py @@ -23,6 +23,7 @@ setup( "fastapi", "uvicorn", "pygame", + "pynput", ], python_requires=">=3.8",