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

from dataclasses import dataclass
import os
from pathlib import Path
from typing import Literal, Optional
import yaml
import decoupled_wbc
from decoupled_wbc.control.main.config_template import ArgsConfig as ArgsConfigTemplate
from decoupled_wbc.control.policy.wbc_policy_factory import WBC_VERSIONS
from decoupled_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 decoupled_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 decoupled_wbc and convert to Path object
package_path = Path(os.path.dirname(decoupled_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)"""