You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
483 lines
16 KiB
483 lines
16 KiB
from dataclasses import dataclass
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Literal, Optional
|
|
|
|
import yaml
|
|
|
|
import gr00t_wbc
|
|
from gr00t_wbc.control.main.config_template import ArgsConfig as ArgsConfigTemplate
|
|
from gr00t_wbc.control.policy.wbc_policy_factory import WBC_VERSIONS
|
|
from gr00t_wbc.control.utils.network_utils import resolve_interface
|
|
|
|
|
|
def override_wbc_config(
|
|
wbc_config: dict, config: "BaseConfig", missed_keys_only: bool = False
|
|
) -> dict:
|
|
"""Override WBC YAML values with dataclass values.
|
|
|
|
Args:
|
|
wbc_config: The loaded WBC YAML configuration dictionary
|
|
config: The BaseConfig dataclass instance with override values
|
|
missed_keys_only: If True, only add keys that don't exist in wbc_config.
|
|
If False, validate all keys exist and override all.
|
|
|
|
Returns:
|
|
Updated wbc_config dictionary with overridden values
|
|
|
|
Raises:
|
|
KeyError: If any required keys are missing from the WBC YAML configuration
|
|
(only when missed_keys_only=False)
|
|
"""
|
|
# Override yaml values with dataclass values
|
|
key_to_value = {
|
|
"INTERFACE": config.interface,
|
|
"ENV_TYPE": config.env_type,
|
|
"VERSION": config.wbc_version,
|
|
"SIMULATOR": config.simulator,
|
|
"SIMULATE_DT": 1 / float(config.sim_frequency),
|
|
"ENABLE_OFFSCREEN": config.enable_offscreen,
|
|
"ENABLE_ONSCREEN": config.enable_onscreen,
|
|
"model_path": config.wbc_model_path,
|
|
"enable_waist": config.enable_waist,
|
|
"with_hands": config.with_hands,
|
|
"verbose": config.verbose,
|
|
"verbose_timing": config.verbose_timing,
|
|
"upper_body_max_joint_speed": config.upper_body_joint_speed,
|
|
"keyboard_dispatcher_type": config.keyboard_dispatcher_type,
|
|
"enable_gravity_compensation": config.enable_gravity_compensation,
|
|
"gravity_compensation_joints": config.gravity_compensation_joints,
|
|
"high_elbow_pose": config.high_elbow_pose,
|
|
}
|
|
|
|
if missed_keys_only:
|
|
# Only add keys that don't exist in wbc_config
|
|
for key in key_to_value:
|
|
if key not in wbc_config:
|
|
wbc_config[key] = key_to_value[key]
|
|
else:
|
|
# Set all keys (overwrite existing)
|
|
for key in key_to_value:
|
|
wbc_config[key] = key_to_value[key]
|
|
|
|
# g1 kp, kd, sim2real gap
|
|
if config.env_type == "real":
|
|
# update waist pitch damping, index 14
|
|
wbc_config["MOTOR_KD"][14] = wbc_config["MOTOR_KD"][14] - 10
|
|
|
|
return wbc_config
|
|
|
|
|
|
@dataclass
|
|
class BaseConfig(ArgsConfigTemplate):
|
|
"""Base config inherited by all G1 control loops"""
|
|
|
|
# WBC Configuration
|
|
wbc_version: Literal[tuple(WBC_VERSIONS)] = "gear_wbc"
|
|
"""Version of the whole body controller."""
|
|
|
|
wbc_model_path: str = (
|
|
"policy/GR00T-WholeBodyControl-Balance.onnx," "policy/GR00T-WholeBodyControl-Walk.onnx"
|
|
)
|
|
"""Path to WBC model file (relative to gr00t_wbc/sim2mujoco/resources/robots/g1)"""
|
|
"""gear_wbc model path: policy/GR00T-WholeBodyControl-Balance.onnx,policy/GR00T-WholeBodyControl-Walk.onnx"""
|
|
|
|
wbc_policy_class: str = "G1DecoupledWholeBodyPolicy"
|
|
"""Whole body policy class."""
|
|
|
|
# System Configuration
|
|
interface: str = "sim"
|
|
"""Interface to use for the control loop. [sim, real, lo, enxe8ea6a9c4e09]"""
|
|
|
|
simulator: str = "mujoco"
|
|
"""Simulator to use."""
|
|
|
|
sim_sync_mode: bool = False
|
|
"""Whether to run the control loop in sync mode."""
|
|
|
|
control_frequency: int = 50
|
|
"""Frequency of the control loop."""
|
|
|
|
sim_frequency: int = 200
|
|
"""Frequency of the simulation loop."""
|
|
|
|
# Robot Configuration
|
|
enable_waist: bool = False
|
|
"""Whether to include waist joints in IK."""
|
|
|
|
with_hands: bool = True
|
|
"""Enable hand functionality. When False, robot operates without hands."""
|
|
|
|
high_elbow_pose: bool = False
|
|
"""Enable high elbow pose configuration for default joint positions."""
|
|
|
|
verbose: bool = True
|
|
"""Whether to print verbose output."""
|
|
|
|
# Additional common fields
|
|
enable_offscreen: bool = False
|
|
"""Whether to enable offscreen rendering."""
|
|
|
|
enable_onscreen: bool = True
|
|
"""Whether to enable onscreen rendering."""
|
|
|
|
upper_body_joint_speed: float = 1000
|
|
"""Upper body joint speed."""
|
|
|
|
env_name: str = "default"
|
|
"""Environment name."""
|
|
|
|
ik_indicator: bool = False
|
|
"""Whether to draw IK indicators."""
|
|
|
|
verbose_timing: bool = False
|
|
"""Enable verbose timing output every iteration."""
|
|
|
|
keyboard_dispatcher_type: str = "raw"
|
|
"""Keyboard dispatcher to use. [raw, ros]"""
|
|
|
|
# Gravity Compensation Configuration
|
|
enable_gravity_compensation: bool = False
|
|
"""Enable gravity compensation using pinocchio dynamics."""
|
|
|
|
gravity_compensation_joints: Optional[list[str]] = None
|
|
"""Joint groups to apply gravity compensation to (e.g., ['arms', 'left_arm', 'right_arm'])."""
|
|
# Teleop/Device Configuration
|
|
body_control_device: str = "dummy"
|
|
"""Device to use for body control. Options: dummy, vive, iphone, leapmotion, joycon."""
|
|
|
|
hand_control_device: Optional[str] = "dummy"
|
|
"""Device to use for hand control. Options: None, manus, joycon, iphone."""
|
|
|
|
body_streamer_ip: str = "10.112.210.229"
|
|
"""IP address for body streamer (vive only)."""
|
|
|
|
body_streamer_keyword: str = "knee"
|
|
"""Body streamer keyword (vive only)."""
|
|
|
|
enable_visualization: bool = False
|
|
"""Whether to enable visualization."""
|
|
|
|
enable_real_device: bool = True
|
|
"""Whether to enable real device."""
|
|
|
|
teleop_frequency: int = 20
|
|
"""Teleoperation frequency (Hz)."""
|
|
|
|
teleop_replay_path: Optional[str] = None
|
|
"""Path to teleop replay data."""
|
|
|
|
# Deployment/Camera Configuration
|
|
robot_ip: str = "192.168.123.164"
|
|
"""Robot IP address"""
|
|
# Data collection settings
|
|
data_collection: bool = True
|
|
"""Enable data collection"""
|
|
|
|
data_collection_frequency: int = 20
|
|
"""Data collection frequency (Hz)"""
|
|
|
|
root_output_dir: str = "outputs"
|
|
"""Root output directory"""
|
|
|
|
# Policy settings
|
|
enable_upper_body_operation: bool = True
|
|
"""Enable upper body operation"""
|
|
|
|
upper_body_operation_mode: Literal["teleop", "inference"] = "teleop"
|
|
"""Upper body operation mode"""
|
|
|
|
def __post_init__(self):
|
|
# Resolve interface (handles sim/real shortcuts, platform differences, and error handling)
|
|
self.interface, self.env_type = resolve_interface(self.interface)
|
|
|
|
def load_wbc_yaml(self) -> dict:
|
|
"""Load and merge wbc yaml with dataclass overrides"""
|
|
# Get the base path to gr00t_wbc and convert to Path object
|
|
package_path = Path(os.path.dirname(gr00t_wbc.__file__))
|
|
|
|
if self.wbc_version == "gear_wbc":
|
|
config_path = str(package_path / "control/main/teleop/configs/g1_29dof_gear_wbc.yaml")
|
|
else:
|
|
raise ValueError(
|
|
f"Invalid wbc_version: {self.wbc_version}, please use one of: " f"gear_wbc"
|
|
)
|
|
|
|
with open(config_path) as file:
|
|
wbc_config = yaml.load(file, Loader=yaml.FullLoader)
|
|
|
|
# Override yaml values with dataclass values
|
|
wbc_config = override_wbc_config(wbc_config, self)
|
|
|
|
return wbc_config
|
|
|
|
|
|
@dataclass
|
|
class ControlLoopConfig(BaseConfig):
|
|
"""Config for running the G1 control loop."""
|
|
|
|
pass
|
|
|
|
|
|
@dataclass
|
|
class TeleopConfig(BaseConfig):
|
|
"""Config for running the G1 teleop policy loop."""
|
|
|
|
robot: Literal["g1"] = "g1"
|
|
"""Name of the robot to use, e.g., 'g1'."""
|
|
|
|
lerobot_replay_path: Optional[str] = None
|
|
"""Path to lerobot replay data."""
|
|
|
|
# Override defaults for teleop-specific values
|
|
body_streamer_ip: str = "10.110.67.24"
|
|
"""IP address for body streamer (vive only)."""
|
|
|
|
body_streamer_keyword: str = "foot"
|
|
"""Keyword for body streamer (vive only)."""
|
|
|
|
teleop_frequency: float = 20 # Override to be float instead of int
|
|
"""Frequency of the teleop loop."""
|
|
|
|
binary_hand_ik: bool = True
|
|
"""Whether to use binary IK."""
|
|
|
|
|
|
@dataclass
|
|
class ComposedCameraClientConfig:
|
|
"""Config for running the composed camera client."""
|
|
|
|
camera_port: int = 5555
|
|
"""Port number"""
|
|
|
|
camera_host: str = "localhost"
|
|
"""Host IP address"""
|
|
|
|
fps: float = 20.0
|
|
"""FPS of the camera viewer"""
|
|
|
|
|
|
@dataclass
|
|
class DataExporterConfig(BaseConfig, ComposedCameraClientConfig):
|
|
"""Config for running the G1 data exporter."""
|
|
|
|
dataset_name: Optional[str] = None
|
|
"""Name of the dataset to save the data to. If the dataset already exists,
|
|
the new episodes will be appended to existing dataset. If the dataset does not exist,
|
|
episodes will be saved under root_output_dir/dataset_name.
|
|
"""
|
|
|
|
task_prompt: str = "demo"
|
|
"""Language Task prompt for the dataset."""
|
|
|
|
state_dim: int = 43
|
|
"""Size of the state."""
|
|
|
|
action_dim: int = 43
|
|
"""Size of the action."""
|
|
|
|
teleoperator_username: Optional[str] = None
|
|
"""Teleoperator username."""
|
|
|
|
support_operator_username: Optional[str] = None
|
|
"""Support operator username."""
|
|
|
|
robot_id: Optional[str] = None
|
|
"""Robot ID."""
|
|
|
|
lower_body_policy: Optional[str] = None
|
|
"""Lower body policy."""
|
|
|
|
img_stream_viewer: bool = False
|
|
"""Whether to open a matplot lib window to view the camera images."""
|
|
|
|
text_to_speech: bool = True
|
|
"""Whether to use text-to-speech for voice feedback."""
|
|
|
|
add_stereo_camera: bool = True
|
|
"""Whether to add stereo camera for data collection. If False, only use a signle ego view camera."""
|
|
|
|
|
|
@dataclass
|
|
class SyncSimDataCollectionConfig(ControlLoopConfig, TeleopConfig):
|
|
"""Args Config for running the data collection loop."""
|
|
|
|
robot: str = "G1"
|
|
"""Name of the robot to collect data for (e.g., G1 variants)."""
|
|
|
|
task_name: str = "GroundOnly"
|
|
"""Name of the task to collect data for. [PnPBottle, GroundOnly, ...]"""
|
|
|
|
body_control_device: str = "dummy"
|
|
"""Device to use for body control. Options: dummy, vive, iphone, leapmotion, joycon."""
|
|
|
|
hand_control_device: Optional[str] = "dummy"
|
|
"""Device to use for hand control. Options: None, manus, joycon, iphone."""
|
|
|
|
remove_existing_dir: bool = False
|
|
"""Whether to remove existing output directory if it exists."""
|
|
|
|
hardcode_teleop_cmd: bool = False
|
|
"""Whether to hardcode the teleop command for testing purposes."""
|
|
|
|
ik_indicator: bool = False
|
|
"""Whether to draw IK indicators."""
|
|
|
|
enable_onscreen: bool = True
|
|
"""Whether to show the onscreen rendering."""
|
|
|
|
save_img_obs: bool = False
|
|
"""Whether to save image observations."""
|
|
|
|
success_hold_steps: int = 50
|
|
"""Number of steps to collect after task completion before saving."""
|
|
|
|
renderer: Literal["mjviewer", "mujoco", "rerun"] = "mjviewer"
|
|
"""Renderer to use for the environment. """
|
|
|
|
replay_data_path: str | None = None
|
|
"""Path to the data (.pkl) to replay. If None, will not replay. Used for CI/CD."""
|
|
|
|
replay_speed: float = 2.5
|
|
"""Speed multiplier for replay data. Higher values make replay slower (e.g., 2.5 for sync sim tests)."""
|
|
|
|
ci_test: bool = False
|
|
"""Whether to run the CI test."""
|
|
|
|
ci_test_mode: Literal["unit", "pre_merge"] = "pre_merge"
|
|
"""'unit' for fast 50-step tests, 'pre_merge' for 500-step test with tracking checks."""
|
|
|
|
manual_control: bool = False
|
|
"""Enable manual control of data collection start/save. When True, use toggle_data_collection
|
|
to manually control episode states (idle -> recording -> need_to_save -> idle).
|
|
When False (default), automatically starts and stops data collection based on task completion."""
|
|
|
|
|
|
@dataclass
|
|
class SyncSimPlaybackConfig(SyncSimDataCollectionConfig):
|
|
"""Configuration class for playback script arguments."""
|
|
|
|
enable_real_device: bool = False
|
|
"""Whether to enable real device"""
|
|
|
|
dataset: str | None = None
|
|
"""Path to the demonstration dataset, either an HDF5 file or a LeRobot folder path."""
|
|
|
|
use_actions: bool = False
|
|
"""Whether to use actions for playback"""
|
|
|
|
use_wbc_goals: bool = False
|
|
"""Whether to use WBC goals for control"""
|
|
|
|
use_teleop_cmd: bool = False
|
|
"""Whether to use teleop IK for action generation"""
|
|
|
|
# Video recording arguments.
|
|
# Warning: enabling this key will leads to divergence between playback and recording.
|
|
save_video: bool = False
|
|
"""Whether to save video of the playback"""
|
|
|
|
# Saving to LeRobot dataset.
|
|
# Warning: enabling this key will leads to divergence between playback and recording.
|
|
save_lerobot: bool = False
|
|
"""Whether to save the playback as a new LeRobot dataset"""
|
|
|
|
video_path: str | None = None
|
|
"""Path to save the output video. If not specified,
|
|
will use the nearest folder to dataset and save as playback_video.mp4"""
|
|
|
|
num_episodes: int = 1
|
|
"""Number of episodes to load and playback/record (loads only the first N episodes from dataset)"""
|
|
|
|
intervention: bool = False
|
|
"""Whether to denote intervention timesteps with colored borders in video frames"""
|
|
|
|
ci_test: bool = False
|
|
"""Whether this is a CI test run, which limits the number of steps for testing purposes"""
|
|
|
|
def validate_args(self):
|
|
# Validate argument combinations
|
|
if self.use_teleop_cmd and not self.use_actions:
|
|
raise ValueError("--use-teleop-cmd requires --use-actions to be set")
|
|
|
|
# Note: using teleop cmd has playback divergence unlike using wbc goals, as TeleopPolicy has a warmup loop
|
|
if self.use_teleop_cmd and self.use_wbc_goals:
|
|
raise ValueError("--use-teleop-cmd and --use-wbc-goals are mutually exclusive")
|
|
|
|
if (self.use_teleop_cmd or self.use_wbc_goals) and not self.use_actions:
|
|
raise ValueError(
|
|
"You are using --use-teleop-cmd or --use-wbc-goals but not --use-actions. "
|
|
"This will not play back actions whether via teleop or wbc goals. "
|
|
"Instead, it'll play back states only."
|
|
)
|
|
|
|
if self.save_img_obs and not self.save_lerobot:
|
|
raise ValueError("--save-img-obs is only supported with --save-lerobot")
|
|
|
|
if self.intervention and not self.save_video:
|
|
raise ValueError("--intervention requires --save-video to be enabled for visualization")
|
|
|
|
|
|
@dataclass
|
|
class WebcamRecorderConfig(BaseConfig):
|
|
"""Config for running the webcam recorder."""
|
|
|
|
output_dir: str = "logs_experiment"
|
|
"""Output directory for webcam recordings"""
|
|
|
|
device_id: int = 0
|
|
"""Camera device ID"""
|
|
|
|
fps: int = 30
|
|
"""Recording frame rate"""
|
|
|
|
duration: Optional[int] = None
|
|
"""Recording duration in seconds (None for continuous)"""
|
|
|
|
|
|
@dataclass
|
|
class SimLoopConfig(BaseConfig):
|
|
"""Config for running the simulation loop."""
|
|
|
|
mp_start_method: str = "spawn"
|
|
"""Multiprocessing start method"""
|
|
|
|
enable_image_publish: bool = False
|
|
"""Enable image publishing in simulation"""
|
|
|
|
camera_port: int = 5555
|
|
"""Camera port for image publishing"""
|
|
|
|
verbose: bool = False
|
|
"""Verbose output, override the base config verbose"""
|
|
|
|
|
|
@dataclass
|
|
class DeploymentConfig(BaseConfig, ComposedCameraClientConfig):
|
|
"""G1 Robot Deployment Configuration
|
|
|
|
Simplified deployment config that inherits all common fields from G1BaseConfig.
|
|
All deployment settings are now available in the base config.
|
|
"""
|
|
|
|
camera_publish_rate: float = 30.0
|
|
"""Camera publish rate (Hz)"""
|
|
|
|
view_camera: bool = True
|
|
"""Enable camera viewer"""
|
|
# Webcam recording settings
|
|
enable_webcam_recording: bool = True
|
|
"""Enable webcam recording for real robot deployment monitoring"""
|
|
|
|
webcam_output_dir: str = "logs_experiment"
|
|
"""Output directory for webcam recordings"""
|
|
|
|
skip_img_transform: bool = False
|
|
"""Skip image transformation in the model (for faster internet)"""
|
|
|
|
sim_in_single_process: bool = False
|
|
"""Run simulator in a separate process. When True, sets simulator to None in main control loop
|
|
and launches run_sim_loop.py separately."""
|
|
|
|
image_publish: bool = False
|
|
"""Enable image publishing in simulation loop (passed to run_sim_loop.py)"""
|