125 lines
4.5 KiB
Python
125 lines
4.5 KiB
Python
# Copyright (c) 2022-2026, The Isaac Lab Project Developers.
|
||
# SPDX-License-Identifier: BSD-3-Clause
|
||
|
||
"""Script to run teleoperation for MindRobot with chassis (wheel) control.
|
||
|
||
Extends teleop_se3_agent.py with differential-drive keyboard control:
|
||
↑ / ↓ — chassis forward / backward
|
||
← / → — chassis left turn / right turn (in-place, skid-steer)
|
||
W/S/A/D/Q/E — left arm IK: translate X/Y/Z
|
||
Z/X T/G C/V — left arm IK: rotate roll/pitch/yaw
|
||
K — left gripper toggle open/close
|
||
R — reset environment
|
||
"""
|
||
|
||
import argparse
|
||
import torch
|
||
|
||
from isaaclab.app import AppLauncher
|
||
|
||
parser = argparse.ArgumentParser(description="MindRobot teleoperation with wheel control.")
|
||
parser.add_argument("--num_envs", type=int, default=1)
|
||
parser.add_argument("--task", type=str, default=None)
|
||
parser.add_argument("--sensitivity", type=float, default=1.0)
|
||
parser.add_argument("--wheel_speed", type=float, default=5.0,
|
||
help="Wheel velocity in rad/s (default: 5.0)")
|
||
AppLauncher.add_app_launcher_args(parser)
|
||
args_cli = parser.parse_args()
|
||
|
||
app_launcher = AppLauncher(args_cli)
|
||
simulation_app = app_launcher.app
|
||
|
||
# ---------- Rest everything follows ----------
|
||
|
||
import logging
|
||
import gymnasium as gym
|
||
|
||
from isaaclab.devices import Se3Keyboard, Se3KeyboardCfg
|
||
from isaaclab.envs import ManagerBasedRLEnvCfg
|
||
|
||
import mindbot.tasks # noqa: F401
|
||
from isaaclab_tasks.utils import parse_env_cfg
|
||
|
||
# mindrobot_keyboard.py lives in the same directory as this script.
|
||
# Python adds the script's own directory to sys.path, so a direct import works.
|
||
from mindrobot_keyboard import WheelKeyboard, MindRobotCombinedKeyboard # noqa: E402
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────
|
||
# main
|
||
# ─────────────────────────────────────────────────────────────
|
||
def main() -> None:
|
||
env_cfg = parse_env_cfg(args_cli.task, device=args_cli.device, num_envs=args_cli.num_envs)
|
||
env_cfg.env_name = args_cli.task
|
||
if not isinstance(env_cfg, ManagerBasedRLEnvCfg):
|
||
raise ValueError(f"Unsupported env config type: {type(env_cfg).__name__}")
|
||
env_cfg.terminations.time_out = None
|
||
|
||
env = gym.make(args_cli.task, cfg=env_cfg).unwrapped
|
||
|
||
should_reset = False
|
||
|
||
def on_reset():
|
||
nonlocal should_reset
|
||
should_reset = True
|
||
print("[teleop] Reset triggered.")
|
||
|
||
# ── 机械臂 IK 控制器（Se3Keyboard，7 维输出）──
|
||
arm_keyboard = Se3Keyboard(
|
||
Se3KeyboardCfg(
|
||
pos_sensitivity=0.05 * args_cli.sensitivity,
|
||
rot_sensitivity=0.05 * args_cli.sensitivity,
|
||
sim_device=args_cli.device,
|
||
)
|
||
)
|
||
arm_keyboard.add_callback("R", on_reset)
|
||
|
||
# ── 底盘轮速控制器（WheelKeyboard，4 维输出）──
|
||
wheel_keyboard = WheelKeyboard(
|
||
wheel_speed=args_cli.wheel_speed,
|
||
sim_device=args_cli.device,
|
||
)
|
||
|
||
print(arm_keyboard)
|
||
print(wheel_keyboard)
|
||
print("\nAction layout: [arm_IK(6) | wheel_vel(4) | gripper(1)] = 11 dims")
|
||
print("Press R to reset. Press Ctrl+C to exit.\n")
|
||
|
||
env.reset()
|
||
arm_keyboard.reset()
|
||
wheel_keyboard.reset()
|
||
|
||
while simulation_app.is_running():
|
||
with torch.inference_mode():
|
||
# ── 1. 获取机械臂 delta pose + gripper (7维) ──
|
||
arm_cmd = arm_keyboard.advance() # [dx,dy,dz,rx,ry,rz, gripper]
|
||
arm_ik = arm_cmd[:6] # (6,)
|
||
gripper = arm_cmd[6:7] # (1,)
|
||
|
||
# ── 2. 获取底盘轮速 (4维) ──
|
||
wheel_vel = wheel_keyboard.advance() # [right_b, left_b, left_f, right_f]
|
||
|
||
# ── 3. 拼接完整 action (11维) ──
|
||
# 顺序与 MindRobotTeleopActionsCfg 中的声明顺序一致:
|
||
# arm_action(6) → wheel_action(4) → gripper_action(1)
|
||
action = torch.cat([arm_ik, wheel_vel, gripper]) # (11,)
|
||
actions = action.unsqueeze(0).repeat(env.num_envs, 1)
|
||
|
||
env.step(actions)
|
||
|
||
if should_reset:
|
||
env.reset()
|
||
arm_keyboard.reset()
|
||
wheel_keyboard.reset()
|
||
should_reset = False
|
||
print("[teleop] Environment reset complete.")
|
||
|
||
env.close()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|
||
simulation_app.close()
|