hangX/mindbot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

记录动作以及回放record_demo,replay_demo

Browse Source
This commit is contained in:
hangX
2026-03-11 16:33:29 +08:00
parent 0c557938a7
commit 08c4cdacb8
21 changed files with 3613 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files

593
scripts/tools/record_demos.py Normal file
View File

@@ -0,0 +1,593 @@
# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause
"""
Script to record demonstrations with Isaac Lab environments using human teleoperation.
This script allows users to record demonstrations operated by human teleoperation for a specified task.
The recorded demonstrations are stored as episodes in a hdf5 file. Users can specify the task, teleoperation
device, dataset directory, and environment stepping rate through command-line arguments.
required arguments:
--task Name of the task.
optional arguments:
-h, --help Show this help message and exit
--teleop_device Device for interacting with environment. (default: keyboard)
--dataset_file File path to export recorded demos. (default: "./datasets/dataset.hdf5")
--step_hz Environment stepping rate in Hz. (default: 30)
--num_demos Number of demonstrations to record. (default: 0)
--num_success_steps Number of continuous steps with task success for concluding a demo as successful.
(default: 10)
"""
"""Launch Isaac Sim Simulator first."""
# Standard library imports
import argparse
import contextlib
# Isaac Lab AppLauncher
from isaaclab.app import AppLauncher
# add argparse arguments
parser = argparse.ArgumentParser(description="Record demonstrations for Isaac Lab environments.")
parser.add_argument("--task", type=str, required=True, help="Name of the task.")
parser.add_argument(
"--teleop_device",
type=str,
default="keyboard",
help=(
"Teleop device. Set here (legacy) or via the environment config. If using the environment config, pass the"
" device key/name defined under 'teleop_devices' (it can be a custom name, not necessarily 'handtracking')."
" Built-ins: keyboard, spacemouse, gamepad. Not all tasks support all built-ins."
),
)
parser.add_argument(
"--dataset_file", type=str, default="./datasets/dataset.hdf5", help="File path to export recorded demos."
)
parser.add_argument("--step_hz", type=int, default=30, help="Environment stepping rate in Hz.")
parser.add_argument(
"--num_demos", type=int, default=0, help="Number of demonstrations to record. Set to 0 for infinite."
)
parser.add_argument(
"--num_success_steps",
type=int,
default=10,
help="Number of continuous steps with task success for concluding a demo as successful. Default is 10.",
)
parser.add_argument(
"--enable_pinocchio",
action="store_true",
default=False,
help="Enable Pinocchio.",
)
# append AppLauncher cli args
AppLauncher.add_app_launcher_args(parser)
# parse the arguments
args_cli = parser.parse_args()
# Validate required arguments
if args_cli.task is None:
parser.error("--task is required")
app_launcher_args = vars(args_cli)
if args_cli.enable_pinocchio:
# Import pinocchio before AppLauncher to force the use of the version
# installed by IsaacLab and not the one installed by Isaac Sim.
# pinocchio is required by the Pink IK controllers and the GR1T2 retargeter
import pinocchio # noqa: F401
if "handtracking" in args_cli.teleop_device.lower():
app_launcher_args["xr"] = True
# launch the simulator
app_launcher = AppLauncher(args_cli)
simulation_app = app_launcher.app
"""Rest everything follows."""
# Third-party imports
import logging
import os
import sys
import time
import gymnasium as gym
import torch
import omni.ui as ui
from isaaclab.devices import Se3Keyboard, Se3KeyboardCfg, Se3SpaceMouse, Se3SpaceMouseCfg
from isaaclab.devices.openxr import remove_camera_configs
from isaaclab.devices.teleop_device_factory import create_teleop_device
import isaaclab_mimic.envs # noqa: F401
from isaaclab_mimic.ui.instruction_display import InstructionDisplay, show_subtask_instructions
if args_cli.enable_pinocchio:
import isaaclab_tasks.manager_based.locomanipulation.pick_place # noqa: F401
import isaaclab_tasks.manager_based.manipulation.pick_place # noqa: F401
from collections.abc import Callable
from isaaclab.envs import DirectRLEnvCfg, ManagerBasedRLEnvCfg
from isaaclab.envs.mdp.recorders.recorders_cfg import ActionStateRecorderManagerCfg
from isaaclab.envs.ui import EmptyWindow
from isaaclab.managers import DatasetExportMode
# Add workspace root to sys.path so mindrobot_keyboard is importable as a package path.
# record_demos.py lives at scripts/tools/, workspace root is three levels up.
_ws_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
if _ws_root not in sys.path:
sys.path.insert(0, _ws_root)
from scripts.environments.teleoperation.mindrobot_keyboard import MindRobotCombinedKeyboard # noqa: E402
import isaaclab_tasks # noqa: F401
import mindbot.tasks # noqa: F401 — registers Isaac-MindRobot-* environments
from isaaclab_tasks.utils.parse_cfg import parse_env_cfg
# import logger
logger = logging.getLogger(__name__)
class RateLimiter:
"""Convenience class for enforcing rates in loops."""
def __init__(self, hz: int):
"""Initialize a RateLimiter with specified frequency.
Args:
hz: Frequency to enforce in Hertz.
"""
self.hz = hz
self.last_time = time.time()
self.sleep_duration = 1.0 / hz
self.render_period = min(0.033, self.sleep_duration)
def sleep(self, env: gym.Env):
"""Attempt to sleep at the specified rate in hz.
Args:
env: Environment to render during sleep periods.
"""
next_wakeup_time = self.last_time + self.sleep_duration
while time.time() < next_wakeup_time:
time.sleep(self.render_period)
env.sim.render()
self.last_time = self.last_time + self.sleep_duration
# detect time jumping forwards (e.g. loop is too slow)
if self.last_time < time.time():
while self.last_time < time.time():
self.last_time += self.sleep_duration
def setup_output_directories() -> tuple[str, str]:
"""Set up output directories for saving demonstrations.
Creates the output directory if it doesn't exist and extracts the file name
from the dataset file path.
Returns:
tuple[str, str]: A tuple containing:
- output_dir: The directory path where the dataset will be saved
- output_file_name: The filename (without extension) for the dataset
"""
# get directory path and file name (without extension) from cli arguments
output_dir = os.path.dirname(args_cli.dataset_file)
output_file_name = os.path.splitext(os.path.basename(args_cli.dataset_file))[0]
# create directory if it does not exist
if not os.path.exists(output_dir):
os.makedirs(output_dir)
print(f"Created output directory: {output_dir}")
return output_dir, output_file_name
def create_environment_config(
output_dir: str, output_file_name: str
) -> tuple[ManagerBasedRLEnvCfg | DirectRLEnvCfg, object | None]:
"""Create and configure the environment configuration.
Parses the environment configuration and makes necessary adjustments for demo recording.
Extracts the success termination function and configures the recorder manager.
Args:
output_dir: Directory where recorded demonstrations will be saved
output_file_name: Name of the file to store the demonstrations
Returns:
tuple[isaaclab_tasks.utils.parse_cfg.EnvCfg, Optional[object]]: A tuple containing:
- env_cfg: The configured environment configuration
- success_term: The success termination object or None if not available
Raises:
Exception: If parsing the environment configuration fails
"""
# parse configuration
try:
env_cfg = parse_env_cfg(args_cli.task, device=args_cli.device, num_envs=1)
env_cfg.env_name = args_cli.task.split(":")[-1]
except Exception as e:
logger.error(f"Failed to parse environment configuration: {e}")
exit(1)
# extract success checking function to invoke in the main loop
success_term = None
if hasattr(env_cfg.terminations, "success"):
success_term = env_cfg.terminations.success
env_cfg.terminations.success = None
else:
logger.warning(
"No success termination term was found in the environment."
" Will not be able to mark recorded demos as successful."
)
if args_cli.xr:
# If cameras are not enabled and XR is enabled, remove camera configs
if not args_cli.enable_cameras:
env_cfg = remove_camera_configs(env_cfg)
env_cfg.sim.render.antialiasing_mode = "DLSS"
# modify configuration such that the environment runs indefinitely until
# the goal is reached or other termination conditions are met
env_cfg.terminations.time_out = None
env_cfg.observations.policy.concatenate_terms = False
env_cfg.recorders: ActionStateRecorderManagerCfg = ActionStateRecorderManagerCfg()
env_cfg.recorders.dataset_export_dir_path = output_dir
env_cfg.recorders.dataset_filename = output_file_name
env_cfg.recorders.dataset_export_mode = DatasetExportMode.EXPORT_SUCCEEDED_ONLY
return env_cfg, success_term
def create_environment(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg) -> gym.Env:
"""Create the environment from the configuration.
Args:
env_cfg: The environment configuration object that defines the environment properties.
This should be an instance of EnvCfg created by parse_env_cfg().
Returns:
gym.Env: A Gymnasium environment instance for the specified task.
Raises:
Exception: If environment creation fails for any reason.
"""
try:
env = gym.make(args_cli.task, cfg=env_cfg).unwrapped
return env
except Exception as e:
logger.error(f"Failed to create environment: {e}")
exit(1)
def setup_teleop_device(callbacks: dict[str, Callable]) -> object:
"""Set up the teleoperation device based on configuration.
Attempts to create a teleoperation device based on the environment configuration.
Falls back to default devices if the specified device is not found in the configuration.
Args:
callbacks: Dictionary mapping callback keys to functions that will be
attached to the teleop device
Returns:
object: The configured teleoperation device interface
Raises:
Exception: If teleop device creation fails
"""
teleop_interface = None
try:
if hasattr(env_cfg, "teleop_devices") and args_cli.teleop_device in env_cfg.teleop_devices.devices:
teleop_interface = create_teleop_device(args_cli.teleop_device, env_cfg.teleop_devices.devices, callbacks)
else:
logger.warning(
f"No teleop device '{args_cli.teleop_device}' found in environment config. Creating default."
)
# Create fallback teleop device
if args_cli.teleop_device.lower() == "keyboard":
teleop_interface = Se3Keyboard(Se3KeyboardCfg(pos_sensitivity=0.2, rot_sensitivity=0.5))
elif args_cli.teleop_device.lower() == "spacemouse":
teleop_interface = Se3SpaceMouse(Se3SpaceMouseCfg(pos_sensitivity=0.2, rot_sensitivity=0.5))
elif args_cli.teleop_device.lower() == "mindrobot_keyboard":
teleop_interface = MindRobotCombinedKeyboard(
pos_sensitivity=0.05,
rot_sensitivity=0.05,
wheel_speed=5.0,
sim_device=args_cli.device,
)
# for key, callback in callbacks.items():
# teleop_interface.add_callback(key, callback)
else:
logger.error(f"Unsupported teleop device: {args_cli.teleop_device}")
logger.error("Supported devices: keyboard, spacemouse, handtracking")
exit(1)
# Add callbacks to fallback device
for key, callback in callbacks.items():
teleop_interface.add_callback(key, callback)
except Exception as e:
logger.error(f"Failed to create teleop device: {e}")
exit(1)
if teleop_interface is None:
logger.error("Failed to create teleop interface")
exit(1)
return teleop_interface
def setup_ui(label_text: str, env: gym.Env) -> InstructionDisplay:
"""Set up the user interface elements.
Creates instruction display and UI window with labels for showing information
to the user during demonstration recording.
Args:
label_text: Text to display showing current recording status
env: The environment instance for which UI is being created
Returns:
InstructionDisplay: The configured instruction display object
"""
instruction_display = InstructionDisplay(args_cli.xr)
if not args_cli.xr:
window = EmptyWindow(env, "Instruction")
with window.ui_window_elements["main_vstack"]:
demo_label = ui.Label(label_text)
subtask_label = ui.Label("")
instruction_display.set_labels(subtask_label, demo_label)
return instruction_display
def process_success_condition(env: gym.Env, success_term: object | None, success_step_count: int) -> tuple[int, bool]:
"""Process the success condition for the current step.
Checks if the environment has met the success condition for the required
number of consecutive steps. Marks the episode as successful if criteria are met.
Args:
env: The environment instance to check
success_term: The success termination object or None if not available
success_step_count: Current count of consecutive successful steps
Returns:
tuple[int, bool]: A tuple containing:
- updated success_step_count: The updated count of consecutive successful steps
- success_reset_needed: Boolean indicating if reset is needed due to success
"""
if success_term is None:
return success_step_count, False
if bool(success_term.func(env, **success_term.params)[0]):
success_step_count += 1
if success_step_count >= args_cli.num_success_steps:
env.recorder_manager.record_pre_reset([0], force_export_or_skip=False)
env.recorder_manager.set_success_to_episodes(
[0], torch.tensor([[True]], dtype=torch.bool, device=env.device)
)
env.recorder_manager.export_episodes([0])
print("Success condition met! Recording completed.")
return success_step_count, True
else:
success_step_count = 0
return success_step_count, False
def handle_reset(
env: gym.Env, success_step_count: int, instruction_display: InstructionDisplay, label_text: str
) -> int:
"""Handle resetting the environment.
Resets the environment, recorder manager, and related state variables.
Updates the instruction display with current status.
Args:
env: The environment instance to reset
success_step_count: Current count of consecutive successful steps
instruction_display: The display object to update
label_text: Text to display showing current recording status
Returns:
int: Reset success step count (0)
"""
print("Resetting environment...")
env.sim.reset()
env.recorder_manager.reset()
env.reset()
success_step_count = 0
instruction_display.show_demo(label_text)
return success_step_count
def run_simulation_loop(
env: gym.Env,
teleop_interface: object | None,
success_term: object | None,
rate_limiter: RateLimiter | None,
) -> int:
"""Run the main simulation loop for collecting demonstrations.
Sets up callback functions for teleop device, initializes the UI,
and runs the main loop that processes user inputs and environment steps.
Records demonstrations when success conditions are met.
Args:
env: The environment instance
teleop_interface: Optional teleop interface (will be created if None)
success_term: The success termination object or None if not available
rate_limiter: Optional rate limiter to control simulation speed
Returns:
int: Number of successful demonstrations recorded
"""
current_recorded_demo_count = 0
success_step_count = 0
should_reset_recording_instance = False
running_recording_instance = not args_cli.xr
# Callback closures for the teleop device
def reset_recording_instance():
nonlocal should_reset_recording_instance
should_reset_recording_instance = True
print("Recording instance reset requested")
def start_recording_instance():
nonlocal running_recording_instance
running_recording_instance = True
print("Recording started")
def stop_recording_instance():
nonlocal running_recording_instance
running_recording_instance = False
print("Recording paused")
def save_success_and_reset():
"""L 键：手动标记当前 episode 为成功并导出，然后重置。"""
nonlocal should_reset_recording_instance
print("[record] Manually marking episode as success and saving...")
env.recorder_manager.record_pre_reset([0], force_export_or_skip=False)
env.recorder_manager.set_success_to_episodes(
[0], torch.tensor([[True]], dtype=torch.bool, device=env.device)
)
env.recorder_manager.export_episodes([0])
should_reset_recording_instance = True
print("[record] Episode saved successfully!")
# Set up teleoperation callbacks
teleoperation_callbacks = {
"R": reset_recording_instance,
"L": save_success_and_reset,
"START": start_recording_instance,
"STOP": stop_recording_instance,
"RESET": reset_recording_instance,
}
teleop_interface = setup_teleop_device(teleoperation_callbacks)
# Reset before starting
env.sim.reset()
env.reset()
teleop_interface.reset()
label_text = f"Recorded {current_recorded_demo_count} successful demonstrations."
instruction_display = setup_ui(label_text, env)
subtasks = {}
with contextlib.suppress(KeyboardInterrupt) and torch.inference_mode():
while simulation_app.is_running():
# Get keyboard command
action = teleop_interface.advance()
# Expand to batch dimension
actions = action.repeat(env.num_envs, 1)
# Perform action on environment
if running_recording_instance:
# Compute actions based on environment
obv = env.step(actions)
if subtasks is not None:
if subtasks == {}:
subtasks = obv[0].get("subtask_terms")
elif subtasks:
show_subtask_instructions(instruction_display, subtasks, obv, env.cfg)
else:
env.sim.render()
# Check for success condition
success_step_count, success_reset_needed = process_success_condition(env, success_term, success_step_count)
if success_reset_needed:
should_reset_recording_instance = True
# Update demo count if it has changed
if env.recorder_manager.exported_successful_episode_count > current_recorded_demo_count:
current_recorded_demo_count = env.recorder_manager.exported_successful_episode_count
label_text = f"Recorded {current_recorded_demo_count} successful demonstrations."
print(label_text)
# Check if we've reached the desired number of demos
if args_cli.num_demos > 0 and env.recorder_manager.exported_successful_episode_count >= args_cli.num_demos:
label_text = f"All {current_recorded_demo_count} demonstrations recorded.\nExiting the app."
instruction_display.show_demo(label_text)
print(label_text)
target_time = time.time() + 0.8
while time.time() < target_time:
if rate_limiter:
rate_limiter.sleep(env)
else:
env.sim.render()
break
# Handle reset if requested
if should_reset_recording_instance:
success_step_count = handle_reset(env, success_step_count, instruction_display, label_text)
should_reset_recording_instance = False
# Check if simulation is stopped
if env.sim.is_stopped():
break
# Rate limiting
if rate_limiter:
rate_limiter.sleep(env)
return current_recorded_demo_count
def main() -> None:
"""Collect demonstrations from the environment using teleop interfaces.
Main function that orchestrates the entire process:
1. Sets up rate limiting based on configuration
2. Creates output directories for saving demonstrations
3. Configures the environment
4. Runs the simulation loop to collect demonstrations
5. Cleans up resources when done
Raises:
Exception: Propagates exceptions from any of the called functions
"""
# if handtracking is selected, rate limiting is achieved via OpenXR
if args_cli.xr:
rate_limiter = None
from isaaclab.ui.xr_widgets import TeleopVisualizationManager, XRVisualization
# Assign the teleop visualization manager to the visualization system
XRVisualization.assign_manager(TeleopVisualizationManager)
else:
rate_limiter = RateLimiter(args_cli.step_hz)
# Set up output directories
output_dir, output_file_name = setup_output_directories()
# Create and configure environment
global env_cfg # Make env_cfg available to setup_teleop_device
env_cfg, success_term = create_environment_config(output_dir, output_file_name)
# Create environment
env = create_environment(env_cfg)
# Run simulation loop
current_recorded_demo_count = run_simulation_loop(env, None, success_term, rate_limiter)
# Clean up
env.close()
print(f"Recording session completed with {current_recorded_demo_count} successful demonstrations")
print(f"Demonstrations saved to: {args_cli.dataset_file}")
if __name__ == "__main__":
# run the main function
main()
# close sim app
simulation_app.close()