150 lines
4.6 KiB
Python
150 lines
4.6 KiB
Python
# -*- 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],
|
|
),
|
|
)
|