diff --git a/GR00T-WholeBodyControl/decoupled_wbc/__init__.py b/GR00T-WholeBodyControl/decoupled_wbc/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3308a84eef2c0aca40b10fc979a62b47ec96beaa --- /dev/null +++ b/GR00T-WholeBodyControl/decoupled_wbc/__init__.py @@ -0,0 +1,3 @@ +from .version import VERSION, VERSION_SHORT # noqa + +__version__ = VERSION # noqa diff --git a/GR00T-WholeBodyControl/decoupled_wbc/version.py b/GR00T-WholeBodyControl/decoupled_wbc/version.py new file mode 100644 index 0000000000000000000000000000000000000000..4e48c3d68b4bba671b9cacf6c176a6f7de7fbc52 --- /dev/null +++ b/GR00T-WholeBodyControl/decoupled_wbc/version.py @@ -0,0 +1,11 @@ +_MAJOR = "0" +_MINOR = "1" +# On main and in a nightly release the patch should be one ahead of the last +# released build. +_PATCH = "0" +# This is mainly for nightly builds which have the suffix ".dev$DATE". See +# https://semver.org/#is-v123-a-semantic-version for the semantics. +_SUFFIX = "" + +VERSION_SHORT = "{0}.{1}".format(_MAJOR, _MINOR) +VERSION = "0.1.0" # or whatever version you want diff --git a/GR00T-WholeBodyControl/docs/Makefile b/GR00T-WholeBodyControl/docs/Makefile new file mode 100644 index 0000000000000000000000000000000000000000..d0c3cbf1020d5c292abdedf27627c6abe25e2293 --- /dev/null +++ b/GR00T-WholeBodyControl/docs/Makefile @@ -0,0 +1,20 @@ +# Minimal makefile for Sphinx documentation +# + +# You can set these variables from the command line, and also +# from the environment for the first two. +SPHINXOPTS ?= +SPHINXBUILD ?= sphinx-build +SOURCEDIR = source +BUILDDIR = build + +# Put it first so that "make" without argument is like "make help". +help: + @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) + +.PHONY: help Makefile + +# Catch-all target: route all unknown targets to Sphinx using the new +# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). +%: Makefile + @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) diff --git a/GR00T-WholeBodyControl/docs/README.md b/GR00T-WholeBodyControl/docs/README.md new file mode 100644 index 0000000000000000000000000000000000000000..0b8f83515c1f2b3dc5a713c06133f667184fd8a8 --- /dev/null +++ b/GR00T-WholeBodyControl/docs/README.md @@ -0,0 +1,74 @@ +# Documentation + +This directory contains the source code for the GR00T-WholeBodyControl documentation website. + +## Building Locally + +### Prerequisites + +Install the required Python packages: + +```bash +pip install sphinx sphinx-book-theme sphinx-design sphinxemoji \ + autodocsumm sphinxcontrib-bibtex myst-parser \ + sphinx-copybutton +``` + +### Build the Documentation + +```bash +cd docs +make html +``` + +The built documentation will be in `build/html/`. Open `build/html/index.html` in your browser. + +### Live Preview + +Start a local web server to preview: + +```bash +cd build/html +python -m http.server 8000 +``` + +Then open http://localhost:8000 + +### Clean Build + +To remove all built files and rebuild from scratch: + +```bash +make clean +make html +``` + +## Deployment + +The documentation is automatically built and deployed to GitHub Pages when changes are pushed to the `main` branch via the GitHub Actions workflow at `.github/workflows/docs.yml`. + +The live documentation will be available at: +**https://nvlabs.github.io/GR00T-WholeBodyControl/** + +## Documentation Structure + +- `source/` - All documentation source files + - `conf.py` - Sphinx configuration + - `index.rst` - Main landing page + - `_static/` - Static assets (CSS, images, logos) + - `tutorials/` - Tutorial pages + - `getting_started/` - Getting started guides + - `user_guide/` - User guide + - `api/` - API reference + - `resources/` - Additional resources + +## Writing Documentation + +- Use Markdown (`.md`) or reStructuredText (`.rst`) files +- Markdown is recommended for simplicity +- Place new files in the appropriate subdirectory +- Update `index.rst` to add new sections to the navigation + +## Theme + +The documentation uses the `sphinx_book_theme` with NVIDIA branding, matching the Isaac Lab documentation style. diff --git a/GR00T-WholeBodyControl/gear_sonic.egg-info/PKG-INFO b/GR00T-WholeBodyControl/gear_sonic.egg-info/PKG-INFO new file mode 100644 index 0000000000000000000000000000000000000000..0d8c36d6a31048b58a80d1290804992ff343cda9 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic.egg-info/PKG-INFO @@ -0,0 +1,72 @@ +Metadata-Version: 2.4 +Name: gear_sonic +Version: 0.1.0 +Author: NVIDIA Gear Lab +License: Apache-2.0 +Classifier: Intended Audience :: Science/Research +Classifier: Development Status :: 3 - Alpha +Classifier: License :: OSI Approved :: Apache Software License +Classifier: Programming Language :: Python :: 3 +Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence +Requires-Python: >=3.10 +Description-Content-Type: text/plain +Requires-Dist: numpy==1.26.4 +Requires-Dist: scipy==1.15.3 +Requires-Dist: torch>=2.4.0 +Requires-Dist: joblib +Requires-Dist: tqdm +Requires-Dist: easydict +Requires-Dist: loguru +Provides-Extra: teleop +Requires-Dist: pyzmq; extra == "teleop" +Requires-Dist: msgpack; extra == "teleop" +Requires-Dist: msgpack-numpy; extra == "teleop" +Requires-Dist: pin; extra == "teleop" +Requires-Dist: pyvista; platform_machine != "aarch64" and extra == "teleop" +Provides-Extra: sim +Requires-Dist: mujoco; extra == "sim" +Requires-Dist: tyro; extra == "sim" +Requires-Dist: pin; extra == "sim" +Requires-Dist: pyyaml; extra == "sim" +Requires-Dist: pyzmq; extra == "sim" +Requires-Dist: msgpack; extra == "sim" +Requires-Dist: msgpack-numpy; extra == "sim" +Requires-Dist: opencv-python; extra == "sim" +Provides-Extra: data-collection +Requires-Dist: pyzmq; extra == "data-collection" +Requires-Dist: msgpack; extra == "data-collection" +Requires-Dist: msgpack-numpy; extra == "data-collection" +Requires-Dist: pin; extra == "data-collection" +Requires-Dist: tyro; extra == "data-collection" +Requires-Dist: pyttsx3==2.90; extra == "data-collection" +Requires-Dist: av>=14.2; extra == "data-collection" +Requires-Dist: opencv-python; extra == "data-collection" +Requires-Dist: lerobot @ git+https://github.com/huggingface/lerobot.git@a445d9c9da6bea99a8972daa4fe1fdd053d711d2 ; extra == "data-collection" +Requires-Dist: datasets==3.6.0; extra == "data-collection" +Provides-Extra: camera +Requires-Dist: pyzmq; extra == "camera" +Requires-Dist: msgpack; extra == "camera" +Requires-Dist: msgpack-numpy; extra == "camera" +Requires-Dist: opencv-python; extra == "camera" +Requires-Dist: tyro; extra == "camera" +Requires-Dist: depthai; extra == "camera" +Requires-Dist: requests; extra == "camera" +Provides-Extra: inference +Requires-Dist: pyzmq; extra == "inference" +Requires-Dist: msgpack; extra == "inference" +Requires-Dist: msgpack-numpy; extra == "inference" +Requires-Dist: pin; extra == "inference" +Requires-Dist: tyro; extra == "inference" +Requires-Dist: opencv-python; extra == "inference" +Requires-Dist: scipy; extra == "inference" +Requires-Dist: Isaac-GR00T @ git+https://github.com/NVIDIA/Isaac-GR00T.git ; extra == "inference" +Provides-Extra: training +Requires-Dist: hydra-core==1.3.2; extra == "training" +Requires-Dist: wandb; extra == "training" +Requires-Dist: trl==0.28.0; extra == "training" +Requires-Dist: transformers>=4.56.2; extra == "training" +Requires-Dist: accelerate>=1.3.0; extra == "training" +Requires-Dist: tensorboard; extra == "training" +Requires-Dist: smpl_sim @ git+https://github.com/ZhengyiLuo/SMPLSim.git ; extra == "training" + +NVIDIA Gear Sonic - Whole Body Control diff --git a/GR00T-WholeBodyControl/gear_sonic.egg-info/SOURCES.txt b/GR00T-WholeBodyControl/gear_sonic.egg-info/SOURCES.txt new file mode 100644 index 0000000000000000000000000000000000000000..a1e6cc43b1b6d080e659ed7af92946a20eb8a78e --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic.egg-info/SOURCES.txt @@ -0,0 +1,276 @@ +pyproject.toml +../gear_sonic/__init__.py +../gear_sonic/eval_agent_trl.py +../gear_sonic/eval_exp.py +../gear_sonic/train_agent_trl.py +../gear_sonic/version.py +../gear_sonic.egg-info/PKG-INFO +../gear_sonic.egg-info/SOURCES.txt +../gear_sonic.egg-info/dependency_links.txt +../gear_sonic.egg-info/requires.txt +../gear_sonic.egg-info/top_level.txt +../gear_sonic/camera/__init__.py +../gear_sonic/camera/composed_camera.py +../gear_sonic/camera/sensor.py +../gear_sonic/camera/sensor_server.py +../gear_sonic/camera/drivers/__init__.py +../gear_sonic/camera/drivers/dummy.py +../gear_sonic/camera/drivers/oak.py +../gear_sonic/camera/drivers/realsense.py +../gear_sonic/camera/drivers/usb_camera.py +../gear_sonic/config/base.yaml +../gear_sonic/config/base_eval.yaml +../gear_sonic/config/eval_exp.yaml +../gear_sonic/config/actor_critic/mlp.yaml +../gear_sonic/config/actor_critic/critics/mlp.yaml +../gear_sonic/config/actor_critic/decoders/g1_dyn_mlp.yaml +../gear_sonic/config/actor_critic/decoders/g1_kin_mf_mlp.yaml +../gear_sonic/config/actor_critic/encoders/g1_mf_mlp.yaml +../gear_sonic/config/actor_critic/encoders/smpl_mlp.yaml +../gear_sonic/config/actor_critic/encoders/soma_mlp.yaml +../gear_sonic/config/actor_critic/encoders/teleop_mlp.yaml +../gear_sonic/config/actor_critic/quantizers/fsq.yaml +../gear_sonic/config/actor_critic/universal_token/all_mlp_v1.yaml +../gear_sonic/config/actor_critic/universal_token/all_mlp_v1_soma.yaml +../gear_sonic/config/algo/ppo_im_phc.yaml +../gear_sonic/config/algo/trl/ppo.yaml +../gear_sonic/config/aux_losses/terms/g1_recon.yaml +../gear_sonic/config/aux_losses/terms/g1_smpl_latent.yaml +../gear_sonic/config/aux_losses/terms/g1_soma_latent.yaml +../gear_sonic/config/aux_losses/terms/g1_teleop_latent.yaml +../gear_sonic/config/aux_losses/terms/reencoded_smpl_g1_latent.yaml +../gear_sonic/config/aux_losses/terms/teleop_smpl_latent.yaml +../gear_sonic/config/aux_losses/universal_token/g1_recon_and_all_latent.yaml +../gear_sonic/config/aux_losses/universal_token/g1_recon_and_all_latent_soma.yaml +../gear_sonic/config/base/hydra.yaml +../gear_sonic/config/base/structure.yaml +../gear_sonic/config/callbacks/im_eval.yaml +../gear_sonic/config/callbacks/im_resample.yaml +../gear_sonic/config/callbacks/model_save.yaml +../gear_sonic/config/callbacks/read_eval.yaml +../gear_sonic/config/callbacks/wandb.yaml +../gear_sonic/config/exp/manager/universal_token/all_modes/sonic_bones_seed.yaml +../gear_sonic/config/exp/manager/universal_token/all_modes/sonic_h2.yaml +../gear_sonic/config/exp/manager/universal_token/all_modes/sonic_release.yaml +../gear_sonic/config/exp/manager/universal_token/all_modes/sonic_v1_1.yaml +../gear_sonic/config/manager_env/base_env.yaml +../gear_sonic/config/manager_env/actions/terms/joint_pos.yaml +../gear_sonic/config/manager_env/actions/tracking/base.yaml +../gear_sonic/config/manager_env/commands/terms/motion.yaml +../gear_sonic/config/manager_env/commands/tracking/base.yaml +../gear_sonic/config/manager_env/curriculum/empty.yaml +../gear_sonic/config/manager_env/events/terms/add_joint_default_pos.yaml +../gear_sonic/config/manager_env/events/terms/base_com.yaml +../gear_sonic/config/manager_env/events/terms/physics_material.yaml +../gear_sonic/config/manager_env/events/terms/push_robot.yaml +../gear_sonic/config/manager_env/events/terms/randomize_rigid_body_mass.yaml +../gear_sonic/config/manager_env/events/tracking/base.yaml +../gear_sonic/config/manager_env/events/tracking/level0_4.yaml +../gear_sonic/config/manager_env/observations/critic/privileged.yaml +../gear_sonic/config/manager_env/observations/critic/privileged_mf_hist.yaml +../gear_sonic/config/manager_env/observations/policy/global.yaml +../gear_sonic/config/manager_env/observations/policy/local_dir_hist.yaml +../gear_sonic/config/manager_env/observations/terms/actions.yaml +../gear_sonic/config/manager_env/observations/terms/base_ang_vel.yaml +../gear_sonic/config/manager_env/observations/terms/base_lin_vel.yaml +../gear_sonic/config/manager_env/observations/terms/body_ori.yaml +../gear_sonic/config/manager_env/observations/terms/body_pos.yaml +../gear_sonic/config/manager_env/observations/terms/command.yaml +../gear_sonic/config/manager_env/observations/terms/command_multi_future.yaml +../gear_sonic/config/manager_env/observations/terms/command_multi_future_lower_body.yaml +../gear_sonic/config/manager_env/observations/terms/command_multi_future_nonflat.yaml +../gear_sonic/config/manager_env/observations/terms/command_z.yaml +../gear_sonic/config/manager_env/observations/terms/command_z_multi_future_nonflat.yaml +../gear_sonic/config/manager_env/observations/terms/encoder_index.yaml +../gear_sonic/config/manager_env/observations/terms/gravity_dir.yaml +../gear_sonic/config/manager_env/observations/terms/joint_pos.yaml +../gear_sonic/config/manager_env/observations/terms/joint_pos_multi_future_wrist_for_smpl.yaml +../gear_sonic/config/manager_env/observations/terms/joint_pos_multi_future_wrist_for_soma.yaml +../gear_sonic/config/manager_env/observations/terms/joint_vel.yaml +../gear_sonic/config/manager_env/observations/terms/motion_anchor_ori_b.yaml +../gear_sonic/config/manager_env/observations/terms/motion_anchor_ori_b_mf_nonflat.yaml +../gear_sonic/config/manager_env/observations/terms/motion_anchor_ori_heading.yaml +../gear_sonic/config/manager_env/observations/terms/motion_anchor_ori_heading_mf_nonflat.yaml +../gear_sonic/config/manager_env/observations/terms/motion_anchor_pos_b.yaml +../gear_sonic/config/manager_env/observations/terms/smpl_joints_multi_future_local_nonflat.yaml +../gear_sonic/config/manager_env/observations/terms/smpl_root_ori_b_multi_future.yaml +../gear_sonic/config/manager_env/observations/terms/smpl_root_ori_heading_multi_future.yaml +../gear_sonic/config/manager_env/observations/terms/soma_joints_multi_future_local_nonflat.yaml +../gear_sonic/config/manager_env/observations/terms/soma_root_ori_b_multi_future.yaml +../gear_sonic/config/manager_env/observations/terms/vr_3point_local_orn_target.yaml +../gear_sonic/config/manager_env/observations/terms/vr_3point_local_target.yaml +../gear_sonic/config/manager_env/observations/tokenizer/unitoken_all_noz.yaml +../gear_sonic/config/manager_env/observations/tokenizer/unitoken_all_noz_heading.yaml +../gear_sonic/config/manager_env/observations/tokenizer/unitoken_all_noz_soma.yaml +../gear_sonic/config/manager_env/recorders/empty.yaml +../gear_sonic/config/manager_env/recorders/render.yaml +../gear_sonic/config/manager_env/rewards/terms/action_rate_l2.yaml +../gear_sonic/config/manager_env/rewards/terms/anti_shake_ang_vel.yaml +../gear_sonic/config/manager_env/rewards/terms/energy_consumption.yaml +../gear_sonic/config/manager_env/rewards/terms/feet_acc.yaml +../gear_sonic/config/manager_env/rewards/terms/joint_limit.yaml +../gear_sonic/config/manager_env/rewards/terms/tracking_anchor_ori.yaml +../gear_sonic/config/manager_env/rewards/terms/tracking_anchor_pos.yaml +../gear_sonic/config/manager_env/rewards/terms/tracking_body_angvel.yaml +../gear_sonic/config/manager_env/rewards/terms/tracking_body_linvel.yaml +../gear_sonic/config/manager_env/rewards/terms/tracking_relative_body_ori.yaml +../gear_sonic/config/manager_env/rewards/terms/tracking_relative_body_pos.yaml +../gear_sonic/config/manager_env/rewards/terms/tracking_vr_2wrists_local_ori.yaml +../gear_sonic/config/manager_env/rewards/terms/tracking_vr_5point_local.yaml +../gear_sonic/config/manager_env/rewards/terms/undesired_contacts.yaml +../gear_sonic/config/manager_env/rewards/tracking/base.yaml +../gear_sonic/config/manager_env/rewards/tracking/base_5point_local_feet_acc.yaml +../gear_sonic/config/manager_env/rewards/tracking/local_feet_acc_energy_5pt.yaml +../gear_sonic/config/manager_env/terminations/terms/anchor_ori_full.yaml +../gear_sonic/config/manager_env/terminations/terms/anchor_pos.yaml +../gear_sonic/config/manager_env/terminations/terms/anchor_pos_adaptive.yaml +../gear_sonic/config/manager_env/terminations/terms/ee_body_pos.yaml +../gear_sonic/config/manager_env/terminations/terms/ee_body_pos_adaptive.yaml +../gear_sonic/config/manager_env/terminations/terms/foot_pos_xyz.yaml +../gear_sonic/config/manager_env/terminations/terms/motion_time_out.yaml +../gear_sonic/config/manager_env/terminations/tracking/base.yaml +../gear_sonic/config/manager_env/terminations/tracking/base_adaptive_strict_ori_foot_xyz.yaml +../gear_sonic/config/manager_env/terminations/tracking/eval.yaml +../gear_sonic/config/opt/wandb.yaml +../gear_sonic/config/trainer/trl.yaml +../gear_sonic/config/trainer/trl_ppo_aux.yaml +../gear_sonic/data/exporter.py +../gear_sonic/data/features_sonic_vla.py +../gear_sonic/data/video_writer.py +../gear_sonic/data/robot_model/__init__.py +../gear_sonic/data/robot_model/robot_model.py +../gear_sonic/data/robot_model/instantiation/__init__.py +../gear_sonic/data/robot_model/instantiation/g1.py +../gear_sonic/data/robot_model/supplemental_info/__init__.py +../gear_sonic/data/robot_model/supplemental_info/robot_supplemental_info.py +../gear_sonic/data/robot_model/supplemental_info/g1/__init__.py +../gear_sonic/data/robot_model/supplemental_info/g1/g1_supplemental_info.py +../gear_sonic/data_process/convert_soma_csv_to_motion_lib.py +../gear_sonic/data_process/extract_soma_joints_from_bvh.py +../gear_sonic/data_process/filter_and_copy_bones_data.py +../gear_sonic/data_process/split_pkl_files.py +../gear_sonic/envs/__init__.py +../gear_sonic/envs/env_utils/__init__.py +../gear_sonic/envs/env_utils/joint_utils.py +../gear_sonic/envs/manager_env/__init__.py +../gear_sonic/envs/manager_env/modular_tracking_env_cfg.py +../gear_sonic/envs/manager_env/mdp/__init__.py +../gear_sonic/envs/manager_env/mdp/actions.py +../gear_sonic/envs/manager_env/mdp/actuators.py +../gear_sonic/envs/manager_env/mdp/commands.py +../gear_sonic/envs/manager_env/mdp/curriculum.py +../gear_sonic/envs/manager_env/mdp/events.py +../gear_sonic/envs/manager_env/mdp/observations.py +../gear_sonic/envs/manager_env/mdp/recorders.py +../gear_sonic/envs/manager_env/mdp/rewards.py +../gear_sonic/envs/manager_env/mdp/terminations.py +../gear_sonic/envs/manager_env/mdp/terrain.py +../gear_sonic/envs/manager_env/mdp/utils.py +../gear_sonic/envs/manager_env/robots/__init__.py +../gear_sonic/envs/manager_env/robots/g1.py +../gear_sonic/envs/manager_env/robots/h2.py +../gear_sonic/envs/wrapper/__init__.py +../gear_sonic/envs/wrapper/manager_env_wrapper.py +../gear_sonic/examples/live_camera_teleop/soma_pt_to_sonic_v3.py +../gear_sonic/examples/live_camera_teleop/soma_to_smpl.py +../gear_sonic/examples/live_camera_teleop/webcam_stream.py +../gear_sonic/isaac_utils/__init__.py +../gear_sonic/isaac_utils/maths.py +../gear_sonic/isaac_utils/rotations.py +../gear_sonic/scripts/launch_data_collection.py +../gear_sonic/scripts/launch_inference.py +../gear_sonic/scripts/pico_manager_thread_server.py +../gear_sonic/scripts/process_dataset.py +../gear_sonic/scripts/run_camera_viewer.py +../gear_sonic/scripts/run_data_exporter.py +../gear_sonic/scripts/run_sim_loop.py +../gear_sonic/scripts/run_vla_inference.py +../gear_sonic/tests/test_input_readers.py +../gear_sonic/trl/__init__.py +../gear_sonic/trl/callbacks/__init__.py +../gear_sonic/trl/callbacks/hv_callback_handler.py +../gear_sonic/trl/callbacks/im_eval_callback.py +../gear_sonic/trl/callbacks/im_resample_callback.py +../gear_sonic/trl/callbacks/model_save_callback.py +../gear_sonic/trl/callbacks/read_eval_callback.py +../gear_sonic/trl/callbacks/wandb_callback.py +../gear_sonic/trl/losses/__init__.py +../gear_sonic/trl/losses/token_losses.py +../gear_sonic/trl/modules/__init__.py +../gear_sonic/trl/modules/actor_critic_modules.py +../gear_sonic/trl/modules/base_module.py +../gear_sonic/trl/modules/data_utils.py +../gear_sonic/trl/modules/universal_token_modules.py +../gear_sonic/trl/trainer/__init__.py +../gear_sonic/trl/trainer/ppo_trainer.py +../gear_sonic/trl/trainer/ppo_trainer_aux_loss.py +../gear_sonic/trl/utils/__init__.py +../gear_sonic/trl/utils/common.py +../gear_sonic/trl/utils/data.py +../gear_sonic/trl/utils/kornia_transform.py +../gear_sonic/trl/utils/math.py +../gear_sonic/trl/utils/order_converter.py +../gear_sonic/trl/utils/rl.py +../gear_sonic/trl/utils/rotation_conversion.py +../gear_sonic/trl/utils/scheduler.py +../gear_sonic/trl/utils/torch_transform.py +../gear_sonic/trl/utils/smplx/smplx_utils.py +../gear_sonic/trl/utils/smplx/body_model/__init__.py +../gear_sonic/trl/utils/smplx/body_model/body_model.py +../gear_sonic/trl/utils/smplx/body_model/body_model_smplh.py +../gear_sonic/trl/utils/smplx/body_model/body_model_smplx.py +../gear_sonic/trl/utils/smplx/body_model/min_lbs.py +../gear_sonic/trl/utils/smplx/body_model/rotation_conversions.py +../gear_sonic/trl/utils/smplx/body_model/smpl_lite.py +../gear_sonic/trl/utils/smplx/body_model/smpl_vert_segmentation.json +../gear_sonic/trl/utils/smplx/body_model/smplx_lite.py +../gear_sonic/trl/utils/smplx/body_model/utils.py +../gear_sonic/utils/__init__.py +../gear_sonic/utils/average_meters.py +../gear_sonic/utils/batch_normalizer.py +../gear_sonic/utils/common.py +../gear_sonic/utils/config_utils.py +../gear_sonic/utils/inference_helpers.py +../gear_sonic/utils/logging.py +../gear_sonic/utils/obs_utils.py +../gear_sonic/utils/running_mean_std.py +../gear_sonic/utils/torch_utils.py +../gear_sonic/utils/data_collection/__init__.py +../gear_sonic/utils/data_collection/episode_state.py +../gear_sonic/utils/data_collection/keyboard_subscriber.py +../gear_sonic/utils/data_collection/telemetry.py +../gear_sonic/utils/data_collection/text_to_speech.py +../gear_sonic/utils/data_collection/transforms.py +../gear_sonic/utils/data_collection/zmq_state_subscriber.py +../gear_sonic/utils/inference/__init__.py +../gear_sonic/utils/inference/initial_poses.py +../gear_sonic/utils/inference/vla_utils.py +../gear_sonic/utils/motion_lib/__init__.py +../gear_sonic/utils/motion_lib/motion_lib_base.py +../gear_sonic/utils/motion_lib/motion_lib_robot.py +../gear_sonic/utils/motion_lib/skeleton.py +../gear_sonic/utils/motion_lib/torch_humanoid_batch.py +../gear_sonic/utils/mujoco_sim/__init__.py +../gear_sonic/utils/mujoco_sim/base_sim.py +../gear_sonic/utils/mujoco_sim/configs.py +../gear_sonic/utils/mujoco_sim/image_publish_utils.py +../gear_sonic/utils/mujoco_sim/metric_utils.py +../gear_sonic/utils/mujoco_sim/robot.py +../gear_sonic/utils/mujoco_sim/sensor_server.py +../gear_sonic/utils/mujoco_sim/sim_utils.py +../gear_sonic/utils/mujoco_sim/simulator_factory.py +../gear_sonic/utils/mujoco_sim/unitree_sdk2py_bridge.py +../gear_sonic/utils/mujoco_sim/wbc_configs/g1_29dof_sonic_model12.yaml +../gear_sonic/utils/network/network_utils.py +../gear_sonic/utils/teleop/input_readers.py +../gear_sonic/utils/teleop/isaac_teleop_client.py +../gear_sonic/utils/teleop/solver/solver.py +../gear_sonic/utils/teleop/solver/hand/g1_gripper_ik_solver.py +../gear_sonic/utils/teleop/vis/vr3pt_pose_visualizer.py +../gear_sonic/utils/teleop/zmq/zmq_planner_sender.py +../gear_sonic/utils/teleop/zmq/zmq_poller.py +../gear_sonic_deploy/visualize_motion.py +../gear_sonic_deploy/reference/convert_motions.py +../gear_sonic_deploy/src/g1/g1_deploy_onnx_ref/tests/pose_estimation_server_onboard_test.py +../gear_sonic_deploy/src/g1/g1_deploy_onnx_ref/tests/test_zmq_manager.py +tests/test_input_readers.py \ No newline at end of file diff --git a/GR00T-WholeBodyControl/gear_sonic.egg-info/dependency_links.txt b/GR00T-WholeBodyControl/gear_sonic.egg-info/dependency_links.txt new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic.egg-info/dependency_links.txt @@ -0,0 +1 @@ + diff --git a/GR00T-WholeBodyControl/gear_sonic.egg-info/requires.txt b/GR00T-WholeBodyControl/gear_sonic.egg-info/requires.txt new file mode 100644 index 0000000000000000000000000000000000000000..cb366ebffaf341a2b415e7050e9a8860ca96ccd7 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic.egg-info/requires.txt @@ -0,0 +1,66 @@ +numpy==1.26.4 +scipy==1.15.3 +torch>=2.4.0 +joblib +tqdm +easydict +loguru + +[camera] +pyzmq +msgpack +msgpack-numpy +opencv-python +tyro +depthai +requests + +[data_collection] +pyzmq +msgpack +msgpack-numpy +pin +tyro +pyttsx3==2.90 +av>=14.2 +opencv-python +lerobot @ git+https://github.com/huggingface/lerobot.git@a445d9c9da6bea99a8972daa4fe1fdd053d711d2 +datasets==3.6.0 + +[inference] +pyzmq +msgpack +msgpack-numpy +pin +tyro +opencv-python +scipy +Isaac-GR00T @ git+https://github.com/NVIDIA/Isaac-GR00T.git + +[sim] +mujoco +tyro +pin +pyyaml +pyzmq +msgpack +msgpack-numpy +opencv-python + +[teleop] +pyzmq +msgpack +msgpack-numpy +pin + +[teleop:platform_machine != "aarch64"] +pyvista + +[training] +hydra-core==1.3.2 +wandb +trl==0.28.0 +transformers>=4.56.2 +accelerate>=1.3.0 +tensorboard +smpl_sim @ git+https://github.com/ZhengyiLuo/SMPLSim.git diff --git a/GR00T-WholeBodyControl/gear_sonic.egg-info/top_level.txt b/GR00T-WholeBodyControl/gear_sonic.egg-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..9216d4effe1c00f116c627085d3565c8ad3dc46d --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic.egg-info/top_level.txt @@ -0,0 +1,2 @@ +gear_sonic +gear_sonic_deploy diff --git a/GR00T-WholeBodyControl/gear_sonic/__init__.py b/GR00T-WholeBodyControl/gear_sonic/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/GR00T-WholeBodyControl/gear_sonic/camera/__init__.py b/GR00T-WholeBodyControl/gear_sonic/camera/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a7c6ea876a0cf9293836084de78a01a1a5a626dd --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/camera/__init__.py @@ -0,0 +1,29 @@ +"""Camera server package for streaming camera images over ZMQ. + +Runs on the robot computer and publishes JPEG-encoded frames that the +data exporter (on the workstation) subscribes to for recording. + +Quickstart (on robot):: + + bash install_scripts/install_camera_server.sh + source .venv_camera/bin/activate + python -m gear_sonic.camera.composed_camera --ego-view-camera oak + +See ``docs/source/tutorials/data_collection.md`` for full setup instructions. +""" + +from gear_sonic.camera.sensor_server import ( + CameraMountPosition, + ImageMessageSchema, + ImageUtils, + SensorClient, + SensorServer, +) + +__all__ = [ + "CameraMountPosition", + "ImageMessageSchema", + "ImageUtils", + "SensorClient", + "SensorServer", +] diff --git a/GR00T-WholeBodyControl/gear_sonic/camera/composed_camera.py b/GR00T-WholeBodyControl/gear_sonic/camera/composed_camera.py new file mode 100644 index 0000000000000000000000000000000000000000..d5f5607c2967ca445001aefe31889fc22ee5c84f --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/camera/composed_camera.py @@ -0,0 +1,692 @@ +"""Composed camera server — orchestrates multiple camera drivers. + +Runs each camera in its own thread with staggered initialization and +automatic reconnection. Publishes all frames as a single merged +``ImageMessageSchema`` payload over ZMQ. + +Usage (on robot):: + + python -m gear_sonic.camera.composed_camera \\ + --ego-view-camera oak \\ + --ego-view-device-id 18443010E1ABC12300 \\ + --port 5555 + +Supported camera types: ``oak``, ``oak_mono``, ``realsense``, +``usb``, or a path to an ``.mp4`` file for replay testing. + +Run ``python -m gear_sonic.camera.composed_camera --help`` for all options. +""" + +from collections import deque +from dataclasses import dataclass +import queue +import threading +import time +from typing import Any + +import cv2 # noqa: F401 — imported early to avoid TSL segfault with camera SDKs +import numpy as np + +from gear_sonic.camera.sensor import Sensor +from gear_sonic.camera.sensor_server import ( + CameraMountPosition, + ImageMessageSchema, + SensorClient, + SensorServer, +) + + +def read_qr_code(data): + """Measure end-to-end latency by decoding QR-code timestamps.""" + current_time = time.monotonic() + detector = cv2.QRCodeDetector() + for key, img in data["images"].items(): + decoded_time, bbox, _ = detector.detectAndDecode(img) + if bbox is not None and decoded_time: + print(f"{key} latency: {(current_time - float(decoded_time)) * 1e3:.1f} ms") + else: + print(f"{key} QR code not detected.") + + +@dataclass +class ComposedCameraConfig: + """Camera configuration for the composed camera server.""" + + ego_view_camera: str | None = "oak" + """Camera type for ego view: oak, oak_mono, realsense, zed, usb, or None.""" + + ego_view_device_id: str | None = None + """Device ID for ego view camera (OAK MxID, RealSense serial, USB /dev/video index).""" + + head_camera: str | None = None + """Camera type for head view.""" + + head_device_id: str | None = None + """Device ID for head camera.""" + + left_wrist_camera: str | None = None + """Camera type for left wrist view.""" + + left_wrist_device_id: str | None = None + """Device ID for left wrist camera.""" + + right_wrist_camera: str | None = None + """Camera type for right wrist view.""" + + right_wrist_device_id: str | None = None + """Device ID for right wrist camera.""" + + fps: int = 30 + """Publish rate. OAK cameras run at 30 FPS; lower values add latency.""" + + run_as_server: bool = True + """Run as ZMQ PUB server (set False for in-process usage).""" + + server: bool = True + """Alias for run_as_server kept for backward compatibility.""" + + port: int = 5555 + """ZMQ port for server / client communication.""" + + test_latency: bool = False + """Decode QR-code timestamps in each frame to measure latency.""" + + queue_size: int = 3 + """Per-camera image queue depth.""" + + use_mjpeg: bool = False + """Use on-device MJPEG encoding on OAK cameras to reduce USB bandwidth.""" + + mjpeg_quality: int = 80 + """MJPEG quality 1-100 (only when use_mjpeg=True).""" + + def __post_init__(self): + self.run_as_server = self.server + + +class ComposedCameraSensor(Sensor, SensorServer): + """Multi-camera orchestrator with per-camera threads and auto-reconnect.""" + + def __init__(self, config: ComposedCameraConfig): + self.config = config + self.camera_queues: dict[str, queue.Queue] = {} + self.camera_threads: dict[str, threading.Thread] = {} + self.shutdown_events: dict[str, threading.Event] = {} + self.error_events: dict[str, threading.Event] = {} + self.error_messages: dict[str, str] = {} + self._observation_spaces: dict[str, Any] = {} + + camera_configs = self._get_camera_configs() + + for _idx, (mount_position, camera_config) in enumerate(camera_configs.items()): + camera_queue = queue.Queue(maxsize=config.queue_size) + shutdown_event = threading.Event() + error_event = threading.Event() + + self.camera_queues[mount_position] = camera_queue + self.shutdown_events[mount_position] = shutdown_event + self.error_events[mount_position] = error_event + + thread = threading.Thread( + target=self._camera_worker_wrapper, + args=( + mount_position, + camera_config["camera_type"], + camera_config["device_id"], + camera_queue, + shutdown_event, + error_event, + ), + ) + thread.start() + self.camera_threads[mount_position] = thread + + # Stagger init to avoid USB bandwidth contention + init_timeout = 15.0 + init_start = time.time() + while time.time() - init_start < init_timeout: + if mount_position in self._observation_spaces: + print(f"[{mount_position}] Camera ready, waiting 3s before next camera...") + time.sleep(3.0) + break + time.sleep(0.5) + else: + print(f"[{mount_position}] Camera init timeout, proceeding anyway...") + + if config.run_as_server: + print("Waiting for all cameras to be ready before starting server...") + self._wait_for_all_cameras_ready(timeout=60.0) + self.start_server(config.port) + + def _get_camera_configs(self) -> dict[str, dict]: + camera_configs = {} + + if self.config.ego_view_camera is not None: + camera_configs[CameraMountPosition.EGO_VIEW.value] = { + "camera_type": self.config.ego_view_camera, + "device_id": self.config.ego_view_device_id, + } + + if self.config.head_camera is not None: + camera_configs[CameraMountPosition.HEAD.value] = { + "camera_type": self.config.head_camera, + "device_id": self.config.head_device_id, + } + + if self.config.left_wrist_camera is not None: + camera_configs[CameraMountPosition.LEFT_WRIST.value] = { + "camera_type": self.config.left_wrist_camera, + "device_id": self.config.left_wrist_device_id, + } + + if self.config.right_wrist_camera is not None: + camera_configs[CameraMountPosition.RIGHT_WRIST.value] = { + "camera_type": self.config.right_wrist_camera, + "device_id": self.config.right_wrist_device_id, + } + + return camera_configs + + def _wait_for_all_cameras_ready(self, timeout: float = 60.0): + expected_cameras = set(self.camera_queues.keys()) + start_time = time.time() + + while time.time() - start_time < timeout: + ready_cameras = set() + for mount_position, camera_queue in self.camera_queues.items(): + if not camera_queue.empty(): + ready_cameras.add(mount_position) + + if ready_cameras == expected_cameras: + print(f"All {len(expected_cameras)} cameras ready: {ready_cameras}") + time.sleep(1.0) + return + + waiting_for = expected_cameras - ready_cameras + print( + f"Waiting for cameras: {waiting_for} " + f"({len(ready_cameras)}/{len(expected_cameras)} ready)" + ) + time.sleep(2.0) + + ready_cameras = set() + for mount_position, camera_queue in self.camera_queues.items(): + if not camera_queue.empty() or mount_position in self._observation_spaces: + ready_cameras.add(mount_position) + missing = expected_cameras - ready_cameras + print( + f"[WARNING] Timeout waiting for all cameras. " + f"Missing: {missing}. Starting anyway with: {ready_cameras}" + ) + + def _camera_worker_wrapper( + self, + mount_position: str, + camera_type: str, + device_id: str | None, + image_queue: queue.Queue, + shutdown_event: threading.Event, + error_event: threading.Event, + ): + """Worker thread with auto-reconnection.""" + max_init_retries = 10 + max_reconnect_attempts = 5 + reconnect_count = 0 + + while not shutdown_event.is_set() and reconnect_count < max_reconnect_attempts: + camera = None + try: + init_retry_delay = 1.0 + + for attempt in range(max_init_retries): + if shutdown_event.is_set(): + return + + try: + if reconnect_count > 0: + print( + f"[{mount_position}] Reconnecting camera " + f"(reconnect {reconnect_count}/{max_reconnect_attempts}, " + f"attempt {attempt + 1}/{max_init_retries})..." + ) + else: + print( + f"[{mount_position}] Initializing camera " + f"(attempt {attempt + 1}/{max_init_retries})..." + ) + camera = self._instantiate_camera(mount_position, camera_type, device_id) + print(f"[{mount_position}] Camera initialized successfully") + break + except Exception as e: + print(f"[{mount_position}] Camera init failed: {e}") + if attempt < max_init_retries - 1: + print(f"[{mount_position}] Retrying in {init_retry_delay:.1f}s...") + time.sleep(init_retry_delay) + init_retry_delay = min(init_retry_delay * 1.5, 10.0) + else: + raise RuntimeError( + f"Camera {mount_position} ({camera_type}) failed to initialize " + f"after {max_init_retries} attempts: {e}" + ) + + obs_space = camera.observation_space() + if obs_space is not None: + self._observation_spaces[mount_position] = obs_space + else: + self._observation_spaces[mount_position] = True + + consecutive_failures = 0 + max_consecutive_failures = 10 + warmup_period = True + warmup_start_time = time.time() + warmup_timeout = 5.0 + + while not shutdown_event.is_set(): + try: + frame = camera.read() + except Exception as e: + print(f"[{mount_position}] Frame read exception: {e}") + frame = None + consecutive_failures = max_consecutive_failures + + if frame: + consecutive_failures = 0 + warmup_period = False + try: + image_queue.put_nowait(frame) + except queue.Full: + try: + image_queue.get_nowait() + image_queue.put_nowait(frame) + except queue.Empty: + pass + else: + if warmup_period: + if time.time() - warmup_start_time > warmup_timeout: + print( + f"[{mount_position}] Warmup timeout — will attempt reconnect" + ) + break + time.sleep(0.1) + else: + consecutive_failures += 1 + if consecutive_failures >= max_consecutive_failures: + print( + f"[{mount_position}] Too many consecutive failures " + f"({consecutive_failures}) — will attempt reconnect" + ) + break + time.sleep(0.01) + + if camera is not None: + try: + camera.close() + except Exception as e: + print(f"[{mount_position}] Error closing camera: {e}") + camera = None + + if not shutdown_event.is_set(): + reconnect_count += 1 + print(f"[{mount_position}] Waiting 5 seconds before reconnect attempt...") + time.sleep(5.0) + + except Exception as e: + print(f"[{mount_position}] Camera error: {e}") + if camera is not None: + try: + camera.close() + except Exception: + pass + camera = None + + if not shutdown_event.is_set(): + reconnect_count += 1 + if reconnect_count < max_reconnect_attempts: + print( + f"[{mount_position}] Waiting 5 seconds before reconnect " + f"attempt {reconnect_count}/{max_reconnect_attempts}..." + ) + time.sleep(5.0) + + if reconnect_count >= max_reconnect_attempts and not shutdown_event.is_set(): + error_msg = ( + f"Camera {mount_position} ({camera_type}) failed " + f"after {max_reconnect_attempts} reconnect attempts" + ) + print(f"[ERROR] {error_msg}") + self.error_messages[mount_position] = error_msg + error_event.set() + + def _instantiate_camera( + self, mount_position: str, camera_type: str, device_id: str | None = None + ) -> Sensor: + """Instantiate a camera sensor based on camera_type (lazy imports).""" + if camera_type in ("oak", "oak_mono"): + from gear_sonic.camera.drivers.oak import OAKConfig, OAKSensor + + oak_config = OAKConfig() + oak_config.use_mjpeg = self.config.use_mjpeg + oak_config.mjpeg_quality = self.config.mjpeg_quality + if camera_type == "oak_mono": + oak_config.enable_mono_cameras = True + print(f"Initializing OAK sensor for camera type: {camera_type}") + return OAKSensor(config=oak_config, mount_position=mount_position, device_id=device_id) + + elif camera_type == "realsense": + from gear_sonic.camera.drivers.realsense import RealSenseSensor + + print(f"Initializing RealSense sensor for camera type: {camera_type}") + return RealSenseSensor(mount_position=mount_position) + + elif camera_type.endswith(".mp4"): + from gear_sonic.camera.drivers.dummy import ReplayDummySensor + + print(f"Initializing Replay Dummy Sensor for camera type: {camera_type}") + return ReplayDummySensor(video_path=camera_type) + + elif camera_type == "usb": + from gear_sonic.camera.drivers.usb_camera import USBCameraConfig, USBCameraSensor + + usb_config = USBCameraConfig() + device_idx = int(device_id) if device_id else 0 + print(f"Initializing USB camera for type: {camera_type}, device: {device_idx}") + return USBCameraSensor( + config=usb_config, mount_position=mount_position, device_index=device_idx + ) + + else: + raise ValueError(f"Unsupported camera type: {camera_type}") + + def _check_for_errors(self): + for mount_position, error_event in self.error_events.items(): + if error_event.is_set(): + error_msg = self.error_messages.get( + mount_position, f"Camera {mount_position} encountered an unknown error" + ) + raise RuntimeError(error_msg) + + def read(self): + """Read frames from all cameras. Returns None unless ALL cameras have frames.""" + self._check_for_errors() + + expected_cameras = set(self.camera_queues.keys()) + message = {} + + for mount_position, camera_queue in self.camera_queues.items(): + frame = self._get_latest_from_queue(camera_queue) + if frame is not None: + message[mount_position] = frame + + if set(message.keys()) == expected_cameras: + return message + return None + + def _get_latest_from_queue(self, camera_queue: queue.Queue) -> dict[str, Any] | None: + latest = None + try: + while True: + latest = camera_queue.get_nowait() + except queue.Empty: + pass + return latest + + def close(self): + for shutdown_event in self.shutdown_events.values(): + shutdown_event.set() + for thread in self.camera_threads.values(): + thread.join(timeout=5.0) + for camera_queue in self.camera_queues.values(): + try: + while True: + camera_queue.get_nowait() + except queue.Empty: + pass + if self.config.run_as_server: + self.stop_server() + + def serialize(self, data: dict[str, Any]) -> dict[str, Any]: + raise NotImplementedError("Use serialize_message() for ComposedCameraSensor") + + def serialize_message(self, message: dict[str, Any]) -> dict[str, Any]: + """Merge per-camera data into a single ImageMessageSchema.""" + all_timestamps = {} + all_images = {} + for _mount, camera_data in message.items(): + all_timestamps.update(camera_data.get("timestamps", {})) + all_images.update(camera_data.get("images", {})) + img_schema = ImageMessageSchema(timestamps=all_timestamps, images=all_images) + return img_schema.serialize() + + def run_server(self): + """Main server loop — reads, serializes and publishes frames.""" + idx = 0 + server_start_time = time.monotonic() + fps_print_time = time.monotonic() + frame_interval = 1.0 / self.config.fps + + while True: + target_time = server_start_time + (idx + 1) * frame_interval + + message = self.read() + if message: + if self.config.test_latency: + read_qr_code(message) + + serialized_message = self.serialize_message(message) + self.send_message(serialized_message) + idx += 1 + + if idx % 10 == 0: + print(f"Image sending FPS: {10 / (time.monotonic() - fps_print_time):.2f}") + fps_print_time = time.monotonic() + + current_time = time.monotonic() + sleep_time = target_time - current_time + if sleep_time > 0: + time.sleep(sleep_time) + else: + if not message: + idx += 1 + + def observation_space(self): + try: + import gymnasium as gym + + return gym.spaces.Dict(self._observation_spaces) + except ImportError: + return None + + +class ComposedCameraClientSensor(Sensor, SensorClient): + """ZMQ client that deserializes merged camera frames from the server.""" + + def __init__(self, server_ip: str = "localhost", port: int = 5555): + self.start_client(server_ip, port) + + self._latest_message = None + self._avg_time_per_frame: deque = deque(maxlen=20) + self._msg_received_time = 0 + self._start_time = 0.0 + self.idx = 0 + + self._last_new_message_time = None + self._last_staleness_warning_time = 0.0 + self._staleness_warning_interval = 2.0 + + print("Initialized composed camera client sensor") + + def read(self, blocking: bool = False, **kwargs) -> dict[str, Any] | None: + self._start_time = time.time() + current_time = time.time() + + if blocking: + message = self.receive_message() + if not message: + return None + else: + message = self.receive_message_nonblocking(timeout_ms=0) + + if message is not None: + self.idx += 1 + self._latest_message = ImageMessageSchema.deserialize(message).asdict() + self._last_new_message_time = current_time + + if self.idx % 10 == 0: + for image_key, image_time in self._latest_message["timestamps"].items(): + image_latency = (time.time() - image_time) * 1000 + print(f"Image latency for {image_key}: {image_latency:.2f} ms") + + self._msg_received_time = time.time() + self._avg_time_per_frame.append(self._msg_received_time - self._start_time) + elif not blocking and self._latest_message is not None: + if self._last_new_message_time is not None: + time_since_last_message = current_time - self._last_new_message_time + if time_since_last_message > 0.1: + if ( + current_time - self._last_staleness_warning_time + >= self._staleness_warning_interval + ): + print( + f"[WARNING] No new image message received for " + f"{time_since_last_message*1000:.1f}ms. " + f"Reusing stale image. Check camera server connection." + ) + self._last_staleness_warning_time = current_time + + return self._latest_message + + def serialize(self, data: dict[str, Any]) -> dict[str, Any]: + raise NotImplementedError("Client does not serialize") + + def close(self): + self.stop_client() + + def fps(self) -> float: + if len(self._avg_time_per_frame) == 0: + return 0.0 + return float(1 / np.mean(self._avg_time_per_frame)) + + +class _MjpegGrabber: + """Background thread that reads an MJPEG stream via raw HTTP.""" + + def __init__(self, url: str): + self.url = url + self.lock = threading.Lock() + self.frame: np.ndarray | None = None + self._running = True + self._thread = threading.Thread(target=self._run, daemon=True) + self._thread.start() + + def _run(self): + import requests + + resp = requests.get(self.url, stream=True, timeout=10) + buf = b"" + for chunk in resp.iter_content(chunk_size=4096): + if not self._running: + break + buf += chunk + while True: + soi = buf.find(b"\xff\xd8") + if soi == -1: + break + eoi = buf.find(b"\xff\xd9", soi + 2) + if eoi == -1: + break + jpeg_bytes = buf[soi : eoi + 2] + buf = buf[eoi + 2 :] + frame = cv2.imdecode( + np.frombuffer(jpeg_bytes, dtype=np.uint8), cv2.IMREAD_COLOR + ) + if frame is not None: + with self.lock: + self.frame = frame + + def get(self) -> np.ndarray | None: + with self.lock: + return self.frame + + def stop(self): + self._running = False + self._thread.join(timeout=2) + + +class ComposedCameraHttpClient: + """Camera client that reads MJPEG streams over HTTP. + + Drop-in replacement for :class:`ComposedCameraClientSensor` when cameras + are served via an HTTP MJPEG server. + + Usage:: + + client = ComposedCameraHttpClient("http://:8000") + data = client.read() # {"images": {"left": ndarray, ...}, "timestamps": {...}} + """ + + DEFAULT_NAME_MAP = { + "center": "ego_view", + "left": "left_wrist", + "right": "right_wrist", + } + + def __init__(self, base_url: str, camera_name_map: dict[str, str] | None = None): + self.base_url = base_url.rstrip("/") + self.camera_name_map = ( + camera_name_map if camera_name_map is not None else self.DEFAULT_NAME_MAP + ) + self.camera_names: list[str] = [] + self._grabbers: dict[str, _MjpegGrabber] = {} + self._connect() + + def _connect(self): + import requests + + resp = requests.get(f"{self.base_url}/cameras", timeout=5) + resp.raise_for_status() + self.camera_names = resp.json() + print(f"HTTP MJPEG: discovered cameras: {self.camera_names}") + for name in self.camera_names: + url = f"{self.base_url}/stream/{name}" + self._grabbers[name] = _MjpegGrabber(url) + + def read(self, blocking: bool = False, **kwargs) -> dict[str, Any] | None: + images = {} + any_ok = False + for name, grabber in self._grabbers.items(): + frame = grabber.get() + mapped_name = self.camera_name_map.get(name, name) + if frame is not None: + images[mapped_name] = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + any_ok = True + else: + images[mapped_name] = None + if not any_ok: + return None + return {"images": images, "timestamps": {n: time.time() for n in images}} + + def close(self): + for grabber in self._grabbers.values(): + grabber.stop() + + +if __name__ == "__main__": + import tyro + + config = tyro.cli(ComposedCameraConfig) + + if config.run_as_server: + composed_camera = ComposedCameraSensor(config) + print("Running composed camera server...") + composed_camera.run_server() + else: + composed_client = ComposedCameraClientSensor(server_ip="localhost", port=config.port) + try: + while True: + data = composed_client.read() + if data is not None: + print(f"FPS: {composed_client.fps():.2f}") + time.sleep(0.1) + except KeyboardInterrupt: + print("Stopping client...") + composed_client.close() diff --git a/GR00T-WholeBodyControl/gear_sonic/config/base.yaml b/GR00T-WholeBodyControl/gear_sonic/config/base.yaml new file mode 100644 index 0000000000000000000000000000000000000000..492df566a06e4c4754f87e2e364dca7945601cd0 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/config/base.yaml @@ -0,0 +1,52 @@ +# First we define the global structures that will be used by all the configs. +defaults: + # - base/fabric + - _self_ + - base/hydra + - base/structure + - callbacks/model_save + - callbacks/wandb + - trainer: trl + - opt/wandb + +num_gpus: 1 +max_retries: 1 + +# These are global variables that all levels of the config can access. +## Experiment setup +seed: 0 +codebase_version: 1.0 # this is recorded to enable auto-conversion of models between different versions of the codebase +headless: True +num_envs: 4096 + +### Checkpoint logic +auto_load_latest: False +checkpoint: null + +### Naming and dir structure +project_name: TEST +experiment_name: TEST + +base_dir: logs_rl +timestamp: ${now:%Y%m%d_%H%M%S} +experiment_dir: ${base_dir}/${project_name}/$${experiment_name}-${timestamp} +save_dir: ${experiment_dir}/.hydra + +force_flat_terrain: False + +use_wandb: false +log_task_name: TEST + +multi_gpu: False +global_rank: 0 + +### Simulation +sim_type: isaacsim +env_spacing: 20 +output_dir: ${experiment_dir}/output + +eval_overrides: + headless: False + num_envs: 1 + auto_load_latest: False + use_wandb: False diff --git a/GR00T-WholeBodyControl/gear_sonic/config/base_eval.yaml b/GR00T-WholeBodyControl/gear_sonic/config/base_eval.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3d21d1b0ac1423765c577ebb1ad4e488f2598335 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/config/base_eval.yaml @@ -0,0 +1,26 @@ +# @package _global_ + +defaults: + - /callbacks/im_eval + - manager_env/recorders: empty + - _self_ + +checkpoint: ??? + +algo: + config: + eval: + num_eval_episodes: 150 + save_videos: false + video_save_prob: 1.0 + save_goal_reached_only: true + save_trajectories: false + num_save_episodes: 200 + +eval_timestamp: ${now:%Y%m%d_%H%M%S} +eval_name: TEST +eval_base_dir: logs_eval +eval_log_dir: ${eval_base_dir}/${eval_timestamp}-${eval_name} +hydra: + run: + dir: ${eval_log_dir} diff --git a/GR00T-WholeBodyControl/gear_sonic/config/eval_exp.yaml b/GR00T-WholeBodyControl/gear_sonic/config/eval_exp.yaml new file mode 100644 index 0000000000000000000000000000000000000000..203f73f9baf4397b119b2f672e20095865083968 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/config/eval_exp.yaml @@ -0,0 +1,34 @@ +# @package _global_ + +# Required parameter: experiment directory to monitor +experiment_dir: ??? +max_train_steps: 1000000 + +# Number of environments for metrics eval (multi-GPU accelerate) +num_eval_envs: 3072 + +# Number of environments for render eval (single GPU) +# VRAM scales ~linearly: 64 envs ≈ 23GB, 32 envs ≈ 12GB +# Keep ≤ 64 on L40 (49GB VRAM) to avoid OOM +num_render_videos: 64 +num_test_render_videos: 32 + +# Monitoring configuration +scan_interval: 0 # seconds between scans for new checkpoints +checkpoint_ready_delay: 60 # seconds to wait after checkpoint modification before evaluating + +eval_frequency: null # null = eval every checkpoint +eval_last_n: null # null = eval all; integer = only eval last N checkpoints +eval_callbacks: im_eval +eval_datasets: null +eval_modes: [null] + +# Single pass mode: evaluate pending checkpoints once and exit +single_pass: false + +# Extra overrides to pass to eval_agent_trl.py (list of strings) +# Example: ["++manager_env.commands.motion.start_from_first_frame=true"] +extra_overrides: [] + +# Set to true to suppress eval subprocess output +capture_output: true diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/__init__.py b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f6c7103708c9c6ca4bdbe270742c88f83c9fdfa3 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/__init__.py @@ -0,0 +1,5 @@ +"""Robot model package: Pinocchio-based FK/IK models for the G1 humanoid.""" + +from .robot_model import ReducedRobotModel, RobotModel + +__all__ = ["RobotModel", "ReducedRobotModel"] diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/__pycache__/__init__.cpython-310.pyc b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a56bcdde10d92bac36d2a54b9b36c8b2f9d8c723 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/__pycache__/__init__.cpython-310.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/__pycache__/robot_model.cpython-310.pyc b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/__pycache__/robot_model.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6f849fba836d663fcdfc37f698ae325633f9476f Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/__pycache__/robot_model.cpython-310.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/instantiation/__init__.py b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/instantiation/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4639c1885aeebda4b5dcf5dce2ffa3eb9626592b --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/instantiation/__init__.py @@ -0,0 +1,5 @@ +"""G1 robot model instantiation helpers.""" + +from .g1 import instantiate_g1_robot_model + +__all__ = ["instantiate_g1_robot_model"] diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/instantiation/__pycache__/__init__.cpython-310.pyc b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/instantiation/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..67d32225dece0096ccc1c2fd842a9c03419017cb Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/instantiation/__pycache__/__init__.cpython-310.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/instantiation/__pycache__/g1.cpython-310.pyc b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/instantiation/__pycache__/g1.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d00c9db56d1ff140d66ae079ecede93b66e528a2 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/instantiation/__pycache__/g1.cpython-310.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/instantiation/g1.py b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/instantiation/g1.py new file mode 100644 index 0000000000000000000000000000000000000000..a5cbaed89e2e8f8e6e7ed8eeaab668f190af1947 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/instantiation/g1.py @@ -0,0 +1,62 @@ +"""Factory function to instantiate a configured G1 RobotModel from URDF.""" + +import os +from pathlib import Path +from typing import Literal + +from gear_sonic.data.robot_model.robot_model import RobotModel +from gear_sonic.data.robot_model.supplemental_info.g1.g1_supplemental_info import ( + ElbowPose, + G1SupplementalInfo, + WaistLocation, +) + + +def instantiate_g1_robot_model( + waist_location: Literal["lower_body", "upper_body", "lower_and_upper_body"] = "lower_body", + high_elbow_pose: bool = False, +): + """ + Instantiate a G1 robot model with configurable waist location and pose. + + Args: + waist_location: Whether to put waist in "lower_body" (default G1 behavior), + "upper_body" (waist controlled with arms/manipulation via IK), + or "lower_and_upper_body" (waist reference from arms/manipulation + via IK then passed to lower body policy) + high_elbow_pose: Whether to use high elbow pose configuration for default joint positions + + Returns: + RobotModel: Configured G1 robot model + """ + model_data_dir = Path(__file__).resolve().parent.parent / "model_data" / "g1" + robot_model_config = { + "asset_path": str(model_data_dir), + "urdf_path": str(model_data_dir / "g1_29dof_with_hand.urdf"), + } + assert waist_location in [ + "lower_body", + "upper_body", + "lower_and_upper_body", + ], f"Invalid waist_location: {waist_location}. Must be 'lower_body' or 'upper_body' or 'lower_and_upper_body'" + + # Map string values to enums + waist_location_enum = { + "lower_body": WaistLocation.LOWER_BODY, + "upper_body": WaistLocation.UPPER_BODY, + "lower_and_upper_body": WaistLocation.LOWER_AND_UPPER_BODY, + }[waist_location] + + elbow_pose_enum = ElbowPose.HIGH if high_elbow_pose else ElbowPose.LOW + + # Create single configurable supplemental info instance + robot_model_supplemental_info = G1SupplementalInfo( + waist_location=waist_location_enum, elbow_pose=elbow_pose_enum + ) + + robot_model = RobotModel( + robot_model_config["urdf_path"], + robot_model_config["asset_path"], + supplemental_info=robot_model_supplemental_info, + ) + return robot_model diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/g1_29dof_with_hand.urdf b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/g1_29dof_with_hand.urdf new file mode 100644 index 0000000000000000000000000000000000000000..e057336b6d0c6f25c8f0084664d28dc1885d0c7f --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/g1_29dof_with_hand.urdf @@ -0,0 +1,1497 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/g1_29dof_with_hand.xml b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/g1_29dof_with_hand.xml new file mode 100644 index 0000000000000000000000000000000000000000..76110b566c95e820ce430a489555b26e2391741a --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/g1_29dof_with_hand.xml @@ -0,0 +1,751 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/left_wrist_pitch_link.STL b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/left_wrist_pitch_link.STL new file mode 100644 index 0000000000000000000000000000000000000000..82cc224a8e41251d879502f9809e31d0988ec7f9 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/meshes/left_wrist_pitch_link.STL differ diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/scene_43dof.xml b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/scene_43dof.xml new file mode 100644 index 0000000000000000000000000000000000000000..d66b23dce5da38aa29ff78d153ac86bf406740e5 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/model_data/g1/scene_43dof.xml @@ -0,0 +1,30 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/robot_model.py b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/robot_model.py new file mode 100644 index 0000000000000000000000000000000000000000..a22cb69df96ef729a5058078e978252fcc445d1a --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/robot_model.py @@ -0,0 +1,817 @@ +"""Pinocchio-based robot model with FK, joint limits, gravity compensation, and reduced-DOF views. + +RobotModel wraps a URDF + optional supplemental info to provide forward kinematics, +Jacobians, gravity torques, and joint-group queries. ReducedRobotModel maps between +a full configuration and an actuated-joint subset. +""" + +from typing import List, Optional, Set, Union + +import numpy as np +import pinocchio as pin + +from gear_sonic.data.robot_model.supplemental_info import RobotSupplementalInfo + + +class RobotModel: + def __init__( + self, + urdf_path, + asset_path, + set_floating_base=False, + supplemental_info: Optional[RobotSupplementalInfo] = None, + ): + self.pinocchio_wrapper = pin.RobotWrapper.BuildFromURDF( + filename=urdf_path, + package_dirs=[asset_path], + root_joint=pin.JointModelFreeFlyer() if set_floating_base else None, + ) + self.is_floating_base_model = set_floating_base + + self.joint_to_dof_index = {} + # Assume we only have single-dof joints + # First two names correspond to universe and floating base joints + names = ( + self.pinocchio_wrapper.model.names[2:] + if set_floating_base + else self.pinocchio_wrapper.model.names[1:] + ) + for name in names: + j_id = self.pinocchio_wrapper.model.getJointId(name) + jmodel = self.pinocchio_wrapper.model.joints[j_id] + self.joint_to_dof_index[name] = jmodel.idx_q + + # Store joint limits only for actual joints (excluding floating base) + # if set floating base is true and the robot can move in the world + # then we don't want to impose joint limits for the 7 dofs corresponding + # to the floating base dofs. + root_nq = 7 if set_floating_base else 0 + self.upper_joint_limits = self.pinocchio_wrapper.model.upperPositionLimit[root_nq:].copy() + self.lower_joint_limits = self.pinocchio_wrapper.model.lowerPositionLimit[root_nq:].copy() + + # Set up supplemental info if provided + self.supplemental_info = supplemental_info + if self.supplemental_info is not None: + # Cache indices for body and hand actuated joints separately + self._body_actuated_joint_indices = [ + self.dof_index(name) for name in self.supplemental_info.body_actuated_joints + ] + self._left_hand_actuated_joint_indices = [ + self.dof_index(name) for name in self.supplemental_info.left_hand_actuated_joints + ] + self._right_hand_actuated_joint_indices = [ + self.dof_index(name) for name in self.supplemental_info.right_hand_actuated_joints + ] + self._hand_actuated_joint_indices = ( + self._left_hand_actuated_joint_indices + self._right_hand_actuated_joint_indices + ) + + # Cache indices for joint groups, handling nested groups + self._joint_group_indices = {} + for group_name, group_info in self.supplemental_info.joint_groups.items(): + indices = [] + # Add indices for direct joints + indices.extend([self.dof_index(name) for name in group_info["joints"]]) + # Add indices from subgroups + for subgroup_name in group_info["groups"]: + indices.extend(self.get_joint_group_indices(subgroup_name)) + self._joint_group_indices[group_name] = sorted(set(indices)) + + # Update joint limits from supplemental info if available + if ( + hasattr(self.supplemental_info, "joint_limits") + and self.supplemental_info.joint_limits + ): + for joint_name, limits in self.supplemental_info.joint_limits.items(): + if joint_name in self.joint_to_dof_index: + # joint_to_dof_index is in full-space (includes floating base DOFs), + # but limits arrays are indexed from 0 starting at the first real joint + idx = self.joint_to_dof_index[joint_name] - root_nq + self.lower_joint_limits[idx] = limits[0] + self.upper_joint_limits[idx] = limits[1] + + # Initialize default body pose + self.default_body_pose = self.q_zero + + # Update with supplemental info if available + if self.supplemental_info is not None: + default_joint_q = self.supplemental_info.default_joint_q + for joint, joint_values in default_joint_q.items(): + # Get the joint name mapping for this type + joint_mapping = self.supplemental_info.joint_name_mapping[joint] + + # Handle both single joint names and left/right mappings + if isinstance(joint_mapping, str): + # Single joint (e.g., waist joints) + if joint_mapping in self.joint_to_dof_index: + joint_idx = self.dof_index(joint_mapping) + self.default_body_pose[joint_idx] = ( + joint_values # joint_values is the value for single joints + ) + else: + # Left/right mapping (e.g., arm joints) + for side, value in joint_values.items(): + if side in joint_mapping and joint_mapping[side] in self.joint_to_dof_index: + joint_idx = self.dof_index(joint_mapping[side]) + self.default_body_pose[joint_idx] = value + + # Initialize initial body pose + self.initial_body_pose = self.default_body_pose.copy() + + @property + def num_dofs(self) -> int: + """Get the number of degrees of freedom of the robot (floating base pose + joints).""" + return self.pinocchio_wrapper.model.nq + + @property + def q_zero(self) -> np.ndarray: + """Get the zero pose of the robot.""" + return self.pinocchio_wrapper.q0.copy() + + @property + def joint_names(self) -> List[str]: + """Get the names of the joints of the robot.""" + return list(self.joint_to_dof_index.keys()) + + @property + def num_joints(self) -> int: + """Get the number of joints of the robot.""" + return len(self.joint_to_dof_index) + + def dof_index(self, joint_name: str) -> int: + """ + Get the index in the degrees of freedom vector corresponding + to the single-DoF joint with name `joint_name`. + """ + if joint_name not in self.joint_to_dof_index: + raise ValueError( + f"Unknown joint name: '{joint_name}'. " + f"Available joints: {list(self.joint_to_dof_index.keys())}" + ) + return self.joint_to_dof_index[joint_name] + + def get_body_actuated_joint_indices(self) -> List[int]: + """ + Get the indices of body actuated joints in the full configuration. + Ordering is that of the actuated joints as defined in the supplemental info. + Requires supplemental_info to be provided. + """ + if self.supplemental_info is None: + raise ValueError("supplemental_info must be provided to use this method") + return self._body_actuated_joint_indices + + def get_hand_actuated_joint_indices(self, side: str = "both") -> List[int]: + """ + Get the indices of hand actuated joints in the full configuration. + Ordering is that of the actuated joints as defined in the supplemental info. + Requires supplemental_info to be provided. + + Args: + side: String specifying which hand to get indices for ('left', 'right', or 'both') + """ + if self.supplemental_info is None: + raise ValueError("supplemental_info must be provided to use this method") + + if side.lower() == "both": + return self._hand_actuated_joint_indices + elif side.lower() == "left": + return self._left_hand_actuated_joint_indices + elif side.lower() == "right": + return self._right_hand_actuated_joint_indices + else: + raise ValueError("side must be 'left', 'right', or 'both'") + + def get_joint_group_indices(self, group_names: Union[str, Set[str]]) -> List[int]: + """ + Get the indices of joints in one or more groups in the full configuration. + Requires supplemental_info to be provided. + The returned indices are sorted in ascending order, so that the joint ordering + of the full model is preserved. + + Args: + group_names: Either a single group name (str) or a set of group names (Set[str]) + + Returns: + List of joint indices in sorted order with no duplicates + """ + if self.supplemental_info is None: + raise ValueError("supplemental_info must be provided to use this method") + + # Convert single string to set for uniform handling + if isinstance(group_names, str): + group_names = {group_names} + + # Collect indices from all groups + all_indices = set() + for group_name in group_names: + if group_name not in self._joint_group_indices: + raise ValueError(f"Unknown joint group: {group_name}") + all_indices.update(self._joint_group_indices[group_name]) + + return sorted(all_indices) + + def cache_forward_kinematics(self, q: np.ndarray, auto_clip=True) -> None: + """ + Perform forward kinematics to update the pose of every joint and frame + in the Pinocchio data structures for the given configuration `q`. + + :param q: A numpy array of shape (num_dofs,) representing the robot configuration. + """ + if q.shape[0] != self.num_dofs: + raise ValueError(f"Expected q of length {self.num_dofs}, got {q.shape[0]} instead.") + + # Apply auto-clip if enabled + if auto_clip: + q = self.clip_configuration(q) + + pin.framesForwardKinematics(self.pinocchio_wrapper.model, self.pinocchio_wrapper.data, q) + + def compute_gravity_compensation_torques( + self, q: np.ndarray, joint_groups: Union[str, List[str], Set[str]] = None, auto_clip=True + ) -> np.ndarray: + """ + Compute gravity compensation torques for specified joint groups using pinocchio. + + :param q: Robot configuration (joint positions) + :param joint_groups: Joint groups to compensate (e.g., "arms", ["left_arm", "waist"], + {"left_arm", "waist"}). If None, compensates all joints + :param auto_clip: Whether to automatically clip joint values to limits + :return: Array of gravity compensation torques for all DOFs (zero for non-compensated joints) + """ + if q.shape[0] != self.num_dofs: + raise ValueError(f"Expected q of length {self.num_dofs}, got {q.shape[0]} instead.") + + # Apply auto-clip if enabled + if auto_clip: + q = self.clip_configuration(q) + + try: + # Cache forward kinematics for the current configuration + self.cache_forward_kinematics(q, auto_clip=False) # Already clipped if needed + + # RNEA with zero velocity and acceleration isolates the gravity term: + # tau = M(q)*0 + C(q,0)*0 + g(q) = g(q), i.e. pure gravity compensation + v = np.zeros(self.num_dofs) + a = np.zeros(self.num_dofs) + + gravity_torques_full = pin.rnea( + self.pinocchio_wrapper.model, self.pinocchio_wrapper.data, q, v, a + ) + + # If no joint groups specified, return full gravity torques + if joint_groups is None: + return gravity_torques_full + + # Convert list to set for get_joint_group_indices compatibility + if isinstance(joint_groups, list): + joint_groups = set(joint_groups) + + # Get joint indices for specified groups - get_joint_group_indices handles str and Set[str] + try: + compensated_joint_indices = self.get_joint_group_indices(joint_groups) + except ValueError as e: + raise ValueError(f"Error resolving joint groups {joint_groups}: {e}") + + # Create mask for joints that should receive gravity compensation + compensation_mask = np.zeros(self.num_dofs, dtype=bool) + for joint_idx in compensated_joint_indices: + if 0 <= joint_idx < len(compensation_mask): + compensation_mask[joint_idx] = True + + # Apply mask to only compensate specified joints + compensated_torques = np.zeros_like(gravity_torques_full) + compensated_torques[compensation_mask] = gravity_torques_full[compensation_mask] + + return compensated_torques + + except Exception as e: + raise RuntimeError(f"Error computing gravity compensation: {e}") + + def clip_configuration(self, q: np.ndarray, margin: float = 1e-6) -> np.ndarray: + """ + Clip the configuration to stay within joint limits with a small tolerance. + + :param q: Configuration to clip + :param margin: Tolerance to keep away from joint limits + :return: Clipped configuration + """ + q_clipped = q.copy() + + # Only clip joint positions, not floating base + root_nq = 7 if self.is_floating_base_model else 0 + q_clipped[root_nq:] = np.clip( + q[root_nq:], self.lower_joint_limits + margin, self.upper_joint_limits - margin + ) + + return q_clipped + + def frame_placement(self, frame_name: str) -> pin.SE3: + """ + Returns the SE3 transform of the specified frame in the world coordinate system. + Note: make sure cache_forward_kinematics() has been previously called. + + :param frame_name: Name of the frame, e.g. "link_elbow_frame", "hand_imu_frame", etc. + :return: A pin.SE3 object representing the pose of the frame. + """ + model = self.pinocchio_wrapper.model + data = self.pinocchio_wrapper.data + + frame_id = model.getFrameId(frame_name) + if frame_id < 0 or frame_id >= len(model.frames): + valid_frames = [f.name for f in model.frames] + raise ValueError(f"Unknown frame '{frame_name}'. Valid frames: {valid_frames}") + + # Pinocchio's data.oMf[frame_id] is a pin.SE3. + return data.oMf[frame_id].copy() + + def frame_jacobian( + self, + frame_name: str, + q: np.ndarray, + reference_frame: pin.ReferenceFrame = pin.LOCAL_WORLD_ALIGNED, + ) -> np.ndarray: + """ + Compute the Jacobian of the specified frame. + + :param frame_name: Name of the frame, e.g. "fingertip_frame", "hand_frame", etc. + :param q: Configuration vector (joint positions) + :param reference_frame: Reference frame for the Jacobian. Options: + - pin.LOCAL: Jacobian expressed in the local frame + - pin.WORLD: Jacobian expressed in the world frame + - pin.LOCAL_WORLD_ALIGNED: Local frame with world orientation (default, best for IK) + :return: A 6xN Jacobian matrix where N is the number of DOFs. + First 3 rows are linear velocity, last 3 rows are angular velocity. + """ + model = self.pinocchio_wrapper.model + data = self.pinocchio_wrapper.data + + # Get frame ID + frame_id = model.getFrameId(frame_name) + if frame_id < 0 or frame_id >= len(model.frames): + valid_frames = [f.name for f in model.frames] + raise ValueError(f"Unknown frame '{frame_name}'. Valid frames: {valid_frames}") + + # Compute the Jacobian + J = pin.computeFrameJacobian(model, data, q, frame_id, reference_frame) + + return J.copy() + + def get_body_actuated_joints(self, q: np.ndarray) -> np.ndarray: + """ + Get the configuration of body actuated joints from a full configuration. + + :param q: Configuration in full space + :return: Configuration of body actuated joints + """ + indices = self.get_body_actuated_joint_indices() + + return q[indices] + + def get_hand_actuated_joints(self, q: np.ndarray, side: str = "both") -> np.ndarray: + """ + Get the configuration of hand actuated joints from a full configuration. + + Args: + q: Configuration in full space + side: String specifying which hand to get joints for ('left', 'right', or 'both') + """ + indices = self.get_hand_actuated_joint_indices(side) + return q[indices] + + def get_configuration_from_actuated_joints( + self, + body_actuated_joint_values: np.ndarray, + hand_actuated_joint_values: Optional[np.ndarray] = None, + left_hand_actuated_joint_values: Optional[np.ndarray] = None, + right_hand_actuated_joint_values: Optional[np.ndarray] = None, + ) -> np.ndarray: + """ + Get the full configuration from the body and hand actuated joint configurations. + Can specify either both hands together or left and right hands separately. + + Args: + body_actuated_joint_values: Configuration of body actuated joints + hand_actuated_joint_values: Configuration of both hands' actuated joints (optional) + left_hand_actuated_joint_values: Configuration of left hand actuated joints (optional) + right_hand_actuated_joint_values: Configuration of right hand actuated joints (optional) + + Returns: + Full configuration including body and hand joints + """ + q = self.pinocchio_wrapper.q0.copy() + q[self.get_body_actuated_joint_indices()] = body_actuated_joint_values + + # Handle hand configurations + if hand_actuated_joint_values is not None: + # Use combined hand configuration + q[self.get_hand_actuated_joint_indices("both")] = hand_actuated_joint_values + else: + # Use separate hand configurations + if left_hand_actuated_joint_values is not None: + q[self.get_hand_actuated_joint_indices("left")] = left_hand_actuated_joint_values + if right_hand_actuated_joint_values is not None: + q[self.get_hand_actuated_joint_indices("right")] = right_hand_actuated_joint_values + + return q + + def reset_forward_kinematics(self) -> None: + """ + Reset the forward kinematics to the initial configuration. + """ + self.cache_forward_kinematics(self.q_zero) + + def get_initial_upper_body_pose(self) -> np.ndarray: + """ + Get the initial upper body pose of the robot. + """ + return self.initial_body_pose[self.get_joint_group_indices("upper_body")] + + def get_default_body_pose(self) -> np.ndarray: + """ + Get the default body pose of the robot. + """ + return self.default_body_pose.copy() + + def set_initial_body_pose(self, q: np.ndarray, q_idx=None) -> None: + """ + Set the initial body pose of the robot. + """ + if q_idx is None: + self.initial_body_pose = q + else: + self.initial_body_pose[q_idx] = q + + +class ReducedRobotModel(RobotModel): + """ + A class that creates a reduced order robot model by fixing certain joints. + This class maintains a mapping between the reduced state space and the full state space. + """ + + def __init__( + self, + full_robot_model: RobotModel, + fixed_joints: List[str], + fixed_values: Optional[List[float]] = None, + ): + """ + Create a reduced order robot model by fixing specified joints. + + :param full_robot_model: The original robot model + :param fixed_joints: List of joint names to fix + :param fixed_values: Optional list of values to fix the joints to. If None, uses the initial + joint positions (q0) from the full robot model. + """ + self.full_robot = full_robot_model + self.supplemental_info = full_robot_model.supplemental_info + + # If fixed_values is None, use q0 from the full robot model + if fixed_values is None: + fixed_values = [] + for joint_name in fixed_joints: + full_idx = full_robot_model.dof_index(joint_name) + fixed_values.append(full_robot_model.pinocchio_wrapper.q0[full_idx]) + elif len(fixed_joints) != len(fixed_values): + raise ValueError("fixed_joints and fixed_values must have the same length") + + # Store fixed joints and their values + self.fixed_joints = fixed_joints + self.fixed_values = fixed_values + + # reduced_to_full[i] = full-space index of the i-th reduced-space DOF + # full_to_reduced[j] = reduced-space index of the j-th full-space DOF (active joints only) + self.reduced_to_full = [] + self.full_to_reduced = {} + + # Initialize with floating base indices if present + if full_robot_model.is_floating_base_model: + self.reduced_to_full.extend(range(7)) # Floating base indices + for i in range(7): + self.full_to_reduced[i] = i + + # Add active joint indices + for joint_name in full_robot_model.joint_names: + if joint_name not in fixed_joints: + full_idx = full_robot_model.dof_index(joint_name) + reduced_idx = len(self.reduced_to_full) + self.reduced_to_full.append(full_idx) + self.full_to_reduced[full_idx] = reduced_idx + + # Create a reduced Pinocchio model using buildReducedModel + # First, get the list of joint IDs to lock + locked_joint_ids = [] + for joint_name in fixed_joints: + joint_id = full_robot_model.pinocchio_wrapper.model.getJointId(joint_name) + # Pinocchio reserves id=0 for "universe" and id=1 for floating base (if present). + # Only lock actual robot joints, not these special entries. + if (full_robot_model.is_floating_base_model and joint_id > 1) or ( + not full_robot_model.is_floating_base_model and joint_id > 0 + ): + locked_joint_ids.append(joint_id) + + # First build the reduced kinematic model + reduced_model = pin.buildReducedModel( + full_robot_model.pinocchio_wrapper.model, + locked_joint_ids, + full_robot_model.pinocchio_wrapper.q0, + ) + + # Then build the reduced geometry models using the reduced kinematic model + self.pinocchio_wrapper = pin.RobotWrapper( + model=reduced_model, + ) + + # Create joint to dof index mapping + self.joint_to_dof_index = {} + # Assume we only have single-dof joints + # First two names correspond to universe and floating base joints + names = ( + self.pinocchio_wrapper.model.names[2:] + if self.full_robot.is_floating_base_model + else self.pinocchio_wrapper.model.names[1:] + ) + for name in names: + j_id = self.pinocchio_wrapper.model.getJointId(name) + jmodel = self.pinocchio_wrapper.model.joints[j_id] + self.joint_to_dof_index[name] = jmodel.idx_q + + # Initialize joint limits + root_nq = 7 if self.full_robot.is_floating_base_model else 0 + self.lower_joint_limits = self.pinocchio_wrapper.model.lowerPositionLimit[root_nq:].copy() + self.upper_joint_limits = self.pinocchio_wrapper.model.upperPositionLimit[root_nq:].copy() + + # Update joint limits from supplemental info if available + if self.supplemental_info is not None: + if ( + hasattr(self.supplemental_info, "joint_limits") + and self.supplemental_info.joint_limits + ): + for joint_name, limits in self.supplemental_info.joint_limits.items(): + if joint_name in self.joint_to_dof_index: + idx = self.joint_to_dof_index[joint_name] - root_nq + self.lower_joint_limits[idx] = limits[0] + self.upper_joint_limits[idx] = limits[1] + + # Get full indices for body and hand actuated joints + full_body_indices = full_robot_model.get_body_actuated_joint_indices() + full_hand_indices = full_robot_model.get_hand_actuated_joint_indices("both") + full_left_hand_indices = full_robot_model.get_hand_actuated_joint_indices("left") + full_right_hand_indices = full_robot_model.get_hand_actuated_joint_indices("right") + + # Map to reduced indices + self._body_actuated_joint_indices = [] + for idx in full_body_indices: + if idx in self.full_to_reduced: + self._body_actuated_joint_indices.append(self.full_to_reduced[idx]) + + self._hand_actuated_joint_indices = [] + for idx in full_hand_indices: + if idx in self.full_to_reduced: + self._hand_actuated_joint_indices.append(self.full_to_reduced[idx]) + + self._left_hand_actuated_joint_indices = [] + for idx in full_left_hand_indices: + if idx in self.full_to_reduced: + self._left_hand_actuated_joint_indices.append(self.full_to_reduced[idx]) + + self._right_hand_actuated_joint_indices = [] + for idx in full_right_hand_indices: + if idx in self.full_to_reduced: + self._right_hand_actuated_joint_indices.append(self.full_to_reduced[idx]) + + # Cache indices for joint groups in reduced space + self._joint_group_indices = {} + for group_name in self.supplemental_info.joint_groups: + full_indices = full_robot_model.get_joint_group_indices(group_name) + reduced_indices = [] + for idx in full_indices: + if idx in self.full_to_reduced: + reduced_indices.append(self.full_to_reduced[idx]) + self._joint_group_indices[group_name] = sorted(set(reduced_indices)) + + # Initialize default body pose in reduced space + self.default_body_pose = self.full_to_reduced_configuration( + full_robot_model.default_body_pose + ) + + # Initialize initial body pose in reduced space + self.initial_body_pose = self.full_to_reduced_configuration( + full_robot_model.initial_body_pose + ) + + @property + def num_joints(self) -> int: + """Get the number of active joints in the reduced model.""" + return len(self.joint_names) + + @property + def joint_names(self) -> List[str]: + """Get the names of the active joints in the reduced model.""" + return [name for name in self.full_robot.joint_names if name not in self.fixed_joints] + + @classmethod + def from_fixed_groups( + cls, + full_robot_model: RobotModel, + fixed_group_names: List[str], + fixed_values: Optional[List[float]] = None, + ) -> "ReducedRobotModel": + """ + Create a reduced order robot model by fixing all joints in specified groups. + + :param full_robot_model: The original robot model + :param fixed_group_names: List of joint group names to fix + :param fixed_values: Optional list of values to fix the joints to. If None, uses the initial + joint positions (q0) from the full robot model. + :return: A ReducedRobotModel instance + """ + if full_robot_model.supplemental_info is None: + raise ValueError("supplemental_info must be provided to use this method") + + # Get all joints in the groups, including those from subgroups + fixed_joints = set() # Use a set to avoid duplicates + + for group_name in fixed_group_names: + if group_name not in full_robot_model.supplemental_info.joint_groups: + raise ValueError(f"Unknown joint group: {group_name}") + + group_info = full_robot_model.supplemental_info.joint_groups[group_name] + + # Add direct joints + fixed_joints.update(group_info["joints"]) + + # Add joints from subgroups + for subgroup_name in group_info["groups"]: + subgroup_joints = full_robot_model.get_joint_group_indices(subgroup_name) + fixed_joints.update([full_robot_model.joint_names[idx] for idx in subgroup_joints]) + + # Convert set back to list for compatibility with the original constructor + return cls(full_robot_model, list(fixed_joints), fixed_values) + + @classmethod + def from_fixed_group( + cls, + full_robot_model: RobotModel, + fixed_group_name: str, + fixed_values: Optional[List[float]] = None, + ) -> "ReducedRobotModel": + """ + Create a reduced order robot model by fixing all joints in a specified group. + This is a convenience method that calls from_fixed_groups with a single group. + + :param full_robot_model: The original robot model + :param fixed_group_name: Name of the joint group to fix + :param fixed_values: Optional list of values to fix the joints to. If None, uses the initial + joint positions (q0) from the full robot model. + :return: A ReducedRobotModel instance + """ + return cls.from_fixed_groups(full_robot_model, [fixed_group_name], fixed_values) + + @classmethod + def from_active_group( + cls, + full_robot_model: RobotModel, + active_group_name: str, + fixed_values: Optional[List[float]] = None, + ) -> "ReducedRobotModel": + """ + Create a reduced order robot model by fixing all joints EXCEPT those in the specified group. + This is a convenience method that calls from_active_groups with a single group. + + :param full_robot_model: The original robot model + :param active_group_name: Name of the joint group to keep active (all other joints will be fixed) + :param fixed_values: Optional list of values to fix the joints to. If None, uses the initial + joint positions (q0) from the full robot model. + :return: A ReducedRobotModel instance + """ + return cls.from_active_groups(full_robot_model, [active_group_name], fixed_values) + + @classmethod + def from_active_groups( + cls, + full_robot_model: RobotModel, + active_group_names: List[str], + fixed_values: Optional[List[float]] = None, + ) -> "ReducedRobotModel": + """ + Create a reduced order robot model by fixing all joints EXCEPT those in the specified groups. + This is useful when you want to keep multiple groups active and fix everything else. + + :param full_robot_model: The original robot model + :param active_group_names: List of joint group names to keep active (all other joints will be fixed) + :param fixed_values: Optional list of values to fix the joints to. If None, uses the initial + joint positions (q0) from the full robot model. + :return: A ReducedRobotModel instance + """ + if full_robot_model.supplemental_info is None: + raise ValueError("supplemental_info must be provided to use this method") + + # Get all joints in the active groups, including those from subgroups + active_joints = set() + + def add_group_joints(group_name: str): + if group_name not in full_robot_model.supplemental_info.joint_groups: + raise ValueError(f"Unknown joint group: {group_name}") + + group_info = full_robot_model.supplemental_info.joint_groups[group_name] + + # Add direct joints + if "joints" in group_info: + active_joints.update(group_info["joints"]) + + # Add joints from subgroups + if "groups" in group_info: + for subgroup_name in group_info["groups"]: + add_group_joints(subgroup_name) + + for group_name in active_group_names: + add_group_joints(group_name) + + # Get all joints from the model + all_joints = set(full_robot_model.joint_names) + + # The fixed joints are all joints minus the active joints + fixed_joints = list(all_joints - active_joints) + + return cls(full_robot_model, fixed_joints, fixed_values) + + def reduced_to_full_configuration(self, q_reduced: np.ndarray) -> np.ndarray: + """ + Convert a reduced configuration to the full configuration space. + + :param q_reduced: Configuration in reduced space + :return: Configuration in full space with fixed joints set to their fixed values + """ + if q_reduced.shape[0] != self.num_dofs: + raise ValueError( + f"Expected q_reduced of length {self.num_dofs}, got {q_reduced.shape[0]} instead" + ) + + q_full = np.zeros(self.full_robot.num_dofs) + + # Set active joints + for reduced_idx, full_idx in enumerate(self.reduced_to_full): + q_full[full_idx] = q_reduced[reduced_idx] + + # Set fixed joints + for joint_name, value in zip(self.fixed_joints, self.fixed_values): + full_idx = self.full_robot.dof_index(joint_name) + q_full[full_idx] = value + + return q_full + + def full_to_reduced_configuration(self, q_full: np.ndarray) -> np.ndarray: + """ + Convert a full configuration to the reduced configuration space. + + :param q_full: Configuration in full space + :return: Configuration in reduced space + """ + if q_full.shape[0] != self.full_robot.num_dofs: + raise ValueError( + f"Expected q_full of length {self.full_robot.num_dofs}, got {q_full.shape[0]} instead" + ) + + q_reduced = np.zeros(self.num_dofs) + + # Copy active joints + for reduced_idx, full_idx in enumerate(self.reduced_to_full): + q_reduced[reduced_idx] = q_full[full_idx] + + return q_reduced + + def cache_forward_kinematics(self, q_reduced: np.ndarray, auto_clip=True) -> None: + """ + Perform forward kinematics using the reduced configuration. + + :param q_reduced: Configuration in reduced space + """ + # First update the full robot's forward kinematics + q_full = self.reduced_to_full_configuration(q_reduced) + self.full_robot.cache_forward_kinematics(q_full, auto_clip) + + # Then update the reduced model's forward kinematics + pin.framesForwardKinematics( + self.pinocchio_wrapper.model, self.pinocchio_wrapper.data, q_reduced + ) + + def clip_configuration(self, q_reduced: np.ndarray, margin: float = 1e-6) -> np.ndarray: + """ + Clip the reduced configuration to stay within joint limits with a small tolerance. + + :param q_reduced: Configuration to clip + :param margin: Tolerance to keep away from joint limits + :return: Clipped configuration + """ + q_full = self.reduced_to_full_configuration(q_reduced) + q_full_clipped = self.full_robot.clip_configuration(q_full, margin) + return self.full_to_reduced_configuration(q_full_clipped) + + def reset_forward_kinematics(self): + """ + Reset the forward kinematics to the initial configuration. + """ + # Reset full robot's forward kinematics + self.full_robot.reset_forward_kinematics() + # Reset reduced model's forward kinematics + self.cache_forward_kinematics(self.q_zero) diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/supplemental_info/__init__.py b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/supplemental_info/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4f71df6ca48c850d3368a63b4337e94c615b0529 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/supplemental_info/__init__.py @@ -0,0 +1,7 @@ +"""Robot supplemental info: non-URDF metadata (joint groups, limits, name maps).""" + +from gear_sonic.data.robot_model.supplemental_info.robot_supplemental_info import ( + RobotSupplementalInfo, +) + +__all__ = ["RobotSupplementalInfo"] diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/supplemental_info/__pycache__/__init__.cpython-310.pyc b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/supplemental_info/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2b90977acb3dd2fe0de6263460e6378679701a74 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/supplemental_info/__pycache__/__init__.cpython-310.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/supplemental_info/__pycache__/robot_supplemental_info.cpython-310.pyc b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/supplemental_info/__pycache__/robot_supplemental_info.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2819b85bb4df0d50b883c1de70e6ab3610129f8b Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/supplemental_info/__pycache__/robot_supplemental_info.cpython-310.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/supplemental_info/g1/__init__.py b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/supplemental_info/g1/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/supplemental_info/g1/__pycache__/__init__.cpython-310.pyc b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/supplemental_info/g1/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..742740bb86446d390fc4212401e7b4220ab9a805 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/supplemental_info/g1/__pycache__/__init__.cpython-310.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/supplemental_info/g1/__pycache__/g1_supplemental_info.cpython-310.pyc b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/supplemental_info/g1/__pycache__/g1_supplemental_info.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e87956141d2fc7705926f656430f9e415112e387 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/supplemental_info/g1/__pycache__/g1_supplemental_info.cpython-310.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/supplemental_info/g1/g1_supplemental_info.py b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/supplemental_info/g1/g1_supplemental_info.py new file mode 100644 index 0000000000000000000000000000000000000000..ba4dd2ea6c60a0b9c3460195034c6a176d8091f6 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/supplemental_info/g1/g1_supplemental_info.py @@ -0,0 +1,334 @@ +"""G1-specific supplemental info: actuated joints, limits, and default poses by waist/elbow config.""" + +from dataclasses import dataclass +from enum import Enum + +import numpy as np + +from gear_sonic.data.robot_model.supplemental_info.robot_supplemental_info import ( + RobotSupplementalInfo, +) + + +class WaistLocation(Enum): + """Enum for waist location configuration.""" + + LOWER_BODY = "lower_body" + UPPER_BODY = "upper_body" + LOWER_AND_UPPER_BODY = "lower_and_upper_body" + + +class ElbowPose(Enum): + """Enum for elbow pose configuration.""" + + LOW = "low" + HIGH = "high" + + +@dataclass +class G1SupplementalInfo(RobotSupplementalInfo): + """ + Supplemental information for the G1 robot. + + Args: + waist_location: Where to place waist joints in the joint groups + elbow_pose: Which elbow pose configuration to use for default joint positions + """ + + def __init__( + self, + waist_location: WaistLocation = WaistLocation.LOWER_BODY, + elbow_pose: ElbowPose = ElbowPose.LOW, + ): + name = "G1_G1ThreeFinger" + + # Define all actuated joints + body_actuated_joints = [ + # Left leg + "left_hip_pitch_joint", + "left_hip_roll_joint", + "left_hip_yaw_joint", + "left_knee_joint", + "left_ankle_pitch_joint", + "left_ankle_roll_joint", + # Right leg + "right_hip_pitch_joint", + "right_hip_roll_joint", + "right_hip_yaw_joint", + "right_knee_joint", + "right_ankle_pitch_joint", + "right_ankle_roll_joint", + # Waist + "waist_yaw_joint", + "waist_roll_joint", + "waist_pitch_joint", + # Left arm + "left_shoulder_pitch_joint", + "left_shoulder_roll_joint", + "left_shoulder_yaw_joint", + "left_elbow_joint", + "left_wrist_roll_joint", + "left_wrist_pitch_joint", + "left_wrist_yaw_joint", + # Right arm + "right_shoulder_pitch_joint", + "right_shoulder_roll_joint", + "right_shoulder_yaw_joint", + "right_elbow_joint", + "right_wrist_roll_joint", + "right_wrist_pitch_joint", + "right_wrist_yaw_joint", + ] + + left_hand_actuated_joints = [ + # Left hand + "left_hand_thumb_0_joint", + "left_hand_thumb_1_joint", + "left_hand_thumb_2_joint", + "left_hand_index_0_joint", + "left_hand_index_1_joint", + "left_hand_middle_0_joint", + "left_hand_middle_1_joint", + ] + + right_hand_actuated_joints = [ + # Right hand + "right_hand_thumb_0_joint", + "right_hand_thumb_1_joint", + "right_hand_thumb_2_joint", + "right_hand_index_0_joint", + "right_hand_index_1_joint", + "right_hand_middle_0_joint", + "right_hand_middle_1_joint", + ] + + # Define joint limits from URDF + joint_limits = { + # Left leg + "left_hip_pitch_joint": [-2.5307, 2.8798], + "left_hip_roll_joint": [-0.5236, 2.9671], + "left_hip_yaw_joint": [-2.7576, 2.7576], + "left_knee_joint": [-0.087267, 2.8798], + "left_ankle_pitch_joint": [-0.87267, 0.5236], + "left_ankle_roll_joint": [-0.2618, 0.2618], + # Right leg + "right_hip_pitch_joint": [-2.5307, 2.8798], + "right_hip_roll_joint": [-2.9671, 0.5236], + "right_hip_yaw_joint": [-2.7576, 2.7576], + "right_knee_joint": [-0.087267, 2.8798], + "right_ankle_pitch_joint": [-0.87267, 0.5236], + "right_ankle_roll_joint": [-0.2618, 0.2618], + # Waist + "waist_yaw_joint": [-2.618, 2.618], + "waist_roll_joint": [-0.52, 0.52], + "waist_pitch_joint": [-0.52, 0.52], + # Left arm + "left_shoulder_pitch_joint": [-3.0892, 2.6704], + "left_shoulder_roll_joint": [0.19, 2.2515], + "left_shoulder_yaw_joint": [-2.618, 2.618], + "left_elbow_joint": [-1.0472, 2.0944], + "left_wrist_roll_joint": [-1.972222054, 1.972222054], + "left_wrist_pitch_joint": [-1.614429558, 1.614429558], + "left_wrist_yaw_joint": [-1.614429558, 1.614429558], + # Right arm + "right_shoulder_pitch_joint": [-3.0892, 2.6704], + "right_shoulder_roll_joint": [-2.2515, -0.19], + "right_shoulder_yaw_joint": [-2.618, 2.618], + "right_elbow_joint": [-1.0472, 2.0944], + "right_wrist_roll_joint": [-1.972222054, 1.972222054], + "right_wrist_pitch_joint": [-1.614429558, 1.614429558], + "right_wrist_yaw_joint": [-1.614429558, 1.614429558], + # Left hand + "left_hand_thumb_0_joint": [-1.04719755, 1.04719755], + "left_hand_thumb_1_joint": [-0.72431163, 1.04719755], + "left_hand_thumb_2_joint": [0, 1.74532925], + "left_hand_index_0_joint": [-1.57079632, 0], + "left_hand_index_1_joint": [-1.74532925, 0], + "left_hand_middle_0_joint": [-1.57079632, 0], + "left_hand_middle_1_joint": [-1.74532925, 0], + # Right hand + "right_hand_thumb_0_joint": [-1.04719755, 1.04719755], + "right_hand_thumb_1_joint": [-0.72431163, 1.04719755], + "right_hand_thumb_2_joint": [0, 1.74532925], + "right_hand_index_0_joint": [-1.57079632, 0], + "right_hand_index_1_joint": [-1.74532925, 0], + "right_hand_middle_0_joint": [-1.57079632, 0], + "right_hand_middle_1_joint": [-1.74532925, 0], + } + + # Define joint groups + joint_groups = { + # Body groups + "waist": { + "joints": ["waist_yaw_joint", "waist_roll_joint", "waist_pitch_joint"], + "groups": [], + }, + # Leg groups + "left_leg": { + "joints": [ + "left_hip_pitch_joint", + "left_hip_roll_joint", + "left_hip_yaw_joint", + "left_knee_joint", + "left_ankle_pitch_joint", + "left_ankle_roll_joint", + ], + "groups": [], + }, + "right_leg": { + "joints": [ + "right_hip_pitch_joint", + "right_hip_roll_joint", + "right_hip_yaw_joint", + "right_knee_joint", + "right_ankle_pitch_joint", + "right_ankle_roll_joint", + ], + "groups": [], + }, + "legs": {"joints": [], "groups": ["left_leg", "right_leg"]}, + # Arm groups + "left_arm": { + "joints": [ + "left_shoulder_pitch_joint", + "left_shoulder_roll_joint", + "left_shoulder_yaw_joint", + "left_elbow_joint", + "left_wrist_roll_joint", + "left_wrist_pitch_joint", + "left_wrist_yaw_joint", + ], + "groups": [], + }, + "right_arm": { + "joints": [ + "right_shoulder_pitch_joint", + "right_shoulder_roll_joint", + "right_shoulder_yaw_joint", + "right_elbow_joint", + "right_wrist_roll_joint", + "right_wrist_pitch_joint", + "right_wrist_yaw_joint", + ], + "groups": [], + }, + "arms": {"joints": [], "groups": ["left_arm", "right_arm"]}, + # Hand groups + "left_hand": { + "joints": [ + "left_hand_index_0_joint", + "left_hand_index_1_joint", + "left_hand_middle_0_joint", + "left_hand_middle_1_joint", + "left_hand_thumb_0_joint", + "left_hand_thumb_1_joint", + "left_hand_thumb_2_joint", + ], + "groups": [], + }, + "right_hand": { + "joints": [ + "right_hand_index_0_joint", + "right_hand_index_1_joint", + "right_hand_middle_0_joint", + "right_hand_middle_1_joint", + "right_hand_thumb_0_joint", + "right_hand_thumb_1_joint", + "right_hand_thumb_2_joint", + ], + "groups": [], + }, + "hands": {"joints": [], "groups": ["left_hand", "right_hand"]}, + # Full body groups + "lower_body": {"joints": [], "groups": ["waist", "legs"]}, + "upper_body_no_hands": {"joints": [], "groups": ["arms"]}, + "body": {"joints": [], "groups": ["lower_body", "upper_body_no_hands"]}, + "upper_body": {"joints": [], "groups": ["upper_body_no_hands", "hands"]}, + } + + # Define joint name mapping from generic types to robot-specific names + joint_name_mapping = { + # Waist joints + "waist_pitch": "waist_pitch_joint", + "waist_roll": "waist_roll_joint", + "waist_yaw": "waist_yaw_joint", + # Shoulder joints + "shoulder_pitch": { + "left": "left_shoulder_pitch_joint", + "right": "right_shoulder_pitch_joint", + }, + "shoulder_roll": { + "left": "left_shoulder_roll_joint", + "right": "right_shoulder_roll_joint", + }, + "shoulder_yaw": { + "left": "left_shoulder_yaw_joint", + "right": "right_shoulder_yaw_joint", + }, + # Elbow joints + "elbow_pitch": {"left": "left_elbow_joint", "right": "right_elbow_joint"}, + # Wrist joints + "wrist_pitch": {"left": "left_wrist_pitch_joint", "right": "right_wrist_pitch_joint"}, + "wrist_roll": {"left": "left_wrist_roll_joint", "right": "right_wrist_roll_joint"}, + "wrist_yaw": {"left": "left_wrist_yaw_joint", "right": "right_wrist_yaw_joint"}, + } + + root_frame_name = "pelvis" + + hand_frame_names = {"left": "left_wrist_yaw_link", "right": "right_wrist_yaw_link"} + + calibration_joint_q = {"elbow_pitch": {"left": 0.0, "right": 0.0}} + + # 90° Y-axis rotation: aligns hand-tracking frame (palm-forward) to robot wrist frame + hand_rotation_correction = np.array([[0, 0, 1], [0, 1, 0], [-1, 0, 0]]) + + # HIGH: arms raised with elbows bent (hands near shoulder height) + # LOW: arms relaxed at sides with slight shoulder roll + if elbow_pose == ElbowPose.HIGH: + default_joint_q = { + "shoulder_roll": {"left": 0.5, "right": -0.5}, + "shoulder_pitch": {"left": -0.2, "right": -0.2}, + "shoulder_yaw": {"left": -0.5, "right": 0.5}, + "wrist_roll": {"left": -0.5, "right": 0.5}, + "wrist_yaw": {"left": 0.5, "right": -0.5}, + "wrist_pitch": {"left": -0.2, "right": -0.2}, + } + else: # ElbowPose.LOW + default_joint_q = { + "shoulder_roll": {"left": 0.2, "right": -0.2}, + } + + teleop_upper_body_motion_scale = 1.0 + + # Configure joint groups based on waist location + modified_joint_groups = joint_groups.copy() + if waist_location == WaistLocation.UPPER_BODY: + # Move waist from lower_body to upper_body_no_hands + modified_joint_groups["lower_body"] = {"joints": [], "groups": ["legs"]} + modified_joint_groups["upper_body_no_hands"] = { + "joints": [], + "groups": ["arms", "waist"], + } + elif waist_location == WaistLocation.LOWER_AND_UPPER_BODY: + # Add waist to upper_body_no_hands while keeping it in lower_body + modified_joint_groups["upper_body_no_hands"] = { + "joints": [], + "groups": ["arms", "waist"], + } + # For LOWER_BODY, keep default joint_groups as is + + super().__init__( + name=name, + body_actuated_joints=body_actuated_joints, + left_hand_actuated_joints=left_hand_actuated_joints, + right_hand_actuated_joints=right_hand_actuated_joints, + joint_limits=joint_limits, + joint_groups=modified_joint_groups, + root_frame_name=root_frame_name, + hand_frame_names=hand_frame_names, + calibration_joint_q=calibration_joint_q, + joint_name_mapping=joint_name_mapping, + hand_rotation_correction=hand_rotation_correction, + default_joint_q=default_joint_q, + teleop_upper_body_motion_scale=teleop_upper_body_motion_scale, + ) diff --git a/GR00T-WholeBodyControl/gear_sonic/data/robot_model/supplemental_info/robot_supplemental_info.py b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/supplemental_info/robot_supplemental_info.py new file mode 100644 index 0000000000000000000000000000000000000000..8f311fdbf526d41306a1faae517a63c1f7b3d8c9 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/data/robot_model/supplemental_info/robot_supplemental_info.py @@ -0,0 +1,93 @@ +"""Base dataclass for robot-specific config not found in URDF (joint groups, limits, names).""" + +from dataclasses import dataclass +from typing import Dict, List, Mapping, Union + +import numpy as np + + +@dataclass +class RobotSupplementalInfo: + """ + Base class for robot-specific information that is not easily extractable from URDF. + This includes information about actuated joints, joint hierarchies, etc. + """ + + name: str + + # List of body actuated joint names (excluding hands) + body_actuated_joints: List[str] + + # List of left hand actuated joint names + left_hand_actuated_joints: List[str] + + # List of right hand actuated joint names + right_hand_actuated_joints: List[str] + + # Dictionary of joint groups, where each group is a dictionary with: + # - "joints": list of joint names + # - "groups": list of subgroup names (optional) + # Example: { + # "right_arm": { + # "joints": ["right_shoulder_pitch_joint", "right_shoulder_roll_joint", "right_elbow_joint"], + # "groups": [] + # }, + # "left_arm": { + # "joints": ["left_shoulder_pitch_joint", "left_shoulder_roll_joint", "left_elbow_joint"], + # "groups": [] + # }, + # "upper_body": { + # "joints": ["torso_pitch_joint", "torso_yaw_joint", "torso_roll_joint"], + # "groups": ["right_arm", "left_arm"] + # } + # } + joint_groups: Dict[str, Dict[str, List[str]]] + + # Name of the root frame + root_frame_name: str + + # Dictionary of hand frame names + # Example: { + # "left": "left_hand_frame", + # "right": "right_hand_frame" + # } + hand_frame_names: Dict[str, str] + + # Dictionary of joint limits + # Example: { + # "left_shoulder_pitch_joint": [-np.pi / 2, np.pi / 2], + # "right_shoulder_pitch_joint": [-np.pi / 2, np.pi / 2] + # } + joint_limits: Dict[str, List[float]] + + # Dictionary of calibration joint positions in radians. + # Structure mirrors default_joint_q for any joints used in calibration. + # Example: { + # "elbow_pitch": {"left": -np.pi / 2, "right": -np.pi / 2} + # } + calibration_joint_q: Mapping[str, Union[float, Mapping[str, float]]] + + # Dictionary of joint name mapping from generic types to robot-specific names + # Example: { + # "waist_pitch": "waist_pitch_joint", + # "shoulder_pitch": { + # "left": "left_shoulder_pitch_joint", + # "right": "right_shoulder_pitch_joint" + # }, + # "elbow_pitch": { + # "left": "left_elbow_pitch_joint", + # "right": "right_elbow_pitch_joint" + # } + # } + joint_name_mapping: Mapping[str, Union[str, Mapping[str, str]]] + + # Maps from generic joint names to robot-specific joint values + # Example: { + # "waist_roll": 0.2, + # "elbow_pitch": {"left": 1.0, "right": 1.0} + # } + default_joint_q: Mapping[str, Union[float, Mapping[str, float]]] + + hand_rotation_correction: np.ndarray + + teleop_upper_body_motion_scale: float diff --git a/GR00T-WholeBodyControl/gear_sonic/eval_agent_trl.py b/GR00T-WholeBodyControl/gear_sonic/eval_agent_trl.py new file mode 100644 index 0000000000000000000000000000000000000000..927e419c48db7d5d39d7219c91ac637225e0f110 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/eval_agent_trl.py @@ -0,0 +1,670 @@ +#!/usr/bin/env python3 +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +try: + import isaaclab # noqa: F401 +except ImportError: + print( + "\n" + "ERROR: Isaac Lab is required for evaluation but not installed.\n" + "\n" + "Isaac Lab is not a pip dependency — it must be installed separately.\n" + "Follow the official guide:\n" + " https://isaac-sim.github.io/IsaacLab/main/source/setup/installation/index.html\n" + "\n" + "After installing, activate the Isaac Lab conda/venv environment\n" + "before running this script.\n" + ) + import sys + sys.exit(1) + +import filelock # noqa: I001 +import json +import os +import shutil +import subprocess +import sys + +sys.path.append(os.getcwd()) +import logging +from pathlib import Path + +import easydict +import hydra +from hydra import utils +from hydra.core import hydra_config +from loguru import logger +import omegaconf +import yaml + +from gear_sonic import train_agent_trl +from gear_sonic.trl.utils import common as trl_utils_common +from gear_sonic.trl.utils import scheduler +from gear_sonic.utils import common as rl_utils_common +from gear_sonic.utils import config_utils, obs_utils + +config_utils.register_rl_resolvers() + + +@hydra.main(config_path="config", config_name="base_eval") +def main(override_config: omegaconf.OmegaConf): + + hydra_log_path = os.path.join(hydra_config.HydraConfig.get().runtime.output_dir, "eval.log") + logger.remove() + logger.add(hydra_log_path, level="DEBUG") + + # Get log level from LOGURU_LEVEL environment variable or use INFO as default + console_log_level = os.environ.get("LOGURU_LEVEL", "INFO").upper() + logger.add(sys.stdout, level=console_log_level, colorize=True) + + from gear_sonic.utils import logging as utils_logging + + logging.basicConfig(level=logging.DEBUG) + logging.getLogger().addHandler(utils_logging.HydraLoggerBridge()) + + os.chdir(hydra.utils.get_original_cwd()) + + if override_config.checkpoint is not None: + has_config = True + checkpoint = Path(override_config.checkpoint) + config_path = checkpoint.parent / "config.yaml" + if not config_path.exists(): + config_path = checkpoint.parent.parent / "config.yaml" + if not config_path.exists(): + has_config = False + logger.error(f"Could not find config path: {config_path}") + + if has_config: + logger.info(f"Loading training config file from {config_path}") + with open(config_path) as file: + raw = file.read() + # Backward compatibility: rewrite internal repo module paths to release repo paths + raw = raw.replace("groot.rl.trl.", "gear_sonic.trl.") + raw = raw.replace("groot.rl.envs.", "gear_sonic.envs.") + raw = raw.replace("groot.rl.utils.", "gear_sonic.utils.") + raw = raw.replace("groot.rl.agents.modules.modules.", "gear_sonic.trl.modules.base_module.") + raw = raw.replace("groot.rl.agents.", "gear_sonic.trl.") + raw = raw.replace("groot/rl/data/", "gear_sonic/data/") + raw = raw.replace("assets/bm/unitree_description/", "assets/robot_description/") + raw = raw.replace("1215_bones_seed_filtered", "bones_seed_smpl") + import io + train_config = omegaconf.OmegaConf.load(io.StringIO(raw)) + + if train_config.eval_overrides is not None: + train_config = omegaconf.OmegaConf.merge(train_config, train_config.eval_overrides) + + config = omegaconf.OmegaConf.merge(train_config, override_config) + else: + config = override_config + + config.experiment_dir = checkpoint.parent + elif override_config.eval_overrides is not None: + config = override_config.copy() + eval_overrides = omegaconf.OmegaConf.to_container(config.eval_overrides, resolve=True) + for arg in sys.argv[1:]: + if not arg.startswith("+"): + key = arg.split("=")[0] + if key in eval_overrides: + del eval_overrides[key] + config.eval_overrides = omegaconf.OmegaConf.create(eval_overrides) + config = omegaconf.OmegaConf.merge(config, eval_overrides) + else: + config = override_config + + meta_path = Path(config.experiment_dir) / "meta.yaml" + if meta_path.exists(): + meta = yaml.safe_load(open(meta_path)) # noqa: SIM115 + if config.get("wandb", None) is not None and meta.get("wandb_run"): + config.wandb.wandb_id = meta["wandb_run"] + print(f"resume wandb from run: {config.wandb.wandb_id}") # noqa: T201 + + with omegaconf.open_dict(config): + for event in config.manager_env.config.get("train_only_events", []): + if event in config.manager_env.events: + config.manager_env.events.pop(event) + remove_schedule_keys = [] + for key in config.trainer.get("schedule_dict", {}): + if event in key: + remove_schedule_keys.append(key) + for key in remove_schedule_keys: + config.trainer.schedule_dict.pop(key) + + for termination in config.manager_env.config.get("train_only_terminations", []): + if termination in config.manager_env.terminations: + config.manager_env.terminations.pop(termination) + + use_encoder = config.get("use_encoder", None) + if use_encoder is not None: + encoder_sample_probs = config.manager_env.commands.motion.encoder_sample_probs + if encoder_sample_probs is not None: + for encoder in encoder_sample_probs: + if encoder != use_encoder: + encoder_sample_probs[encoder] = 0.0 + print(f"Using encoder: {use_encoder}") # noqa: T201 + print(f"Encoder sample probs: {encoder_sample_probs}") # noqa: T201 + + simulator_type = "IsaacSim" + env_config = config.manager_env + + import datetime as dt + + import accelerate + import torch # noqa: E402, RUF100 + + kwargs = accelerate.InitProcessGroupKwargs(timeout=dt.timedelta(seconds=6000)) + accelerator = accelerate.Accelerator(kwargs_handlers=[kwargs]) + + device = str(accelerator.device) + if accelerator.device.type == "cuda": + try: + torch.cuda.set_device(accelerator.local_process_index) + except Exception: # noqa: S110, BLE001 + pass + + device = str(accelerator.device) + config.multi_gpu = accelerator.num_processes > 1 + if config.multi_gpu: + config.global_rank = accelerator.process_index + config.seed += accelerator.process_index + config.algo.config.global_rank = accelerator.process_index + config.algo.config.world_size = accelerator.num_processes + rl_utils_common.seeding(config.seed) + + def _pick_display_gpu_index(default_idx: int = 0) -> int: + try: + out = subprocess.check_output( + ["nvidia-smi", "--query-gpu=index,display_active,name", "--format=csv,noheader"], + text=True, + ) + for line in out.strip().splitlines(): + parts = [p.strip() for p in line.split(",")] + if len(parts) >= 2: + idx, active = int(parts[0]), parts[1].lower() + if active.startswith("enabled") or active.startswith("on"): + return idx + except Exception: # noqa: S110, BLE001 + pass + return default_idx + + render_gpu_idx = _pick_display_gpu_index(default_idx=0) + + if simulator_type == "IsaacSim": + try: + with open("./rl/simulator/isaacsim/.isaacsim_version", encoding="utf-8") as f: + DEFAULT_ISAACSIM_VERSION = f.read().strip() + except FileNotFoundError: + DEFAULT_ISAACSIM_VERSION = "4.5" + + if DEFAULT_ISAACSIM_VERSION == "4.5": + from isaaclab.app import AppLauncher + elif DEFAULT_ISAACSIM_VERSION == "4.2": + logger.warning("Using IsaacSim 4.2, replacing isaaclab with omni.isaac.lab") + from omni.isaac.lab.app import AppLauncher # 4.2 + import argparse + + parser = argparse.ArgumentParser(description="Evaluate an RL agent with TRL.") + AppLauncher.add_app_launcher_args(parser) + + args_cli, hydra_args = parser.parse_known_args() + sys.argv = [sys.argv[0]] + hydra_args # noqa: RUF005 + args_cli.num_envs = config.num_envs + args_cli.seed = config.seed + args_cli.env_spacing = env_config.config.env_spacing + args_cli.output_dir = config.output_dir + args_cli.enable_cameras = env_config.config.get( + "render_results", False + ) or env_config.config.get("enable_cameras", False) + + args_cli.headless = config.headless + args_cli.multi_gpu = config.multi_gpu + args_cli.distributed = config.multi_gpu + args_cli.device = device + + base_kit_args = ( + "--/log/level=error --/log/fileLogLevel=error --/log/outputStreamLevel=error" + ) + if args_cli.headless: + args_cli.kit_args = base_kit_args + " --no-window" + else: + args_cli.kit_args = base_kit_args + f" --/renderer/activeGpu={render_gpu_idx}" + + # Allow air-gapped machines to use an experience file with online + # extension registries disabled, while preserving the default behavior. + offline_experience = os.environ.get("ISAACLAB_EXPERIENCE") + if offline_experience: + args_cli.experience = offline_experience + + _lock_path = "/tmp/isaaclab_app_launcher.lock" # noqa: S108 + with filelock.FileLock(_lock_path): + app_launcher = AppLauncher(args_cli) + simulation_app = app_launcher.app # noqa: F841 + + import torch + + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + torch.backends.cudnn.deterministic = False + torch.backends.cudnn.benchmark = False + + unresolved_conf = omegaconf.OmegaConf.to_container(config, resolve=False) # noqa: F841 + os.chdir(hydra.utils.get_original_cwd()) + + ckpt_num = config.checkpoint.split("/")[-1].split("_")[-1].split(".")[0] + + if env_config.config.get("save_rendering_dir", None) is None: + env_config.config.save_rendering_dir = str( + checkpoint.parent / "renderings" / f"ckpt_{ckpt_num}" + ) + + metrics_file = config.get("metrics_file", None) + if metrics_file is not None: + metrics_file = Path(metrics_file) + assert metrics_file.exists(), f"Metrics file {metrics_file} does not exist" + if metrics_file.exists(): + metrics = json.load(open(metrics_file)) # noqa: SIM115 + all_dict = metrics["eval/all_metrics_dict"] + + # Check if this is grab evaluation (has success_lift) + has_obj_metrics = "obj_pos_error" in all_dict + if "success_lift" in all_dict: + # Grab evaluation: prioritize failed grasps (not lifted) and terminated trajectories + motion_keys = all_dict["motion_keys"] + terminated = all_dict["terminated"] + success_lift = all_dict["success_lift"] + progress = all_dict.get("progress", [1.0] * len(motion_keys)) + obj_pos_errors = all_dict.get("obj_pos_error", [0.0] * len(motion_keys)) + + pairs = [] + for i in range(len(motion_keys)): + term = bool(terminated[i]) if i < len(terminated) else False + lifted = bool(success_lift[i]) if i < len(success_lift) else False + prog = progress[i] if i < len(progress) else 1.0 + obj_err = obj_pos_errors[i] if i < len(obj_pos_errors) else 0.0 + priority = 0 if not lifted else (1 if term else 2) + pairs.append((motion_keys[i], term, lifted, prog, obj_err, priority)) + + pairs_sorted = sorted(pairs, key=lambda x: (x[5], x[3])) + if len(pairs_sorted) > config.num_envs: + pairs_sorted = pairs_sorted[: config.num_envs] + + render_info = [] + for pair in pairs_sorted: + motion_key, term, lifted, prog, obj_err, _ = pair + status = "FAILED" if not lifted else ("TERMINATED" if term else "SUCCESS") + info = [ + f"{motion_key}", + f"lifted: {lifted}", + f"progress: {prog:.3f}", + f"status: {status}", + ] + if has_obj_metrics: + info.append(f"obj_pos_err: {obj_err:.4f}m") + render_info.append(tuple(info)) + + filter_keys = [pair[0] for pair in pairs_sorted] + + with omegaconf.open_dict(env_config.config): + env_config.config.render_info = render_info + env_config.config.max_render_envs = len(render_info) + with omegaconf.open_dict(env_config.commands.motion): + env_config.commands.motion.filter_motion_keys = filter_keys + if "motion_lib_cfg" in env_config.commands.motion: + env_config.commands.motion.motion_lib_cfg.filter_motion_keys = filter_keys + else: + # Imitation evaluation: use MPJPE-based sorting + obj_pos_errors = all_dict.get("obj_pos_error", None) + success_pair = [ + ( + all_dict["motion_keys"][i], + all_dict["mpjpe_l"][i], + all_dict["mpjpe_g"][i], + True, + obj_pos_errors[i] if obj_pos_errors is not None else 0.0, + ) + for i in range(len(all_dict["motion_keys"])) + if not all_dict["terminated"][i] + ] + render_sort_by = config.get("render_sort_by", "mpjpe_l") + sort_idx = 4 if render_sort_by == "obj_pos_error" else 1 + success_pair_sorted = sorted(success_pair, key=lambda x: x[sort_idx], reverse=True) + failed_pair = [ + ( + all_dict["motion_keys"][i], + all_dict["mpjpe_l"][i], + all_dict["mpjpe_g"][i], + False, + obj_pos_errors[i] if obj_pos_errors is not None else 0.0, + ) + for i in range(len(all_dict["motion_keys"])) + if all_dict["terminated"][i] + ] + failed_pair_sorted = sorted(failed_pair, key=lambda x: x[sort_idx], reverse=True) + all_pair = failed_pair_sorted + success_pair_sorted + if len(all_pair) > config.num_envs: + all_pair = all_pair[: config.num_envs] + render_info = [] + for pair in all_pair: + info = [ + f"{pair[0]}", + f"mpjpe_l: {pair[1]:.2f}", + f"mpjpe_g: {pair[2]:.2f}", + f"success: {pair[3]}", + ] + if has_obj_metrics: + info.append(f"obj_pos_err: {pair[4]:.4f}m") + render_info.append(tuple(info)) + with omegaconf.open_dict(env_config.config): + env_config.config.render_info = render_info + env_config.config.max_render_envs = len(all_pair) + filter_keys = [pair[0] for pair in all_pair] + with omegaconf.open_dict(env_config.commands.motion): + env_config.commands.motion.filter_motion_keys = filter_keys + if "motion_lib_cfg" in env_config.commands.motion: + env_config.commands.motion.motion_lib_cfg.filter_motion_keys = filter_keys + + env = train_agent_trl.create_manager_env(config, device, args_cli) + + module_dim_dict = getattr(config.algo.config, "module_dim", {}) + policy_backbone_kwargs = {} + critic_backbone_kwargs = {} + env.config["obs"]["obs_dims"]["actor_obs"] = env.env.observation_space["policy"].shape[-1] + env.config["obs"]["obs_dims"]["critic_obs"] = env.env.observation_space["critic"].shape[-1] + env.config["robot"]["algo_obs_dim_dict"]["actor_obs"] = env.env.observation_space[ + "policy" + ].shape[-1] + env.config["robot"]["algo_obs_dim_dict"]["critic_obs"] = env.env.observation_space[ + "critic" + ].shape[-1] + example_obs = env.reset(flatten_dict_obs=False) + for key in env.env.observation_space: + if key not in ["policy", "critic"]: + group_obs_dims, group_obs_names, group_obs_total_dim = ( + obs_utils.get_group_term_obs_shape(example_obs, key) + ) + env.config["obs"]["group_obs_dims"][key] = group_obs_dims + env.config["obs"]["group_obs_names"][key] = group_obs_names + env.config["obs"]["obs_dims"][key] = group_obs_total_dim + env.config["robot"]["algo_obs_dim_dict"][key] = group_obs_total_dim + + meta_action_dim = env.config.get("meta_action_dim", None) + if meta_action_dim is not None and meta_action_dim > 0: + env.config["robot"]["actions_dim"] = meta_action_dim + else: + env.config["robot"]["actions_dim"] = env.env.action_space.shape[-1] + + policy = trl_utils_common.custom_instantiate( + config.algo.config.actor, + env_config=env.config, + algo_config=config.algo.config, + module_dim_dict=module_dim_dict, + backbone_kwargs=policy_backbone_kwargs, + _resolve=False, + ).to(device) + + if not getattr(config.algo.config, "distill_only", False): + value_model = trl_utils_common.custom_instantiate( + config.algo.config.critic, + env_config=env.config, + algo_config=config.algo.config, + module_dim_dict=module_dim_dict, + backbone_kwargs=critic_backbone_kwargs, + _resolve=False, + ).to(device) + + accelerator.wait_for_everyone() + + args = easydict.EasyDict() + args.is_main_process = accelerator.is_main_process + args.global_rank = accelerator.process_index + args.world_size = accelerator.num_processes + state = easydict.EasyDict() + + from gear_sonic.trl.trainer import ppo_trainer + + model = ppo_trainer.PolicyAndValueWrapper(policy, value_model) + + checkpoint_path = str(config.checkpoint) + logger.info(f"Loading checkpoint from {checkpoint_path}") + checkpoint = torch.load(checkpoint_path, map_location=accelerator.device, weights_only=False) + + # Load policy state dict with backward compatibility for std/log_std + if "actor_model_state_dict" in checkpoint: + state_dict = checkpoint["actor_model_state_dict"] + elif "policy_state_dict" in checkpoint: + state_dict = checkpoint["policy_state_dict"] + else: + state_dict = None + + if state_dict is not None: + model_uses_std = "std" in model.policy.state_dict() + checkpoint_has_std = "std" in state_dict + checkpoint_has_log_std = "log_std" in state_dict + + logger.info(f"Model parameterization: {'std' if model_uses_std else 'log_std'}") + logger.info( + f"Checkpoint parameterization: {'std' if checkpoint_has_std else 'log_std' if checkpoint_has_log_std else 'unknown'}" # noqa: E501 + ) + + if model_uses_std and checkpoint_has_log_std and not checkpoint_has_std: + logger.info("Transforming 'log_std' -> 'std' (applying exp) for backward compatibility") + state_dict["std"] = torch.exp(state_dict.pop("log_std")) + elif not model_uses_std and checkpoint_has_std and not checkpoint_has_log_std: + logger.info("Transforming 'std' -> 'log_std' (applying log) for backward compatibility") + state_dict["log_std"] = torch.log(state_dict.pop("std")) + + model.policy.load_state_dict(state_dict) + logger.info("Successfully loaded policy state dict") + + state.global_step = checkpoint["state"].global_step + + schedule_wrapper = easydict.EasyDict(env=env, model=model) + if "schedule_dict" in config.trainer: + scheduled_params_dict = scheduler.update_scheduled_params( # noqa: F841 + schedule_wrapper, config.trainer.schedule_dict, state.global_step + ) + env.reinit_dr() + + global_step = checkpoint["state"].global_step + exported_policy_path = os.path.join(config.experiment_dir, "exported") + os.makedirs(exported_policy_path, exist_ok=True) + exported_onnx_name = f"model_step_{global_step:06d}.onnx" + new_cp_path = f"{os.path.dirname(config.checkpoint)}/model_step_{global_step:06d}.pt" + if not os.path.exists(new_cp_path): + shutil.copy(checkpoint_path, new_cp_path) + + if config.get("export_onnx_only", False): + + def get_example_obs(): + obs_dict = env.reset_all() + for k in obs_dict: + obs_dict[k] = obs_dict[k].cpu() + return obs_dict + + assert config.num_envs == 1, "num_envs must be 1 for exporting onnx" + from gear_sonic.utils import inference_helpers + + example_obs_dict = get_example_obs() + + # Check if actor has universal-token encoder structure + has_actor_module = hasattr(model.policy, "actor_module") + has_encoders = has_actor_module and hasattr( + model.policy.actor_module, "encoders_to_iterate" + ) + + if "tokenizer" in example_obs_dict and has_encoders: + + inference_helpers.export_universal_token_module_as_onnx( + model.policy.actor_module, + encoder_name="smpl", + decoder_name="g1_dyn", + path=exported_policy_path, + exported_model_name=exported_onnx_name.replace(".onnx", "_smpl.onnx"), + batch_size=1, + ) + inference_helpers.export_universal_token_module_as_onnx( + model.policy.actor_module, + encoder_name="g1", + decoder_name="g1_dyn", + path=exported_policy_path, + exported_model_name=exported_onnx_name.replace(".onnx", "_g1.onnx"), + batch_size=1, + ) + inference_helpers.export_universal_token_module_as_onnx( + model.policy.actor_module, + encoder_name="teleop", + decoder_name="g1_dyn", + path=exported_policy_path, + exported_model_name=exported_onnx_name.replace(".onnx", "_teleop.onnx"), + batch_size=1, + ) + + inference_helpers.export_universal_token_encoders_as_onnx( + model.policy.actor_module, + path=exported_policy_path, + exported_model_name=exported_onnx_name.replace(".onnx", "_encoder.onnx"), + batch_size=1, + ) + inference_helpers.export_universal_token_decoder_as_onnx( + model.policy.actor_module, + decoder_name="g1_dyn", + path=exported_policy_path, + exported_model_name=exported_onnx_name.replace(".onnx", "_decoder.onnx"), + batch_size=1, + ) + print( # noqa: T201 + f'Exported encoders ONNX to {os.path.join(exported_policy_path, exported_onnx_name.replace(".onnx", "_encoder.onnx"))}' # noqa: E501 + ) + print( # noqa: T201 + f'Exported decoder ONNX to {os.path.join(exported_policy_path, exported_onnx_name.replace(".onnx", "_decoder.onnx"))}' # noqa: E501 + ) + + else: + inference_helpers.export_policy_as_onnx( + {"actor": model.policy}, exported_policy_path, exported_onnx_name, example_obs_dict + ) + + logger.info(f"Exported policy as onnx to: {os.path.join(exported_policy_path)}") + + # Export configs to YAML + export_config = { + "env_config": omegaconf.OmegaConf.to_container(env.config, resolve=True), + "algo_config": omegaconf.OmegaConf.to_container(config.algo.config, resolve=True), + } + config_yaml_path = os.path.join(os.path.dirname(config.checkpoint), "model_config.yaml") + with open(config_yaml_path, "w") as f: + yaml.dump(export_config, f, default_flow_style=False) + logger.info(f"Exported config to: {config_yaml_path}") + exit() # noqa: PLR1722 + + eval_callbacks = config.get("eval_callbacks", []) + if isinstance(eval_callbacks, str): + eval_callbacks = [eval_callbacks] + + callbacks = {} + for callback_name in eval_callbacks: + if callback_name == "im_eval": + with omegaconf.open_dict(config.callbacks.im_eval): + config.callbacks.im_eval.eval_only = True + config.callbacks.im_eval.eval_frequency = 1 + config.callbacks.im_eval.output_dir = config.get("eval_output_dir", None) + config.callbacks.im_eval.log_keys = config.get("log_keys", None) + if callback_name not in config.callbacks: + raise ValueError(f"Callback {callback_name} not found") + callbacks[callback_name] = utils.instantiate(config.callbacks[callback_name]) + + for callback_name, callback in callbacks.items(): # noqa: B007 + if hasattr(callback, "model") and callback.model is None: + callback.model = model + + for callback_name, callback in callbacks.items(): # noqa: B007 + callback.on_step_end(args, state, None, env=env, model=model, accelerator=accelerator) + + if config.get("run_eval_loop", True): + env.set_is_evaluating(True) + obs_dict = env.reset_all() + model.eval() + for obs_key in obs_dict: + obs_dict[obs_key] = obs_dict[obs_key].to(device) + + eval_step_callbacks = { + name: cb + for name, cb in callbacks.items() + if hasattr(cb, "eval_step") and callable(getattr(cb, "eval_step")) # noqa: B009 + } + if eval_step_callbacks: + logger.info(f"Eval step callbacks enabled: {list(eval_step_callbacks.keys())}") + + step_count = 0 + max_render_steps = config.get("max_render_steps", 0) + + run_once = config.get("run_once", False) + envs_completed = torch.zeros(config.num_envs, dtype=torch.bool, device=device) + + with torch.no_grad(): + while True: + policy_model = model.policy + value_model = model.value_model + policy_model.init_rollout() + + actor_state = {} + actions = policy_model.rollout(obs_dict=obs_dict) + actor_state["actions"] = policy_model.action_mean.detach() + actor_state["obs_dict"] = actions["obs_dict"] + + step_count += 1 + + if max_render_steps > 0 and step_count >= max_render_steps: + logger.info(f"Reached max_render_steps={max_render_steps}. Exiting.") + if hasattr(env, "end_render_results"): + env.end_render_results() + break + + results = env.step(actor_state) + obs_dict, rewards, dones, infos = ( + results[0], + results[1], + results[2], + results[3], + ) # noqa: F841 + + if eval_step_callbacks: + all_want_exit = all( + cb.eval_step(env, results) for cb in eval_step_callbacks.values() + ) + if all_want_exit: + logger.info("All eval step callbacks signaled exit. Exiting evaluation loop.") + break + + if run_once: + envs_completed = ( + envs_completed | dones.squeeze(-1) + if dones.dim() > 1 + else envs_completed | dones + ) + if envs_completed.all(): + logger.info("All environments completed one episode. Exiting (run_once=True).") + if hasattr(env, "end_render_results"): + env.end_render_results() + break + + for obs_key in obs_dict.keys(): # noqa: SIM118 + obs_dict[obs_key] = obs_dict[obs_key].to(device) + + if simulator_type == "IsaacSim": + os._exit(0) + + +if __name__ == "__main__": + main() diff --git a/GR00T-WholeBodyControl/gear_sonic/eval_exp.py b/GR00T-WholeBodyControl/gear_sonic/eval_exp.py new file mode 100644 index 0000000000000000000000000000000000000000..92b7a72109310410f2cb7d627e1a7c87bfaa328a --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/eval_exp.py @@ -0,0 +1,679 @@ +#!/usr/bin/env python3 # noqa: EXE001 +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import glob +import itertools +import json +import os +from pathlib import Path +import signal +import subprocess +import sys +import time +import hydra +from loguru import logger +import omegaconf +import wandb +import yaml + +from gear_sonic.trl.callbacks import im_eval_callback +from gear_sonic.utils import config_utils + +config_utils.register_rl_resolvers() + + +class CheckpointEvaluator: + """Continuously monitors an experiment directory for new checkpoints and evaluates them sequentially.""" + + def __init__(self, config): + self.config = config + self.experiment_dir = Path(config.experiment_dir) + self.evaluated_checkpoints: set[str] = set() + self.shutdown_flag = False + self.last_evaluation_time = time.time() + self.evaluation_timeout = config.get("evaluation_timeout", 24 * 3600) + self.eval_frequency = config.get("eval_frequency", None) + self.eval_last_n = config.get("eval_last_n", None) + + if not self.experiment_dir.exists(): + raise ValueError(f"Experiment directory does not exist: {self.experiment_dir}") + + self.find_evaluated_checkpoints() + + self.wandb_run_id = None + self.wandb_project = None + self.wandb_entity = None + self._load_wandb_config() + + logger.info(f"Monitoring experiment directory: {self.experiment_dir}") + logger.info(f"Scan interval: {config.scan_interval} seconds") + logger.info(f"Evaluation timeout: {self.evaluation_timeout / 3600:.1f} hours") + if self.wandb_run_id: + logger.info(f"Wandb logging enabled: run_id={self.wandb_run_id}") + + self._backfill_wandb() + + signal.signal(signal.SIGINT, self._signal_handler) + signal.signal(signal.SIGTERM, self._signal_handler) + + def _signal_handler(self, signum, frame): # noqa: ARG002 + logger.info("Received shutdown signal. Stopping checkpoint monitoring...") + self.shutdown_flag = True + + def _load_wandb_config(self): + meta_path = self.experiment_dir / "meta.yaml" + if meta_path.exists(): + with open(meta_path) as f: + meta = yaml.safe_load(f) + self.wandb_run_id = meta.get("wandb_run") + + config_path = self.experiment_dir / ".hydra" / "config.yaml" + if not config_path.exists(): + config_path = self.experiment_dir / "config.yaml" + if config_path.exists(): + try: + with open(config_path) as f: + train_config = yaml.safe_load(f) + wandb_cfg = train_config.get("wandb", {}) + self.wandb_project = train_config.get("project_name", "TRL_G1_Track") + self.wandb_entity = wandb_cfg.get("wandb_entity", None) + except Exception as e: # noqa: BLE001 + logger.warning(f"Could not load training config for wandb: {e}") + + if self.wandb_project is None: + self.wandb_project = "TRL_G1_Track" + if self.wandb_entity is None: + self.wandb_entity = None # uses wandb default entity + + def _get_wandb_logged_steps(self) -> set[int]: + logged_steps = set() + if not self.wandb_run_id: + return logged_steps + + try: + api = wandb.Api(timeout=30) + run = api.run(f"{self.wandb_entity}/{self.wandb_project}/{self.wandb_run_id}") + hist = run.scan_history( + keys=["eval/success/success_rate", "eval_step"], + min_step=0, + page_size=10000, + ) + for row in hist: + if row.get("eval/success/success_rate") is not None: + step = row.get("eval_step") + if step is not None: + logged_steps.add(int(step)) + except Exception as e: # noqa: BLE001 + logger.warning(f"Could not query wandb for logged eval steps: {e}") + + return logged_steps + + def _backfill_wandb(self): + if not self.wandb_run_id: + self._load_wandb_config() + if not self.wandb_run_id: + logger.info("No wandb run ID available, skipping backfill") + return + + eval_dir = self.experiment_dir / "eval" + if not eval_dir.exists(): + return + + completed_steps = [] + for eval_subdir in sorted(eval_dir.iterdir()): + if not eval_subdir.is_dir(): + continue + try: + step_num = int(eval_subdir.name) + except ValueError: + continue + if (eval_subdir / "all_eval_finish.txt").exists(): + completed_steps.append((step_num, str(eval_subdir))) + + if not completed_steps: + logger.info("No completed eval steps found on disk, nothing to backfill") + return + + logged_steps = self._get_wandb_logged_steps() + missing = [(step, path) for step, path in completed_steps if step not in logged_steps] + + if not missing: + logger.info( + f"All {len(completed_steps)} eval steps already logged to wandb, no backfill needed" + ) + return + + logger.info(f"Backfilling {len(missing)}/{len(completed_steps)} eval steps to wandb") + for eval_step, checkpoint_work_dir in missing: + self._log_eval_to_wandb(eval_step, checkpoint_work_dir) + + logger.info("Backfill complete") + + def _log_eval_to_wandb(self, eval_step: int, checkpoint_work_dir: str): + if not self.wandb_run_id: + self._load_wandb_config() + if not self.wandb_run_id: + logger.warning("No wandb run ID found, skipping wandb logging") + return + + eval_dir = Path(checkpoint_work_dir) + if not eval_dir.exists(): + return + + try: + wandb.init( + id=self.wandb_run_id, + project=self.wandb_project, + entity=self.wandb_entity, + resume="allow", + ) + + wandb.define_metric("eval_step") + wandb.define_metric("eval/*", step_metric="eval_step") + wandb.define_metric("videos_hard*", step_metric="eval_step") + for subdir in sorted(eval_dir.iterdir()): + if subdir.is_dir() and subdir.name != "train": + wandb.define_metric(f"{subdir.name}/*", step_metric="eval_step") + + all_metrics = {"eval_step": eval_step} + + for subdir in sorted(eval_dir.iterdir()): + if not subdir.is_dir(): + continue + + try: + metrics_file = subdir / "metrics_eval.json" + metrics_finish = subdir / "metrics_finish.txt" + if metrics_finish.exists() and metrics_file.exists(): + self._log_metrics(eval_step, metrics_file) + + render_finish = subdir / "render_finish.txt" + video_dir = subdir / "render_results" + if render_finish.exists() and video_dir.exists(): + self._log_videos(eval_step, metrics_file, video_dir) + except Exception as subdir_e: # noqa: BLE001 + logger.error( + f"Failed to log subdir {subdir.name} for step {eval_step}: {subdir_e}" + ) + + wandb.log(all_metrics) + wandb.finish() + logger.info(f"Logged eval results to wandb for step {eval_step}") + + except Exception as e: # noqa: BLE001 + logger.error(f"Failed to log to wandb for step {eval_step}: {e}") + try: + wandb.finish() + except Exception as e: # noqa: BLE001 + logger.error(f"Error finishing wandb: {e}") + + def _load_metrics(self, eval_step: int, metrics_file: Path) -> dict | None: + try: + with open(metrics_file) as f: + metrics_eval = json.load(f) + except json.JSONDecodeError: + logger.error(f"Error loading {metrics_file}") + return None + + log_keys = metrics_eval.pop("log_keys", None) + + file_size_mb = metrics_file.stat().st_size / 1024 / 1024 + if file_size_mb > 20: + metrics_eval.pop("eval/all_metrics_dict", None) + metrics_eval.pop("eval/failed_metrics_dict", None) + logger.info( + f"Skipping per-motion dicts for {metrics_file.parent.name} ({file_size_mb:.0f} MB > 20 MB)" + ) + else: + if "eval/all_metrics_dict" in metrics_eval: + metrics_eval["eval/all_metrics_dict"] = im_eval_callback.create_html_table( + metrics_eval["eval/all_metrics_dict"] + ) + if "eval/failed_metrics_dict" in metrics_eval: + metrics_eval["eval/failed_metrics_dict"] = im_eval_callback.create_html_table( + metrics_eval["eval/failed_metrics_dict"] + ) + + for key in ["failed_keys", "failed_idxes"]: + metrics_eval.pop(key, None) + + metrics_eval["eval_step"] = eval_step + + if log_keys is not None: + metrics_eval = {f"{log_keys}/{k}": v for k, v in metrics_eval.items()} + metrics_eval["eval_step"] = eval_step + + return metrics_eval + + def _log_metrics(self, eval_step: int, metrics_file: Path): + metrics = self._load_metrics(eval_step, metrics_file) + if metrics: + wandb.log(metrics) + + def _log_videos(self, eval_step: int, metrics_file: Path, video_dir: Path): + if not video_dir.exists(): + return + + log_keys = None + if metrics_file.exists(): + try: + with open(metrics_file) as f: + metrics = json.load(f) + log_keys = metrics.get("log_keys") + except Exception as e: # noqa: BLE001 + logger.error(f"Error getting log_keys from metrics file: {e}") + + video_files = sorted( + [ + (i, f) + for i, f in enumerate(sorted(video_dir.iterdir())) + if f.is_file() and f.name.endswith(".mp4") + ] + ) + + if not video_files: + return + + prefix = f"videos_hard_{log_keys}" if log_keys else "videos_hard" + wandb_videos = { + f"{prefix}/{i:04d}": wandb.Video(str(video_file), format="mp4") + for i, video_file in reversed(video_files) + } + wandb_videos["eval_step"] = eval_step + wandb.log(wandb_videos) + + def find_evaluated_checkpoints(self): + """Find all checkpoints that have been successfully evaluated.""" + eval_dir = self.experiment_dir / "eval" + + if not eval_dir.exists(): + logger.info("No eval directory found") + return + + for eval_subdir in sorted(eval_dir.iterdir()): + if eval_subdir.is_dir(): + metrics_finish_file = eval_subdir / "metrics_finish.txt" + metrics_file = eval_subdir / "metrics_eval.json" + render_finish_file = eval_subdir / "render_finish.txt" + if ( + metrics_finish_file.exists() + and metrics_file.exists() + and render_finish_file.exists() + ): + try: + step_num = int(eval_subdir.name) + checkpoint_path = ( + self.experiment_dir / f"model_step_{step_num:06d}.pt" + ) + if checkpoint_path.exists(): + self.evaluated_checkpoints.add(str(checkpoint_path)) + except ValueError: + pass + + logger.info(f"Found {len(self.evaluated_checkpoints)} already evaluated checkpoints") + + def find_checkpoints(self) -> list[Path]: + """Find all checkpoint files in the experiment directory.""" + checkpoint_pattern = str(self.experiment_dir / "model_step_*.pt") + checkpoints = sorted( + [Path(p) for p in glob.glob(checkpoint_pattern)], + key=lambda p: int(p.stem.split("_")[-1]), + ) + return checkpoints + + def is_checkpoint_ready(self, checkpoint_path: Path) -> bool: + """Check if a checkpoint is ready for evaluation (not being written).""" + checkpoint_ready_delay = self.config.get("checkpoint_ready_delay", 60) + mtime = checkpoint_path.stat().st_mtime + age = time.time() - mtime + return age > checkpoint_ready_delay + + def evaluate_checkpoint( + self, + checkpoint_path: Path, + mode: str = "metrics", + work_dir: str = None, + num_render_videos: int = None, + eval_step: int = None, # noqa: ARG002 + eval_dataset: str = None, + eval_mode: str = None, + ): + """Evaluate a single checkpoint using eval_agent_trl.py.""" + checkpoint_str = str(checkpoint_path) + success = False + + mode_finish_file = os.path.join(work_dir, f"{mode}_finish.txt") + metrics_file = os.path.join(work_dir, "metrics_eval.json") + skip = os.path.exists(mode_finish_file) + if skip and mode == "metrics" and not os.path.exists(metrics_file): + logger.info(f"[{mode}] Not skipping since metrics file not found: {metrics_file}") + skip = False + if skip: + logger.info( + f"[{mode}] Skipping evaluation for checkpoint: {checkpoint_path} because it has already been evaluated" # noqa: E501 + ) + return True + + try: + logger.info(f"[{mode}] Starting evaluation for checkpoint: {checkpoint_path}") + + eval_callbacks = self.config.get("eval_callbacks", "im_eval") + + if mode == "metrics": + cmd = f"accelerate launch gear_sonic/eval_agent_trl.py +checkpoint={checkpoint_str} +headless=True ++eval_callbacks={eval_callbacks} ++run_eval_loop=False" # noqa: E501 + cmd += f" ++num_envs={self.config.num_eval_envs}" + cmd += f" ++eval_output_dir={work_dir}" + if eval_mode is not None: + cmd += f" ++use_encoder={eval_mode}" + cmd += " ++manager_env.commands.motion.motion_lib_cfg.multi_thread=False" + cmd += " +manager_env/terminations=tracking/eval" + if eval_dataset is not None: + cmd += ( + f" +manager_env.commands.motion.motion_lib_cfg.motion_file={eval_dataset}" + ) + cmd += f" +log_keys={Path(eval_dataset).name}_{eval_mode if eval_mode is not None else 'all'}" + + elif mode == "render": + cmd = f"python -u gear_sonic/eval_agent_trl.py +checkpoint={checkpoint_str} +headless=True ++eval_callbacks={eval_callbacks} ++run_eval_loop=False" # noqa: E501 + cmd += f" ++num_envs={num_render_videos}" + cmd += f" ++metrics_file={metrics_file}" + render_sort_by = self.config.get("render_sort_by", None) + if render_sort_by is not None: + cmd += f" ++render_sort_by={render_sort_by}" + cmd += f" ++manager_env.config.save_rendering_dir={work_dir}/render_results" + cmd += " ++manager_env.config.render_results=True" + cmd += " ++manager_env.config.env_spacing=10.0" + cmd += " +manager_env/recorders=render" + cmd += " ++manager_env.commands.motion.motion_lib_cfg.multi_thread=False" + if eval_mode is not None: + cmd += f" ++use_encoder={eval_mode}" + + if eval_dataset is not None: + cmd += ( + f" +manager_env.commands.motion.motion_lib_cfg.motion_file={eval_dataset}" + ) + + extra_overrides = self.config.get("extra_overrides", []) + for override in extra_overrides: + cmd += f" {override}" + + logger.info(f"Running command: {cmd}") + capture_output = self.config.get("capture_output", True) + timeout_seconds = self.config.get("render_timeout", 3600) if mode == "render" else 21600 + proc = subprocess.Popen( + cmd, + shell=True, + preexec_fn=os.setsid, + stdout=subprocess.PIPE if capture_output else None, + stderr=subprocess.PIPE if capture_output else None, + text=True, + ) + try: + stdout_data, stderr_data = proc.communicate(timeout=timeout_seconds) + except subprocess.TimeoutExpired: + os.killpg(os.getpgid(proc.pid), signal.SIGKILL) + proc.communicate() + logger.error(f"Evaluation timeout for checkpoint: {checkpoint_path}") + return False + result_returncode = proc.returncode + result_stdout = stdout_data or "" + result_stderr = stderr_data or "" + + success = result_returncode == 0 and os.path.exists(metrics_file) + + if mode == "render": + found_videos = len(glob.glob(os.path.join(work_dir, "render_results", "*.mp4"))) + expected_videos = num_render_videos or self.config.get("num_render_videos", 64) + if result_returncode == 0 and found_videos < expected_videos: + logger.warning( + f"[{mode}] Fewer videos than requested: found {found_videos}/{expected_videos} " + f"(OK if dataset has fewer unique motions)" + ) + elif found_videos == 0: + logger.error(f"[{mode}] No videos produced") + success = False + + if success: + logger.info(f"[{mode}] Successfully evaluated checkpoint: {checkpoint_path}") + self.last_evaluation_time = time.time() + else: + logger.error(f"[{mode}] Evaluation failed for checkpoint {checkpoint_path}") + logger.error("=" * 20 + " stdout " + "=" * 20) + logger.error(result_stdout) + logger.error("=" * 20 + " stderr " + "=" * 20) + logger.error(result_stderr) + logger.error("=" * 20 + " end " + "=" * 20 + "\n") + if not os.path.exists(metrics_file): + logger.error(f"[{mode}] Metrics file not found: {metrics_file}") + + except Exception as e: # noqa: BLE001 + logger.error(f"Error evaluating checkpoint {checkpoint_path}: {e}") + return False + + return success + + def run(self): + """Main monitoring loop.""" + single_pass = self.config.get("single_pass", False) + if single_pass: + logger.info("Running in single-pass mode...") + else: + logger.info("Starting checkpoint monitoring loop...") + eval_datasets = self.config.get("eval_datasets", None) + eval_modes = self.config.get("eval_modes", [None]) + num_render_videos = self.config.get("num_render_videos", 64) + num_test_render_videos = self.config.get("num_test_render_videos", 32) + while not self.shutdown_flag: + try: + checkpoints = self.find_checkpoints() + + new_checkpoints = [] + for cp in checkpoints: + cp_str = str(cp) + if ( + self.eval_frequency is not None + and int(cp.stem.split("_")[-1]) % self.eval_frequency != 0 + ): + continue + if cp_str not in self.evaluated_checkpoints and self.is_checkpoint_ready(cp): + new_checkpoints.append(cp) + + if self.eval_last_n is not None and len(new_checkpoints) > self.eval_last_n: + skipped = len(new_checkpoints) - self.eval_last_n + new_checkpoints = new_checkpoints[-self.eval_last_n :] + logger.info( + f"eval_last_n={self.eval_last_n}: skipping {skipped} earlier checkpoints" + ) + + if single_pass and not new_checkpoints: + logger.info("Single-pass mode: no new checkpoints to evaluate, exiting") + break + + evaluation_success_count = 0 + for checkpoint in new_checkpoints: + if self.shutdown_flag: + break + + eval_step = int(checkpoint.stem.split("_")[-1]) + checkpoint_work_dir = os.path.join( + self.experiment_dir, "eval", f"{eval_step:06d}" + ) + os.makedirs(checkpoint_work_dir, exist_ok=True) + logger.info(f"Found new checkpoint: {checkpoint}") + + success = True + metrics_success = True + + for mode in ["metrics", "render"]: + mode_work_dir = checkpoint_work_dir + "/train" + mode_success = self.evaluate_checkpoint( + checkpoint, + mode=mode, + work_dir=mode_work_dir, + eval_step=eval_step, + num_render_videos=num_render_videos, + ) + if mode_success: + with open(os.path.join(mode_work_dir, f"{mode}_finish.txt"), "w") as f: + f.write(f"{mode}_finish") + success = success and mode_success + if mode == "metrics": + metrics_success = metrics_success and mode_success + + if eval_datasets is not None: + for eval_dataset, eval_mode in itertools.product(eval_datasets, eval_modes): + for mode in ["metrics", "render"]: + mode_work_dir = ( + checkpoint_work_dir + + f"/{Path(eval_dataset).name}_{eval_mode if eval_mode is not None else 'all'}" + ) + mode_success = self.evaluate_checkpoint( + checkpoint, + mode=mode, + work_dir=mode_work_dir, + eval_step=eval_step, + eval_dataset=eval_dataset, + num_render_videos=num_test_render_videos, + eval_mode=eval_mode, + ) + if mode_success: + with open( + os.path.join(mode_work_dir, f"{mode}_finish.txt"), "w" + ) as f: + f.write(f"{mode}_finish") + success = success and mode_success + if mode == "metrics": + metrics_success = metrics_success and mode_success + + if success: + with open( + os.path.join(checkpoint_work_dir, "all_eval_finish.txt"), "w" + ) as f: + f.write("all_eval_finish") + self._log_eval_to_wandb(eval_step, checkpoint_work_dir) + elif metrics_success: + logger.warning( + f"Render failed for step {eval_step}, logging metrics-only to W&B" + ) + self._log_eval_to_wandb(eval_step, checkpoint_work_dir) + + if success: + self.evaluated_checkpoints.add(str(checkpoint)) + evaluation_success_count += 1 + if eval_step >= self.config.max_train_steps: + logger.info( + f"Reached max train steps: {eval_step} >= {self.config.max_train_steps}. Shutting down..." # noqa: E501 + ) + self.shutdown_flag = True + break + + if new_checkpoints: + logger.info(f"Evaluated {evaluation_success_count} new checkpoints") + logger.info(f"Total evaluated checkpoints: {len(self.evaluated_checkpoints)}") + + if single_pass: + logger.info( + f"Single-pass mode: evaluated {evaluation_success_count} checkpoint(s), exiting" + ) + break + + time_since_last_eval = time.time() - self.last_evaluation_time + if time_since_last_eval > self.evaluation_timeout: + logger.info( + f"No checkpoints evaluated in {time_since_last_eval / 3600:.1f} hours. Shutting down..." + ) + self.shutdown_flag = True + break + + time.sleep(self.config.scan_interval) + + except KeyboardInterrupt: + logger.info("Received keyboard interrupt. Shutting down...") + break + except Exception as e: # noqa: BLE001 + logger.error(f"Error in monitoring loop: {e}") + time.sleep(self.config.scan_interval) + + logger.info("Checkpoint monitoring stopped.") + + +@hydra.main(config_path="config", config_name="eval_exp", version_base="1.1") +def main(config: omegaconf.OmegaConf) -> None: + """Main function to start checkpoint monitoring and evaluation.""" + os.chdir(hydra.utils.get_original_cwd()) + + single_pass = config.get("single_pass", False) + + experiment_dir = Path(config.experiment_dir) + if not experiment_dir.exists(): + parent_dir = experiment_dir.parent + prefix = experiment_dir.name + logger.info( + f"Experiment directory doesn't exist, looking for prefix match: {prefix}* in {parent_dir}" + ) + + while True: + if parent_dir.exists(): + matches = sorted( + [d for d in parent_dir.iterdir() if d.is_dir() and d.name.startswith(prefix)] + )[::-1] + if matches: + experiment_dir = None + for match in matches: + if (match / "meta.yaml").exists(): + experiment_dir = match + logger.info( + f"Found matching directory with meta.yaml: {experiment_dir}" + ) + break + if experiment_dir is None: + experiment_dir = matches[-1] + logger.info( + f"Found matching directory (no meta.yaml yet): {experiment_dir}" + ) + config.experiment_dir = str(experiment_dir) + break + + if single_pass: + logger.info("Single-pass mode: no matching directory found, exiting") + return + + logger.info("No match found yet, waiting...") + time.sleep(5) + + meta_file = os.path.join(config.experiment_dir, "meta.yaml") + logger.info(f"Waiting for meta.yaml to exist: {meta_file}") + + if single_pass and not os.path.exists(meta_file): + logger.info("Single-pass mode: meta.yaml not found, exiting") + return + while not os.path.exists(meta_file): + time.sleep(1) + meta = yaml.safe_load(open(meta_file)) # noqa: SIM115 + config.max_train_steps = meta["max_train_steps"] + logger.info(f"Loaded meta: {meta}") + + hydra_log_path = os.path.join(config.experiment_dir, "eval_exp.log") + logger.remove() + logger.add(hydra_log_path, level="DEBUG") + console_log_level = os.environ.get("LOGURU_LEVEL", "INFO").upper() + logger.add(sys.stdout, level=console_log_level, colorize=True) + + evaluator = CheckpointEvaluator(config) + evaluator.run() + + +if __name__ == "__main__": + main() diff --git a/GR00T-WholeBodyControl/gear_sonic/pyproject.toml b/GR00T-WholeBodyControl/gear_sonic/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..471ce161652956992a61d0190a4e3bfecafe2ee1 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/pyproject.toml @@ -0,0 +1,116 @@ +[build-system] +requires = ["setuptools>=67", "wheel", "pip"] +build-backend = "setuptools.build_meta" + +[project] +name = "gear_sonic" +dynamic = ["version"] +readme = {text = "NVIDIA Gear Sonic - Whole Body Control", content-type = "text/plain"} +classifiers = [ + "Intended Audience :: Science/Research", + "Development Status :: 3 - Alpha", + "License :: OSI Approved :: Apache Software License", + "Programming Language :: Python :: 3", + "Topic :: Scientific/Engineering :: Artificial Intelligence", +] +authors = [ + {name = "NVIDIA Gear Lab"} +] +requires-python = ">=3.10" +dependencies = [ + "numpy==1.26.4", + "scipy==1.15.3", + "torch>=2.4.0", + "joblib", + "tqdm", + "easydict", + "loguru", +] +license = {text = "Apache-2.0"} + +[project.optional-dependencies] +# Teleop: minimal deps for running ZMQ-based teleop +# pico_manager_thread_server.py + visualization (pyzmq, pyvista, pinocchio FK) +# Usage: pip install -e "gear_sonic[teleop]" +teleop = [ + "pyzmq", + "msgpack", + "msgpack-numpy", + "pin", + "pyvista; platform_machine != 'aarch64'", +] +# MuJoCo simulation (run_sim_loop.py and related scripts) +# Usage: pip install -e "gear_sonic[sim]" +sim = [ + "mujoco", + "tyro", + "pin", + "pyyaml", + "pyzmq", + "msgpack", + "msgpack-numpy", + "opencv-python", +] +# Data collection: Sonic VLA data exporter with LeRobot dataset output +# Usage: pip install -e "gear_sonic[data_collection]" +data_collection = [ + "pyzmq", + "msgpack", + "msgpack-numpy", + "pin", + "tyro", + "pyttsx3==2.90", + "av>=14.2", + "opencv-python", + "lerobot @ git+https://github.com/huggingface/lerobot.git@a445d9c9da6bea99a8972daa4fe1fdd053d711d2", + "datasets==3.6.0", +] +# Camera server: runs on the robot to publish camera frames over ZMQ +# Includes depthai (OAK cameras) by default. For other SDKs (pyrealsense2), +# install them separately into the .venv_camera venv. +# Usage: pip install -e "gear_sonic[camera]" +camera = [ + "pyzmq", + "msgpack", + "msgpack-numpy", + "opencv-python", + "tyro", + "depthai", + "requests", +] +# Inference: VLA inference client for running Isaac-GR00T policies +# Usage: pip install -e "gear_sonic[inference]" +inference = [ + "pyzmq", + "msgpack", + "msgpack-numpy", + "pin", + "tyro", + "opencv-python", + "scipy", + "Isaac-GR00T @ git+https://github.com/NVIDIA/Isaac-GR00T.git", +] +# Training: full RL training stack (Isaac Lab must be installed separately) +# Usage: pip install -e "gear_sonic[training]" +training = [ + "hydra-core==1.3.2", + "wandb", + "trl==0.28.0", + "transformers>=4.56.2", + "accelerate>=1.3.0", + "tensorboard", + "smpl_sim @ git+https://github.com/ZhengyiLuo/SMPLSim.git", +] + +[tool.setuptools.packages.find] +where = [".."] +include = ["gear_sonic*"] + +[tool.setuptools] +include-package-data = true + +[tool.setuptools.package-data] +gear_sonic = ["py.typed", "**/*.json", "**/*.yaml"] + +[tool.setuptools.dynamic] +version = {attr = "gear_sonic.version.VERSION"} \ No newline at end of file diff --git a/GR00T-WholeBodyControl/gear_sonic/scripts/launch_data_collection.py b/GR00T-WholeBodyControl/gear_sonic/scripts/launch_data_collection.py new file mode 100644 index 0000000000000000000000000000000000000000..1a1fec4daae020817de2fdac1b5f201e625a5e4a --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/scripts/launch_data_collection.py @@ -0,0 +1,473 @@ +""" +All-in-one tmux launcher for SONIC data collection. + +Starts the full data collection stack in a single tmux session: + + Window 0 — data_collection (4 panes): + ┌───────────────────────┬───────────────────────┐ + │ Pane 0: C++ Deploy │ Pane 2: Data Exporter │ + │ (gear_sonic_deploy) │ (.venv_data_collection)│ + ├───────────────────────┼───────────────────────┤ + │ Pane 1: Teleop │ Pane 3: Camera Viewer │ + │ (.venv_teleop) │ (.venv_data_collection)│ + └───────────────────────┴───────────────────────┘ + + Window 1 — sim (only when --sim is passed): + ┌─────────────────────────────────────────────────┐ + │ MuJoCo Simulator (run_sim_loop.py) │ + │ (.venv_sim) │ + └─────────────────────────────────────────────────┘ + +Prerequisites: + - tmux installed (sudo apt install tmux) + - Virtual environments set up: + bash install_scripts/install_pico.sh -> .venv_teleop + bash install_scripts/install_data_collection.sh -> .venv_data_collection + - gear_sonic_deploy built (see docs) + - For sim: .venv_sim must exist (see install instructions) + +Usage (from repo root — no venv activation needed): + python gear_sonic/scripts/launch_data_collection.py # real robot (default) + python gear_sonic/scripts/launch_data_collection.py --sim # MuJoCo sim + python gear_sonic/scripts/launch_data_collection.py --no-camera-viewer # skip viewer + python gear_sonic/scripts/launch_data_collection.py --pico-input-source isaac-teleop # in-process CloudXR / DeviceIO +""" + +from dataclasses import dataclass +from pathlib import Path +import os +import shutil +import signal +import socket +import subprocess +import sys +import time + + +def _bootstrap_venv(): + """Re-exec with the .venv_data_collection Python if tyro is not available.""" + try: + import tyro # noqa: F401 + return + except ImportError: + pass + + repo_root = Path(__file__).resolve().parent.parent.parent + venv_python = repo_root / ".venv_data_collection" / "bin" / "python" + if not venv_python.exists(): + print( + "ERROR: tyro is not installed and .venv_data_collection not found.\n" + " Run: bash install_scripts/install_data_collection.sh" + ) + sys.exit(1) + + print(f"Re-launching with {venv_python} ...") + os.execv(str(venv_python), [str(venv_python)] + sys.argv) + + +_bootstrap_venv() + +import tyro + + +def _get_local_ip() -> str: + """Best-effort detection of the PC's LAN IP address.""" + try: + s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + s.connect(("8.8.8.8", 80)) + ip = s.getsockname()[0] + s.close() + return ip + except Exception: + return "unknown" + + +@dataclass +class DataCollectionLaunchConfig: + """CLI config for the all-in-one data collection tmux launcher.""" + + # Deployment mode + sim: bool = False + """Run against MuJoCo sim (deploy.sh sim) instead of real robot.""" + + # C++ deploy options + deploy_input_type: str = "zmq_manager" + """Input type for the C++ deploy (zmq_manager, keyboard, etc.).""" + + deploy_zmq_host: str = "localhost" + """ZMQ host for the C++ deploy to listen on.""" + + deploy_checkpoint: str = "" + """Checkpoint path for deploy.sh (e.g., 'policy/checkpoints/my_model/model_step_100000'). + Leave empty to use the deploy.sh default.""" + + deploy_obs_config: str = "" + """Observation config file for deploy.sh. Leave empty for default.""" + + deploy_planner: str = "" + """Planner model path for deploy.sh. Leave empty for default.""" + + deploy_motion_data: str = "" + """Motion data path for deploy.sh. Leave empty for default.""" + + deploy_output_type: str = "" + """Output type for deploy.sh. Leave empty for default.""" + + # Teleop streamer options + pico_manager: bool = True + """Run pico_manager_thread_server with --manager flag.""" + + pico_input_source: str = "xrt" + """Teleop input source for pico_manager_thread_server.py (xrt or isaac-teleop).""" + + pico_vis_vr3pt: bool = False + """Enable VR 3-point visualization on the teleop streamer.""" + + pico_vis_smpl: bool = False + """Enable SMPL visualization on the teleop streamer.""" + + pico_waist_tracking: bool = False + """Enable waist tracking on the teleop streamer.""" + + # Data exporter options + task_prompt: str = "demo" + """Language task prompt for the data exporter.""" + + dataset_name: str = "" + """Dataset name for the data exporter. Leave empty to auto-generate from timestamp.""" + + data_exporter_frequency: int = 50 + """Data collection frequency (Hz) for the data exporter.""" + + record_wrist_cameras: bool = False + """Record wrist camera streams (left_wrist, right_wrist) in the dataset.""" + + text_to_speech: bool = True + """Enable voice feedback via espeak (data exporter).""" + + # Camera viewer + camera_viewer: bool = True + """Start the camera viewer pane.""" + + camera_host: str = "localhost" + """Camera server host (shared by data exporter and viewer).""" + + camera_port: int = 5555 + """Camera server port (shared by data exporter and viewer).""" + + +SESSION_NAME = "sonic_data_collection" + + +def _check_prerequisites(config: DataCollectionLaunchConfig): + """Verify that required tools and venvs exist.""" + errors = [] + + if not shutil.which("tmux"): + errors.append("tmux is not installed. Install with: sudo apt install tmux") + + repo_root = Path(__file__).resolve().parent.parent.parent + + if not (repo_root / ".venv_teleop" / "bin" / "activate").exists(): + errors.append( + ".venv_teleop not found. Run: bash install_scripts/install_pico.sh" + ) + + if not (repo_root / ".venv_data_collection" / "bin" / "activate").exists(): + errors.append( + ".venv_data_collection not found. Run: " + "bash install_scripts/install_data_collection.sh" + ) + + deploy_dir = repo_root / "gear_sonic_deploy" + if not (deploy_dir / "deploy.sh").exists(): + errors.append( + f"gear_sonic_deploy/deploy.sh not found at {deploy_dir}. " + "Ensure the deploy directory is set up." + ) + + if config.sim and not (repo_root / ".venv_sim" / "bin" / "activate").exists(): + errors.append( + ".venv_sim not found. Set up the simulation venv first " + "(see install instructions)." + ) + + if config.pico_input_source not in {"xrt", "isaac-teleop"}: + errors.append("--pico-input-source must be one of: xrt, isaac-teleop") + + if errors: + print("ERROR: Prerequisites not met:\n") + for e in errors: + print(f" - {e}") + print() + sys.exit(1) + + +def _kill_existing_session(): + """Kill any existing tmux session with our name.""" + subprocess.run( + ["tmux", "kill-session", "-t", SESSION_NAME], + capture_output=True, + ) + + +def _create_tmux_session(): + """Create a 4-pane tmux layout.""" + # Create detached session + subprocess.run( + ["tmux", "new-session", "-d", "-s", SESSION_NAME], + check=True, + ) + + # Enable mouse support (click panes, scroll, resize) + subprocess.run( + ["tmux", "set-option", "-t", SESSION_NAME, "-g", "mouse", "on"], + ) + + # Bind Ctrl+\ to kill the entire session (no prefix needed) + subprocess.run( + ["tmux", "bind-key", "-T", "root", "C-\\", "kill-session"], + ) + + # Rename default window + subprocess.run( + ["tmux", "rename-window", "-t", f"{SESSION_NAME}:0", "data_collection"], + ) + + # Split into 4 panes: + # 0 | 1 + # ----- + # 2 | 3 + + # Split horizontally: pane 0 (left) and pane 1 (right) + subprocess.run( + ["tmux", "split-window", "-t", f"{SESSION_NAME}:0", "-h"], + ) + + # Split left pane vertically: pane 0 (top-left) and pane 2 (bottom-left) + subprocess.run( + ["tmux", "split-window", "-t", f"{SESSION_NAME}:0.0", "-v"], + ) + + # Split right pane vertically: pane 1 becomes top-right, new pane 3 bottom-right + subprocess.run( + ["tmux", "split-window", "-t", f"{SESSION_NAME}:0.2", "-v"], + ) + + # Let all pane shells finish initialization (.bashrc, conda, etc.) + time.sleep(5) + + +def _send_to_pane(pane_index: int, cmd: str, wait: float = 1.0): + """Send a command string to a tmux pane.""" + target = f"{SESSION_NAME}:0.{pane_index}" + + subprocess.run( + ["tmux", "send-keys", "-t", target, cmd, "C-m"], + ) + time.sleep(wait) + + +def _check_pane_alive(pane_index: int) -> bool: + """Check if a tmux pane's process is still running.""" + target = f"{SESSION_NAME}:0.{pane_index}" + result = subprocess.run( + ["tmux", "list-panes", "-t", target, "-F", "#{pane_dead}"], + capture_output=True, + text=True, + ) + return result.stdout.strip() != "1" + + +def main(config: DataCollectionLaunchConfig): + repo_root = Path(__file__).resolve().parent.parent.parent + + _check_prerequisites(config) + _kill_existing_session() + + print("=" * 60) + print(" SONIC Data Collection Launcher") + print("=" * 60) + print(f" Mode: {'Simulation' if config.sim else 'Real Robot'}") + print(f" Task prompt: {config.task_prompt}") + print(f" Dataset name: {config.dataset_name or '(auto)'}") + print(f" Deploy input: {config.deploy_input_type}") + print(f" Teleop input: {config.pico_input_source}") + if config.deploy_checkpoint: + print(f" Checkpoint: {config.deploy_checkpoint}") + print(f" Camera: {config.camera_host}:{config.camera_port}") + print(f" DC frequency: {config.data_exporter_frequency} Hz") + print(f" Camera viewer: {'Yes' if config.camera_viewer else 'No'}") + print(f" Wrist cameras: {'Yes' if config.record_wrist_cameras else 'No'}") + print(f" Text-to-speech: {'Yes' if config.text_to_speech else 'No'}") + print(f" PC IP (for PICO): {_get_local_ip()}") + print(f" Teleop vis: vr3pt={config.pico_vis_vr3pt} smpl={config.pico_vis_smpl}") + print("=" * 60) + + _create_tmux_session() + print(f"Created tmux session: {SESSION_NAME}") + + # --- Window 1 (sim only): MuJoCo Simulator --- + if config.sim: + subprocess.run( + ["tmux", "new-window", "-t", SESSION_NAME, "-n", "sim"], + ) + sim_cmd = ( + f"cd {repo_root} && " + f"source .venv_sim/bin/activate && " + f"python gear_sonic/scripts/run_sim_loop.py " + f"--enable-image-publish --enable-offscreen " + f"--camera-port {config.camera_port}" + ) + sim_target = f"{SESSION_NAME}:sim" + subprocess.run( + ["tmux", "send-keys", "-t", sim_target, sim_cmd, "C-m"], + ) + print("Starting MuJoCo simulator (window: sim)...") + time.sleep(3.0) + + # Switch back to the data_collection window for the remaining panes + subprocess.run( + ["tmux", "select-window", "-t", f"{SESSION_NAME}:data_collection"], + ) + + # --- Pane 0 (top-left): C++ Deploy --- + deploy_mode = "sim" if config.sim else "real" + deploy_cmd = ( + f"cd {repo_root / 'gear_sonic_deploy'} && " + f"./deploy.sh " + f"--input-type {config.deploy_input_type} " + f"--zmq-host {config.deploy_zmq_host} " + ) + if config.deploy_checkpoint: + deploy_cmd += f"--cp {config.deploy_checkpoint} " + if config.deploy_obs_config: + deploy_cmd += f"--obs-config {config.deploy_obs_config} " + if config.deploy_planner: + deploy_cmd += f"--planner {config.deploy_planner} " + if config.deploy_motion_data: + deploy_cmd += f"--motion-data {config.deploy_motion_data} " + if config.deploy_output_type: + deploy_cmd += f"--output-type {config.deploy_output_type} " + deploy_cmd += deploy_mode + + print("Starting C++ deploy (pane 0)...") + _send_to_pane(0, deploy_cmd, wait=3.0) + + if not _check_pane_alive(0): + print("WARNING: C++ deploy pane may have failed to start.") + + # --- Pane 2 (bottom-left): Teleop Streamer --- + pico_cmd = ( + f"cd {repo_root} && " + f"source .venv_teleop/bin/activate && " + f"python gear_sonic/scripts/pico_manager_thread_server.py " + f"--input-source {config.pico_input_source}" + ) + if config.pico_manager: + pico_cmd += " --manager" + if config.pico_vis_vr3pt: + pico_cmd += " --vis_vr3pt" + if config.pico_vis_smpl: + pico_cmd += " --vis_smpl" + if config.pico_waist_tracking: + pico_cmd += " --waist_tracking" + + print("Starting teleop streamer (pane 2)...") + _send_to_pane(1, pico_cmd, wait=2.0) + + # --- Pane 3 (bottom-right): Camera Viewer --- + if config.camera_viewer: + viewer_cmd = ( + f"cd {repo_root} && " + f"source .venv_data_collection/bin/activate && " + f"python gear_sonic/scripts/run_camera_viewer.py " + f"--camera-host {config.camera_host} " + f"--camera-port {config.camera_port}" + ) + print("Starting camera viewer (pane 3)...") + _send_to_pane(3, viewer_cmd, wait=2.0) + + # --- Pane 1 (top-right): Data Exporter --- + exporter_cmd = ( + f"cd {repo_root} && " + f"source .venv_data_collection/bin/activate && " + f"python gear_sonic/scripts/run_data_exporter.py " + f"--task-prompt '{config.task_prompt}' " + f"--data-collection-frequency {config.data_exporter_frequency} " + f"--camera-host {config.camera_host} " + f"--camera-port {config.camera_port}" + ) + if config.dataset_name: + exporter_cmd += f" --dataset-name '{config.dataset_name}'" + if config.record_wrist_cameras: + exporter_cmd += " --record-wrist-cameras" + if not config.text_to_speech: + exporter_cmd += " --no-text-to-speech" + + print("Starting data exporter (pane 1)...") + _send_to_pane(2, exporter_cmd, wait=1.0) + + # Select the data exporter pane so the user lands there for interactive input + subprocess.run( + ["tmux", "select-pane", "-t", f"{SESSION_NAME}:0.2"], + ) + + print() + print("=" * 60) + print(" All components launched!") + print() + print(f" tmux session: {SESSION_NAME}") + print() + if config.sim: + print(" Window 'sim':") + print(" MuJoCo Simulator (.venv_sim)") + print() + print(" Window 'data_collection':") + print(" Pane 0 (top-left): C++ Deploy") + print(" Pane 1 (bottom-left): Teleop Streamer") + print(" Pane 2 (top-right): Data Exporter <-- you are here") + if config.camera_viewer: + print(" Pane 3 (bottom-right): Camera Viewer") + print() + print(" ** deploy.sh (pane 0) is waiting for confirmation —") + print(" click on pane 0 and press Enter to proceed **") + print() + print(" Controls:") + print(" Ctrl+b, arrow keys - Switch between panes") + if config.sim: + print(" Ctrl+b, n / p - Next / previous window") + print(" Ctrl+b, d - Detach from session") + print(" Ctrl+\\ - Kill entire session") + print("=" * 60) + + # Attach to the session + try: + subprocess.run(["tmux", "attach", "-t", SESSION_NAME]) + except KeyboardInterrupt: + pass + + # After detach/exit, offer cleanup + result = subprocess.run( + ["tmux", "has-session", "-t", SESSION_NAME], + capture_output=True, + ) + if result.returncode == 0: + print(f"\nSession '{SESSION_NAME}' is still running.") + print(f" Reattach: tmux attach -t {SESSION_NAME}") + print(f" Kill: tmux kill-session -t {SESSION_NAME}") + + +def _signal_handler(sig, frame): + print("\nShutdown requested...") + subprocess.run( + ["tmux", "kill-session", "-t", SESSION_NAME], + capture_output=True, + ) + sys.exit(0) + + +if __name__ == "__main__": + signal.signal(signal.SIGINT, _signal_handler) + config = tyro.cli(DataCollectionLaunchConfig) + main(config) diff --git a/GR00T-WholeBodyControl/gear_sonic/scripts/launch_inference.py b/GR00T-WholeBodyControl/gear_sonic/scripts/launch_inference.py new file mode 100644 index 0000000000000000000000000000000000000000..4b7a65b741c876ad2c9b7bea2a00716f78e757a0 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/scripts/launch_inference.py @@ -0,0 +1,465 @@ +""" +All-in-one tmux launcher for SONIC VLA inference. + +Starts the inference stack in a single tmux session: + + Window 0 — inference (4 panes): + ┌───────────────────────┬───────────────────────┐ + │ Pane 0: C++ Deploy │ Pane 1: VLA Inference │ + │ (gear_sonic_deploy) │ (.venv_inference) │ + ├───────────────────────┼───────────────────────┤ + │ Pane 2: Keyboard Pub │ Pane 3: Data Exporter │ + │ (.venv_inference) │ (.venv_data_collection)│ + └───────────────────────┴───────────────────────┘ + + Window 1 — sim (only when --sim is passed): + ┌─────────────────────────────────────────────────┐ + │ MuJoCo Simulator (run_sim_loop.py) │ + │ (.venv_sim) │ + └─────────────────────────────────────────────────┘ + +Prerequisites: + - tmux installed (sudo apt install tmux) + - Virtual environments set up: + bash install_scripts/install_inference.sh -> .venv_inference + bash install_scripts/install_data_collection.sh -> .venv_data_collection (optional, for recording) + - gear_sonic_deploy built (see docs) + - Isaac-GR00T PolicyServer running separately + +Usage (from repo root — no venv activation needed): + python gear_sonic/scripts/launch_inference.py # real robot + python gear_sonic/scripts/launch_inference.py --sim # MuJoCo sim + python gear_sonic/scripts/launch_inference.py --no-data-exporter # no recording pane +""" + +from dataclasses import dataclass +from pathlib import Path +import os +import shutil +import signal +import socket +import base64 +import subprocess +import sys +import textwrap +import time + + +def _bootstrap_venv(): + """Re-exec with the .venv_inference Python if tyro is not available.""" + try: + import tyro # noqa: F401 + return + except ImportError: + pass + + repo_root = Path(__file__).resolve().parent.parent.parent + venv_python = repo_root / ".venv_inference" / "bin" / "python" + if not venv_python.exists(): + print( + "ERROR: tyro is not installed and .venv_inference not found.\n" + " Run: bash install_scripts/install_inference.sh" + ) + sys.exit(1) + + print(f"Re-launching with {venv_python} ...") + os.execv(str(venv_python), [str(venv_python)] + sys.argv) + + +_bootstrap_venv() + +import tyro + + +def _get_local_ip() -> str: + """Best-effort detection of the PC's LAN IP address.""" + try: + s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + s.connect(("8.8.8.8", 80)) + ip = s.getsockname()[0] + s.close() + return ip + except Exception: + return "unknown" + + +@dataclass +class InferenceLaunchConfig: + """CLI config for the all-in-one VLA inference tmux launcher.""" + + # Deployment mode + sim: bool = False + """Run against MuJoCo sim instead of real robot.""" + + # C++ deploy options + deploy_input_type: str = "zmq_manager" + """Input type for the C++ deploy.""" + + deploy_zmq_host: str = "localhost" + """ZMQ host for the C++ deploy to listen on.""" + + deploy_checkpoint: str = "" + """Checkpoint path for deploy.sh. Leave empty for default.""" + + deploy_obs_config: str = "" + """Observation config file for deploy.sh. Leave empty for default.""" + + deploy_planner: str = "" + """Planner model path for deploy.sh. Leave empty for default.""" + + deploy_motion_data: str = "" + """Motion data path for deploy.sh. Leave empty for default.""" + + deploy_output_type: str = "" + """Output type for deploy.sh. Leave empty for default.""" + + # VLA inference options + policy_host: str = "localhost" + """Isaac-GR00T PolicyServer host.""" + + policy_port: int = 5550 + """Isaac-GR00T PolicyServer port.""" + + embodiment_tag: str = "unitree_g1_sonic" + """Embodiment tag for policy inference.""" + + prompt: str = "demo" + """Language prompt for inference.""" + + action_publish_rate: int = 50 + """Rate at which individual actions are published to the C++ control loop (Hz).""" + + action_horizon: int = 40 + """Action horizon of the VLA policy.""" + + # Camera + camera_host: str = "localhost" + """Camera server host.""" + + camera_port: int = 5555 + """Camera server port.""" + + # Data exporter (optional recording during inference) + data_exporter: bool = True + """Start the data exporter pane for recording during inference.""" + + data_exporter_frequency: int = 50 + """Data collection frequency (Hz) for the data exporter.""" + + task_prompt: str = "" + """Task prompt for the data exporter. Defaults to the inference prompt if empty.""" + + dataset_name: str = "" + """Dataset name for the data exporter. Leave empty to auto-generate.""" + + +SESSION_NAME = "sonic_inference" + + +def _check_prerequisites(config: InferenceLaunchConfig): + """Verify that required tools and venvs exist.""" + errors = [] + + if not shutil.which("tmux"): + errors.append("tmux is not installed. Install with: sudo apt install tmux") + + repo_root = Path(__file__).resolve().parent.parent.parent + + if not (repo_root / ".venv_inference" / "bin" / "activate").exists(): + errors.append( + ".venv_inference not found. Run: bash install_scripts/install_inference.sh" + ) + + deploy_dir = repo_root / "gear_sonic_deploy" + if not (deploy_dir / "deploy.sh").exists(): + errors.append( + f"gear_sonic_deploy/deploy.sh not found at {deploy_dir}. " + "Ensure the deploy directory is set up." + ) + + if config.data_exporter: + if not (repo_root / ".venv_data_collection" / "bin" / "activate").exists(): + errors.append( + ".venv_data_collection not found (needed for data exporter). Run: " + "bash install_scripts/install_data_collection.sh" + ) + + if config.sim and not (repo_root / ".venv_sim" / "bin" / "activate").exists(): + errors.append( + ".venv_sim not found. Set up the simulation venv first." + ) + + if errors: + print("ERROR: Prerequisites not met:\n") + for e in errors: + print(f" - {e}") + print() + sys.exit(1) + + +def _kill_existing_session(): + subprocess.run( + ["tmux", "kill-session", "-t", SESSION_NAME], + capture_output=True, + ) + + +def _create_tmux_session(): + subprocess.run( + ["tmux", "new-session", "-d", "-s", SESSION_NAME], + check=True, + ) + subprocess.run( + ["tmux", "set-option", "-t", SESSION_NAME, "-g", "mouse", "on"], + ) + subprocess.run( + ["tmux", "bind-key", "-T", "root", "C-\\", "kill-session"], + ) + subprocess.run( + ["tmux", "rename-window", "-t", f"{SESSION_NAME}:0", "inference"], + ) + + # Split into 4 panes: 0|1 / 2|3 + subprocess.run( + ["tmux", "split-window", "-t", f"{SESSION_NAME}:0", "-h"], + ) + subprocess.run( + ["tmux", "split-window", "-t", f"{SESSION_NAME}:0.0", "-v"], + ) + subprocess.run( + ["tmux", "split-window", "-t", f"{SESSION_NAME}:0.2", "-v"], + ) + + time.sleep(5) + + +def _send_to_pane(pane_index: int, cmd: str, wait: float = 1.0): + target = f"{SESSION_NAME}:0.{pane_index}" + subprocess.run( + ["tmux", "send-keys", "-t", target, cmd, "C-m"], + ) + time.sleep(wait) + + +def _check_pane_alive(pane_index: int) -> bool: + target = f"{SESSION_NAME}:0.{pane_index}" + result = subprocess.run( + ["tmux", "list-panes", "-t", target, "-F", "#{pane_dead}"], + capture_output=True, + text=True, + ) + return result.stdout.strip() != "1" + + +def main(config: InferenceLaunchConfig): + repo_root = Path(__file__).resolve().parent.parent.parent + + _check_prerequisites(config) + _kill_existing_session() + + exporter_prompt = config.task_prompt if config.task_prompt else config.prompt + + print("=" * 60) + print(" SONIC VLA Inference Launcher") + print("=" * 60) + print(f" Mode: {'Simulation' if config.sim else 'Real Robot'}") + print(f" PolicyServer: {config.policy_host}:{config.policy_port}") + print(f" Embodiment: {config.embodiment_tag}") + print(f" Prompt: {config.prompt}") + print(f" Action rate: {config.action_publish_rate} Hz") + print(f" Action horizon: {config.action_horizon}") + print(f" Camera: {config.camera_host}:{config.camera_port}") + print(f" Data exporter: {'Yes' if config.data_exporter else 'No'}") + if config.data_exporter: + print(f" DC frequency: {config.data_exporter_frequency} Hz") + print(f" Task prompt: {exporter_prompt}") + print(f" PC IP: {_get_local_ip()}") + print("=" * 60) + + _create_tmux_session() + print(f"Created tmux session: {SESSION_NAME}") + + # --- Window 1 (sim only): MuJoCo Simulator --- + if config.sim: + subprocess.run( + ["tmux", "new-window", "-t", SESSION_NAME, "-n", "sim"], + ) + sim_cmd = ( + f"cd {repo_root} && " + f"source .venv_sim/bin/activate && " + f"python gear_sonic/scripts/run_sim_loop.py " + f"--enable-image-publish --enable-offscreen " + f"--camera-port {config.camera_port}" + ) + sim_target = f"{SESSION_NAME}:sim" + subprocess.run( + ["tmux", "send-keys", "-t", sim_target, sim_cmd, "C-m"], + ) + print("Starting MuJoCo simulator (window: sim)...") + time.sleep(3.0) + + subprocess.run( + ["tmux", "select-window", "-t", f"{SESSION_NAME}:inference"], + ) + + # --- Pane 0 (top-left): C++ Deploy --- + deploy_mode = "sim" if config.sim else "real" + deploy_cmd = ( + f"cd {repo_root / 'gear_sonic_deploy'} && " + f"./deploy.sh " + f"--input-type {config.deploy_input_type} " + f"--zmq-host {config.deploy_zmq_host} " + ) + if config.deploy_checkpoint: + deploy_cmd += f"--cp {config.deploy_checkpoint} " + if config.deploy_obs_config: + deploy_cmd += f"--obs-config {config.deploy_obs_config} " + if config.deploy_planner: + deploy_cmd += f"--planner {config.deploy_planner} " + if config.deploy_motion_data: + deploy_cmd += f"--motion-data {config.deploy_motion_data} " + if config.deploy_output_type: + deploy_cmd += f"--output-type {config.deploy_output_type} " + deploy_cmd += deploy_mode + + print("Starting C++ deploy (pane 0)...") + _send_to_pane(0, deploy_cmd, wait=3.0) + + if not _check_pane_alive(0): + print("WARNING: C++ deploy pane may have failed to start.") + + # --- Pane 2 (bottom-left): Keyboard Publisher --- + keyboard_script = textwrap.dedent("""\ + import zmq, time + ctx = zmq.Context() + pub = ctx.socket(zmq.PUB) + pub.bind('tcp://localhost:5580') + time.sleep(0.5) + print('Keyboard publisher ready. Keys: p=pause, k=start/stop, i=init pose, [/]=toggle hands, t=prompt') + while True: + key = input() + if key.startswith('t '): + pub.send_string('prompt:' + key[2:]) + print('Sent prompt: ' + key[2:]) + else: + pub.send_string(key) + print('Sent: ' + key) + """) + encoded = base64.b64encode(keyboard_script.encode()).decode() + keyboard_cmd = ( + f"cd {repo_root} && " + f"source .venv_inference/bin/activate && " + f"python -c \"import base64;exec(base64.b64decode('{encoded}'))\"" + ) + + print("Starting keyboard publisher (pane 2)...") + _send_to_pane(1, keyboard_cmd, wait=2.0) + + # --- Pane 3 (bottom-right): Data Exporter (optional) --- + if config.data_exporter: + exporter_cmd = ( + f"cd {repo_root} && " + f"source .venv_data_collection/bin/activate && " + f"python gear_sonic/scripts/run_data_exporter.py " + f"--task-prompt '{exporter_prompt}' " + f"--data-collection-frequency {config.data_exporter_frequency} " + f"--camera-host {config.camera_host} " + f"--camera-port {config.camera_port}" + ) + if config.dataset_name: + exporter_cmd += f" --dataset-name '{config.dataset_name}'" + + print("Starting data exporter (pane 3)...") + _send_to_pane(3, exporter_cmd, wait=2.0) + + # --- Pane 1 (top-right): VLA Inference --- + inference_cmd = ( + f"cd {repo_root} && " + f"source .venv_inference/bin/activate && " + f"python gear_sonic/scripts/run_vla_inference.py " + f"--host {config.policy_host} " + f"--port {config.policy_port} " + f"--embodiment-tag {config.embodiment_tag} " + f"--prompt '{config.prompt}' " + f"--action-publish-rate {config.action_publish_rate} " + f"--action-horizon {config.action_horizon} " + f"--camera-host {config.camera_host} " + f"--camera-port {config.camera_port}" + ) + + print("Starting VLA inference (pane 1)...") + _send_to_pane(2, inference_cmd, wait=1.0) + + # Select the VLA inference pane + subprocess.run( + ["tmux", "select-pane", "-t", f"{SESSION_NAME}:0.2"], + ) + + print() + print("=" * 60) + print(" All components launched!") + print() + print(f" tmux session: {SESSION_NAME}") + print() + if config.sim: + print(" Window 'sim':") + print(" MuJoCo Simulator (.venv_sim)") + print() + print(" Window 'inference':") + print(" Pane 0 (top-left): C++ Deploy") + print(" Pane 1 (bottom-left): Keyboard Publisher") + print(" Pane 2 (top-right): VLA Inference <-- you are here") + if config.data_exporter: + print(" Pane 3 (bottom-right): Data Exporter") + print() + print(" ** deploy.sh (pane 0) is waiting for confirmation --") + print(" click on pane 0 and press Enter to proceed **") + print() + print(" Keyboard controls (type in pane 1):") + print(" p - Pause / resume inference") + print(" k - Start / stop C++ control loop") + print(" i - Send initial pose") + print(" [ - Toggle left hand open/closed (initial pose)") + print(" ] - Toggle right hand open/closed (initial pose)") + print(" t - Change inference prompt") + if config.data_exporter: + print(" c - Start recording episode") + print(" s - Stop recording (success)") + print(" f - Stop recording (failure)") + print() + print(" Navigation:") + print(" Ctrl+b, arrow keys - Switch between panes") + if config.sim: + print(" Ctrl+b, n / p - Next / previous window") + print(" Ctrl+b, d - Detach from session") + print(" Ctrl+\\ - Kill entire session") + print("=" * 60) + + try: + subprocess.run(["tmux", "attach", "-t", SESSION_NAME]) + except KeyboardInterrupt: + pass + + result = subprocess.run( + ["tmux", "has-session", "-t", SESSION_NAME], + capture_output=True, + ) + if result.returncode == 0: + print(f"\nSession '{SESSION_NAME}' is still running.") + print(f" Reattach: tmux attach -t {SESSION_NAME}") + print(f" Kill: tmux kill-session -t {SESSION_NAME}") + + +def _signal_handler(_sig, _frame): + print("\nShutdown requested...") + subprocess.run( + ["tmux", "kill-session", "-t", SESSION_NAME], + capture_output=True, + ) + sys.exit(0) + + +if __name__ == "__main__": + signal.signal(signal.SIGINT, _signal_handler) + config = tyro.cli(InferenceLaunchConfig) + main(config) diff --git a/GR00T-WholeBodyControl/gear_sonic/scripts/pico_manager_thread_server.py b/GR00T-WholeBodyControl/gear_sonic/scripts/pico_manager_thread_server.py new file mode 100644 index 0000000000000000000000000000000000000000..39835307071e1233e99cf7c5660c7cd055360969 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/scripts/pico_manager_thread_server.py @@ -0,0 +1,2311 @@ +# Pico SMPL stream server for body tracking visualization + +""" + +# Recommended Command Line Arguments: + # With VR3 PT visualization (by --vis_vr3pt) and optional SMPL body visualization (by --vis_smpl) + # If you want to enable waist tracking in the VR3 PT visualization, please add --waist_tracking + python pico_manager_thread_server.py --manager \ + --vis_vr3pt --vis_smpl \ + --waist_tracking + + # VR3 PT visualization only (without SMPL body) — lower latency + python pico_manager_thread_server.py --manager --vis_vr3pt + +# DEBUG VR3 PT VISUALIZATION: + # A standalone test mode that captures one live frame and visualizes it. + python pico_manager_thread_server.py --vr3pt_live + +# TIMING COMPARISON: + # The visualizer automatically reports timing every 5 seconds when running: + # [Vis Timing] vr3pt: X.XXms | smpl: X.XXms | render: X.XXms | vr3pt_only: X.XXms | both(vr3pt+smpl): X.XXms + +""" + +from collections import defaultdict, deque +from enum import Enum, IntEnum +import os +import subprocess +import threading +import time + +import msgpack +import numpy as np +from scipy.spatial.transform import Rotation as R, Rotation as sRot +import torch +import zmq + +from gear_sonic.utils.teleop import input_readers +from gear_sonic.utils.teleop.zmq.zmq_poller import ZMQPoller +from gear_sonic.trl.utils.rotation_conversion import decompose_rotation_aa +from gear_sonic.trl.utils.torch_transform import ( + angle_axis_to_quaternion, + compute_human_joints, + quat_apply, + quat_inv, + quaternion_to_angle_axis, + quaternion_to_rotation_matrix, +) + +try: + from gear_sonic.utils.teleop.zmq.zmq_planner_sender import ( + build_command_message, + build_planner_message, + pack_pose_message, + ) +except ImportError: + + def build_command_message(*args, **kwargs) -> bytes: + raise RuntimeError("build_command_message unavailable") + + def build_planner_message(*args, **kwargs) -> bytes: + raise RuntimeError("build_planner_message unavailable") + + def pack_pose_message(*args, **kwargs) -> bytes: + raise RuntimeError("pack_pose_message unavailable") + + +try: + from gear_sonic.isaac_utils.rotations import remove_smpl_base_rot, smpl_root_ytoz_up +except ImportError: + print("Warning: gear_sonic.isaac_utils.rotations not available.") + remove_smpl_base_rot = None + smpl_root_ytoz_up = None + +try: + import xrobotoolkit_sdk as xrt +except ImportError: + xrt = None + +try: + from gear_sonic.utils.teleop.solver.hand.g1_gripper_ik_solver import ( + G1GripperInverseKinematicsSolver, + ) +except ImportError: + print("Warning: G1GripperInverseKinematicsSolver not available.") + G1GripperInverseKinematicsSolver = None + +try: + from gear_sonic.utils.teleop.vis.vr3pt_pose_visualizer import VR3PtPoseVisualizer +except ImportError: + print("Warning: VR3PtPoseVisualizer not available (pyvista may not be installed).") + VR3PtPoseVisualizer = None + +try: + from gear_sonic.utils.teleop.vis.vr3pt_pose_visualizer import get_g1_key_frame_poses +except ImportError: + print("Warning: get_g1_key_frame_poses not available (pyvista may not be installed).") + get_g1_key_frame_poses = None + + +class LocomotionMode(IntEnum): + """Locomotion mode enum for robot movement.""" + + IDLE = 0 + SLOW_WALK = 1 + WALK = 2 + RUN = 3 + IDLE_SQUAT = 4 + IDLE_KNEEL_TWO_LEGS = 5 + IDLE_KNEEL = 6 + IDLE_LYING_FACE_DOWN = 7 + CRAWLING = 8 + IDLE_BOXING = 9 + WALK_BOXING = 10 + LEFT_PUNCH = 11 + RIGHT_PUNCH = 12 + RANDOM_PUNCH = 13 + ELBOW_CRAWLING = 14 + LEFT_HOOK = 15 + RIGHT_HOOK = 16 + FORWARD_JUMP = 17 + STEALTH_WALK = 18 + INJURED_WALK = 19 + + +class StreamMode(Enum): + OFF = 0 + POSE = 1 + PLANNER = 2 + PLANNER_FROZEN_UPPER_BODY = 3 + POSE_PAUSE = 4 + PLANNER_VR_3PT = 5 + + +### Parse 3 point pose from SMPL +# +# OFFSETS: Rotation corrections applied to each keypoint to align SMPL joint frames +# with the desired robot/visualization coordinate convention. +# +# Index mapping (based on [0, 22, 23, 12].index(joint_id)): +# - OFFSETS[0]: Root/Pelvis (joint 0) +# - OFFSETS[1]: Left Wrist (joint 22) +# - OFFSETS[2]: Right Wrist (joint 23) +# - OFFSETS[3]: Neck (joint 12) - more stable than Head (joint 15) for body tracking +# +# Scipy euler rotation convention: +# - Lowercase "xyz" = EXTRINSIC rotations (about the FIXED/ORIGINAL frame's axes) +# - Uppercase "XYZ" = INTRINSIC rotations (about the ROTATING body's axes) +# +# For EXTRINSIC "xyz" with angles [a, b, c]: +# All rotations are about the ORIGINAL frame's axes (before any rotation): +# R_total = R_z(c) @ R_y(b) @ R_x(a) (matrix multiplication order) +# Applied as: first rotate 'a' about original X, then 'b' about original Y, then 'c' about original Z +# +# For INTRINSIC "XYZ" with angles [a, b, c]: +# Each rotation is about the CURRENT (rotated) frame's axis: +# R_total = R_x(a) @ R_y(b) @ R_z(c) (matrix multiplication order) +# Applied as: first rotate 'a' about X, then 'b' about NEW Y, then 'c' about NEW Z +# +OFFSETS = [ + sRot.from_euler("xyz", [0, 0, -90], degrees=True), # Root: yaw -90° about fixed Z + sRot.from_euler("xyz", [90, 0, 0], degrees=True), # L-Wrist: roll +90° about fixed X + sRot.from_euler( + "xyz", [-90, 0, 180], degrees=True + ), # R-Wrist: roll -90° about fixed X, then yaw 180° about fixed Z + sRot.from_euler("xyz", [0, 0, -90], degrees=True), # Neck: yaw -90° about fixed Z +] + + +def _compute_rel_transform(pose, world_frame, scalar_first=True): + """ + Transform a pose from Unity coordinate frame to robot coordinate frame. + + Args: + pose: np.ndarray shape (7,) - [x, y, z, qx, qy, qz, qw] in Unity frame + world_frame: np.ndarray shape (7,) - reference frame to compute relative transform + scalar_first: bool - if True, quaternion is [qw, qx, qy, qz]; if False, [qx, qy, qz, qw] + + Returns: + rel_pos: np.ndarray (3,) - position in robot frame + rel_rot: np.ndarray (4,) - quaternion [qw, qx, qy, qz] in robot frame + + Coordinate transform matrix Q converts Unity (Y-up, left-handed) to Robot (Z-up, right-handed): + Unity: X-right, Y-up, Z-forward + Robot: X-forward, Y-left, Z-up + """ + world_frame = world_frame.copy() + + # Q transforms Unity coordinates to Robot coordinates + # Unity [x, y, z] -> Robot [-x, z, y] + Q = np.array([[-1, 0, 0], [0, 0, 1], [0, 1, 0.0]]) + pose[:3] = Q @ pose[:3] + world_frame[:3] = Q @ world_frame[:3] + rot_base = sRot.from_quat(world_frame[3:], scalar_first=scalar_first).as_matrix() + rot = sRot.from_quat(pose[3:], scalar_first=scalar_first).as_matrix() + rel_rot = sRot.from_matrix(Q @ (rot_base.T @ rot) @ Q.T) + rel_pos = sRot.from_matrix(Q @ rot_base.T @ Q.T).apply(pose[:3] - world_frame[:3]) + return rel_pos, rel_rot.as_quat(scalar_first=True) + + +def _process_3pt_pose(smpl_pose_np): + """ + Extract 3-point VR pose (L-Wrist, R-Wrist, Neck) from full SMPL body joint poses. + + NOTE: We use Neck (joint 12) instead of Head (joint 15) because: + - Neck is more rigidly coupled to the torso + - Head has high DoF (looking around) which doesn't reflect body pose + - Neck provides more stable tracking for upper body orientation + + Args: + smpl_pose_np: np.ndarray shape (24, 7) - 24 SMPL joints, each [x, y, z, qx, qy, qz, qw] + in Unity frame (scalar-last quaternion format) + + Returns: + vr_3pt_pose: np.ndarray shape (3, 7) - 3 keypoints in robot frame + Each row is [x, y, z, qw, qx, qy, qz] (scalar-FIRST quaternion format) + Row 0: Left Wrist (SMPL joint 22) + Row 1: Right Wrist (SMPL joint 23) + Row 2: Neck (SMPL joint 12) + + IMPORTANT: Positions and orientations are RELATIVE TO ROOT (pelvis). + + Processing Steps: + 1. Transform all 24 joints from Unity frame to robot frame + 2. Extract 4 keypoints: Root(0), L-Wrist(22), R-Wrist(23), Neck(12) + 3. Apply per-joint rotation OFFSETS to align joint frames + 4. Make L-Wrist, R-Wrist, Neck relative to Root (both position and orientation) + 5. Return only the 3 non-root keypoints + + Note: Position calibration (wrist offsets, neck kinematic chain) is done in + ThreePointPose.apply_calibration() to ensure consistency with calibrated + orientations. + """ + + # Defensive copy: _compute_rel_transform modifies pose[:3] in-place, which would + # corrupt the caller's array (e.g. PicoReader._latest) and cause wrong results + # if the same sample is processed more than once. + smpl_pose_np = smpl_pose_np.copy() + + # ========================================================================= + # STEP 1: Transform all joints from Unity frame to robot frame + # ========================================================================= + # Input: smpl_pose_np[i] = [x, y, z, qx, qy, qz, qw] in Unity frame (scalar-last) + # Output: body_poses[i] = [x, y, z, qw, qx, qy, qz] in robot frame (scalar-first) + body_poses = np.zeros((smpl_pose_np.shape[0], 7), dtype=np.float32) + for i in range(smpl_pose_np.shape[0]): + pos, orn = _compute_rel_transform( + smpl_pose_np[i], [0, 0, 0, 0, 0, 0, 1], scalar_first=False + ) + body_poses[i, :3] = pos # Position in robot frame + body_poses[i, 3:] = orn # Quaternion [qw, qx, qy, qz] in robot frame + + # ========================================================================= + # STEP 2 & 3: Extract 4 keypoints and apply rotation OFFSETS + # ========================================================================= + # We only care about these SMPL joint indices: + # - Joint 0: Root/Pelvis (reference frame) + # - Joint 22: Left Wrist + # - Joint 23: Right Wrist + # - Joint 12: Neck (more stable than Head joint 15) + # + # kp_poses maps these to indices 0, 1, 2, 3 respectively + positions = np.array([[p[0], p[1], p[2]] for p in body_poses]) + kp_poses = np.zeros((4, 7), dtype=np.float32) + + for i, pose in enumerate(body_poses): + if i not in [0, 22, 23, 12]: + continue # Skip joints we don't care about + + pos = positions[i] + + # Map SMPL joint index to our keypoint index (0-3) + # rel_i: 0=Root, 1=L-Wrist, 2=R-Wrist, 3=Neck + rel_i = [0, 22, 23, 12].index(i) + + # Extract quaternion and apply rotation offset + # pose[3:7] is [qw, qx, qy, qz] (scalar-first from _compute_rel_transform) + quat = np.array([pose[3], pose[4], pose[5], pose[6]]) + + # Apply offset: new_rotation = original_rotation * OFFSET + # This post-multiplies the offset (intrinsic rotation) + rot_quat = (sRot.from_quat(quat, scalar_first=True) * OFFSETS[rel_i]).as_quat( + scalar_first=False + ) + + kp_poses[rel_i, 3:] = rot_quat # Store as scalar-last temporarily for scipy compatibility + kp_poses[rel_i, :3] = pos + + # ========================================================================= + # STEP 4: Make positions and orientations RELATIVE TO ROOT + # ========================================================================= + # This transforms everything into the root's local coordinate frame. + # After this step: + # - Root's position would be (0,0,0) and orientation identity (but we don't return root) + # - Other keypoints are expressed relative to root + root_pos = kp_poses[0, :3].copy() + root_quat = kp_poses[0, 3:].copy() # Still scalar-last for scipy + + for i in range(1, 4): + # Position: subtract root position, then rotate by inverse of root orientation + kp_poses[i, :3] = sRot.from_quat(root_quat).inv().apply(kp_poses[i, :3] - root_pos) + + # Orientation: compute relative rotation (root_inv * keypoint_rot) + # Result stored as scalar-FIRST [qw, qx, qy, qz] + kp_poses[i, 3:] = ( + sRot.from_quat(root_quat).inv() * sRot.from_quat(kp_poses[i, 3:]) + ).as_quat(scalar_first=True) + + # ========================================================================= + # STEP 5: Return only L-Wrist, R-Wrist, Neck (skip Root) + # ========================================================================= + # NOTE: Position and orientation calibration (including neck position via kinematic + # chain) is done in ThreePointPose.apply_calibration() to ensure consistency + # between calibrated orientation and computed neck position. + # kp_poses[1:] = indices 1, 2, 3 = L-Wrist, R-Wrist, Neck + # Each row: [x, y, z, qw, qx, qy, qz] relative to root, scalar-first quaternion + return kp_poses[1:] + + +# ============================================================================= +# VR 3-Point Pose Visualization Functions +# ============================================================================= + + +def run_vr3pt_visualizer_test(): + """ + Standalone test for VR 3-point pose visualizer using PyVista. + Run this to verify the reference frames are displayed correctly. + """ + if VR3PtPoseVisualizer is None: + raise ImportError("VR3PtPoseVisualizer not available. Install pyvista: pip install pyvista") + + print("=" * 60) + print("VR 3-Point Pose Visualizer Test (PyVista)") + print("=" * 60) + print("\nExpected reference frames (all with RGB axes for XYZ):") + print(" 1. WHITE ball at origin (0, 0, 0) - World frame") + print(" 2. CYAN ball at (0, 0, 0.35) - Looking forward (identity)") + print(" 3. MAGENTA ball at (0, 0.4, 0.25) - Looking left (yaw +90°)") + print(" 4. YELLOW ball at (0.4, 0, 0.15) - Looking down (pitch +90°)") + print("\nClose the window to exit.") + print("=" * 60) + + visualizer = VR3PtPoseVisualizer(axis_length=0.08, ball_radius=0.015, with_g1_robot=True) + visualizer.show_static() + + +def run_vr3pt_live_visualizer(): + """ + Live visualizer for real VR 3-point pose data from Pico. + Captures one frame from Pico and displays it alongside reference frames. + """ + if xrt is None: + raise ImportError( + "XRoboToolkit SDK not available. Install xrobotoolkit_sdk to use live visualizer." + ) + + if VR3PtPoseVisualizer is None: + raise ImportError("VR3PtPoseVisualizer not available. Install pyvista: pip install pyvista") + + print("=" * 60) + print("VR 3-Point Pose Live Visualizer (PyVista)") + print("=" * 60) + + # Initialize XRT + subprocess.Popen(["bash", "/opt/apps/roboticsservice/runService.sh"]) + xrt.init() + print("Waiting for body tracking data...") + while not xrt.is_body_data_available(): + print("waiting for body data...") + time.sleep(1) + + print("Body data available! Capturing VR 3-point pose...") + + # Capture body poses and compute vr_3pt_pose + body_poses = xrt.get_body_joints_pose() + body_poses_np = np.array(body_poses) + + # Process to get 3-point pose (L-Wrist, R-Wrist, Neck) + vr_3pt_pose = _process_3pt_pose(body_poses_np) + + print(f"\nCaptured vr_3pt_pose shape: {vr_3pt_pose.shape}") + print(f" L-Wrist: pos={vr_3pt_pose[0, :3]}, quat_wxyz={vr_3pt_pose[0, 3:]}") + print(f" R-Wrist: pos={vr_3pt_pose[1, :3]}, quat_wxyz={vr_3pt_pose[1, 3:]}") + print(f" Neck: pos={vr_3pt_pose[2, :3]}, quat_wxyz={vr_3pt_pose[2, 3:]}") + + print("\nDisplaying visualization...") + print("Close the window to exit.") + print("=" * 60) + + visualizer = VR3PtPoseVisualizer(axis_length=0.08, ball_radius=0.015, with_g1_robot=True) + visualizer.show_with_vr_pose(vr_3pt_pose) + + +def run_vr3pt_realtime_visualizer(update_hz: int = 10): + """ + Real-time visualizer for VR 3-point pose data from Pico. + Continuously updates the visualization with live data. + + Args: + update_hz: Update rate in Hz (default 10) + """ + if xrt is None: + raise ImportError( + "XRoboToolkit SDK not available. Install xrobotoolkit_sdk to use realtime visualizer." + ) + + if VR3PtPoseVisualizer is None: + raise ImportError("VR3PtPoseVisualizer not available. Install pyvista: pip install pyvista") + + print("=" * 60) + print("VR 3-Point Pose Real-time Visualizer (PyVista)") + print("=" * 60) + + # Initialize XRT + subprocess.Popen(["bash", "/opt/apps/roboticsservice/runService.sh"]) + xrt.init() + print("Waiting for body tracking data...") + while not xrt.is_body_data_available(): + print("waiting for body data...") + time.sleep(1) + + print("Body data available! Starting real-time visualization...") + print(f"Update rate: {update_hz} Hz") + print("Close the window or press 'q' to exit.") + print("=" * 60) + + # Use the VR3PtPoseVisualizer for real-time visualization with G1 robot + visualizer = VR3PtPoseVisualizer(axis_length=0.08, ball_radius=0.015, with_g1_robot=True) + visualizer.create_realtime_plotter(interactive=True) + + try: + while visualizer.is_open: + # Get new data from Pico + body_poses = xrt.get_body_joints_pose() + body_poses_np = np.array(body_poses) + vr_3pt_pose = _process_3pt_pose(body_poses_np) + + # Update visualization + visualizer.update_vr_poses(vr_3pt_pose) + visualizer.render() + + time.sleep(1.0 / update_hz) + except KeyboardInterrupt: + print("\nInterrupted by user") + finally: + visualizer.close() + + +def process_smpl_joints(body_pose, global_orient, transl): + """Process SMPL parameters to compute local joints. + + Args: + body_pose: Body pose tensor, shape (T, 69) + global_orient: Global orientation tensor, shape (T, 3) + transl: Translation tensor, shape (T, 3) + + Returns: + Dictionary with processed joints and parameters + """ + # Convert global_orient to quaternion and apply transformations (robust if utils missing) + global_orient_quat = angle_axis_to_quaternion(global_orient) + if smpl_root_ytoz_up is not None: + global_orient_quat = smpl_root_ytoz_up(global_orient_quat) + global_orient_new = quaternion_to_angle_axis(global_orient_quat) + + # Compute joints and vertices using SMPL model (single forward pass) + joints = compute_human_joints( + body_pose=body_pose[..., :63], + global_orient=global_orient_new, + ) # (*, 24, 3) + + # Apply base rotation removal and compute local joints + if remove_smpl_base_rot is not None: + global_orient_quat = remove_smpl_base_rot(global_orient_quat, w_last=False) + + global_orient_quat_inv = quat_inv(global_orient_quat).unsqueeze(1).repeat(1, joints.shape[1], 1) + smpl_joints_local = quat_apply(global_orient_quat_inv, joints) + global_orient_mat = quaternion_to_rotation_matrix(global_orient_quat) + global_orient_6d = global_orient_mat[..., :2].reshape(1, 6) + + return { + "smpl_pose": body_pose, + "joints": joints, + "smpl_joints_local": smpl_joints_local, + "global_orient_quat": global_orient_quat, + "global_orient_6d": global_orient_6d, + "adjusted_transl": transl, + } + + +def generate_finger_data(hand: str, trigger: float, grip: float) -> np.ndarray: + """ + Generate finger position data from Pico controller button states. + + Args: + hand: "left" or "right" + trigger: Trigger button value (0-1) + grip: Grip button value (0-1) + + Returns: + Array of shape [25, 4, 4] representing fingertip positions + """ + fingertips = np.zeros([25, 4, 4]) + + thumb = 0 + middle = 10 + # Control thumb based on shoulder button state (index 4 is thumb tip) + fingertips[4 + thumb, 0, 3] = 1.0 # open thumb + if trigger > 0.5: + fingertips[4 + middle, 0, 3] = 1.0 # close middle + + return fingertips + + +# Joystick deadzone threshold +JOYSTICK_DEADZONE = 0.15 + + +class YawAccumulator: + """Accumulates yaw heading angle based on joystick input.""" + + def __init__(self, yaw_gain: float = 1.5, deadzone: float = JOYSTICK_DEADZONE): + self.yaw_gain = yaw_gain + self.deadzone = deadzone + self.reset() + + def reset(self): + """Reset facing direction to default (1,0,0).""" + self.heading = [1.0, 0.0, 0.0] + self.yaw_angle_rad = 0.0 + self.dyaw = 0.0 + print("YawAccumulator: reset yaw angle to 0.0") + + def yaw_angle(self) -> float: + """Get current yaw angle in radians.""" + return self.yaw_angle_rad + + def yaw_angle_change(self) -> float: + """Get current yaw angle change in radians.""" + return self.dyaw + + def update(self, rx: float, dt: float) -> list[float]: + """ + Update facing direction based on right stick x-axis input. + + Args: + rx: Right stick x-axis value (-1 to 1) + dt: Time delta in seconds + + Returns: + Facing direction as [x, y, 0.0] + """ + self.dyaw = self.yaw_gain * (-rx) * dt + if abs(rx) >= self.deadzone: + self.yaw_angle_rad += self.dyaw + self.heading = [np.cos(self.yaw_angle_rad), np.sin(self.yaw_angle_rad), 0.0] + return self.heading + + +def compute_from_body_poses(parent_indices: list, device, body_poses_np: np.ndarray): + """ + Compute local joints and body orientation from provided body_poses_np. + """ + positions = body_poses_np[:, :3] + global_quats = body_poses_np[:, [6, 3, 4, 5]] + + # Convert to local rotations + global_rots = sRot.from_quat(global_quats, scalar_first=True) + global_rots = global_rots * sRot.from_euler("y", 180, degrees=True) + + local_rots = [] + for i in range(24): + if parent_indices[i] == -1: + local_rots.append(global_rots[i]) + else: + local_rot = global_rots[parent_indices[i]].inv() * global_rots[i] + local_rots.append(local_rot) + + pose_aa = np.array([rot.as_rotvec() for rot in local_rots]) + + body_pose = torch.from_numpy(pose_aa[1:].flatten()).float().to(device).unsqueeze(0) + global_orient = torch.from_numpy(pose_aa[0]).float().to(device).unsqueeze(0) + transl = torch.from_numpy(positions[0]).float().to(device).unsqueeze(0) + + return process_smpl_joints(body_pose, global_orient, transl) + + +# def compute_latest_frame(parent_indices: list, device) -> tuple[np.ndarray, np.ndarray]: +# """ +# Pull body data from XRoboToolkit, compute local SMPL joints and body orientation. +# Returns (smpl_joints_local_np [24,3], global_orient_quat_np [4,]) +# """ +# body_poses = xrt.get_body_joints_pose() +# body_poses_np = np.array(body_poses) +# return compute_from_body_poses(parent_indices, device, body_poses_np) + + +def init_hand_ik_solvers(): + """Initialize hand IK solvers if available.""" + if G1GripperInverseKinematicsSolver is not None: + left_solver = G1GripperInverseKinematicsSolver(side="left") + right_solver = G1GripperInverseKinematicsSolver(side="right") + print("Hand IK solvers initialized") + return left_solver, right_solver + print("Warning: Hand IK solvers not available") + return None, None + + +# Readers that expose `get_controller_data()` returning the IsaacTeleop +# controller_data dict schema (left/right trigger/squeeze, thumbstick, clicks). +# Tuple form keeps the dispatch sites uniform if/when a second reader speaks +# the same schema. +_ISAAC_TELEOP_READERS = (input_readers.IsaacTeleopReader,) + + +def get_controller_inputs(reader=None): + """Fetch controller button/trigger states from XRoboToolkit or IsaacTeleop.""" + if isinstance(reader, _ISAAC_TELEOP_READERS): + ctrl = reader.get_controller_data() + if ctrl is None: + return False, 0.0, 0.0, 0.0, 0.0 + return ( + False, + float(ctrl.get("left_trigger_value", 0.0)), + float(ctrl.get("right_trigger_value", 0.0)), + float(ctrl.get("left_squeeze_value", 0.0)), + float(ctrl.get("right_squeeze_value", 0.0)), + ) + left_trigger = xrt.get_left_trigger() + right_trigger = xrt.get_right_trigger() + left_grip = xrt.get_left_grip() + right_grip = xrt.get_right_grip() + left_menu_button = xrt.get_left_menu_button() + return left_menu_button, left_trigger, right_trigger, left_grip, right_grip + + +def get_controller_axes(reader=None): + """Fetch joystick axes (lx, ly, rx, ry). Falls back to zeros if not available.""" + if isinstance(reader, _ISAAC_TELEOP_READERS): + ctrl = reader.get_controller_data() + if ctrl is None: + return 0.0, 0.0, 0.0, 0.0 + left_thumbstick = ctrl.get("left_thumbstick", [0.0, 0.0]) + right_thumbstick = ctrl.get("right_thumbstick", [0.0, 0.0]) + return ( + float(left_thumbstick[0]), + float(left_thumbstick[1]), + float(right_thumbstick[0]), + float(right_thumbstick[1]), + ) + if xrt is None: + return 0.0, 0.0, 0.0, 0.0 + try: + left_axis = xrt.get_left_axis() # expected [x, y] + right_axis = xrt.get_right_axis() # expected [x, y] + lx = float(left_axis[0]) if len(left_axis) >= 1 else 0.0 + ly = float(left_axis[1]) if len(left_axis) >= 2 else 0.0 + rx = float(right_axis[0]) if len(right_axis) >= 1 else 0.0 + ry = float(right_axis[1]) if len(right_axis) >= 2 else 0.0 + return lx, ly, rx, ry + except Exception: + return 0.0, 0.0, 0.0, 0.0 + + +def get_menu_buttons(reader=None): + """Fetch both menu buttons (left, right). Falls back to False if not available.""" + if isinstance(reader, _ISAAC_TELEOP_READERS): + return False, False + if xrt is None: + return False, False + + def _safe_btn(attr): + try: + fn = getattr(xrt, attr) + return bool(fn()) + except Exception: + return False + + left = _safe_btn("get_left_menu_button") + right = _safe_btn("get_right_menu_button") + return left, right + + +def get_axis_clicks(reader=None): + """Fetch both axis click buttons (left, right). Falls back to False if not available.""" + if isinstance(reader, _ISAAC_TELEOP_READERS): + ctrl = reader.get_controller_data() + if ctrl is None: + return False, False + return ( + float(ctrl.get("left_thumbstick_click", 0.0)) > 0.5, + float(ctrl.get("right_thumbstick_click", 0.0)) > 0.5, + ) + if xrt is None: + return False, False + + def _safe_btn(attr): + try: + fn = getattr(xrt, attr) + return bool(fn()) + except Exception: + return False + + left = _safe_btn("get_left_axis_click") + right = _safe_btn("get_right_axis_click") + return left, right + + +def get_face_buttons(reader=None): + """Fetch primary face buttons A and X. Returns (a_pressed, x_pressed).""" + if isinstance(reader, _ISAAC_TELEOP_READERS): + ctrl = reader.get_controller_data() + if ctrl is None: + return False, False + return ( + float(ctrl.get("right_primary_click", 0.0)) > 0.5, + float(ctrl.get("left_primary_click", 0.0)) > 0.5, + ) + if xrt is None: + return False, False + try: + a_pressed = bool(xrt.get_A_button()) + x_pressed = bool(xrt.get_X_button()) + return a_pressed, x_pressed + except Exception: + return False, False + + +def get_abxy_buttons(reader=None): + """Fetch A,B,X,Y face buttons as booleans (a,b,x,y).""" + if isinstance(reader, _ISAAC_TELEOP_READERS): + ctrl = reader.get_controller_data() + if ctrl is None: + return False, False, False, False + return ( + float(ctrl.get("right_primary_click", 0.0)) > 0.5, + float(ctrl.get("right_secondary_click", 0.0)) > 0.5, + float(ctrl.get("left_primary_click", 0.0)) > 0.5, + float(ctrl.get("left_secondary_click", 0.0)) > 0.5, + ) + if xrt is None: + return False, False, False, False + try: + a_pressed = bool(xrt.get_A_button()) + b_pressed = bool(xrt.get_B_button()) + x_pressed = bool(xrt.get_X_button()) + y_pressed = bool(xrt.get_Y_button()) + return a_pressed, b_pressed, x_pressed, y_pressed + except Exception: + return False, False, False, False + + +def compute_hand_joints_from_inputs( + left_solver, right_solver, left_trigger, left_grip, right_trigger, right_grip +) -> tuple[np.ndarray, np.ndarray]: + """Compute left/right hand joints using IK solvers, or zeros if unavailable.""" + if left_solver is not None and right_solver is not None: + left_finger_data = generate_finger_data("left", left_trigger, left_grip) + right_finger_data = generate_finger_data("right", right_trigger, right_grip) + left_hand_joints = left_solver({"position": left_finger_data}) + right_hand_joints = right_solver({"position": right_finger_data}) + else: + left_hand_joints = np.zeros((1, 7), dtype=np.float32) + right_hand_joints = np.zeros((1, 7), dtype=np.float32) + return left_hand_joints, right_hand_joints + + +def _quat_lerp_normalized(q0: np.ndarray, q1: np.ndarray, alpha: float) -> np.ndarray: + """ + Linear interpolate two quaternions and renormalize. Input shape (4,), xyzw order. + Ensures shortest path by flipping sign if dot < 0. + """ + dot = float(np.dot(q0, q1)) + if dot < 0.0: + q1 = -q1 + q = (1.0 - alpha) * q0 + alpha * q1 + norm = np.linalg.norm(q) + if norm > 0: + q = q / norm + return q + + +def _interp_pose_axis_angle( + prev_pose: np.ndarray, curr_pose: np.ndarray, alpha: float +) -> np.ndarray: + """ + Interpolate axis-angle joint poses by converting to quats, lerp-normalize, then back. + prev_pose, curr_pose: (21,3) axis-angle (rotvec) + Returns (21,3) axis-angle. + """ + prev_quats = sRot.from_rotvec(prev_pose.reshape(-1, 3)).as_quat() # (N,4) xyzw + curr_quats = sRot.from_rotvec(curr_pose.reshape(-1, 3)).as_quat() + out_quats = np.empty_like(prev_quats) + for i in range(prev_quats.shape[0]): + out_quats[i] = _quat_lerp_normalized(prev_quats[i], curr_quats[i], alpha) + out_pose = sRot.from_quat(out_quats).as_rotvec().reshape(prev_pose.shape) + return out_pose + + +class PicoReader: + """ + Background reader that pulls Pico/XRT data as fast as possible and computes dt/FPS. + """ + + def __init__(self, max_queue_size: int = 15): + self._stop = threading.Event() + self._thread = threading.Thread(target=self._run, daemon=True) + self._last_t = None + self._fps_ema = 0.0 + self._last_stamp_ns = None + self._latest = None + self._lock = threading.Lock() + + def start(self): + self._thread.start() + + def stop(self): + self._stop.set() + self._thread.join(timeout=1.0) + + def get_latest(self): + with self._lock: + return self._latest + + @property + def disconnected(self) -> bool: + return False + + def clear_disconnect(self): + pass + + def get_timestamp_ns(self) -> int: + if xrt is None: + return 0 + return int(xrt.get_time_stamp_ns()) + + def _run(self): + last_report = time.time() + while not self._stop.is_set(): + if not xrt.is_body_data_available(): + time.sleep(0.001) + continue + stamp_ns = xrt.get_time_stamp_ns() + prev_stamp_ns = self._last_stamp_ns + if prev_stamp_ns is not None and stamp_ns == prev_stamp_ns: + time.sleep(0.000001) + continue + # Compute device-based dt/fps using timestamp deltas (ns -> s) + device_dt = ((stamp_ns - prev_stamp_ns) * 1e-9) if prev_stamp_ns is not None else 0.0 + if device_dt > 0.0: + inst = 1.0 / device_dt + self._fps_ema = inst if self._fps_ema == 0.0 else (0.9 * self._fps_ema + 0.1 * inst) + self._last_stamp_ns = stamp_ns + t_realtime = time.time() + t_monotonic = time.monotonic() + try: + body_poses = xrt.get_body_joints_pose() + + sample = { + "body_poses_np": np.array(body_poses), + "timestamp_realtime": t_realtime, + "timestamp_monotonic": t_monotonic, + "timestamp_ns": stamp_ns, + "dt": device_dt, + "fps": self._fps_ema, + } + with self._lock: + self._latest = sample + now = time.time() + if now - last_report >= 5.0: + print( + f"[PicoReader] dt_ts: {device_dt*1000.0:.2f} ms, fps: {self._fps_ema:.2f}" + ) + last_report = now + except Exception as e: + print(f"[PicoReader] read error: {e}") + + +def _pose_stream_common( + socket, + buffer_size: int, + num_frames_to_send: int, + target_fps: int, + use_cuda: bool, + record_dir: str, + record_format: str, + stop_event: threading.Event | None = None, + log_prefix: str = "PoseLoop", + enable_vis_vr3pt: bool = False, + with_g1_robot: bool = True, + enable_waist_tracking: bool = False, + enable_smpl_vis: bool = False, + reader=None, +): + """Shared pose streaming loop used by run_pico.""" + if reader is None: + if xrt is None: + raise ImportError( + "XRoboToolkit SDK not available. Install xrobotoolkit_sdk to run pose streaming." + ) + + # Create reader and start it + reader = PicoReader(max_queue_size=buffer_size) + reader.start() + + # Create 3-point pose processor with visualization settings + three_point = ThreePointPose( + enable_vis_vr3pt=enable_vis_vr3pt, + with_g1_robot=with_g1_robot, + enable_waist_tracking=enable_waist_tracking, + enable_smpl_vis=enable_smpl_vis, + log_prefix=log_prefix, + ) + + streamer = PoseStreamer( + socket=socket, + reader=reader, + three_point=three_point, + num_frames_to_send=num_frames_to_send, + target_fps=target_fps, + use_cuda=use_cuda, + record_dir=record_dir, + record_format=record_format, + log_prefix=log_prefix, + ) + + if stop_event is None: + stop_event = threading.Event() + + try: + while not stop_event.is_set(): + streamer.run_once() + except KeyboardInterrupt: + pass + finally: + # Cleanup resources + reader.stop() + three_point.close() + + +class ThreePointPose: + """ + Encapsulates everything around calculating 3-point pose from SMPL input. + + This includes: + - Processing SMPL poses to extract 3-point VR pose (L-Wrist, R-Wrist, Neck) + - Calibration logic to align VR poses with G1 robot + - Optional visualization of 3-point poses + + Calibration is done in two steps: + 1. Neck orientation: Captures initial neck orientation to align subsequent poses as upright + 2. Wrist positions: Aligns wrist positions to match G1 robot key frame positions + """ + + # Kinematic chain constants for neck position (matches VR3PtPoseVisualizer) + TORSO_LINK_OFFSET_Z = 0.05 # meters from root to torso_link + NECK_LINK_LENGTH = 0.35 # meters from torso_link to neck along neck's local Z + + def __init__( + self, + enable_vis_vr3pt: bool = False, + with_g1_robot: bool = True, + enable_waist_tracking: bool = False, + enable_smpl_vis: bool = False, + log_prefix: str = "ThreePointPose", + robot_model=None, + ): + """ + Initialize 3-point pose processor. + + Args: + enable_vis_vr3pt: Whether to enable VR 3pt pose visualization (requires display) + with_g1_robot: Whether to include G1 robot in visualization + enable_waist_tracking: Whether to enable waist tracking in visualization + enable_smpl_vis: Whether to render SMPL body joints in the VR3pt visualizer + log_prefix: Prefix for log messages + robot_model: Optional pre-instantiated RobotModel. If None, will create one. + Used for FK-based calibration (no display required). + """ + self.log_prefix = log_prefix + self.with_g1_robot = with_g1_robot + self.enable_waist_tracking = enable_waist_tracking + self.enable_smpl_vis = enable_smpl_vis + + # Robot model for FK-based calibration (headless, no display required) + self._robot_model = robot_model + if self._robot_model is None: + from gear_sonic.data.robot_model.instantiation.g1 import ( + instantiate_g1_robot_model, + ) + + self._robot_model = instantiate_g1_robot_model() + print(f"[{log_prefix}] Robot model loaded for FK calibration") + + # Optional visualization (requires display + PyVista) + self.vr3pt_visualizer = None + if enable_vis_vr3pt: + if VR3PtPoseVisualizer is None: + raise ImportError( + "VR3PtPoseVisualizer could not be imported but --vis_vr3pt was requested. " + "Ensure pyvista is installed: pip install pyvista" + ) + self.vr3pt_visualizer = VR3PtPoseVisualizer( + axis_length=0.08, + ball_radius=0.015, + with_g1_robot=with_g1_robot, + robot_model=self._robot_model, + enable_waist_tracking=enable_waist_tracking, + enable_smpl_vis=enable_smpl_vis, + ) + self.vr3pt_visualizer.create_realtime_plotter(interactive=True) + g1_str = " with G1 robot" if with_g1_robot else "" + waist_str = " + waist tracking" if enable_waist_tracking else "" + smpl_str = " + SMPL body" if enable_smpl_vis else "" + print(f"[{log_prefix}] VR 3pt pose visualization enabled{g1_str}{waist_str}{smpl_str}") + + # Calibration state — triggered explicitly by calibrate_now() or reset_with_measured_q() + self._calibration_pending = False + self._calibration_neck_quat_inv: np.ndarray | None = None # inv(initial neck quat) + self._calibration_lwrist_offset: np.ndarray | None = None # position offset + self._calibration_rwrist_offset: np.ndarray | None = None + self._calibration_lwrist_rot_offset: sRot | None = None # orientation offset + self._calibration_rwrist_rot_offset: sRot | None = None + # Override robot q for FK during recalibration (e.g. measured joints for VR 3PT) + self._override_robot_q: np.ndarray | None = None + + @property + def is_pending(self) -> bool: + """Check if calibration is pending.""" + return self._calibration_pending + + @property + def is_calibrated(self) -> bool: + """Check if calibration has been captured.""" + return self._calibration_neck_quat_inv is not None + + def process_smpl_pose( + self, + smpl_pose_np: np.ndarray, + smpl_joints_local: np.ndarray | None = None, + ) -> np.ndarray: + """ + Process SMPL pose to extract and calibrate 3-point VR pose. + + Args: + smpl_pose_np: np.ndarray shape (24, 7) - 24 SMPL joints + smpl_joints_local: Optional np.ndarray shape (24, 3) - SMPL local joint + positions for body visualization. If provided and SMPL + visualization is enabled, the joint spheres are updated. + + Returns: + vr_3pt_pose: np.ndarray shape (3, 7) - Calibrated 3-point pose + [L-Wrist, R-Wrist, Neck], each row [x, y, z, qw, qx, qy, qz] + """ + # Extract raw 3-point pose from SMPL + vr_3pt_pose_raw = _process_3pt_pose(smpl_pose_np) + + # Capture calibration on first valid frame (or after reset) + if self._calibration_pending: + self._capture_calibration(vr_3pt_pose_raw) + + # Apply calibration to get the final pose + vr_3pt_pose = self._apply_calibration(vr_3pt_pose_raw) + + if self.vr3pt_visualizer is not None: + self.vr3pt_visualizer.update_from_vr_pose(vr_3pt_pose, waist_scale=1.0) + if smpl_joints_local is not None: + self.vr3pt_visualizer.update_smpl_joints(smpl_joints_local) + self.vr3pt_visualizer.render() + + return vr_3pt_pose + + def close(self) -> None: + """Close and cleanup visualizer resources.""" + if self.vr3pt_visualizer is not None: + try: + self.vr3pt_visualizer.close() + except Exception as e: + print(f"[{self.log_prefix}] Warning: Error closing VR3pt visualizer: {e}") + + def calibrate_now(self, body_poses_np: np.ndarray) -> bool: + """Calibrate using current SMPL frame against FK of all-zero body joints. + Operator should be in zero-reference pose when calling this.""" + try: + vr_3pt_pose_raw = _process_3pt_pose(body_poses_np) + self._override_robot_q = np.zeros(29, dtype=np.float64) + self._capture_calibration(vr_3pt_pose_raw) + print(f"[{self.log_prefix}] Calibration completed (zero-pose reference)") + return True + except Exception as e: + print(f"[{self.log_prefix}] Calibration failed: {e}") + import traceback + + traceback.print_exc() + return False + + def _capture_calibration(self, vr_3pt_pose: np.ndarray) -> None: + """Capture calibration offsets from vr_3pt_pose against G1 FK reference. + If neck calibration already exists (e.g. from calibrate_now), it is preserved + to avoid jumps from SMPL noise during recalibration.""" + + # Step 1: Neck orientation — only capture if not already set + if self._calibration_neck_quat_inv is None: + neck_quat_wxyz = vr_3pt_pose[2, 3:].copy() + neck_rot = sRot.from_quat(neck_quat_wxyz, scalar_first=True) + self._calibration_neck_quat_inv = neck_rot.inv().as_quat(scalar_first=True) + calib_inv_rot = sRot.from_quat(self._calibration_neck_quat_inv, scalar_first=True) + + # Step 2: Rotate VR wrist positions/orientations by neck inverse + lwrist_pos_corrected = calib_inv_rot.apply(vr_3pt_pose[0, :3].copy()) + rwrist_pos_corrected = calib_inv_rot.apply(vr_3pt_pose[1, :3].copy()) + lwrist_rot_corrected = calib_inv_rot * sRot.from_quat(vr_3pt_pose[0, 3:], scalar_first=True) + rwrist_rot_corrected = calib_inv_rot * sRot.from_quat(vr_3pt_pose[1, 3:], scalar_first=True) + + # Step 3: Get G1 FK reference poses + if self._robot_model is None: + raise RuntimeError( + "Robot model is required for calibration but was not loaded. " + "Ensure the G1 robot model and URDF are available." + ) + if get_g1_key_frame_poses is None: + raise RuntimeError( + "get_g1_key_frame_poses could not be imported. " + "Ensure gear_sonic.utils.teleop.vis.vr3pt_pose_visualizer is available." + ) + + # Convert 29-DOF override to full model config if needed + if self._override_robot_q is not None: + robot_q = self._robot_model.get_configuration_from_actuated_joints( + body_actuated_joint_values=self._override_robot_q[:29] + ) + else: + robot_q = None + g1_poses = get_g1_key_frame_poses(self._robot_model, q=robot_q) + + g1_lwrist_pos = g1_poses["left_wrist"]["position"] + g1_rwrist_pos = g1_poses["right_wrist"]["position"] + g1_lwrist_rot = sRot.from_quat( + g1_poses["left_wrist"]["orientation_wxyz"], scalar_first=True + ) + g1_rwrist_rot = sRot.from_quat( + g1_poses["right_wrist"]["orientation_wxyz"], scalar_first=True + ) + + # Compute position offsets: calibrated = neck_corrected - offset + self._calibration_lwrist_offset = lwrist_pos_corrected - g1_lwrist_pos + self._calibration_rwrist_offset = rwrist_pos_corrected - g1_rwrist_pos + + # Compute orientation offsets: calibrated = rot_offset * neck_corrected + self._calibration_lwrist_rot_offset = g1_lwrist_rot * lwrist_rot_corrected.inv() + self._calibration_rwrist_rot_offset = g1_rwrist_rot * rwrist_rot_corrected.inv() + + self._calibration_pending = False + self._override_robot_q = None + + # Log summary + source = "override q" if g1_lwrist_pos.any() else "default/zero" + print( + f"[{self.log_prefix}] Calibration captured (FK ref: {source}):\n" + f" L-Wrist pos offset: [{self._calibration_lwrist_offset[0]:.4f}, " + f"{self._calibration_lwrist_offset[1]:.4f}, {self._calibration_lwrist_offset[2]:.4f}]\n" + f" R-Wrist pos offset: [{self._calibration_rwrist_offset[0]:.4f}, " + f"{self._calibration_rwrist_offset[1]:.4f}, {self._calibration_rwrist_offset[2]:.4f}]" + ) + + def _apply_calibration(self, vr_3pt_pose: np.ndarray) -> np.ndarray: + """Apply stored calibration offsets to raw VR 3-point pose.""" + if self._calibration_neck_quat_inv is None: + return vr_3pt_pose + + calibrated = vr_3pt_pose.copy() + calib_inv_rot = sRot.from_quat(self._calibration_neck_quat_inv, scalar_first=True) + + # Neck orientation: calibrated = inv(initial) * current + neck_rot = sRot.from_quat(vr_3pt_pose[2, 3:], scalar_first=True) + calibrated[2, 3:] = (calib_inv_rot * neck_rot).as_quat(scalar_first=True) + + # Wrist positions: rotate by neck inverse, then subtract offset + if self._calibration_lwrist_offset is not None: + calibrated[0, :3] = ( + calib_inv_rot.apply(vr_3pt_pose[0, :3]) - self._calibration_lwrist_offset + ) + if self._calibration_rwrist_offset is not None: + calibrated[1, :3] = ( + calib_inv_rot.apply(vr_3pt_pose[1, :3]) - self._calibration_rwrist_offset + ) + + # Wrist orientations: rot_offset * (neck_inv * current) + if self._calibration_lwrist_rot_offset is not None: + lw_corrected = calib_inv_rot * sRot.from_quat(vr_3pt_pose[0, 3:], scalar_first=True) + calibrated[0, 3:] = (self._calibration_lwrist_rot_offset * lw_corrected).as_quat( + scalar_first=True + ) + if self._calibration_rwrist_rot_offset is not None: + rw_corrected = calib_inv_rot * sRot.from_quat(vr_3pt_pose[1, 3:], scalar_first=True) + calibrated[1, 3:] = (self._calibration_rwrist_rot_offset * rw_corrected).as_quat( + scalar_first=True + ) + + # Neck position via kinematic chain: root → torso_link (+Z) → neck (along calibrated Z) + neck_z = sRot.from_quat(calibrated[2, 3:], scalar_first=True).apply([0, 0, 1]) + calibrated[2, :3] = ( + np.array([0, 0, self.TORSO_LINK_OFFSET_Z]) + self.NECK_LINK_LENGTH * neck_z + ).astype(np.float32) + + return calibrated + + def _clear_calibration(self): + """Clear all calibration state.""" + self._calibration_neck_quat_inv = None + self._calibration_lwrist_offset = None + self._calibration_rwrist_offset = None + self._calibration_lwrist_rot_offset = None + self._calibration_rwrist_rot_offset = None + self._override_robot_q = None + + def reset(self) -> None: + """Reset calibration. Next process_smpl_pose() call will recalibrate.""" + self._clear_calibration() + self._calibration_pending = True + print(f"[{self.log_prefix}] Calibration reset, will re-calibrate on next frame") + + def reset_with_measured_q(self, body_q_measured: np.ndarray) -> None: + """Recalibrate wrist offsets using measured robot joints (29 DOFs). + Preserves neck calibration to avoid jumps from SMPL noise. + Next process_smpl_pose() will recompute wrist offsets against FK of these joints.""" + # Preserve neck calibration — only clear wrist offsets + self._calibration_lwrist_offset = None + self._calibration_rwrist_offset = None + self._calibration_lwrist_rot_offset = None + self._calibration_rwrist_rot_offset = None + self._override_robot_q = body_q_measured.copy() + self._calibration_pending = True + print(f"[{self.log_prefix}] Wrist recalibration pending (neck preserved, measured q)") + + +class PoseStreamer: + """Encapsulates the pose streaming loop state and logic.""" + + def __init__( + self, + socket, + reader: "PicoReader | input_readers.IsaacTeleopReader", + three_point: ThreePointPose, + num_frames_to_send: int, + target_fps: int, + use_cuda: bool, + record_dir: str, + record_format: str, + log_prefix: str = "PoseLoop", + ): + self.socket = socket + self.reader = reader + self.num_frames_to_send = num_frames_to_send + self.target_fps = target_fps + self.record_dir = record_dir + self.log_prefix = log_prefix + + # Injected dependencies + self.reader = reader + self.three_point = three_point + + self.device = ( + torch.device("cuda") if use_cuda and torch.cuda.is_available() else torch.device("cpu") + ) + + if record_dir: + os.makedirs(record_dir, exist_ok=True) + self.record_idx = 0 + + self.left_hand_ik_solver, self.right_hand_ik_solver = init_hand_ik_solvers() + self.parent_indices = [ + -1, + 0, + 0, + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 9, + 9, + 12, + 13, + 14, + 16, + 17, + 18, + 19, + 20, + 22, + 23, + ][:24] + + self.step = 0 + self.last_fps_report = time.time() + self.fps_counter = 0 + # NOTE: Sleep budget set to 95% of the ideal frame period so that the actual + # FPS lands closer to target_fps despite per-frame processing overhead. + self.frame_time = 0.95 / max(1, target_fps) + self.frame_buffer = defaultdict(lambda: deque(maxlen=num_frames_to_send)) + + self.prev_stamp_ns = None + self.prev_smpl_pose_np = None + self.prev_smpl_joints_np = None + self.prev_body_quat_np = None + self.next_target_ns = None + self.frame_start = time.time() + + # Data collection button state tracking (edge-triggered) + self.toggle_data_collection_last = False + self.toggle_data_abort_last = False + + self.buffer_cleared = ( + True # Start with buffer cleared - wait for full buffer before first send + ) + self.yaw_accumulator = YawAccumulator() + + def reset_yaw(self): + """Called when entering pose mode. Resets yaw only. + Calibration is triggered separately by the operator (A+B+X+Y → calibrate_now).""" + self.yaw_accumulator.reset() + + def on_mode_exit(self): + self.frame_buffer.clear() + self.prev_stamp_ns = None + self.prev_smpl_pose_np = None + self.prev_smpl_joints_np = None + self.prev_body_quat_np = None + self.next_target_ns = None + self.buffer_cleared = True + self.step = 0 + + def run_once(self): + """Execute one iteration of the pose streaming loop.""" + sample = self.reader.get_latest() + + if sample is None: + time.sleep(0.005) + return + + latest_data = compute_from_body_poses( + self.parent_indices, self.device, sample["body_poses_np"] + ) + left_menu_button, left_trigger, right_trigger, left_grip, right_grip = get_controller_inputs( + self.reader + ) + # Get A and B button states for data collection control + a_pressed, b_pressed, x_pressed, y_pressed = get_abxy_buttons(self.reader) + + # Data collection toggle logic (edge-triggered) + # Left grip + A = toggle_data_collection + # Left grip + B = toggle_data_abort + toggle_data_collection_tmp = a_pressed and left_grip > 0.5 + toggle_data_abort_tmp = b_pressed and left_grip > 0.5 + + # Detect rising edge + toggle_data_collection = toggle_data_collection_tmp and not self.toggle_data_collection_last + toggle_data_abort = toggle_data_abort_tmp and not self.toggle_data_abort_last + self.toggle_data_collection_last = toggle_data_collection_tmp + self.toggle_data_abort_last = toggle_data_abort_tmp + + left_hand_joints, right_hand_joints = compute_hand_joints_from_inputs( + self.left_hand_ik_solver, + self.right_hand_ik_solver, + left_trigger, + left_grip, + right_trigger, + right_grip, + ) + smpl_pose_np = ( + latest_data["smpl_pose"].detach().cpu().numpy()[:, :63].reshape(-1, 21, 3)[0] + ).astype(np.float32) + smpl_joints_np = ( + latest_data["smpl_joints_local"].detach().cpu().numpy()[0].astype(np.float32) + ) + body_quat_np = ( + latest_data["global_orient_quat"].detach().cpu().numpy()[0].astype(np.float32) + ) + curr_stamp_ns = int(sample.get("timestamp_ns", 0)) + step_ns = int(1e9 / max(1, self.target_fps)) + if self.prev_stamp_ns is None: + self.prev_stamp_ns = curr_stamp_ns + self.prev_smpl_pose_np = smpl_pose_np + self.prev_smpl_joints_np = smpl_joints_np + self.prev_body_quat_np = body_quat_np + self.next_target_ns = curr_stamp_ns + return + if curr_stamp_ns <= self.prev_stamp_ns: + return + if self.next_target_ns is None: + self.next_target_ns = self.prev_stamp_ns + step_ns + if self.next_target_ns < self.prev_stamp_ns: + self.next_target_ns = self.prev_stamp_ns + if self.next_target_ns > curr_stamp_ns: + return + denom = float(curr_stamp_ns - self.prev_stamp_ns) + alpha = float(self.next_target_ns - self.prev_stamp_ns) / denom if denom > 0.0 else 1.0 + if alpha < 0.0: + alpha = 0.0 + elif alpha > 1.0: + alpha = 1.0 + use_joints = (1.0 - alpha) * self.prev_smpl_joints_np + alpha * smpl_joints_np + use_pose = _interp_pose_axis_angle(self.prev_smpl_pose_np, smpl_pose_np, alpha).astype( + np.float32 + ) + use_body_quat = _quat_lerp_normalized(self.prev_body_quat_np, body_quat_np, alpha).astype( + np.float32 + ) + N = len(self.frame_buffer["frame_index"]) + + ##### From @Jiefeng for directly setting the joint position ###### + joint_pos = np.zeros(29) + body_pose = use_pose.reshape(-1, 21, 3) + + SMPL_L_ELBOW_IDX = 17 + SMPL_L_WRIST_IDX = 19 + SMPL_R_ELBOW_IDX = 18 + SMPL_R_WRIST_IDX = 20 + + # G1_L_ELBOW_IDX = 0 + G1_L_WRIST_ROLL_IDX = 23 + G1_L_WRIST_PITCH_IDX = 25 + G1_L_WRIST_YAW_IDX = 27 + + # G1_R_ELBOW_IDX = 0 + G1_R_WRIST_ROLL_IDX = 24 # Done + G1_R_WRIST_PITCH_IDX = 26 + G1_R_WRIST_YAW_IDX = 28 + smpl_l_elbow_aa = body_pose[:, SMPL_L_ELBOW_IDX] + smpl_l_wrist_aa = body_pose[:, SMPL_L_WRIST_IDX] + smpl_r_elbow_aa = body_pose[:, SMPL_R_ELBOW_IDX] + smpl_r_wrist_aa = body_pose[:, SMPL_R_WRIST_IDX] + + g1_l_elbow_axis = np.array([0, 1, 0]) + g1_l_elbow_q_twist, g1_l_elbow_q_swing = decompose_rotation_aa( + smpl_l_elbow_aa, g1_l_elbow_axis + ) + + g1_r_elbow_axis = np.array([0, 1, 0]) + g1_r_elbow_q_twist, g1_r_elbow_q_swing = decompose_rotation_aa( + smpl_r_elbow_aa, g1_r_elbow_axis + ) + + # Move elbow roll/yaw into wrist while preserving wrist pitch from SMPL + l_elbow_swing_euler = R.from_quat(g1_l_elbow_q_swing[:, [1, 2, 3, 0]]).as_euler( + "XYZ", degrees=False + ) + r_elbow_swing_euler = R.from_quat(g1_r_elbow_q_swing[:, [1, 2, 3, 0]]).as_euler( + "XYZ", degrees=False + ) + + l_wrist_euler = R.from_rotvec(smpl_l_wrist_aa).as_euler("XYZ", degrees=False) + r_wrist_euler = R.from_rotvec(smpl_r_wrist_aa).as_euler("XYZ", degrees=False) + + g1_l_wrist_roll = l_elbow_swing_euler[:, 0] + l_wrist_euler[:, 0] + g1_l_wrist_pitch = -l_wrist_euler[:, 1] + g1_l_wrist_yaw = l_elbow_swing_euler[:, 2] + l_wrist_euler[:, 2] + + g1_r_wrist_roll = -(r_elbow_swing_euler[:, 0] + r_wrist_euler[:, 0]) + g1_r_wrist_pitch = -r_wrist_euler[:, 1] + g1_r_wrist_yaw = r_elbow_swing_euler[:, 2] + r_wrist_euler[:, 2] + + joint_pos[G1_L_WRIST_ROLL_IDX] = g1_l_wrist_roll[0] + joint_pos[G1_L_WRIST_PITCH_IDX] = -g1_l_wrist_pitch[0] + joint_pos[G1_L_WRIST_YAW_IDX] = g1_l_wrist_yaw[0] + + joint_pos[G1_R_WRIST_ROLL_IDX] = g1_r_wrist_roll[0] + joint_pos[G1_R_WRIST_PITCH_IDX] = g1_r_wrist_pitch[0] + joint_pos[G1_R_WRIST_YAW_IDX] = g1_r_wrist_yaw[0] + + # Process SMPL pose to get calibrated 3-point VR pose and update visualization + # Pass SMPL local joints for optional body visualization in the VR3Pt viewer + smpl_joints_for_vis = ( + latest_data["smpl_joints_local"].detach().cpu().numpy()[0] + if self.three_point.enable_smpl_vis + else None + ) + vr_3pt_pose = self.three_point.process_smpl_pose( + sample["body_poses_np"], smpl_joints_local=smpl_joints_for_vis + ) + ##### From @Jiefeng for directly setting the joint position ###### + + self.frame_buffer["smpl_pose"].append(use_pose) + self.frame_buffer["smpl_joints"].append(use_joints) + self.frame_buffer["body_quat_w"].append(use_body_quat) + self.frame_buffer["frame_index"].append(int(self.step)) + self.frame_buffer["joint_pos"].append(joint_pos) + pico_dt = float(sample.get("dt", 0.0)) + pico_fps = float(sample.get("fps", 0.0)) + N = len(self.frame_buffer["frame_index"]) + + # Wait for buffer to be completely filled before sending first message after clearing + buffer_is_full = len(self.frame_buffer["frame_index"]) >= self.num_frames_to_send + if buffer_is_full and self.buffer_cleared: + # Buffer is now full with fresh data, can start sending + self.buffer_cleared = False + + # Get joystick axes for yaw accumulation + _, _, rx, _ = get_controller_axes(self.reader) + self.yaw_accumulator.update(rx, self.frame_time) + + # Only send if buffer is full and we're not waiting for fresh data + if buffer_is_full and not self.buffer_cleared: + numpy_data = { + "smpl_pose": np.stack((self.frame_buffer["smpl_pose"]), axis=0), + "smpl_joints": np.stack((self.frame_buffer["smpl_joints"]), axis=0), + "body_quat_w": np.stack((self.frame_buffer["body_quat_w"]), axis=0), + "joint_pos": np.stack((self.frame_buffer["joint_pos"]), axis=0), + "joint_vel": np.zeros((N, 29)), + "vr_position": vr_3pt_pose[:, :3].flatten(), + "vr_orientation": vr_3pt_pose[:, 3:].flatten(), + "frame_index": np.array((self.frame_buffer["frame_index"]), dtype=np.int64), + "left_trigger": np.array([left_trigger], dtype=np.float32), + "right_trigger": np.array([right_trigger], dtype=np.float32), + "left_grip": np.array([left_grip], dtype=np.float32), + "right_grip": np.array([right_grip], dtype=np.float32), + "pico_dt": np.array([pico_dt], dtype=np.float32), + "pico_fps": np.array([pico_fps], dtype=np.float32), + "timestamp_realtime": np.array( + [sample.get("timestamp_realtime", 0.0)], dtype=np.float64 + ), + "timestamp_monotonic": np.array( + [sample.get("timestamp_monotonic", 0.0)], dtype=np.float64 + ), + "left_hand_joints": left_hand_joints.reshape(-1).astype(np.float32), + "right_hand_joints": right_hand_joints.reshape(-1).astype(np.float32), + "toggle_data_collection": np.array([toggle_data_collection], dtype=bool), + "toggle_data_abort": np.array([toggle_data_abort], dtype=bool), + "heading_increment": np.array( + [self.yaw_accumulator.yaw_angle_change()], dtype=np.float32 + ), + } + + packed_message = pack_pose_message(numpy_data, topic="pose") + self.socket.send(packed_message) + + if self.record_dir: + out_path = os.path.join(self.record_dir, f"pose_{self.record_idx:06d}.npz") + np.savez_compressed(out_path, **numpy_data) + self.record_idx += 1 + + self.step += 1 + self.next_target_ns += step_ns + self.prev_stamp_ns = curr_stamp_ns + self.prev_smpl_pose_np = smpl_pose_np + self.prev_smpl_joints_np = smpl_joints_np + self.prev_body_quat_np = body_quat_np + self.fps_counter += 1 + current_time = time.time() + if current_time - self.last_fps_report >= 5.0: + fps = self.fps_counter / (current_time - self.last_fps_report) + print(f"[{self.log_prefix}] FPS: {fps:.2f}, Step: {self.step}") + self.fps_counter = 0 + self.last_fps_report = current_time + elapsed = time.time() - self.frame_start + if elapsed < self.frame_time: + time.sleep(self.frame_time - elapsed) + self.frame_start = time.time() + + +def _init_input_source( + input_source: str, + buffer_size: int, +) -> "PicoReader | input_readers.IsaacTeleopReader": + """Create, start, and wait for readiness of the requested teleop input source.""" + if input_source == "isaac-teleop": + reader = input_readers.IsaacTeleopReader(max_queue_size=buffer_size) + reader.start() + print("Using Isaac Teleop (in-process CloudXR / DeviceIO), waiting for data...") + while reader.get_latest() is None: + print("waiting for Isaac Teleop body data (connect the headset to CloudXR)...") + time.sleep(1) + return reader + + if xrt is None: + raise ImportError( + "XRoboToolkit SDK not available. Install xrobotoolkit_sdk to run Pico streaming." + ) + + subprocess.Popen(["bash", "/opt/apps/roboticsservice/runService.sh"]) + xrt.init() + print("Waiting for body tracking data...") + while not xrt.is_body_data_available(): + print("waiting for body data...") + time.sleep(1) + + reader = PicoReader(max_queue_size=buffer_size) + reader.start() + return reader + + +def run_pico( + buffer_size: int = 15, + port: int = 5556, + num_frames_to_send: int = 5, + target_fps: int = 50, + use_cuda: bool = False, + record_dir: str = "", + record_format: str = "npz", + enable_vis_vr3pt: bool = False, + with_g1_robot: bool = True, + enable_waist_tracking: bool = False, + enable_smpl_vis: bool = False, + input_source: str = "xrt", +): + """Run body tracking with real-time visualization and ZMQ streaming.""" + reader = _init_input_source(input_source, buffer_size) + context = zmq.Context() + socket = context.socket(zmq.PUB) + socket.bind(f"tcp://*:{port}") + time.sleep(0.1) + print(f"ZMQ socket bound to port {port}") + if build_command_message is not None and build_planner_message is not None: + try: + socket.send(build_command_message(start=False, stop=False, planner=False)) + socket.send(build_planner_message(0, [0.0, 0.0, 0.0], [1.0, 0.0, 0.0], -1.0, -1.0)) + except Exception as e: + print(f"Warning: failed to send initial command/planner messages: {e}") + try: + _pose_stream_common( + socket=socket, + buffer_size=buffer_size, + num_frames_to_send=num_frames_to_send, + target_fps=target_fps, + use_cuda=use_cuda, + record_dir=record_dir, + record_format=record_format, + stop_event=None, + log_prefix="Main", + enable_vis_vr3pt=enable_vis_vr3pt, + with_g1_robot=with_g1_robot, + enable_waist_tracking=enable_waist_tracking, + enable_smpl_vis=enable_smpl_vis, + reader=reader, + ) + finally: + socket.close() + context.term() + print("Threads stopped, ZMQ socket closed") + + +class FeedbackReader: + """Reads feedback from robot via ZMQ and processes measured upper body position to use as frozen targets.""" + + def __init__(self, zmq_feedback_host: str = "localhost", zmq_feedback_port: int = 5557): + self.poller = ZMQPoller(host=zmq_feedback_host, port=zmq_feedback_port, topic="g1_debug") + + self.upper_body_joint_indices = self._get_upper_body_joint_indices() + + self.upper_body_position_target = None + self.left_hand_position_target = None + self.right_hand_position_target = None + # Full body joint configuration (29 DOFs) as measured from robot, + # used for FK when recalibrating VR 3PT tracking against actual robot pose + self.full_body_q_measured: np.ndarray | None = None + + def _get_upper_body_joint_indices(self) -> list[int]: + # TODO: get from robot model, not hardcoded + # robot_model = instantiate_g1_robot_model() + # return robot_model.get_joint_group_indices("upper_body") + return [12, 13, 14, 15, 22, 16, 23, 17, 24, 18, 25, 19, 26, 20, 27, 21, 28] + + def poll_feedback(self): + """Poll for feedback once, and update internal state.""" + ( + self.upper_body_position_target, + self.left_hand_position_target, + self.right_hand_position_target, + self.full_body_q_measured, + ) = self._process_upper_body_position_targets() + print("[PlannerLoop] Saved upper body position target:", self.upper_body_position_target) + + def _process_upper_body_position_targets( + self, + ) -> tuple[np.ndarray | None, np.ndarray | None, np.ndarray | None, np.ndarray | None]: + data = self.poller.get_data() + + if data is None: + print("[PlannerLoop] No feedback data received") + return None, None, None, None + + unpacked = msgpack.unpackb(data, raw=False) + full_body_q = None + if "body_q_measured" in unpacked: + body_q_swizzled = unpacked["body_q_measured"] + full_body_q = np.array(body_q_swizzled, dtype=np.float64) + body_q = [body_q_swizzled[i] for i in self.upper_body_joint_indices] + else: + print("[PlannerLoop] body_q_measured not in feedback data") + body_q = None + + if "left_hand_q_measured" in unpacked: + left_hand_q = unpacked["left_hand_q_measured"] + else: + print("[PlannerLoop] left_hand_q_measured not in feedback data") + left_hand_q = None + + if "right_hand_q_measured" in unpacked: + right_hand_q = unpacked["right_hand_q_measured"] + else: + print("[PlannerLoop] right_hand_q_measured not in feedback data") + right_hand_q = None + + return body_q, left_hand_q, right_hand_q, full_body_q + + +class PlannerStreamer: + """Encapsulates the planner control loop state and logic.""" + + def __init__( + self, + socket, + reader: "PicoReader | input_readers.IsaacTeleopReader", + three_point: ThreePointPose, + poll_hz: int = 20, + zmq_feedback_host: str = "localhost", + zmq_feedback_port: int = 5557, + ): + self.socket = socket + self.reader = reader + self.three_point = three_point + self.feedback_reader = FeedbackReader( + zmq_feedback_host=zmq_feedback_host, zmq_feedback_port=zmq_feedback_port + ) + + self.dt = 1.0 / max(1, poll_hz) + # Current locomotion mode, default IDLE + self.mode = LocomotionMode.IDLE + self.prev_ab = False + self.prev_xy = False + # Persistent facing buffer (unit vector on XY plane) + self.yaw_accumulator = YawAccumulator() + self.last_send = time.time() + self.last_xrt_timestamp = None + + # Hand IK solvers for trigger-controlled hand open/close in VR 3PT mode + self.left_hand_ik_solver, self.right_hand_ik_solver = init_hand_ik_solvers() + + def reset_yaw(self): + """Called when entering planner mode. Resets state for fresh start.""" + self.yaw_accumulator.reset() + + def save_upper_body_position_target(self): + """Poll feedback and save upper body position target.""" + self.feedback_reader.poll_feedback() + + def recalibrate_for_vr3pt(self): + """ + Recalibrate VR 3-point pose tracking using the robot's current measured joints. + + Polls the g1_debug feedback to get the robot's actual joint state, then + schedules recalibration so VR tracking aligns with the robot's current pose. + This prevents sudden jumps when entering VR 3PT mode from PLANNER mode. + """ + self.feedback_reader.poll_feedback() + if self.feedback_reader.full_body_q_measured is not None: + self.three_point.reset_with_measured_q(self.feedback_reader.full_body_q_measured) + print("[PlannerLoop] VR 3PT recalibration scheduled with measured robot pose") + else: + # Fallback: use zeros if no feedback available + print( + "[PlannerLoop] WARNING: No feedback data for VR 3PT recalibration, " + "using zero body_q as fallback" + ) + self.three_point.reset_with_measured_q(np.zeros(29, dtype=np.float64)) + + def run_once(self, stream_mode: StreamMode): + """Execute one iteration of the planner control loop.""" + try: + # Avoid sending old commands if XRT timestamp hasn't advanced, in case of headset disconnect + xrt_timestamp = self.reader.get_timestamp_ns() + if xrt_timestamp == self.last_xrt_timestamp: + return + self.last_xrt_timestamp = xrt_timestamp + + # A+B => next mode; X+Y => previous mode (rising edges) + a_pressed, b_pressed, x_pressed, y_pressed = get_abxy_buttons(self.reader) + ab_now = bool(a_pressed) and bool(b_pressed) + xy_now = bool(x_pressed) and bool(y_pressed) + if ab_now and not self.prev_ab: + self.mode = LocomotionMode(min(LocomotionMode.INJURED_WALK, self.mode + 1)) + print(f"[PlannerLoop] Mode -> {self.mode.value}: {self.mode.name}") + if xy_now and not self.prev_xy: + self.mode = LocomotionMode(max(LocomotionMode.IDLE, self.mode - 1)) + print(f"[PlannerLoop] Mode -> {self.mode.value}: {self.mode.name}") + self.prev_ab = ab_now + self.prev_xy = xy_now + + # Read axes/joysticks to control movement, facing, speed and mode + lx, ly, rx, ry = get_controller_axes(self.reader) + + # Facing from RIGHT stick: continuous yaw based on rx (right = turn right, left = turn left) + facing = self.yaw_accumulator.update(rx, self.dt) + + raw_mag = np.hypot(lx, ly) + raw_mag = np.clip(raw_mag, 0.0, 1.0) + if np.abs(raw_mag) < JOYSTICK_DEADZONE: + mag = 0.0 + speed = -1.0 + mode_to_send = LocomotionMode.IDLE + else: + mag = (raw_mag - JOYSTICK_DEADZONE) / (1.0 - JOYSTICK_DEADZONE) + if mag > 1.0: + mag = 1.0 + mode_to_send = self.mode + + if self.mode == LocomotionMode.SLOW_WALK: + speed = 0.1 + 0.5 * mag # 0.1 .. 0.6 + elif self.mode == LocomotionMode.WALK: + speed = -1.0 + elif self.mode == LocomotionMode.RUN: + speed = 1.5 + 3 * mag # 1.5 .. 4.5 + else: + speed = mag # default 0 .. 1.0 + + denom = raw_mag if raw_mag > 0.0 else 1.0 + scale = mag / denom + movement_local = np.array([-lx, ly]) * scale + perp_x, perp_y = -facing[1], facing[0] + rotation_facing = np.array([[perp_x, perp_y], [facing[0], facing[1]]]) + movement_global = rotation_facing @ movement_local + + movement = [movement_global[0], movement_global[1], 0.0] + + upper_body_position = None + left_hand_position = None + right_hand_position = None + if stream_mode == StreamMode.PLANNER_FROZEN_UPPER_BODY: + upper_body_position = self.feedback_reader.upper_body_position_target + left_hand_position = self.feedback_reader.left_hand_position_target + right_hand_position = self.feedback_reader.right_hand_position_target + + vr_3pt_position = None + vr_3pt_orientation = None + vr_3pt_compliance = None + if stream_mode == StreamMode.PLANNER_VR_3PT: + sample = self.reader.get_latest() + if sample is not None: + print("[PlannerLoop] Sending VR 3-point pose as target") + vr_3pt_pose = self.three_point.process_smpl_pose(sample["body_poses_np"]) + vr_3pt_position = (vr_3pt_pose[:, :3].flatten()).tolist() + vr_3pt_orientation = vr_3pt_pose[:, 3:].flatten().tolist() + + # Compute hand joints from trigger/grip inputs so operator can + # control hand open/close while in VR 3PT mode + ( + left_menu_button, + left_trigger, + right_trigger, + left_grip, + right_grip, + ) = get_controller_inputs(self.reader) + lh_joints, rh_joints = compute_hand_joints_from_inputs( + self.left_hand_ik_solver, + self.right_hand_ik_solver, + left_trigger, + left_grip, + right_trigger, + right_grip, + ) + left_hand_position = lh_joints.reshape(-1).astype(np.float32).tolist() + right_hand_position = rh_joints.reshape(-1).astype(np.float32).tolist() + + msg = build_planner_message( + mode_to_send.value, + movement, + facing, + speed=speed, + height=-1.0, + upper_body_position=upper_body_position, + left_hand_position=left_hand_position, + right_hand_position=right_hand_position, + vr_3pt_position=vr_3pt_position, + vr_3pt_orientation=vr_3pt_orientation, + vr_3pt_compliance=vr_3pt_compliance, + ) + self.socket.send(msg) + except Exception as e: + import traceback + + print(f"[PlannerLoop] error: {e}") + traceback.print_exc() + raise + + # pacing + now = time.time() + sleep_t = self.dt - (now - self.last_send) + if sleep_t > 0: + time.sleep(sleep_t) + self.last_send = time.time() + + +def run_pico_manager( + port: int = 5556, + buffer_size: int = 15, + num_frames_to_send: int = 5, + target_fps: int = 50, + use_cuda: bool = False, + record_dir: str = "", + record_format: str = "npz", + zmq_feedback_host: str = "localhost", + zmq_feedback_port: int = 5557, + enable_vis_vr3pt: bool = False, + with_g1_robot: bool = True, + enable_waist_tracking: bool = False, + enable_smpl_vis: bool = False, + input_source: str = "xrt", +): + """ + Manager: creates shared PUB socket and runs pose/planner streamers based on current mode. + Controller input: + A+X: Toggle between planner and pose mode + A+B+X+Y: Toggle policy start/stop + """ + reader = _init_input_source(input_source, buffer_size) + + context = zmq.Context() + socket = context.socket(zmq.PUB) + socket.bind(f"tcp://*:{port}") + time.sleep(0.1) + print(f"[Manager] ZMQ socket bound to port {port}") + + # Print available locomotion modes + try: + print("[Manager] Available modes:") + for mode in LocomotionMode: + print(f" {mode.value}: {mode.name}") + except Exception: + pass + + three_point = ThreePointPose( + enable_vis_vr3pt=enable_vis_vr3pt, + with_g1_robot=with_g1_robot, + enable_waist_tracking=enable_waist_tracking, + enable_smpl_vis=enable_smpl_vis, + log_prefix="PoseLoop", + ) + + pose_streamer = PoseStreamer( + socket=socket, + reader=reader, + three_point=three_point, + num_frames_to_send=num_frames_to_send, + target_fps=target_fps, + use_cuda=use_cuda, + record_dir=record_dir, + record_format=record_format, + log_prefix="PoseLoop", + ) + planner_streamer = PlannerStreamer( + socket=socket, + reader=reader, + three_point=three_point, + poll_hz=20, + zmq_feedback_host=zmq_feedback_host, + zmq_feedback_port=zmq_feedback_port, + ) + + # State machine diagram: + # + # Chain 1 (by_pressed enters/exits, left_axis_click toggles sub-mode): + # POSE <--(by)--> PLANNER_FROZEN_UPPER_BODY <--(left_axis_click)--> PLANNER_VR_3PT + # | + # (by)--> POSE + # + # Chain 2 (ax_pressed enters/exits, left_axis_click toggles sub-mode): + # POSE <--(ax)--> PLANNER <--(left_axis_click)--> PLANNER_VR_3PT + # | + # (ax)--> POSE + # + # Emergency stop from any mode: A+B+X+Y (start_combo) --> OFF + # POSE_PAUSE: left_menu_button held --> POSE_PAUSE, released --> POSE + # + print("Manager controls: A+X=toggle mode, A+B+X+Y=start/stop policy") + current_mode = StreamMode.OFF + # Track which mode VR_3PT was entered from, so left_axis_click returns to it. + # Will be either PLANNER or PLANNER_FROZEN_UPPER_BODY. + vr3pt_parent_mode = StreamMode.PLANNER + prev_toggle_dc = False + prev_toggle_da = False + try: + prev_ax_pressed = False + prev_by_pressed = False + prev_start_combo = False + prev_left_axis_click = False + while True: + # Poll Pico controller for buttons/axes + a_pressed, b_pressed, x_pressed, y_pressed = get_abxy_buttons(reader) + + left_menu_button, _, _, left_grip_mgr, _ = get_controller_inputs(reader) + + left_axis_click, _ = get_axis_clicks(reader) + + # Rising edge: A+X pressed together -> toggle POSE/PLANNER mode + ax_pressed = (a_pressed) and (x_pressed) + + # Rising edge: B+Y pressed together -> toggle POSE/PLANNER_FROZEN_UPPER_BODY mode + by_pressed = (b_pressed) and (y_pressed) + + # Rising edge: A+B+X+Y pressed together -> toggle policy start/stop (planner=True) + start_combo = (a_pressed) and (b_pressed) and (x_pressed) and (y_pressed) + + new_mode = current_mode + if current_mode == StreamMode.OFF: + if start_combo and not prev_start_combo: + new_mode = StreamMode.PLANNER + # Calibrate VR 3pt tracking NOW: operator should be in zero-ref pose. + # Uses the current Pico SMPL frame + FK of all-zero body joints. + sample = reader.get_latest() + if sample is not None: + three_point.calibrate_now(sample["body_poses_np"]) + else: + print("[Manager] WARNING: No SMPL data available for calibration") + + elif current_mode == StreamMode.PLANNER: + # Chain 2: POSE <--(ax)--> PLANNER <--(left_axis_click)--> VR_3PT + if start_combo and not prev_start_combo: + new_mode = StreamMode.OFF + elif ax_pressed and not prev_ax_pressed: + new_mode = StreamMode.POSE + elif left_axis_click and not prev_left_axis_click: + new_mode = StreamMode.PLANNER_VR_3PT + + elif current_mode == StreamMode.POSE: + if start_combo and not prev_start_combo: + new_mode = StreamMode.OFF + elif ax_pressed and not prev_ax_pressed: + new_mode = StreamMode.PLANNER # Enter chain 2 + elif by_pressed and not prev_by_pressed: + new_mode = StreamMode.PLANNER_FROZEN_UPPER_BODY # Enter chain 1 + elif left_menu_button: + new_mode = StreamMode.POSE_PAUSE + + elif current_mode == StreamMode.PLANNER_FROZEN_UPPER_BODY: + # Chain 1: POSE <--(by)--> FROZEN <--(left_axis_click)--> VR_3PT + if start_combo and not prev_start_combo: + new_mode = StreamMode.OFF + elif by_pressed and not prev_by_pressed: + new_mode = StreamMode.POSE + elif left_axis_click and not prev_left_axis_click: + new_mode = StreamMode.PLANNER_VR_3PT + + elif current_mode == StreamMode.POSE_PAUSE: + if start_combo and not prev_start_combo: + new_mode = StreamMode.OFF + elif not left_menu_button: + new_mode = StreamMode.POSE + + elif current_mode == StreamMode.PLANNER_VR_3PT: + # VR_3PT is reachable from both chains: + # left_axis_click → return to parent (PLANNER or FROZEN) + # ax_pressed → POSE (chain 2 exit) + # by_pressed → POSE (chain 1 exit) + if start_combo and not prev_start_combo: + new_mode = StreamMode.OFF + elif left_axis_click and not prev_left_axis_click: + new_mode = vr3pt_parent_mode # Return to parent mode + elif ax_pressed and not prev_ax_pressed: + new_mode = StreamMode.POSE + elif by_pressed and not prev_by_pressed: + new_mode = StreamMode.POSE + + # Handle mode transitions before running loop + if new_mode != current_mode: + if current_mode == StreamMode.POSE: + pose_streamer.on_mode_exit() + + # Track parent when entering VR_3PT + if new_mode == StreamMode.PLANNER_VR_3PT: + vr3pt_parent_mode = current_mode + print(f"[Manager] VR_3PT parent: {vr3pt_parent_mode.name}") + + if new_mode == StreamMode.POSE: + pose_streamer.reset_yaw() + elif new_mode == StreamMode.PLANNER and current_mode != StreamMode.PLANNER_VR_3PT: + # Only reset yaw when freshly entering PLANNER from POSE, + # not when returning from VR_3PT sub-mode + planner_streamer.reset_yaw() + elif new_mode == StreamMode.PLANNER_FROZEN_UPPER_BODY: + if current_mode != StreamMode.PLANNER_VR_3PT: + # Freshly entering from POSE: reset yaw and grab initial targets + planner_streamer.reset_yaw() + # Always re-grab the latest robot state as frozen targets, + # whether entering from POSE or returning from VR_3PT + # (the old targets are stale after VR_3PT moved the arms) + planner_streamer.save_upper_body_position_target() + elif new_mode == StreamMode.PLANNER_VR_3PT: + # Recalibrate VR tracking against the robot's actual current pose + # (read via g1_debug feedback + FK) to prevent sudden jumps + planner_streamer.recalibrate_for_vr3pt() + + # Run one iteration of the new mode + if new_mode == StreamMode.POSE: + pose_streamer.run_once() + elif ( + new_mode == StreamMode.PLANNER + or new_mode == StreamMode.PLANNER_FROZEN_UPPER_BODY + or new_mode == StreamMode.PLANNER_VR_3PT + ): + planner_streamer.run_once(new_mode) + + # Make sure to send command messages after loop iteration to ensure data arrives before mode switch + if new_mode != current_mode: + if new_mode == StreamMode.OFF: + socket.send(build_command_message(start=False, stop=True, planner=True)) + exit() + elif ( + new_mode == StreamMode.PLANNER + or new_mode == StreamMode.PLANNER_FROZEN_UPPER_BODY + or new_mode == StreamMode.PLANNER_VR_3PT + ): + socket.send(build_command_message(start=True, stop=False, planner=True)) + elif new_mode == StreamMode.POSE: + socket.send(build_command_message(start=True, stop=False, planner=False)) + + print(f"[Manager] StreamMode switch: {current_mode.name} -> {new_mode.name}") + current_mode = new_mode + + # Mode-independent: send manager_state for data exporter + toggle_dc_tmp = bool(a_pressed) and left_grip_mgr > 0.5 + toggle_da_tmp = bool(b_pressed) and left_grip_mgr > 0.5 + toggle_dc = toggle_dc_tmp and not prev_toggle_dc + toggle_da = toggle_da_tmp and not prev_toggle_da + prev_toggle_dc = toggle_dc_tmp + prev_toggle_da = toggle_da_tmp + socket.send( + pack_pose_message( + { + "stream_mode": np.array([current_mode.value], dtype=np.int32), + "toggle_data_collection": np.array([toggle_dc], dtype=bool), + "toggle_data_abort": np.array([toggle_da], dtype=bool), + }, + topic="manager_state", + ) + ) + + prev_ax_pressed = ax_pressed + prev_by_pressed = by_pressed + prev_start_combo = start_combo + prev_left_axis_click = left_axis_click + + except KeyboardInterrupt: + print("\nStopping manager...") + finally: + # Cleanup resources + reader.stop() + three_point.close() + socket.close() + context.term() + print("[Manager] Shutdown complete") + + +if __name__ == "__main__": + + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument("--buffer_size", type=int, default=15, help="Sliding window buffer size") + parser.add_argument("--port", type=int, default=5556, help="ZMQ server port (default: 5556)") + parser.add_argument( + "--num_frames_to_send", type=int, default=5, help="Number of frames to send (default: 200)" + ) + parser.add_argument("--target_fps", type=int, default=50, help="Target loop FPS (default: 50)") + parser.add_argument( + "--cuda", action="store_true", help="Use CUDA for tensors and model (default: CPU)" + ) + parser.add_argument( + "--record_dir", + type=str, + default="", + help="Directory to save sent batches (default: disabled)", + ) + parser.add_argument( + "--record_format", + type=str, + default="npz", + help="Recording format: 'npz' or 'bin' (default: npz)", + ) + parser.add_argument( + "--manager", + action="store_true", + help="Run manager with planner and pose threads (interactive)", + ) + parser.add_argument( + "--zmq_feedback_host", + type=str, + default="localhost", + help="ZMQ feedback host (default: localhost)", + ) + parser.add_argument( + "--zmq_feedback_port", + type=int, + default=5557, + help="ZMQ feedback port (default: 5557)", + ) + parser.add_argument( + "--vr3pt_test", + action="store_true", + help="Run VR 3-point pose visualizer test (reference frames only)", + ) + parser.add_argument( + "--vr3pt_live", + action="store_true", + help="Capture one frame of VR 3-point pose and visualize with reference frames", + ) + parser.add_argument( + "--vr3pt_realtime", + action="store_true", + help="Run standalone real-time VR 3-point pose visualizer", + ) + parser.add_argument( + "--vis_vr3pt", + action="store_true", + help="Enable inline VR 3-point pose visualization in pose streaming mode", + ) + parser.add_argument( + "--vr3pt_hz", + type=int, + default=10, + help="Update rate for real-time VR visualization in Hz (default: 10)", + ) + parser.add_argument( + "--no_g1", + action="store_true", + help="Disable G1 robot visualization in VR 3pt pose view (G1 is shown by default)", + ) + parser.add_argument( + "--waist_tracking", + action="store_true", + help="Enable G1 robot waist to follow VR head orientation (disabled by default for performance)", + ) + parser.add_argument( + "--vis_smpl", + action="store_true", + help="Enable SMPL body joint visualization (24 joint spheres) in the VR3pt viewer", + ) + parser.add_argument( + "--input-source", + type=str, + default="xrt", + choices=["xrt", "isaac-teleop"], + help=( + "Input source: 'xrt' for XRoboToolkit SDK (default), " + "'isaac-teleop' for in-process IsaacTeleop / CloudXR DeviceIO" + ), + ) + args = parser.parse_args() + + # Standalone VR3Pt test modes (exit after finishing) + if args.vr3pt_test: + print("Running VR 3-point pose visualizer test...") + run_vr3pt_visualizer_test() + print("VR 3-point pose visualizer test completed") + exit(0) + + if args.vr3pt_live: + print("Running VR 3-point pose live capture...") + run_vr3pt_live_visualizer() + print("VR 3-point pose live visualizer completed") + exit(0) + + if args.vr3pt_realtime: + print("Running VR 3-point pose real-time visualizer...") + run_vr3pt_realtime_visualizer(update_hz=args.vr3pt_hz) + print("VR 3-point pose real-time visualizer completed") + exit(0) + + # Main execution modes + # G1 robot visualization is enabled by default when vis_vr3pt is used + with_g1_robot = not args.no_g1 + + if args.manager: + run_pico_manager( + port=args.port, + buffer_size=args.buffer_size, + num_frames_to_send=args.num_frames_to_send, + target_fps=args.target_fps, + use_cuda=args.cuda, + record_dir=args.record_dir, + record_format=args.record_format, + zmq_feedback_host=args.zmq_feedback_host, + zmq_feedback_port=args.zmq_feedback_port, + enable_vis_vr3pt=args.vis_vr3pt, + with_g1_robot=with_g1_robot, + enable_waist_tracking=args.waist_tracking, + enable_smpl_vis=args.vis_smpl, + input_source=args.input_source, + ) + else: + # Run legacy single-thread pose streaming + run_pico( + buffer_size=args.buffer_size, + port=args.port, + num_frames_to_send=args.num_frames_to_send, + target_fps=args.target_fps, + use_cuda=args.cuda, + record_dir=args.record_dir, + record_format=args.record_format, + enable_vis_vr3pt=args.vis_vr3pt, + with_g1_robot=with_g1_robot, + enable_waist_tracking=args.waist_tracking, + enable_smpl_vis=args.vis_smpl, + input_source=args.input_source, + ) diff --git a/GR00T-WholeBodyControl/gear_sonic/scripts/process_dataset.py b/GR00T-WholeBodyControl/gear_sonic/scripts/process_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..2f3ad46914401531ec05aa44db2d85330a1f3273 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/scripts/process_dataset.py @@ -0,0 +1,634 @@ +""" +Post-process a LeRobot dataset recorded by the data exporter. + +Removes discarded episodes (flagged during collection) and stale SMPL frames +(all-zero teleop.smpl_pose and frozen lead-in frames that precede them) which +occur during teleop pauses or ZMQ frame drops. Can also merge multiple +recording sessions into a single dataset. + +The script operates directly on the LeRobot v2.1 on-disk format +(parquet + mp4) without any external training framework dependencies. + +Usage: + + # Clean a single dataset in-place + python gear_sonic/scripts/process_dataset.py \\ + --dataset-path outputs/my_dataset + + # Clean and write to a new directory (non-destructive) + python gear_sonic/scripts/process_dataset.py \\ + --dataset-path outputs/my_dataset \\ + --output-path outputs/my_dataset_cleaned + + # Merge multiple datasets into one (validates matching script_config) + python gear_sonic/scripts/process_dataset.py \\ + --dataset-path outputs/session1 outputs/session2 outputs/session3 \\ + --output-path outputs/merged_dataset + + # Merge from a text file listing dataset paths (one per line) + python gear_sonic/scripts/process_dataset.py \\ + --dataset-list datasets.txt \\ + --output-path outputs/merged_dataset + + # Skip SMPL cleaning (merge only) + python gear_sonic/scripts/process_dataset.py \\ + --dataset-path outputs/session1 outputs/session2 \\ + --output-path outputs/merged \\ + --no-remove-stale-smpl + + # Remove discarded episodes (flagged during collection via 'x' key) + python gear_sonic/scripts/process_dataset.py \\ + --dataset-path outputs/my_dataset \\ + --output-path outputs/my_dataset_cleaned \\ + --remove-discarded +""" + +from dataclasses import dataclass, field +import json +from pathlib import Path +import shutil +from typing import Optional + +import av +import numpy as np +import pandas as pd +import tyro + + +SMPL_POSE_COLUMN = "teleop.smpl_pose" + + +# --------------------------------------------------------------------------- +# Stale SMPL frame detection +# --------------------------------------------------------------------------- + +def build_stale_mask(smpl_arr: np.ndarray) -> np.ndarray: + """Return a boolean mask where True = frame should be removed. + + Marks all-zero rows AND any consecutive frozen (identical-to-next) rows + that immediately precede a zero row. Frozen runs that do NOT lead into + a zero row are left untouched — those occur naturally when the SMPL + stream publishes at a slightly lower rate than the collection loop. + """ + n = len(smpl_arr) + is_zero = np.all(smpl_arr == 0, axis=1) + remove = is_zero.copy() + + diffs = np.zeros(n) + diffs[1:] = np.sum(np.abs(smpl_arr[1:] - smpl_arr[:-1]), axis=1) + + for i in range(n): + if is_zero[i]: + j = i - 1 + while j >= 0 and diffs[j] == 0.0 and not is_zero[j]: + remove[j] = True + j -= 1 + + return remove + + +# --------------------------------------------------------------------------- +# LeRobot on-disk helpers +# --------------------------------------------------------------------------- + +def load_info(dataset_path: Path) -> dict: + info_path = dataset_path / "meta" / "info.json" + with open(info_path, encoding="utf-8") as f: + return json.load(f) + + +def load_episodes_meta(dataset_path: Path) -> list[dict]: + episodes_path = dataset_path / "meta" / "episodes.jsonl" + episodes = [] + with open(episodes_path, encoding="utf-8") as f: + for line in f: + line = line.strip() + if line: + episodes.append(json.loads(line)) + return episodes + + +def load_tasks_meta(dataset_path: Path) -> list[dict]: + tasks_path = dataset_path / "meta" / "tasks.jsonl" + tasks = [] + if tasks_path.exists(): + with open(tasks_path, encoding="utf-8") as f: + for line in f: + line = line.strip() + if line: + tasks.append(json.loads(line)) + return tasks + + +def get_parquet_path(dataset_path: Path, info: dict, episode_index: int) -> Path: + data_path_pattern = info.get("data_path", "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet") + chunks_size = info.get("chunks_size", 1000) + episode_chunk = episode_index // chunks_size + return dataset_path / data_path_pattern.format( + episode_chunk=episode_chunk, episode_index=episode_index, + ) + + +def get_video_keys(info: dict) -> list[str]: + """Extract video keys from info.json, falling back to features if needed.""" + keys = info.get("video_keys", []) + if not keys: + keys = [ + k for k, v in info.get("features", {}).items() + if v.get("dtype") in ("video", "image") + ] + return keys + + +def get_video_paths(dataset_path: Path, info: dict, episode_index: int) -> dict[str, Path]: + video_path_pattern = info.get( + "video_path", + "videos/{video_key}/episode_{episode_index:06d}.mp4", + ) + video_keys = get_video_keys(info) + chunks_size = info.get("chunks_size", 1000) + episode_chunk = episode_index // chunks_size + paths = {} + for key in video_keys: + paths[key] = dataset_path / video_path_pattern.format( + video_key=key, episode_index=episode_index, + episode_chunk=episode_chunk, + ) + return paths + + +def filter_video_frames(video_path: Path, valid_indices: np.ndarray, fps: int): + """Re-encode a video keeping only frames at valid_indices.""" + input_container = av.open(str(video_path)) + input_stream = input_container.streams.video[0] + + all_frames = [] + for frame in input_container.decode(input_stream): + all_frames.append(frame.to_ndarray(format="rgb24")) + input_container.close() + + if len(all_frames) == 0: + return + + filtered = [all_frames[i] for i in valid_indices if i < len(all_frames)] + if len(filtered) == 0: + return + + tmp_path = video_path.with_suffix(".tmp.mp4") + output_container = av.open(str(tmp_path), mode="w") + output_stream = output_container.add_stream("h264", rate=fps) + h, w = filtered[0].shape[:2] + output_stream.width = w + output_stream.height = h + output_stream.pix_fmt = "yuv420p" + + for img in filtered: + frame = av.VideoFrame.from_ndarray(img, format="rgb24") + for packet in output_stream.encode(frame): + output_container.mux(packet) + for packet in output_stream.encode(): + output_container.mux(packet) + output_container.close() + + tmp_path.replace(video_path) + + +# --------------------------------------------------------------------------- +# Script config validation +# --------------------------------------------------------------------------- + +def validate_script_configs(dataset_paths: list[Path]) -> dict | None: + """Check that all datasets share the same script_config. + + Returns the common config if they match, or raises an error with + details about which datasets differ. + """ + configs = {} + for ds_path in dataset_paths: + info = load_info(ds_path) + sc = info.get("script_config") + if sc is not None: + configs[ds_path.name] = sc + + if not configs: + return None + + canonical = json.dumps(next(iter(configs.values())), sort_keys=True) + mismatched = [] + for name, cfg in configs.items(): + if json.dumps(cfg, sort_keys=True) != canonical: + mismatched.append(name) + + if mismatched: + print("\nERROR: script_config mismatch across datasets.") + print("The following datasets have different robot configurations:\n") + ref_name = next(iter(configs.keys())) + print(f" Reference: {ref_name}") + for name in mismatched: + print(f" Differs: {name}") + print( + "\nDatasets recorded with different robot configurations cannot be " + "merged. Verify that all sessions used the same robot setup." + ) + raise SystemExit(1) + + return next(iter(configs.values())) + + +# --------------------------------------------------------------------------- +# Core processing +# --------------------------------------------------------------------------- + +def process_single_dataset( + dataset_path: Path, + remove_stale_smpl: bool, + remove_discarded: bool = False, + episode_index_offset: int = 0, +) -> dict: + """Process one dataset: optionally clean stale SMPL frames. + + Returns stats dict and the list of (parquet_df, video_paths, episode_meta) + tuples for merging. + """ + info = load_info(dataset_path) + episodes_meta = load_episodes_meta(dataset_path) + fps = info.get("fps", 50) + + discarded_indices = set(info.get("discarded_episode_indices", [])) if remove_discarded else set() + + stats = { + "total_episodes": len(episodes_meta), + "episodes_with_stale": 0, + "total_frames": 0, + "frames_removed": 0, + "zero_frames": 0, + "frozen_leadin_frames": 0, + "episodes_dropped": 0, + "episodes_discarded": 0, + } + processed_episodes = [] + + for ep_meta in episodes_meta: + ep_idx = ep_meta["episode_index"] + + if ep_idx in discarded_indices: + stats["episodes_discarded"] += 1 + print(f" Episode {ep_idx}: discarded during collection — removing") + continue + + parquet_path = get_parquet_path(dataset_path, info, ep_idx) + video_paths = get_video_paths(dataset_path, info, ep_idx) + + if not parquet_path.exists(): + print(f" WARNING: Missing parquet for episode {ep_idx}, skipping") + continue + + df = pd.read_parquet(parquet_path) + ep_len = len(df) + stats["total_frames"] += ep_len + + valid_indices = None + + if remove_stale_smpl and SMPL_POSE_COLUMN in df.columns: + smpl_arr = np.vstack( + [np.asarray(x, dtype=np.float32) for x in df[SMPL_POSE_COLUMN]] + ) + mask = build_stale_mask(smpl_arr) + n_remove = int(mask.sum()) + n_zero = int(np.all(smpl_arr == 0, axis=1).sum()) + n_frozen = n_remove - n_zero + + if n_remove > 0: + stats["episodes_with_stale"] += 1 + stats["frames_removed"] += n_remove + stats["zero_frames"] += n_zero + stats["frozen_leadin_frames"] += n_frozen + pct = 100.0 * n_remove / ep_len + print( + f" Episode {ep_idx}: removing {n_remove}/{ep_len} frames " + f"({pct:.1f}%) — {n_zero} zero + {n_frozen} frozen lead-in" + ) + + if n_remove == ep_len: + print(f" Episode {ep_idx}: ALL frames stale — dropping episode") + stats["episodes_dropped"] += 1 + continue + + valid_indices = np.where(~mask)[0] + df = df.iloc[valid_indices].copy().reset_index(drop=True) + if "timestamp" in df.columns: + df["timestamp"] -= df["timestamp"].iloc[0] + + new_ep_idx = ep_idx + episode_index_offset + processed_episodes.append({ + "df": df, + "source_video_paths": video_paths, + "valid_indices": valid_indices, + "episode_meta": ep_meta, + "new_episode_index": new_ep_idx, + "fps": fps, + }) + + return stats, processed_episodes, info + + +def write_output_dataset( + dest_path: Path, + all_episodes: list[dict], + reference_info: dict, + tasks_meta: list[dict], + script_config: dict | None, +): + """Write processed episodes to a new LeRobot dataset directory.""" + dest_path.mkdir(parents=True, exist_ok=True) + meta_dir = dest_path / "meta" + meta_dir.mkdir(exist_ok=True) + + info = reference_info.copy() + fps = info.get("fps", 50) + chunks_size = info.get("chunks_size", 1000) + + if script_config is not None: + info["script_config"] = script_config + + total_frames = 0 + episodes_jsonl = [] + + for i, ep in enumerate(all_episodes): + df = ep["df"] + ep_len = len(df) + + df["episode_index"] = i + df["index"] = range(total_frames, total_frames + ep_len) + df["frame_index"] = range(ep_len) + if "timestamp" in df.columns: + df["timestamp"] = [j / fps for j in range(ep_len)] + + episode_chunk = i // chunks_size + data_path_pattern = info.get( + "data_path", + "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet", + ) + parquet_rel = data_path_pattern.format(episode_chunk=episode_chunk, episode_index=i) + parquet_path = dest_path / parquet_rel + parquet_path.parent.mkdir(parents=True, exist_ok=True) + df.to_parquet(parquet_path) + + video_keys = get_video_keys(info) + video_path_pattern = info.get( + "video_path", + "videos/{video_key}/episode_{episode_index:06d}.mp4", + ) + for vkey in video_keys: + src_video = ep["source_video_paths"].get(vkey) + dst_rel = video_path_pattern.format( + video_key=vkey, episode_index=i, episode_chunk=episode_chunk, + ) + dst_video = dest_path / dst_rel + dst_video.parent.mkdir(parents=True, exist_ok=True) + + if src_video and src_video.exists(): + if ep["valid_indices"] is not None: + shutil.copy2(src_video, dst_video) + filter_video_frames(dst_video, ep["valid_indices"], fps) + else: + shutil.copy2(src_video, dst_video) + + ep_meta = { + "episode_index": i, + "tasks": ep["episode_meta"].get("tasks", []), + "length": ep_len, + } + episodes_jsonl.append(ep_meta) + + total_frames += ep_len + + info["total_episodes"] = len(all_episodes) + info["total_frames"] = total_frames + info.pop("discarded_episode_indices", None) + + with open(meta_dir / "info.json", "w", encoding="utf-8") as f: + json.dump(info, f, indent=4) + + with open(meta_dir / "episodes.jsonl", "w", encoding="utf-8") as f: + for ep in episodes_jsonl: + f.write(json.dumps(ep) + "\n") + + if tasks_meta: + with open(meta_dir / "tasks.jsonl", "w", encoding="utf-8") as f: + for task in tasks_meta: + f.write(json.dumps(task) + "\n") + + return total_frames + + +def copy_modality_json(source_paths: list[Path], output_path: Path): + """Copy modality.json from the first source that has one.""" + for src in source_paths: + modality_path = src / "meta" / "modality.json" + if modality_path.exists(): + dst = output_path / "meta" / "modality.json" + dst.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(modality_path, dst) + return + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + +@dataclass +class ProcessDatasetConfig: + """Post-process LeRobot datasets: clean stale SMPL frames and/or merge.""" + + dataset_path: list[str] = field(default_factory=list) + """One or more dataset directories to process.""" + + dataset_list: Optional[str] = None + """Path to a text file listing dataset directories (one per line). + Can be used instead of or in addition to --dataset-path.""" + + output_path: Optional[str] = None + """Output directory for the processed dataset. If not specified and a + single dataset is given, the dataset is modified in-place.""" + + remove_stale_smpl: bool = True + """Remove frames where teleop.smpl_pose is all zeros (stale/dropped + SMPL data) and frozen lead-in frames that precede them.""" + + remove_discarded: bool = True + """Remove episodes that were flagged as discarded during data collection + (stored in meta/info.json under discarded_episode_indices).""" + + +def main(cfg: ProcessDatasetConfig): + dataset_paths = [Path(p) for p in cfg.dataset_path] + + if cfg.dataset_list: + list_file = Path(cfg.dataset_list) + with open(list_file, encoding="utf-8") as f: + for line in f: + line = line.strip() + if line and not line.startswith("#"): + dataset_paths.append(Path(line)) + + if not dataset_paths: + print("ERROR: No dataset paths provided. Use --dataset-path or --dataset-list.") + raise SystemExit(1) + + for ds in dataset_paths: + if not ds.exists(): + print(f"ERROR: Dataset path does not exist: {ds}") + raise SystemExit(1) + if not (ds / "meta" / "info.json").exists(): + print(f"ERROR: Not a valid LeRobot dataset (missing meta/info.json): {ds}") + raise SystemExit(1) + + merging = len(dataset_paths) > 1 + in_place = cfg.output_path is None + + if merging and in_place: + print("ERROR: --output-path is required when merging multiple datasets.") + raise SystemExit(1) + + output_path = Path(cfg.output_path) if cfg.output_path else dataset_paths[0] + + print("=" * 70) + print(" LeRobot Dataset Processor") + print("=" * 70) + print(f" Input datasets: {len(dataset_paths)}") + for ds in dataset_paths: + print(f" - {ds}") + print(f" Output: {output_path}{' (in-place)' if in_place else ''}") + print(f" Remove stale SMPL: {cfg.remove_stale_smpl}") + print(f" Remove discarded: {cfg.remove_discarded}") + print("=" * 70) + + # Validate script configs match across all datasets + if merging: + print("\nValidating script_config across datasets...") + script_config = validate_script_configs(dataset_paths) + print(" All datasets have matching robot configurations.\n") + else: + info = load_info(dataset_paths[0]) + script_config = info.get("script_config") + + # Collect tasks from all datasets (deduplicated) + all_tasks_meta: list[dict] = [] + seen_task_ids: set = set() + for ds in dataset_paths: + for task in load_tasks_meta(ds): + tid = task.get("task_index", id(task)) + if tid not in seen_task_ids: + all_tasks_meta.append(task) + seen_task_ids.add(tid) + + # Process each dataset + all_episodes = [] + total_stats = { + "total_episodes": 0, + "episodes_with_stale": 0, + "total_frames": 0, + "frames_removed": 0, + "zero_frames": 0, + "frozen_leadin_frames": 0, + "episodes_dropped": 0, + "episodes_discarded": 0, + } + reference_info = None + + for ds_path in dataset_paths: + print(f"\nProcessing: {ds_path}") + stats, episodes, info = process_single_dataset( + ds_path, + remove_stale_smpl=cfg.remove_stale_smpl, + remove_discarded=cfg.remove_discarded, + episode_index_offset=len(all_episodes), + ) + + if reference_info is None: + reference_info = info + + all_episodes.extend(episodes) + for key in total_stats: + total_stats[key] += stats[key] + + if not all_episodes: + print("\nERROR: No valid episodes after processing.") + raise SystemExit(1) + + # Write output + if in_place: + # In-place: rewrite parquet files and re-encode videos + print(f"\nRewriting dataset in-place at {output_path}...") + ds_info = load_info(output_path) + fps = ds_info.get("fps", 50) + + # Delete files for discarded episodes + if cfg.remove_discarded: + discarded_indices = set(ds_info.get("discarded_episode_indices", [])) + for ep_idx in discarded_indices: + parquet_path = get_parquet_path(output_path, ds_info, ep_idx) + if parquet_path.exists(): + parquet_path.unlink() + video_paths = get_video_paths(output_path, ds_info, ep_idx) + for _vkey, vpath in video_paths.items(): + if vpath.exists(): + vpath.unlink() + + for ep in all_episodes: + ep_idx = ep["episode_meta"]["episode_index"] + parquet_path = get_parquet_path(output_path, ds_info, ep_idx) + ep["df"].to_parquet(parquet_path) + + if ep["valid_indices"] is not None: + video_paths = get_video_paths(output_path, ds_info, ep_idx) + for _vkey, vpath in video_paths.items(): + if vpath.exists(): + filter_video_frames(vpath, ep["valid_indices"], fps) + + # Update episode metadata + episodes_meta = [] + for ep in all_episodes: + meta = ep["episode_meta"].copy() + meta["length"] = len(ep["df"]) + episodes_meta.append(meta) + + with open(output_path / "meta" / "episodes.jsonl", "w", encoding="utf-8") as f: + for em in episodes_meta: + f.write(json.dumps(em) + "\n") + + ds_info["total_frames"] = sum(len(ep["df"]) for ep in all_episodes) + ds_info["total_episodes"] = len(all_episodes) + if cfg.remove_discarded: + ds_info.pop("discarded_episode_indices", None) + with open(output_path / "meta" / "info.json", "w", encoding="utf-8") as f: + json.dump(ds_info, f, indent=4) + else: + print(f"\nWriting output dataset to {output_path}...") + write_output_dataset( + output_path, all_episodes, reference_info, all_tasks_meta, script_config, + ) + copy_modality_json(dataset_paths, output_path) + + # Print summary + kept = total_stats["total_frames"] - total_stats["frames_removed"] + kept_episodes = total_stats["total_episodes"] - total_stats["episodes_dropped"] - total_stats["episodes_discarded"] + + print("\n" + "=" * 70) + print(" Processing complete!") + print("=" * 70) + print(f" Episodes: {kept_episodes} kept / {total_stats['total_episodes']} total" + f" ({total_stats['episodes_dropped']} dropped, {total_stats['episodes_discarded']} discarded)") + print(f" Frames: {kept} kept / {total_stats['total_frames']} total" + f" ({total_stats['frames_removed']} removed)") + if total_stats["frames_removed"] > 0: + print(f" Zero SMPL: {total_stats['zero_frames']}") + print(f" Frozen lead-in: {total_stats['frozen_leadin_frames']}") + print(f" Episodes affected: {total_stats['episodes_with_stale']}") + print(f" Output: {output_path}") + print("=" * 70) + + +if __name__ == "__main__": + main(tyro.cli(ProcessDatasetConfig)) diff --git a/GR00T-WholeBodyControl/gear_sonic/scripts/run_camera_viewer.py b/GR00T-WholeBodyControl/gear_sonic/scripts/run_camera_viewer.py new file mode 100644 index 0000000000000000000000000000000000000000..17e5dc06de7342ae246a9beec632c351281c1977 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/scripts/run_camera_viewer.py @@ -0,0 +1,264 @@ +""" +ROS-free camera viewer with optional recording. + +Connects to a ZMQ camera server (MuJoCo sim SensorServer or real robot camera) +and displays live camera feeds using OpenCV. Supports recording to MP4. + +Virtual environment setup (run from repo root): + bash install_scripts/install_data_collection.sh + source .venv_data_collection/bin/activate + +Usage: + python gear_sonic/scripts/run_camera_viewer.py --camera-host localhost --camera-port 5555 + +Headless recording: + python gear_sonic/scripts/run_camera_viewer.py --no-display --record --duration 10 + +Controls (OpenCV window must be focused): + R - Start/stop recording + Q - Quit + +Output structure: + camera_recordings/ + └── rec_20260403_143052/ + ├── ego_view.mp4 + └── head_left_color_image.mp4 +""" + +from dataclasses import dataclass +from pathlib import Path +import time +from typing import Optional + +import cv2 +import numpy as np +import tyro + +from gear_sonic.camera.composed_camera import ComposedCameraClientSensor + + +@dataclass +class CameraViewerConfig: + """CLI config for the ROS-free camera viewer.""" + + camera_host: str = "localhost" + """Camera server hostname.""" + + camera_port: int = 5555 + """Camera server port.""" + + fps: int = 30 + """Target display refresh rate (Hz).""" + + output_path: Optional[str] = None + """Output directory for recordings. Auto-creates 'camera_recordings/' if not set.""" + + codec: str = "mp4v" + """Video codec for recording (e.g., 'mp4v', 'XVID').""" + + max_display_width: int = 640 + """Max width per camera tile in the display window.""" + + display: bool = True + """Show the OpenCV preview window. Disable this on headless servers.""" + + record: bool = False + """Start recording immediately instead of waiting for the R key.""" + + duration: Optional[float] = None + """Stop automatically after this many seconds of recording.""" + + +def main(config: CameraViewerConfig): + if not config.display and not config.record: + raise ValueError("--no-display requires --record") + if config.duration is not None and not config.record: + raise ValueError("--duration requires --record") + if config.duration is not None and config.duration <= 0: + raise ValueError("--duration must be greater than zero") + + client = ComposedCameraClientSensor(server_ip=config.camera_host, port=config.camera_port) + + print("Waiting for first camera frame...") + sample = None + for _ in range(100): + sample = client.read(blocking=False) + if sample and sample.get("images"): + break + time.sleep(0.1) + + if sample is None or not sample.get("images"): + print("ERROR: No camera frames received after 10s. Check the camera server.") + return + + camera_names = sorted(sample["images"].keys()) + print(f"Detected {len(camera_names)} camera stream(s): {camera_names}") + + output_dir = Path(config.output_path) if config.output_path else Path("camera_recordings") + + is_recording = False + video_writers: dict[str, cv2.VideoWriter] = {} + frame_count = 0 + recording_start_time = 0.0 + recording_dir = Path(".") + loop_period = 1.0 / config.fps + + window_name = "SONIC Camera Viewer" + + print(f"Target FPS: {config.fps}") + print(f"Recordings will be saved to: {output_dir}") + if config.display: + print("Controls: R = start/stop recording, Q = quit") + + if config.record: + recording_dir = output_dir / f"rec_{time.strftime('%Y%m%d_%H%M%S')}" + recording_dir.mkdir(parents=True, exist_ok=True) + + fourcc = cv2.VideoWriter_fourcc(*config.codec) + for name in camera_names: + img = sample["images"].get(name) + if img is not None: + h, w = img.shape[:2] + path = recording_dir / f"{name}.mp4" + writer = cv2.VideoWriter(str(path), fourcc, config.fps, (w, h)) + if not writer.isOpened(): + raise RuntimeError(f"Failed to open video writer: {path}") + video_writers[name] = writer + + is_recording = True + recording_start_time = time.time() + print(f"Recording started: {recording_dir}") + + try: + while True: + t_start = time.monotonic() + + image_data = client.read(blocking=False) + if image_data is None or not image_data.get("images"): + elapsed = time.monotonic() - t_start + remaining = loop_period - elapsed + if remaining > 0: + time.sleep(remaining) + continue + + tiles = [] + for name in camera_names: + img = image_data["images"].get(name) + if img is None: + continue + + if img.shape[2] == 3: + img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) + else: + img_bgr = img + + if is_recording and name in video_writers: + video_writers[name].write(img_bgr) + + h, w = img_bgr.shape[:2] + if w > config.max_display_width: + scale = config.max_display_width / w + img_bgr = cv2.resize( + img_bgr, (config.max_display_width, int(h * scale)) + ) + + label = f"{name}" + if is_recording: + label = f"[REC] {name}" + cv2.putText( + img_bgr, label, (10, 25), + cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2, + ) + tiles.append(img_bgr) + + if tiles: + max_h = max(t.shape[0] for t in tiles) + padded = [] + for t in tiles: + if t.shape[0] < max_h: + pad = np.zeros( + (max_h - t.shape[0], t.shape[1], 3), dtype=np.uint8 + ) + t = np.vstack([t, pad]) + padded.append(t) + canvas = np.hstack(padded) + + if is_recording: + frame_count += 1 + elapsed_rec = time.time() - recording_start_time + status = f"REC {frame_count}f / {elapsed_rec:.1f}s" + cv2.putText( + canvas, status, (canvas.shape[1] - 300, 25), + cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2, + ) + + if config.display: + cv2.imshow(window_name, canvas) + + key = cv2.waitKey(1) & 0xFF if config.display else 0xFF + + if key == ord("q"): + print("Quit requested.") + break + elif key == ord("r"): + if not is_recording: + recording_dir = output_dir / f"rec_{time.strftime('%Y%m%d_%H%M%S')}" + recording_dir.mkdir(parents=True, exist_ok=True) + + fourcc = cv2.VideoWriter_fourcc(*config.codec) + video_writers = {} + for name in camera_names: + img = image_data["images"].get(name) + if img is not None: + h, w = img.shape[:2] + path = recording_dir / f"{name}.mp4" + video_writers[name] = cv2.VideoWriter( + str(path), fourcc, config.fps, (w, h) + ) + + is_recording = True + recording_start_time = time.time() + frame_count = 0 + print(f"Recording started: {recording_dir}") + else: + is_recording = False + for writer in video_writers.values(): + writer.release() + video_writers = {} + duration = time.time() - recording_start_time + print( + f"Recording stopped - {duration:.1f}s, {frame_count} frames " + f"-> {recording_dir}" + ) + + if ( + is_recording + and config.duration is not None + and time.time() - recording_start_time >= config.duration + ): + print(f"Recording duration reached: {config.duration:.1f}s") + break + + elapsed = time.monotonic() - t_start + remaining = loop_period - elapsed + if remaining > 0: + time.sleep(remaining) + + except KeyboardInterrupt: + print("\nExiting...") + finally: + if video_writers: + for writer in video_writers.values(): + writer.release() + if is_recording: + duration = time.time() - recording_start_time + print(f"Final recording: {duration:.1f}s, {frame_count} frames") + + client.close() + if config.display: + cv2.destroyAllWindows() + + +if __name__ == "__main__": + config = tyro.cli(CameraViewerConfig) + main(config) diff --git a/GR00T-WholeBodyControl/gear_sonic/scripts/run_data_exporter.py b/GR00T-WholeBodyControl/gear_sonic/scripts/run_data_exporter.py new file mode 100644 index 0000000000000000000000000000000000000000..4737bcb0b8b891ff1d7d9223898eb31913d7e0d9 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/scripts/run_data_exporter.py @@ -0,0 +1,963 @@ +""" +Sonic VLA data exporter for G1 -- NO ROS 2 DEPENDENCY. + +All data sources use ZMQ: + 1. Robot state -> ZMQ SUB on ``g1_debug`` topic (port 5557, from C++ zmq_output_handler) + 2. SMPL pose -> ZMQ SUB on ``pose`` topic (port 5556, from pico_manager_thread_server) + 3. Camera -> ZMQ/TCP via ComposedCameraClientSensor + +Robot config (``script_config`` in info.json) is read from the ``robot_config`` +ZMQ topic re-published every ~2 s by the C++ process. If the config is not +received within the timeout the exporter exits with an error. + +Virtual environment setup (run from repo root): + bash install_scripts/install_data_collection.sh + source .venv_data_collection/bin/activate + +Usage (from repo root): + python gear_sonic/scripts/run_data_exporter.py --task-prompt "pick up the cup" + python gear_sonic/scripts/run_data_exporter.py --task-prompt "walk forward" --dataset-name my_session +""" + +from collections import deque +from dataclasses import dataclass +from datetime import datetime +import json +import time + +import numpy as np +from scipy.spatial.transform import Rotation as R +import tyro +import zmq + +from gear_sonic.data.exporter import Gr00tDataExporter +from gear_sonic.data.features_sonic_vla import ( + get_features_sonic_vla, + get_g1_robot_model, + get_modality_config_sonic_vla, + get_wrist_camera_features, + get_wrist_camera_modality_config, +) +from gear_sonic.camera.composed_camera import ComposedCameraClientSensor +from gear_sonic.utils.data_collection.episode_state import EpisodeState +from gear_sonic.utils.data_collection.keyboard_subscriber import ZMQKeyboardSubscriber +from gear_sonic.utils.data_collection.telemetry import Telemetry +from gear_sonic.utils.data_collection.text_to_speech import TextToSpeech +from gear_sonic.utils.data_collection.transforms import compute_projected_gravity, quat_to_rot6d +from gear_sonic.utils.data_collection.zmq_state_subscriber import ( + ZMQStateSubscriber, + poll_robot_config_zmq, +) + +# --------------------------------------------------------------------------- +# Config +# --------------------------------------------------------------------------- + + +@dataclass +class SonicDataExporterConfig: + """CLI config for the ROS-free Sonic data exporter.""" + + # Dataset + dataset_name: str | None = None + """Dataset name (auto-generated if creating new).""" + + task_prompt: str = "demo" + """Language task prompt.""" + + root_output_dir: str = "outputs" + """Root output directory.""" + + data_collection_frequency: int = 50 + """Data collection frequency (Hz).""" + + + # Camera + camera_host: str = "localhost" + """Camera server host.""" + + camera_port: int = 5555 + """Camera server port.""" + + # ZMQ: Sonic / SMPL pose (from pico_manager_thread_server) + sonic_zmq_host: str = "localhost" + """ZMQ host for Sonic SMPL pose messages.""" + + sonic_zmq_port: int = 5556 + """ZMQ port for Sonic SMPL pose messages.""" + + # ZMQ: Robot state (from C++ zmq_output_handler, g1_debug topic) + state_zmq_host: str = "localhost" + """ZMQ host for robot state (g1_debug topic from C++ deploy).""" + + state_zmq_port: int = 5557 + """ZMQ port for robot state (same socket as robot_config topic).""" + + # Robot config + robot_config_timeout: float = 0 + """Seconds to wait for the ZMQ robot_config message at startup (0 = wait forever).""" + + record_wrist_cameras: bool = False + """Record wrist camera streams (left_wrist, right_wrist). Requires cameras to be available.""" + + text_to_speech: bool = True + """Use text-to-speech voice feedback.""" + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +class TimeDeltaException(Exception): + def __init__(self, failure_count: int, reset_timeout_sec: float): + self.failure_count = failure_count + self.reset_timeout_sec = reset_timeout_sec + self.message = f"{self.failure_count} failures in {self.reset_timeout_sec} seconds" + super().__init__(self.message) + + +def unpack_pose_message(packed_data: bytes, topic: str = "pose") -> dict: + """Unpack a single-frame packed message from pico_manager_thread_server. + + Wire format: [topic_prefix][1280-byte JSON header][concatenated binary fields] + """ + HEADER_SIZE = 1280 + + topic_bytes = topic.encode("utf-8") + if not packed_data.startswith(topic_bytes): + raise ValueError(f"Message does not start with expected topic '{topic}'") + + offset = len(topic_bytes) + if len(packed_data) < offset + HEADER_SIZE: + raise ValueError(f"Packed data too small: {len(packed_data)} < {offset + HEADER_SIZE}") + + header_bytes = packed_data[offset : offset + HEADER_SIZE] + null_idx = header_bytes.find(b"\x00") + if null_idx > 0: + header_bytes = header_bytes[:null_idx] + + header = json.loads(header_bytes.decode("utf-8")) + fields = header.get("fields", []) + + result = {"version": header.get("v", 0), "endian": header.get("endian", "le")} + current_offset = offset + HEADER_SIZE + dtype_map = { + "f32": np.float32, + "f64": np.float64, + "i32": np.int32, + "i64": np.int64, + "bool": bool, + } + + for field in fields: + dtype = dtype_map.get(field["dtype"], np.float32) + shape = tuple(field["shape"]) + n_bytes = int(np.prod(shape)) * np.dtype(dtype).itemsize + result[field["name"]] = ( + np.frombuffer(packed_data[current_offset : current_offset + n_bytes], dtype=dtype) + .reshape(shape) + .copy() + ) + current_offset += n_bytes + + return result + + +class TimingThresholdMonitor: + def __init__(self, max_failures=3, reset_timeout_sec=5, time_delta=0.2, raise_exception=False): + self.max_failures = max_failures + self.reset_timeout_sec = reset_timeout_sec + self.failure_count = 0 + self.last_failure_time = 0 + self.time_delta = time_delta + self.raise_exception = raise_exception + + def reset(self): + self.failure_count = 0 + self.last_failure_time = 0 + + def log_time_delta(self, time_delta_sec: float): + time_delta = abs(time_delta_sec) + if time_delta > self.time_delta: + self.failure_count += 1 + self.last_failure_time = time.monotonic() + + if self.is_threshold_exceeded(): + print( + f"Time delta exception: {self.failure_count} failures in " + f"{self.reset_timeout_sec} seconds, time delta: {time_delta}" + ) + if self.raise_exception: + raise TimeDeltaException(self.failure_count, self.reset_timeout_sec) + + def is_threshold_exceeded(self): + if self.failure_count >= self.max_failures: + return True + if time.monotonic() - self.last_failure_time > self.reset_timeout_sec: + self.reset() + return False + + +# --------------------------------------------------------------------------- +# Data Collector +# --------------------------------------------------------------------------- + + +class GrootDataCollector: + """Collects data from G1 robot in Sonic CPP + SMPL mode -- no ROS 2. + + Data sources (all ZMQ): + - ``g1_debug`` topic -> proprio (body_q, hand_q, actions, base_quat, ...) + - ``pose`` topic -> SMPL pose (smpl_joints, body_quat_w, hand_joints, ...) + - ``planner`` topic -> planner commands (vr_position, vr_orientation, ...) + - ``manager_state`` topic -> current stream mode + toggle flags + - Camera client -> ego-view images + """ + + def __init__( + self, + camera_host: str, + camera_port: int, + data_exporter: Gr00tDataExporter, + robot_model, + text_to_speech=None, + frequency: int = 20, + sonic_data_zmq_host: str = "localhost", + sonic_data_zmq_port: int = 5556, + state_zmq_host: str = "localhost", + state_zmq_port: int = 5557, + ): + self.text_to_speech = text_to_speech + self.frequency = frequency + self.loop_period = 1.0 / frequency + self.data_exporter = data_exporter + self.robot_model = robot_model + + self._episode_state = EpisodeState() + self._keyboard_listener = ZMQKeyboardSubscriber() + + self._image_subscriber = ComposedCameraClientSensor(server_ip=camera_host, port=camera_port) + + self.obs_act_buffer = deque(maxlen=100) + self.latest_image_msg = None + self.latest_proprio_msg = None + self.latest_sonic_msg = None + self.latest_planner_msg = None + + self.current_stream_mode = 0 + + self._manager_toggle_dc = False + self._manager_toggle_da = False + + self._state_subscriber = ZMQStateSubscriber( + host=state_zmq_host, + port=state_zmq_port, + ) + + self._sonic_zmq_ctx = None + self._sonic_zmq_socket = None + try: + self._sonic_zmq_ctx = zmq.Context() + self._sonic_zmq_socket = self._sonic_zmq_ctx.socket(zmq.SUB) + self._sonic_zmq_socket.connect(f"tcp://{sonic_data_zmq_host}:{sonic_data_zmq_port}") + self._sonic_zmq_socket.setsockopt(zmq.RCVTIMEO, 100) + self._sonic_zmq_socket.setsockopt(zmq.CONFLATE, 0) + self._sonic_zmq_socket.setsockopt(zmq.RCVHWM, 20) + self._sonic_zmq_socket.setsockopt_string(zmq.SUBSCRIBE, "pose") + self._sonic_zmq_socket.setsockopt_string(zmq.SUBSCRIBE, "planner") + self._sonic_zmq_socket.setsockopt_string(zmq.SUBSCRIBE, "manager_state") + time.sleep(0.5) + print(f"[Sonic] Connected to ZMQ at {sonic_data_zmq_host}:{sonic_data_zmq_port}") + print("[Sonic] Subscribed to: pose, planner, manager_state") + except Exception as e: + print(f"[Sonic] Warning: Failed to initialize ZMQ subscriber: {e}") + self._sonic_zmq_socket = None + + self.telemetry = Telemetry(window_size=100) + self.sonic_timing_monitor = TimingThresholdMonitor( + max_failures=3, reset_timeout_sec=5, time_delta=0.1 + ) + + self._last_latency_log_time = 0.0 + self._initial_yaw = None + + print(f"Recording to {self.data_exporter.meta.root}") + + @property + def current_episode_index(self): + return self.data_exporter.episode_buffer["episode_index"] + + def _print_and_say(self, message: str, say: bool = True, blocking: bool = False): + if self.text_to_speech is not None: + self.text_to_speech.print_and_say(message, say, blocking=blocking) + else: + print(message) + + def _poll_state_zmq(self): + """Poll the ``g1_debug`` ZMQ topic for robot state (non-blocking).""" + msg = self._state_subscriber.get_msg(clear=True) + if msg is None: + return + + if msg.get("ros_timestamp", 0.0) == 0.0: + msg["ros_timestamp"] = time.time() + + self.latest_proprio_msg = msg + + def _check_recording_commands(self): + """Check keyboard + ZMQ toggle flags for recording commands.""" + key = self._keyboard_listener.read_msg() + + if self._manager_toggle_da: + key = "x" + self._manager_toggle_da = False + elif self._manager_toggle_dc: + key = "c" + self._manager_toggle_dc = False + + if key == "c": + self._episode_state.change_state() + if self._episode_state.get_state() == self._episode_state.RECORDING: + self._initial_yaw = None + self._print_and_say( + f"Started recording {self.current_episode_index}", blocking=False + ) + elif self._episode_state.get_state() == self._episode_state.NEED_TO_SAVE: + self._print_and_say("Stopping recording, preparing to save", blocking=False) + elif self._episode_state.get_state() == self._episode_state.IDLE: + self._print_and_say("Saved episode and back to idle state", blocking=False) + elif key == "x": + if self._episode_state.get_state() == self._episode_state.RECORDING: + self.data_exporter.save_episode_as_discarded() + self._episode_state.reset_state() + self._initial_yaw = None + self._print_and_say("Discarded episode", blocking=False) + + def _poll_sonic_zmq_messages(self): + """Poll ZMQ for pose, planner, and manager_state messages (non-blocking).""" + if self._sonic_zmq_socket is None: + return + + max_polls = 20 + for _ in range(max_polls): + try: + raw = self._sonic_zmq_socket.recv(zmq.NOBLOCK) + except zmq.Again: + break + + if raw.startswith(b"manager_state"): + self._handle_manager_state(raw) + elif raw.startswith(b"planner"): + self._handle_planner_message(raw) + elif raw.startswith(b"pose"): + self._handle_pose_message(raw) + + def _handle_manager_state(self, raw: bytes) -> None: + try: + data = unpack_pose_message(raw, topic="manager_state") + except Exception: + return + + if "stream_mode" in data: + self.current_stream_mode = int(data["stream_mode"].flat[0]) + + if self._extract_bool(data, "toggle_data_collection"): + self._manager_toggle_dc = True + if self._extract_bool(data, "toggle_data_abort"): + self._manager_toggle_da = True + + def _handle_planner_message(self, raw: bytes) -> None: + try: + data = unpack_pose_message(raw, topic="planner") + except Exception: + return + + planner_mode = int(data["mode"].flat[0]) if "mode" in data else 0 + planner_movement = ( + data["movement"].flatten().astype(np.float32) + if "movement" in data and data["movement"].size == 3 + else np.zeros(3, dtype=np.float32) + ) + planner_facing = ( + data["facing"].flatten().astype(np.float32) + if "facing" in data and data["facing"].size == 3 + else np.array([1.0, 0.0, 0.0], dtype=np.float32) + ) + planner_speed = float(data["speed"].flat[0]) if "speed" in data else -1.0 + planner_height = float(data["height"].flat[0]) if "height" in data else -1.0 + + vr_3pt_position = None + if "vr_position" in data and data["vr_position"].size == 9: + vr_3pt_position = data["vr_position"].flatten().astype(np.float32) + vr_3pt_orientation = None + if "vr_orientation" in data and data["vr_orientation"].size == 12: + vr_3pt_orientation = data["vr_orientation"].flatten().astype(np.float32) + + self.latest_planner_msg = { + "planner_mode": planner_mode, + "planner_movement": planner_movement, + "planner_facing": planner_facing, + "planner_speed": planner_speed, + "planner_height": planner_height, + "vr_3pt_position": vr_3pt_position, + "vr_3pt_orientation": vr_3pt_orientation, + "left_hand_joints": self._extract_hand_joints(data, "left_hand_joints"), + "right_hand_joints": self._extract_hand_joints(data, "right_hand_joints"), + "receive_timestamp": time.time(), + } + + def _handle_pose_message(self, raw: bytes) -> None: + G1_L_WRIST_ROLL_IDX = 23 + G1_L_WRIST_PITCH_IDX = 25 + G1_L_WRIST_YAW_IDX = 27 + G1_R_WRIST_ROLL_IDX = 24 + G1_R_WRIST_PITCH_IDX = 26 + G1_R_WRIST_YAW_IDX = 28 + + try: + pose_data = unpack_pose_message(raw, topic="pose") + except Exception as e: + print(f"[Sonic] Error unpacking pose message: {e}") + return + + try: + if "smpl_joints" not in pose_data or len(pose_data["smpl_joints"].shape) != 3: + return + + left_wrist_joints = None + right_wrist_joints = None + if "joint_pos" in pose_data and len(pose_data["joint_pos"].shape) == 2: + joint_pos = pose_data["joint_pos"][0] + left_wrist_joints = np.array( + [ + joint_pos[G1_L_WRIST_ROLL_IDX], + joint_pos[G1_L_WRIST_PITCH_IDX], + joint_pos[G1_L_WRIST_YAW_IDX], + ], + dtype=np.float32, + ) + right_wrist_joints = np.array( + [ + joint_pos[G1_R_WRIST_ROLL_IDX], + joint_pos[G1_R_WRIST_PITCH_IDX], + joint_pos[G1_R_WRIST_YAW_IDX], + ], + dtype=np.float32, + ) + + frame_index = None + if "frame_index" in pose_data: + frame_index = np.array([pose_data["frame_index"].flat[0]], dtype=np.int64) + + smpl_pose = np.zeros(63, dtype=np.float32) + if "smpl_pose" in pose_data: + raw_pose = pose_data["smpl_pose"] + if raw_pose.ndim == 3: + smpl_pose = raw_pose[0].flatten().astype(np.float32) + elif raw_pose.ndim == 2: + smpl_pose = raw_pose.flatten().astype(np.float32) + elif raw_pose.ndim == 1 and raw_pose.size == 63: + smpl_pose = raw_pose.astype(np.float32) + + left_hand_joints = self._extract_hand_joints(pose_data, "left_hand_joints") + right_hand_joints = self._extract_hand_joints(pose_data, "right_hand_joints") + + vr_3pt_position = None + if "vr_position" in pose_data and pose_data["vr_position"].size == 9: + vr_3pt_position = pose_data["vr_position"].flatten().astype(np.float32) + vr_3pt_orientation = None + if "vr_orientation" in pose_data and pose_data["vr_orientation"].size == 12: + vr_3pt_orientation = pose_data["vr_orientation"].flatten().astype(np.float32) + + self.latest_sonic_msg = { + "smpl_joints": pose_data["smpl_joints"][0], + "smpl_pose": smpl_pose, + "body_quat_w": ( + pose_data["body_quat_w"][0] if "body_quat_w" in pose_data else None + ), + "left_hand_joints": left_hand_joints, + "right_hand_joints": right_hand_joints, + "left_wrist_joints": left_wrist_joints, + "right_wrist_joints": right_wrist_joints, + "vr_3pt_position": vr_3pt_position, + "vr_3pt_orientation": vr_3pt_orientation, + "frame_index": frame_index, + "receive_timestamp": time.time(), + } + except Exception as e: + if not hasattr(self, "_sonic_error_count"): + self._sonic_error_count = 0 + self._sonic_error_count += 1 + if self._sonic_error_count == 1 or self._sonic_error_count % 100 == 0: + print(f"[Sonic] Error processing pose message: {e}") + + @staticmethod + def _extract_hand_joints(pose_data: dict, key: str) -> np.ndarray: + arr = pose_data.get(key) + if arr is not None: + if arr.ndim > 1: + arr = arr[0] + return arr.astype(np.float32) + return np.zeros(7, dtype=np.float32) + + @staticmethod + def _extract_bool(pose_data: dict, key: str) -> bool: + val = pose_data.get(key) + if val is None: + return False + if isinstance(val, np.ndarray): + return bool(val.flat[0]) + return bool(val) + + def _log_latency_periodic( + self, + sonic_latency_ms: float | None = None, + ): + current_time = time.time() + if current_time - self._last_latency_log_time >= 1.0: + self._last_latency_log_time = current_time + parts = [] + if sonic_latency_ms is not None: + parts.append(f"Sonic Pose: {sonic_latency_ms:.1f}ms") + if parts: + print(f"[Latency] {', '.join(parts)}") + + def _add_images_to_frame_data(self, frame_data: dict) -> None: + if self.latest_image_msg is None: + return + images = self.latest_image_msg["images"] + for feature_name, feature_info in self.data_exporter.features.items(): + if feature_info.get("dtype") in ["image", "video"]: + image_key = feature_name.split(".")[-1] + if image_key not in images: + raise ValueError( + f"Required image '{image_key}' for feature '{feature_name}' " + f"not found in image message. Available: {list(images.keys())}" + ) + frame_data[feature_name] = images[image_key] + + def _finalize_frame(self, t_start: float) -> bool: + t_end = time.monotonic() + if t_end - t_start > (1 / self.frequency): + print(f"DataExporter Missed: {t_end - t_start} sec") + + if self._episode_state.get_state() == self._episode_state.NEED_TO_SAVE: + buffer_size = self.data_exporter.episode_buffer.get("size", 0) + if buffer_size > 0: + self.data_exporter.save_episode() + self.sonic_timing_monitor.reset() + self._initial_yaw = None + self._print_and_say("Finished saving episode") + else: + self._print_and_say("Skipping save: no frames collected", say=False) + self._episode_state.change_state() + return True + + def _add_data_frame(self): + t_start = time.monotonic() + + if self.latest_proprio_msg is None or self.latest_image_msg is None: + self._print_and_say( + f"Waiting for message. " + f"Avail msg: proprio {self.latest_proprio_msg is not None} | " + f"image {self.latest_image_msg is not None}", + say=False, + ) + return False + + if self._episode_state.get_state() != self._episode_state.RECORDING: + return self._finalize_frame(t_start) + + return self._add_data_frame_sonic(t_start) + + def _add_data_frame_sonic(self, t_start: float) -> bool: + """Build one data frame in Sonic CPP + SMPL mode.""" + assert self.latest_proprio_msg is not None + proprio = self.latest_proprio_msg + + whole_q = self.robot_model.get_configuration_from_actuated_joints( + body_actuated_joint_values=proprio["body_q"], + left_hand_actuated_joint_values=proprio["left_hand_q"], + right_hand_actuated_joint_values=proprio["right_hand_q"], + ) + whole_action_wbc = self.robot_model.get_configuration_from_actuated_joints( + body_actuated_joint_values=proprio["last_action"], + left_hand_actuated_joint_values=proprio["last_left_hand_action"], + right_hand_actuated_joint_values=proprio["last_right_hand_action"], + ) + + self.robot_model.cache_forward_kinematics(whole_q) + eef_parts = [] + for side in ["left", "right"]: + placement = self.robot_model.frame_placement( + self.robot_model.supplemental_info.hand_frame_names[side] + ) + pos = placement.translation[:3] + quat = R.from_matrix(placement.rotation).as_quat(scalar_first=True) + eef_parts.append(np.concatenate([pos, quat])) + observation_eef_state = np.concatenate(eef_parts) + + frame_data: dict = { + "observation.state": whole_q, + "observation.eef_state": observation_eef_state, + "action.wbc": whole_action_wbc, + } + + self._add_cpp_state_features(frame_data, proprio) + + sonic_latency_ms = self._add_sonic_pose_features(frame_data) + + self._add_images_to_frame_data(frame_data) + + self._log_latency_periodic(sonic_latency_ms) + + self.data_exporter.add_frame(frame_data) + return self._finalize_frame(t_start) + + def _add_cpp_state_features(self, frame_data: dict, proprio: dict) -> None: + if "base_quat" in proprio: + base_quat = np.asarray(proprio["base_quat"], dtype=np.float64) + frame_data["observation.root_orientation"] = base_quat + frame_data["observation.projected_gravity"] = compute_projected_gravity( + base_quat + ).astype(np.float64) + + if "init_ref_data_root_rot_array" in proprio: + frame_data["observation.cpp_rotation_offset"] = np.asarray( + proprio["init_ref_data_root_rot_array"], dtype=np.float64 + ) + else: + frame_data["observation.cpp_rotation_offset"] = np.array( + [1.0, 0.0, 0.0, 0.0], dtype=np.float64 + ) + else: + frame_data["observation.root_orientation"] = np.array( + [1.0, 0.0, 0.0, 0.0], dtype=np.float64 + ) + frame_data["observation.projected_gravity"] = np.array( + [0.0, 0.0, -1.0], dtype=np.float64 + ) + frame_data["observation.cpp_rotation_offset"] = np.array( + [1.0, 0.0, 0.0, 0.0], dtype=np.float64 + ) + + if "init_base_quat" in proprio: + frame_data["observation.init_base_quat"] = np.asarray( + proprio["init_base_quat"], dtype=np.float64 + ) + else: + frame_data["observation.init_base_quat"] = np.array( + [1.0, 0.0, 0.0, 0.0], dtype=np.float64 + ) + + if "delta_heading" in proprio: + dh = proprio["delta_heading"] + if isinstance(dh, np.ndarray): + dh = dh.item() if dh.size == 1 else dh[0] + frame_data["teleop.delta_heading"] = np.array([float(dh)], dtype=np.float64) + else: + frame_data["teleop.delta_heading"] = np.zeros(1, dtype=np.float64) + + if "token_state" in proprio: + frame_data["action.motion_token"] = np.asarray(proprio["token_state"], dtype=np.float64) + else: + frame_data["action.motion_token"] = np.zeros(64, dtype=np.float64) + + def _add_sonic_pose_features(self, frame_data: dict) -> float | None: + """Add teleop features based on current stream mode.""" + sonic_latency_ms = None + + frame_data["teleop.stream_mode"] = np.array([self.current_stream_mode], dtype=np.int32) + + smpl_msg = self.latest_sonic_msg + use_smpl = False + if self.current_stream_mode in (1, 4) and smpl_msg is not None: + receive_ts = smpl_msg.get("receive_timestamp") + if receive_ts is not None: + age_sec = time.time() - receive_ts + sonic_latency_ms = age_sec * 1000 + self.sonic_timing_monitor.log_time_delta(age_sec) + if sonic_latency_ms <= 100.0: + use_smpl = True + elif (self.sonic_timing_monitor.failure_count + 1) % 10 == 0: + self._print_and_say( + f"Sonic pose stale ({sonic_latency_ms:.1f}ms old), using zeros", + say=False, + ) + else: + use_smpl = True + + planner_msg = self.latest_planner_msg + use_planner = False + if self.current_stream_mode == 5 and planner_msg is not None: + receive_ts = planner_msg.get("receive_timestamp") + if receive_ts is not None: + age_sec = time.time() - receive_ts + planner_latency_ms = age_sec * 1000 + if sonic_latency_ms is None: + sonic_latency_ms = planner_latency_ms + if planner_latency_ms <= 200.0: + use_planner = True + else: + use_planner = True + + # SMPL features + if use_smpl and smpl_msg.get("smpl_joints") is not None: + joints = np.asarray(smpl_msg["smpl_joints"], dtype=np.float32) + if joints.ndim == 2: + joints = joints.flatten() + frame_data["teleop.smpl_joints"] = np.ascontiguousarray(joints, dtype=np.float32) + else: + frame_data["teleop.smpl_joints"] = np.zeros(72, dtype=np.float32) + + if use_smpl and smpl_msg.get("smpl_pose") is not None: + pose = np.asarray(smpl_msg["smpl_pose"], dtype=np.float32) + if pose.ndim > 1: + pose = pose.flatten() + frame_data["teleop.smpl_pose"] = np.ascontiguousarray(pose, dtype=np.float32) + else: + frame_data["teleop.smpl_pose"] = np.zeros(63, dtype=np.float32) + + if use_smpl and smpl_msg.get("body_quat_w") is not None: + body_quat_w = smpl_msg["body_quat_w"].astype(np.float32) + frame_data["teleop.body_quat_w"] = body_quat_w + frame_data["teleop.target_body_orientation"] = self._compute_target_body_orientation( + body_quat_w, frame_data + ) + else: + frame_data["teleop.body_quat_w"] = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32) + frame_data["teleop.target_body_orientation"] = quat_to_rot6d( + np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32) + ) + + frame_data["teleop.left_wrist_joints"] = ( + smpl_msg["left_wrist_joints"].astype(np.float32) + if use_smpl and smpl_msg.get("left_wrist_joints") is not None + else np.zeros(3, dtype=np.float32) + ) + frame_data["teleop.right_wrist_joints"] = ( + smpl_msg["right_wrist_joints"].astype(np.float32) + if use_smpl and smpl_msg.get("right_wrist_joints") is not None + else np.zeros(3, dtype=np.float32) + ) + + frame_data["teleop.smpl_frame_index"] = ( + smpl_msg["frame_index"].astype(np.int64) + if use_smpl and smpl_msg is not None and smpl_msg.get("frame_index") is not None + else np.array([0], dtype=np.int64) + ) + + hand_msg = ( + smpl_msg if self.current_stream_mode in (1, 4) and smpl_msg is not None + else planner_msg if planner_msg is not None + else smpl_msg + ) + frame_data["teleop.left_hand_joints"] = ( + hand_msg["left_hand_joints"].astype(np.float32) + if hand_msg is not None + and hand_msg.get("left_hand_joints") is not None + else np.zeros(7, dtype=np.float32) + ) + frame_data["teleop.right_hand_joints"] = ( + hand_msg["right_hand_joints"].astype(np.float32) + if hand_msg is not None + and hand_msg.get("right_hand_joints") is not None + else np.zeros(7, dtype=np.float32) + ) + + # Planner command fields + frame_data["teleop.planner_mode"] = np.array( + [planner_msg["planner_mode"]] if use_planner else [0], + dtype=np.int32, + ) + frame_data["teleop.planner_movement"] = ( + planner_msg["planner_movement"].copy() + if use_planner and planner_msg.get("planner_movement") is not None + else np.zeros(3, dtype=np.float32) + ) + frame_data["teleop.planner_facing"] = ( + planner_msg["planner_facing"].copy() + if use_planner and planner_msg.get("planner_facing") is not None + else np.array([1.0, 0.0, 0.0], dtype=np.float32) + ) + frame_data["teleop.planner_speed"] = np.array( + [planner_msg["planner_speed"]] if use_planner else [-1.0], + dtype=np.float32, + ) + frame_data["teleop.planner_height"] = np.array( + [planner_msg["planner_height"]] if use_planner else [-1.0], + dtype=np.float32, + ) + + # VR 3-point pose + frame_data["teleop.vr_3pt_position"] = ( + planner_msg["vr_3pt_position"].astype(np.float32) + if use_planner and planner_msg.get("vr_3pt_position") is not None + else np.zeros(9, dtype=np.float32) + ) + if use_planner and planner_msg.get("vr_3pt_orientation") is not None: + frame_data["teleop.vr_3pt_orientation"] = quat_to_rot6d( + planner_msg["vr_3pt_orientation"].astype(np.float32) + ) + else: + frame_data["teleop.vr_3pt_orientation"] = np.zeros(18, dtype=np.float32) + + return sonic_latency_ms + + def _compute_target_body_orientation( + self, body_quat_w: np.ndarray, frame_data: dict + ) -> np.ndarray: + """Compute yaw-normalised target body orientation as rot6d (6-dim).""" + delta_heading = float(frame_data.get("teleop.delta_heading", [0.0])[0]) + + body_rot = R.from_quat(body_quat_w, scalar_first=True) + target_rot = R.from_euler("z", delta_heading, degrees=False) * body_rot + + euler = target_rot.as_euler("ZYX", degrees=False) + current_yaw = euler[0] + + if self._initial_yaw is None: + self._initial_yaw = current_yaw + + normalised_euler = np.array([current_yaw - self._initial_yaw, euler[1], euler[2]]) + target_quat = ( + R.from_euler("ZYX", normalised_euler, degrees=False) + .as_quat(scalar_first=True) + .astype(np.float32) + ) + return quat_to_rot6d(target_quat) + + def save_and_cleanup(self): + try: + self._print_and_say("saving episode done", blocking=False) + buffer_size = self.data_exporter.episode_buffer.get("size", 0) + if buffer_size > 0: + self.data_exporter.save_episode() + self._print_and_say( + f"Recording complete: {self.data_exporter.meta.root}", say=False, blocking=True + ) + except Exception as e: + self._print_and_say(f"Error saving episode: {e}", blocking=True) + + try: + self._state_subscriber.close() + except Exception: + pass + for sock in [self._sonic_zmq_socket]: + if sock is not None: + try: + sock.close() + except Exception: + pass + for ctx in [self._sonic_zmq_ctx]: + if ctx is not None: + try: + ctx.term() + except Exception: + pass + + self._print_and_say("Shutting down data exporter...", say=False) + + def run(self): + try: + while True: + t_start = time.monotonic() + with self.telemetry.timer("total_loop"): + with self.telemetry.timer("poll_state"): + self._poll_state_zmq() + + with self.telemetry.timer("poll_sonic"): + self._poll_sonic_zmq_messages() + + with self.telemetry.timer("poll_image"): + img_msg = self._image_subscriber.read() + if img_msg is not None: + self.latest_image_msg = img_msg + + with self.telemetry.timer("add_frame"): + self._add_data_frame() + + with self.telemetry.timer("check_recording_commands"): + self._check_recording_commands() + + end_time = time.monotonic() + + elapsed = time.monotonic() - t_start + sleep_time = self.loop_period - elapsed + if sleep_time > 0: + time.sleep(sleep_time) + + if (end_time - t_start) > self.loop_period: + self.telemetry.log_timing_info( + context="Data Exporter Loop Missed", threshold=0.001 + ) + + except KeyboardInterrupt: + print("Data exporter terminated by user") + buffer_size = self.data_exporter.episode_buffer.get("size", 0) + if buffer_size > 0: + self.data_exporter.save_episode_as_discarded() + + finally: + self.save_and_cleanup() + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- + + +def main(config: SonicDataExporterConfig): + g1_rm = get_g1_robot_model() + + dataset_features = get_features_sonic_vla(g1_rm) + modality_config = get_modality_config_sonic_vla(g1_rm) + + if config.record_wrist_cameras: + print("[Camera] Wrist cameras enabled — adding to dataset schema") + dataset_features.update(get_wrist_camera_features()) + wrist_modality = get_wrist_camera_modality_config() + for key, value in wrist_modality.items(): + if key in modality_config: + modality_config[key].update(value) + else: + modality_config[key] = value + + text_to_speech = TextToSpeech() if config.text_to_speech else None + + robot_config = poll_robot_config_zmq( + config.state_zmq_host, config.state_zmq_port, config.robot_config_timeout + ) + + data_exporter = Gr00tDataExporter.create( + save_root=f"{config.root_output_dir}/{config.dataset_name}", + fps=config.data_collection_frequency, + features=dataset_features, + modality_config=modality_config, + task=config.task_prompt, + script_config={**robot_config, "record_wrist_cameras": config.record_wrist_cameras}, + ) + + data_collector = GrootDataCollector( + frequency=config.data_collection_frequency, + data_exporter=data_exporter, + robot_model=g1_rm, + camera_host=config.camera_host, + camera_port=config.camera_port, + text_to_speech=text_to_speech, + sonic_data_zmq_host=config.sonic_zmq_host, + sonic_data_zmq_port=config.sonic_zmq_port, + state_zmq_host=config.state_zmq_host, + state_zmq_port=config.state_zmq_port, + ) + data_collector.run() + + +if __name__ == "__main__": + config = tyro.cli(SonicDataExporterConfig) + + if config.dataset_name is None: + config.dataset_name = datetime.now().strftime("%Y-%m-%d-%H-%M-%S") + + main(config) diff --git a/GR00T-WholeBodyControl/gear_sonic/scripts/run_sim_loop.py b/GR00T-WholeBodyControl/gear_sonic/scripts/run_sim_loop.py new file mode 100644 index 0000000000000000000000000000000000000000..9115e77616232868fbaa6c906fa7748505ca704f --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/scripts/run_sim_loop.py @@ -0,0 +1,68 @@ +"""Entry point for running a MuJoCo simulation loop with the G1 robot model. + +Parses a YAML-based WBC config via tyro CLI, instantiates the G1 robot model, +and launches the simulator (optionally with offscreen image publishing). +""" + +from typing import Dict + +import tyro + +from gear_sonic.utils.mujoco_sim.simulator_factory import SimulatorFactory, init_channel +from gear_sonic.utils.mujoco_sim.configs import SimLoopConfig +from gear_sonic.data.robot_model.instantiation.g1 import ( + instantiate_g1_robot_model, +) +from gear_sonic.data.robot_model.robot_model import RobotModel + +ArgsConfig = SimLoopConfig + + +class SimWrapper: + def __init__(self, robot_model: RobotModel, env_name: str, config: Dict[str, any], **kwargs): + self.robot_model = robot_model + self.config = config + + init_channel(config=self.config) + + # Create simulator using factory + self.sim = SimulatorFactory.create_simulator( + config=self.config, + env_name=env_name, + **kwargs, + ) + + +def main(config: ArgsConfig): + wbc_config = config.load_wbc_yaml() + # NOTE: we will override the interface to local if it is not specified + wbc_config["ENV_NAME"] = config.env_name + + if config.enable_image_publish: + assert ( + config.enable_offscreen + ), "enable_offscreen must be True when enable_image_publish is True" + + robot_model = instantiate_g1_robot_model() + + sim_wrapper = SimWrapper( + robot_model=robot_model, + env_name=config.env_name, + config=wbc_config, + onscreen=wbc_config.get("ENABLE_ONSCREEN", True), + offscreen=wbc_config.get("ENABLE_OFFSCREEN", False), + enable_image_publish=config.enable_image_publish, + ) + # Start simulator as independent process + SimulatorFactory.start_simulator( + sim_wrapper.sim, + as_thread=False, + enable_image_publish=config.enable_image_publish, + mp_start_method=config.mp_start_method, + camera_port=config.camera_port, + ) + + +if __name__ == "__main__": + config = tyro.cli(ArgsConfig) + main(config) \ No newline at end of file diff --git a/GR00T-WholeBodyControl/gear_sonic/scripts/run_vla_inference.py b/GR00T-WholeBodyControl/gear_sonic/scripts/run_vla_inference.py new file mode 100644 index 0000000000000000000000000000000000000000..7a5792d9818622049df5a67abfa8c3649e11f940 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/scripts/run_vla_inference.py @@ -0,0 +1,777 @@ +""" +VLA inference runner — NO ROS 2 DEPENDENCY. + +Runs an Isaac-GR00T VLA policy against the Sonic whole-body control stack. +All communication uses ZMQ: + 1. Robot state -> ZMQ SUB on ``g1_debug`` topic (from C++ zmq_output_handler) + 2. Actions out -> ZMQ PUB (latent protocol v4: motion token + hand joints) + 3. Camera -> ZMQ/TCP via ComposedCameraClientSensor + 4. Keyboard -> ZMQ SUB via ZMQKeyboardSubscriber + +Uses the Isaac-GR00T PolicyClient (ZMQ REQ/REP) to communicate with a +running PolicyServer. + +Keyboard commands (received via ZMQ from the standalone keyboard publisher): + p -> pause / resume the policy loop + k -> start / stop the C++ control loop + i -> blend smoothly to initial pose (or snap if no prior token) and switch to POSE mode + t -> change prompt at runtime (publisher sends ``prompt:``) + [ -> toggle left hand open/closed for initial pose + ] -> toggle right hand open/closed for initial pose + c -> start recording (handled by data exporter if running) + s -> stop recording success (handled by data exporter) + f -> stop recording failure (handled by data exporter) +""" + +from dataclasses import dataclass +import queue +import threading +import time + +import numpy as np +import tyro +import zmq + +from gear_sonic.camera.composed_camera import ComposedCameraClientSensor +from gear_sonic.data.robot_model.instantiation.g1 import instantiate_g1_robot_model +from gear_sonic.utils.data_collection.keyboard_subscriber import ( + DEFAULT_ZMQ_KEYBOARD_PORT, + ZMQKeyboardSubscriber, +) +from gear_sonic.utils.data_collection.telemetry import Telemetry +from gear_sonic.utils.data_collection.transforms import compute_projected_gravity +from gear_sonic.utils.data_collection.zmq_state_subscriber import ZMQStateSubscriber +from gear_sonic.utils.inference.initial_poses import LATENT_INITIAL_MOTION_TOKEN +from gear_sonic.utils.inference.vla_utils import ( + calculate_latency_compensated_index, + concat_action, + prepare_observation_for_eval, + should_trigger_new_inference, +) +from gear_sonic.utils.teleop.solver.hand.g1_gripper_ik_solver import ( + G1GripperInverseKinematicsSolver, +) +from gear_sonic.utils.teleop.zmq.zmq_planner_sender import ( + build_command_message, + pack_pose_message, +) + + +@dataclass +class InferenceConfig: + """CLI config for the VLA inference runner.""" + + # Policy server (Isaac-GR00T PolicyServer) + host: str = "localhost" + """The host address of the Isaac-GR00T PolicyServer.""" + + port: int = 5550 + """The port of the Isaac-GR00T PolicyServer.""" + + # Control + action_publish_rate: int = 50 + """Rate at which individual actions are published to the C++ control loop (Hz).""" + + action_horizon: int = 40 + """Action horizon of the VLA policy (number of future actions per inference).""" + + rate: float = 1 / 0.4 + """Rate at which we run the forward pass of the VLA policy (Hz).""" + + # Camera + camera_host: str = "localhost" + """Camera server host.""" + + camera_port: int = 5555 + """Camera server port.""" + + # ZMQ: Robot state (from C++ zmq_output_handler, g1_debug topic) + state_zmq_host: str = "localhost" + """ZMQ host for robot state (g1_debug topic from C++ deploy).""" + + state_zmq_port: int = 5557 + """ZMQ port for robot state (same socket as robot_config topic).""" + + # ZMQ: Action output (latent actions to C++ control loop) + action_zmq_host: str = "localhost" + """ZMQ host for action output (PUB socket).""" + + action_zmq_port: int = 5556 + """ZMQ port for action output.""" + + # ZMQ: Keyboard input + keyboard_zmq_host: str = "localhost" + """ZMQ host for keyboard input.""" + + keyboard_zmq_port: int = DEFAULT_ZMQ_KEYBOARD_PORT + """ZMQ port for keyboard input.""" + + # Embodiment + embodiment_tag: str = "unitree_g1_sonic" + """Embodiment tag for policy inference.""" + + # Prompt / eval + prompt: str = "demo" + """The language prompt for the VLA policy.""" + + # Initial pose + initial_pose_blend_duration: float = 1.0 + """Duration (seconds) for smooth interpolation to initial pose. The robot + blends from its current motion token to the initial pose token over this + period. Set to 0 to snap instantly (no blend).""" + + # Debug + verbose_timing: bool = False + """Whether to always print timing info (not just when loop is slow).""" + + +def print_green(x): + print(f"\033[92m{x}\033[0m") + + +# --------------------------------------------------------------------------- +# Action packing (latent protocol v4) +# --------------------------------------------------------------------------- + + +def pack_latent_action_message( + motion_token: np.ndarray, + frame_index: np.ndarray, + left_hand_joints: np.ndarray = None, + right_hand_joints: np.ndarray = None, +) -> bytes: + """Pack a single motion-token action into a ZMQ message (Protocol v4). + + Args: + motion_token: Shape ``[64]`` (flat) or ``[1, 64]``. + frame_index: Shape ``[1]``. + left_hand_joints: Shape ``[7]`` or ``[1, 7]``, optional. + right_hand_joints: Shape ``[7]`` or ``[1, 7]``, optional. + + Returns: + Packed ZMQ message bytes. + """ + motion_token = np.asarray(motion_token, dtype=np.float32) + frame_index = np.asarray(frame_index, dtype=np.int64) + + if frame_index.ndim == 0: + frame_index = np.array([frame_index], dtype=np.int64) + elif frame_index.shape[0] != 1: + frame_index = frame_index[:1] + + if motion_token.ndim == 1: + motion_token = motion_token.reshape(1, -1) + + pose_data = { + "token_state": motion_token, + "frame_index": frame_index, + } + + if left_hand_joints is not None: + left_hand_joints = np.asarray(left_hand_joints, dtype=np.float32) + if left_hand_joints.ndim == 1: + if left_hand_joints.shape[0] != 7: + raise ValueError( + f"left_hand_joints must have shape [7], got {left_hand_joints.shape}" + ) + left_hand_joints = left_hand_joints.reshape(1, 7) + pose_data["left_hand_joints"] = left_hand_joints + + if right_hand_joints is not None: + right_hand_joints = np.asarray(right_hand_joints, dtype=np.float32) + if right_hand_joints.ndim == 1: + if right_hand_joints.shape[0] != 7: + raise ValueError( + f"right_hand_joints must have shape [7], got {right_hand_joints.shape}" + ) + right_hand_joints = right_hand_joints.reshape(1, 7) + pose_data["right_hand_joints"] = right_hand_joints + + return pack_pose_message(pose_data, topic="pose", version=4) + + +def get_action_field(action_dict: dict, key: str): + """Get action field from dict, checking both with and without 'action.' prefix.""" + value = action_dict.get(key) + if value is not None: + return value + value = action_dict.get(f"action.{key}") + if value is not None: + return value + raise AssertionError( + f"Required action field '{key}' (or 'action.{key}') not found in processed_action. " + f"Available keys: {list(action_dict.keys())}" + ) + + +# --------------------------------------------------------------------------- +# Observation / inference helpers +# --------------------------------------------------------------------------- + + +def prepare_observation_from_sensors( + camera_subscriber, + state_subscriber, + robot_model, + language_prompt: str, + log_errors: bool = False, +): + """Read sensors and prepare observation for the VLA policy. + + Returns: + observation dict, or None if sensor data not yet available. + """ + camera_msg = camera_subscriber.read() + if camera_msg is None: + if log_errors: + print("[DEBUG] prepare_observation: waiting for camera msg..", flush=True) + return None + + state_msg = state_subscriber.get_msg() + if state_msg is None: + if log_errors: + print("[DEBUG] prepare_observation: waiting for state msg..", flush=True) + return None + + cam_img = camera_msg["images"]["ego_view"] + + # Copy index finger data to middle finger (hardware coupling) + state_msg["left_hand_q"][5] = state_msg["left_hand_q"][3] + state_msg["left_hand_q"][6] = state_msg["left_hand_q"][4] + + qpos = robot_model.get_configuration_from_actuated_joints( + body_actuated_joint_values=state_msg["body_q"], + left_hand_actuated_joint_values=state_msg["left_hand_q"], + right_hand_actuated_joint_values=state_msg["right_hand_q"], + ) + + video = {"ego_view": cam_img[np.newaxis, np.newaxis]} + if "left_wrist" in camera_msg["images"]: + video["left_wrist"] = camera_msg["images"]["left_wrist"][np.newaxis, np.newaxis] + if "right_wrist" in camera_msg["images"]: + video["wrist_view"] = camera_msg["images"]["right_wrist"][np.newaxis, np.newaxis] + + observation = { + "video": video, + "state": {}, + "language": { + "annotation.human.task_description": [[language_prompt]], + }, + "q": np.asarray(qpos, dtype=np.float32)[np.newaxis, np.newaxis], + "timestamps": camera_msg["timestamps"]["ego_view"], + } + + observation = prepare_observation_for_eval(robot_model, observation) + + # Projected gravity for Sonic latent embodiment + assert "base_quat" in state_msg, "base_quat not found in state_msg" + base_quat = np.asarray(state_msg["base_quat"], dtype=np.float64) + assert base_quat.shape == (4,), "base_quat must have shape (4,)" + projected_gravity = compute_projected_gravity(base_quat) + observation["state"]["projected_gravity"] = np.asarray( + projected_gravity, dtype=np.float32 + )[np.newaxis, np.newaxis] + + return observation + + +def run_policy_inference_and_process(policy, observation, robot_model): + """Run policy inference via Isaac-GR00T PolicyClient and process results. + + Returns: + processed_action dict or None on error. + """ + try: + action, _info = policy.get_action(observation) + + action.pop("task_progress", None) + action.pop("action.task_progress", None) + + motion_key = "motion_token" if "motion_token" in action else "action.motion_token" + if np.abs(action[motion_key]).max() > 1.25: + print( + f"[Warning] action['{motion_key}'] max " + f"({np.abs(action[motion_key]).max():.4f}) > 1.25. " + "Exceeds action bound, skipping." + ) + return None + + processed_action = concat_action(robot_model, action) + return processed_action + except Exception as e: + print(f"Error in inference: {e}") + import traceback + + traceback.print_exc() + return None + + +def _inference_worker_loop( + inference_queue: queue.Queue, + result_queue: queue.Queue, + stop_event: threading.Event, + busy_event: threading.Event, + prepare_obs_fn, + inference_fn, +): + """Persistent worker thread for async inference.""" + while not stop_event.is_set(): + try: + try: + inference_queue.get(timeout=0.1) + except queue.Empty: + continue + + busy_event.set() + try: + observation = prepare_obs_fn() + if observation is None: + print("[DEBUG] Worker thread: Observation is None, skipping", flush=True) + continue + + inference_start_time = time.monotonic() + processed_action = inference_fn(observation) + + if processed_action is not None: + try: + result_queue.put_nowait((processed_action, inference_start_time)) + except queue.Full: + try: + result_queue.get_nowait() + result_queue.put_nowait((processed_action, inference_start_time)) + except queue.Empty: + result_queue.put_nowait((processed_action, inference_start_time)) + finally: + busy_event.clear() + except Exception as e: + print(f"Error in inference worker thread: {e}") + import traceback + + traceback.print_exc() + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def _compute_closed_hand_joints(side: str) -> np.ndarray: + """Compute closed hand joint positions using G1GripperInverseKinematicsSolver.""" + side_str = "left" if side.upper() == "L" else "right" + solver = G1GripperInverseKinematicsSolver(side=side_str) + return solver._get_middle_close_q_desired().astype(np.float32) + + +def main(config: InferenceConfig): + pause_loop = True + + robot_model = instantiate_g1_robot_model(waist_location="lower_and_upper_body") + + # Isaac-GR00T PolicyClient + from gr00t.policy.server_client import PolicyClient + + n1_policy = PolicyClient(host=config.host, port=config.port) + + print(f"Connecting to PolicyServer at {config.host}:{config.port}...") + if n1_policy.ping(): + print_green("PolicyServer is reachable.") + else: + print("WARNING: PolicyServer not reachable. Inference will fail until server is up.") + + state_subscriber = ZMQStateSubscriber( + host=config.state_zmq_host, + port=config.state_zmq_port, + ) + + camera_subscriber = ComposedCameraClientSensor( + server_ip=config.camera_host, port=config.camera_port + ) + + zmq_context = zmq.Context() + zmq_socket = zmq_context.socket(zmq.PUB) + zmq_socket.bind(f"tcp://{config.action_zmq_host}:{config.action_zmq_port}") + time.sleep(0.1) + print_green( + f"ZMQ action socket bound to tcp://{config.action_zmq_host}:{config.action_zmq_port}" + ) + print_green(f"Using embodiment tag: {config.embodiment_tag}") + + keyboard_listener = ZMQKeyboardSubscriber( + port=config.keyboard_zmq_port, host=config.keyboard_zmq_host + ) + + telemetry = Telemetry(window_size=100) + + loop_rate = config.action_publish_rate + loop_period = 1.0 / loop_rate + + # Track C++ control loop state + cpp_loop_running = False + cpp_mode = "OFF" # "OFF", "PLANNER", or "POSE" + + # Track initial pose hand states + initial_pose_left_hand_closed = False + initial_pose_right_hand_closed = False + + def publish_initial_pose(): + """Publish initial pose command to move robot to starting position.""" + print("Moving to initial pose") + left_hand = ( + _compute_closed_hand_joints("L") + if initial_pose_left_hand_closed + else np.zeros(7, dtype=np.float32) + ) + right_hand = ( + _compute_closed_hand_joints("R") + if initial_pose_right_hand_closed + else np.zeros(7, dtype=np.float32) + ) + zmq_message = pack_latent_action_message( + motion_token=LATENT_INITIAL_MOTION_TOKEN, + frame_index=np.array([0], dtype=np.int64), + left_hand_joints=left_hand, + right_hand_joints=right_hand, + ) + zmq_socket.send(zmq_message) + print_green("Sent latent initial pose via ZMQ") + time.sleep(1.0) + print("Initial pose done.") + + def blend_to_initial_pose(duration_s: float) -> bool: + """Smoothly interpolate from the last sent motion token to the initial pose. + + Linearly blends over ``duration_s`` seconds at the action publish rate, + sending intermediate tokens each loop iteration. Returns True if blend + was performed, False if skipped (no previous token available). + """ + nonlocal last_sent_motion_token + if last_sent_motion_token is None: + print("No previous motion token — snapping to initial pose instead.") + publish_initial_pose() + return False + + start_token = last_sent_motion_token.copy() + target_token = LATENT_INITIAL_MOTION_TOKEN.copy() + num_steps = max(1, round(config.action_publish_rate * duration_s)) + step_period = 1.0 / config.action_publish_rate + + left_hand = ( + _compute_closed_hand_joints("L") + if initial_pose_left_hand_closed + else np.zeros(7, dtype=np.float32) + ) + right_hand = ( + _compute_closed_hand_joints("R") + if initial_pose_right_hand_closed + else np.zeros(7, dtype=np.float32) + ) + + print( + f"Blending to initial pose over {duration_s:.2f}s " + f"({num_steps} steps at {config.action_publish_rate} Hz)" + ) + + for step in range(num_steps): + t_step_start = time.monotonic() + alpha = (step + 1) / num_steps + blended_token = ((1.0 - alpha) * start_token + alpha * target_token).astype( + np.float32 + ) + zmq_message = pack_latent_action_message( + motion_token=blended_token, + frame_index=np.array([0], dtype=np.int64), + left_hand_joints=left_hand, + right_hand_joints=right_hand, + ) + zmq_socket.send(zmq_message) + last_sent_motion_token = blended_token.copy() + + elapsed = time.monotonic() - t_step_start + remaining = step_period - elapsed + if remaining > 0: + time.sleep(remaining) + + print_green("Initial pose blend complete.") + return True + + def send_cpp_control_command(start: bool, planner: bool = False): + """Send C++ control loop start/stop commands via ZMQ.""" + nonlocal cpp_loop_running, cpp_mode + try: + cmd_msg = build_command_message(start=start, stop=not start, planner=planner) + zmq_socket.send(cmd_msg) + time.sleep(0.01) + action_str = "start" if start else "stop" + mode_str = "planner" if planner else "pose" + cpp_loop_running = start + if start: + cpp_mode = "PLANNER" if planner else "POSE" + else: + cpp_mode = "OFF" + print_green(f"Sent ZMQ command: {action_str} control loop ({mode_str} mode)") + return True + except Exception as e: + action_str = "start" if start else "stop" + print(f"Warning: Failed to send {action_str} command message: {e}") + return False + + # Async inference state + cached_action_chunk = None + action_chunk_index = 0 + last_inference_time = 0.0 + inference_interval = 1.0 / config.rate + + zmq_frame_counter = 0 + last_sent_motion_token: np.ndarray | None = None + + PROMPT_MSG_PREFIX = "prompt:" + + def check_keyboard_input(): + nonlocal pause_loop, cpp_loop_running, cpp_mode + nonlocal initial_pose_left_hand_closed, initial_pose_right_hand_closed + nonlocal cached_action_chunk, action_chunk_index, last_inference_time + nonlocal zmq_frame_counter, last_sent_motion_token + + key = keyboard_listener.read_msg() + if key is None: + return + + if key.startswith(PROMPT_MSG_PREFIX): + new_prompt = key[len(PROMPT_MSG_PREFIX):] + if new_prompt: + old_prompt = language_prompt_ref[0] + language_prompt_ref[0] = new_prompt + print_green(f'Inference prompt changed: "{old_prompt}" -> "{new_prompt}"') + else: + print("Received empty prompt change -- ignoring.") + return + + if key == "c": + print("Keyboard: 'c' (start recording -- handled by data exporter)") + elif key == "s": + print("Keyboard: 's' (stop recording success -- handled by data exporter)") + elif key == "f": + print("Keyboard: 'f' (stop recording failure -- handled by data exporter)") + elif key == "i": + if cpp_loop_running and cpp_mode == "PLANNER": + if send_cpp_control_command(start=True, planner=False): + print("Switched to POSE mode (from PLANNER mode)") + else: + print("Warning: Failed to switch to POSE mode") + elif not cpp_loop_running: + print("Note: C++ loop not running - press 'k' to start") + + pause_loop = True + if config.initial_pose_blend_duration > 0 and last_sent_motion_token is not None: + blend_to_initial_pose(config.initial_pose_blend_duration) + else: + publish_initial_pose() + + zmq_frame_counter = 0 + cached_action_chunk = None + action_chunk_index = 0 + print("Cleared cached action chunk, reset frame counter") + elif key == "p": + pause_loop = not pause_loop + print(f"{'Paused' if pause_loop else 'Resumed'} policy loop") + if pause_loop: + print("Policy loop paused (C++ loop still running - press 'k' to stop)") + else: + print("Policy loop resumed") + elif key == "k": + if cpp_loop_running: + current_planner = cpp_mode == "PLANNER" + print(f"Stopping C++ control loop (from {cpp_mode} mode)...") + if send_cpp_control_command(start=False, planner=current_planner): + print("Stopped C++ control loop") + else: + print("Starting C++ control loop in PLANNER mode...") + if send_cpp_control_command(start=True, planner=True): + print("Started C++ control loop in PLANNER mode") + print("Press 'i' to send initial pose and switch to POSE mode") + if pause_loop: + print("Note: Policy loop is paused - press 'p' to resume") + elif key == "[": + initial_pose_left_hand_closed = not initial_pose_left_hand_closed + print( + f"Initial pose left hand: {'closed' if initial_pose_left_hand_closed else 'open'}" + ) + elif key == "]": + initial_pose_right_hand_closed = not initial_pose_right_hand_closed + print( + f"Initial pose right hand: " + f"{'closed' if initial_pose_right_hand_closed else 'open'}" + ) + + # Mutable prompt container (single-writer from keyboard, single-reader from inference) + language_prompt_ref: list[str] = [config.prompt] + print(f"Starting the policy loop with language prompt: {language_prompt_ref[0]}") + + inference_queue = queue.Queue(maxsize=1) + result_queue = queue.Queue(maxsize=1) + inference_stop_event = threading.Event() + inference_busy_event = threading.Event() + + inference_worker_thread = threading.Thread( + target=_inference_worker_loop, + args=( + inference_queue, + result_queue, + inference_stop_event, + inference_busy_event, + lambda: prepare_observation_from_sensors( + camera_subscriber=camera_subscriber, + state_subscriber=state_subscriber, + robot_model=robot_model, + language_prompt=language_prompt_ref[0], + log_errors=True, + ), + lambda obs: run_policy_inference_and_process( + policy=n1_policy, + observation=obs, + robot_model=robot_model, + ), + ), + daemon=True, + ) + inference_worker_thread.start() + + try: + while True: + t_start = time.monotonic() + check_keyboard_input() + + # Consume result first so last_inference_time is fresh before trigger check + try: + processed_action, inference_start_time = result_queue.get_nowait() + inference_delay = time.monotonic() - inference_start_time + action_chunk_index = calculate_latency_compensated_index( + inference_delay, config.action_publish_rate, config.action_horizon + ) + cached_action_chunk = processed_action + last_inference_time = time.monotonic() + print_green( + f'New action chunk (prompt: "{language_prompt_ref[0]}", ' + f"latency: {inference_delay:.3f}s)" + ) + except queue.Empty: + pass + + worker_is_busy = inference_busy_event.is_set() + should_start = should_trigger_new_inference( + cached_chunk_exists=(cached_action_chunk is not None), + inference_thread_running=worker_is_busy, + time_since_last_inference=(time.monotonic() - last_inference_time), + inference_interval=inference_interval, + ) + + if should_start: + try: + inference_queue.put_nowait(None) + except queue.Full: + pass + + if pause_loop: + print("Pausing...", end="", flush=True) + time.sleep(0.2) + print(".", end="", flush=True) + continue + + with telemetry.timer("total_loop"): + if cached_action_chunk is None: + print("[DEBUG] No cached chunk yet, waiting...", flush=True) + _sleep_remaining(t_start, loop_period) + continue + + processed_action = cached_action_chunk + + if processed_action is None or not processed_action: + print("[DEBUG] processed_action is None or empty, skipping", flush=True) + else: + motion_token = np.asarray( + get_action_field(processed_action, "motion_token"), + dtype=np.float32, + ) + left_hand_joints = np.asarray( + get_action_field(processed_action, "left_hand_joints"), + dtype=np.float32, + ) + right_hand_joints = np.asarray( + get_action_field(processed_action, "right_hand_joints"), + dtype=np.float32, + ) + + # Action arrays arrive as (B, T, D) from the model. + # Squeeze batch dim to get (T, D), then index by time step. + if motion_token.ndim == 3: + motion_token = motion_token[0] + if left_hand_joints.ndim == 3: + left_hand_joints = left_hand_joints[0] + if right_hand_joints.ndim == 3: + right_hand_joints = right_hand_joints[0] + + horizon = motion_token.shape[0] if motion_token.ndim == 2 else 1 + current_idx = min(action_chunk_index, horizon - 1) + + if motion_token.ndim == 2: + motion_token = motion_token[current_idx] + if left_hand_joints.ndim == 2: + left_hand_joints = left_hand_joints[current_idx] + if right_hand_joints.ndim == 2: + right_hand_joints = right_hand_joints[current_idx] + + frame_index = np.array([zmq_frame_counter], dtype=np.int64) + zmq_frame_counter += 1 + + zmq_message = pack_latent_action_message( + motion_token, + frame_index, + left_hand_joints=left_hand_joints, + right_hand_joints=right_hand_joints, + ) + zmq_socket.send(zmq_message) + last_sent_motion_token = motion_token.copy() + if zmq_frame_counter % 50 == 0: + print_green( + f"ZMQ: Sent latent action - " + f"frame: {frame_index[0]}, " + f"token shape: {motion_token.shape}" + ) + + action_chunk_index = min(action_chunk_index + 1, config.action_horizon - 1) + + end_time = time.monotonic() + + if config.verbose_timing: + telemetry.log_timing_info(context="VLA Inference Loop", threshold=0.0) + elif (end_time - t_start) > (1 / config.rate): + telemetry.log_timing_info( + context="VLA Inference Loop Missed", threshold=0.001 + ) + + _sleep_remaining(t_start, loop_period) + + except KeyboardInterrupt: + print("VLA inference loop terminated by user") + + finally: + inference_stop_event.set() + inference_worker_thread.join(timeout=1.0) + zmq_socket.close() + zmq_context.term() + state_subscriber.close() + keyboard_listener.close() + print("Shutdown complete.") + + +def _sleep_remaining(t_start: float, loop_period: float): + """Sleep for the remainder of the loop period.""" + elapsed = time.monotonic() - t_start + remaining = loop_period - elapsed + if remaining > 0: + time.sleep(remaining) + + +if __name__ == "__main__": + config = tyro.cli(InferenceConfig) + main(config) diff --git a/GR00T-WholeBodyControl/gear_sonic/scripts/run_vla_inference_dex1_head.py b/GR00T-WholeBodyControl/gear_sonic/scripts/run_vla_inference_dex1_head.py new file mode 100644 index 0000000000000000000000000000000000000000..5249bddaab37f9c31debc016c77bc6e6f8f078b0 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/scripts/run_vla_inference_dex1_head.py @@ -0,0 +1,697 @@ +"""Run standard GR00T N1.7 inference for SONIC + Dex1 + a two-axis head. + +This client is intentionally *not* an RTC client. Each policy request calls +``PolicyClient.get_action(observation)`` with no previous action chunk and no +RTC options. It retains the official GR00T-WholeBodyControl asynchronous +inference cadence and latency-compensated chunk index. + +Checkpoint contract used by this entry point: + +* video: ``ego_view``, ``left_wrist``, ``right_wrist`` as separate tensors; +* state: G1 body groups, two scalar grippers, measured yaw/pitch, gravity; +* action: 64-D SONIC token, two scalar grippers, 2-D absolute head target; +* horizon: 40 steps. + +The program starts paused. In the terminal, enter ``k`` (start SONIC planner), +``i`` (initialize and switch to pose mode), then ``p`` (run the policy). +""" + +# ruff: noqa: E402 -- remove ROS Python 3.10 paths before importing Pinocchio users + +from __future__ import annotations + +from dataclasses import dataclass +import queue +import sys +import threading +import time +import traceback +from typing import Any + +# ROS Humble can inject incompatible Python 3.10 packages into a Python 3.12 +# inference environment before cmeel/Pinocchio is imported. +sys.path[:] = [ + path + for path in sys.path + if not ("/opt/ros/humble/" in path and "python3.10" in path) +] + +import numpy as np +import tyro +import zmq + +from gear_sonic.camera.composed_camera import ComposedCameraClientSensor +from gear_sonic.data.robot_model.instantiation.g1 import instantiate_g1_robot_model +from gear_sonic.scripts.run_vla_inference import pack_latent_action_message +from gear_sonic.utils.data_collection.keyboard_subscriber import ( + DEFAULT_ZMQ_KEYBOARD_PORT, + ZMQKeyboardSubscriber, +) +from gear_sonic.utils.data_collection.transforms import compute_projected_gravity +from gear_sonic.utils.data_collection.zmq_state_subscriber import ZMQStateSubscriber +from gear_sonic.utils.inference.dex1_head import ( + Dex1CommandFilter, + HeadCommandFilter, + LatestMsgpackSubscriber, + pack_gripper_command, + pack_head_command, + parse_dex1_state, + parse_head_state, + validate_action_chunk, +) +from gear_sonic.utils.inference.initial_poses import LATENT_INITIAL_MOTION_TOKEN +from gear_sonic.utils.inference.vla_utils import ( + calculate_latency_compensated_index, + prepare_observation_for_eval, + should_trigger_new_inference, +) +from gear_sonic.utils.teleop.zmq.zmq_planner_sender import build_command_message + + +BODY_STATE_KEYS = ("left_leg", "right_leg", "waist", "left_arm", "right_arm") +EXPECTED_STATE_DIMS = { + "left_leg": 6, + "right_leg": 6, + "waist": 3, + "left_arm": 7, + "right_arm": 7, + "left_gripper": 1, + "right_gripper": 1, + "head_joints": 2, + "projected_gravity": 3, +} +EXPECTED_VIDEO_KEYS = ("ego_view", "left_wrist", "right_wrist") + + +@dataclass +class InferenceConfig: + """Command-line configuration for the real-robot policy client.""" + + host: str = "localhost" + """GR00T PolicyServer host.""" + + port: int = 5550 + """GR00T PolicyServer port.""" + + policy_timeout_ms: int = 30000 + """Policy request timeout in milliseconds.""" + + embodiment_tag: str = "unitree_g1_sonic" + """Expected checkpoint embodiment (logged as a deployment assertion).""" + + prompt: str = "Pick up the bottle and put it in the box" + """Initial language instruction.""" + + action_publish_rate: int = 50 + """SONIC/Dex1/head command rate in Hz.""" + + action_horizon: int = 40 + """Checkpoint action horizon.""" + + rate: float = 2.5 + """Minimum completed-request cadence in Hz, matching the official client.""" + + camera_host: str = "192.168.123.164" + camera_port: int = 5555 + state_zmq_host: str = "localhost" + state_zmq_port: int = 5557 + hand_state_host: str = "192.168.123.164" + hand_state_port: int = 5559 + head_state_host: str = "192.168.123.164" + head_state_port: int = 5561 + + output_bind_host: str = "0.0.0.0" + action_zmq_port: int = 5556 + hand_command_port: int = 5558 + head_command_port: int = 5560 + + use_zmq_keyboard: bool = False + """Read commands from the legacy ZMQ keyboard publisher instead of stdin.""" + + keyboard_zmq_host: str = "localhost" + keyboard_zmq_port: int = DEFAULT_ZMQ_KEYBOARD_PORT + + state_timeout: float = 0.5 + """Maximum age of Dex1/head feedback in seconds.""" + + max_gripper_step: float = 0.08 + head_yaw_limits: tuple[float, float] = (-1.2, 1.2) + head_pitch_limits: tuple[float, float] = (-0.6, 0.6) + max_yaw_step: float = 0.08 + max_pitch_step: float = 0.06 + + dry_run: bool = False + """Run sensing and inference but never publish hardware commands.""" + + debug: bool = False + verbose_timing: bool = False + + +def _endpoint(host: str, port: int) -> str: + return f"tcp://{host}:{port}" + + +def _green(message: str) -> None: + print(f"\033[92m{message}\033[0m", flush=True) + + +class StdinCommandReader: + """Read terminal commands on a daemon thread without blocking the control loop.""" + + def __init__(self) -> None: + self._commands: queue.Queue[str] = queue.Queue() + self._thread = threading.Thread(target=self._read_loop, daemon=True) + self._thread.start() + + def _read_loop(self) -> None: + while True: + try: + value = input().strip() + except EOFError: + return + if value: + self._commands.put(value) + + def read_msg(self) -> str | None: + try: + return self._commands.get_nowait() + except queue.Empty: + return None + + def close(self) -> None: + return + + +class DiagnosticLogger: + def __init__(self, enabled: bool, interval: float = 1.0): + self.enabled = enabled + self.interval = interval + self._last: dict[str, float] = {} + + def log(self, stage: str, message: str, *, force: bool = False) -> None: + if not self.enabled: + return + now = time.monotonic() + if not force and now - self._last.get(stage, float("-inf")) < self.interval: + return + self._last[stage] = now + print(f"[debug][{stage}] {message}", flush=True) + + def exception(self, stage: str, error: Exception) -> None: + self.log(stage, f"{type(error).__name__}: {error}", force=True) + if self.enabled: + traceback.print_exc() + + +def _validate_image(name: str, image: Any) -> np.ndarray: + value = np.asarray(image) + if value.dtype != np.uint8 or value.ndim != 3 or value.shape[-1] != 3: + raise ValueError( + f"Camera {name!r} must be uint8 HxWx3, got dtype={value.dtype}, shape={value.shape}" + ) + return np.ascontiguousarray(value) + + +def validate_observation_schema(observation: dict[str, Any]) -> None: + """Fail locally before a malformed observation crosses the policy network.""" + if set(observation.get("video", {})) != set(EXPECTED_VIDEO_KEYS): + raise ValueError( + f"Video keys must be exactly {EXPECTED_VIDEO_KEYS}, got " + f"{sorted(observation.get('video', {}))}" + ) + for key in EXPECTED_VIDEO_KEYS: + value = observation["video"][key] + if value.dtype != np.uint8 or value.ndim != 5 or value.shape[:2] != (1, 1): + raise ValueError(f"video.{key} must be uint8 [1,1,H,W,3], got {value.shape}") + + if set(observation.get("state", {})) != set(EXPECTED_STATE_DIMS): + raise ValueError( + f"State keys must be exactly {tuple(EXPECTED_STATE_DIMS)}, got " + f"{sorted(observation.get('state', {}))}" + ) + for key, width in EXPECTED_STATE_DIMS.items(): + value = observation["state"][key] + if value.dtype != np.float32 or value.shape != (1, 1, width): + raise ValueError(f"state.{key} must be float32 [1,1,{width}], got {value.shape}") + + language = observation.get("language", {}).get("annotation.human.task_description") + if not isinstance(language, list) or len(language) != 1: + raise ValueError("Language prompt must have a single batch entry") + + +def prepare_observation( + camera_subscriber, + state_subscriber, + hand_subscriber: LatestMsgpackSubscriber, + head_subscriber: LatestMsgpackSubscriber, + robot_model, + language_prompt: str, + state_timeout: float, + diagnostics: DiagnosticLogger | None = None, +) -> dict[str, Any] | None: + """Build the exact observation schema saved in checkpoint-20000.""" + camera_msg = camera_subscriber.read() + state_msg = state_subscriber.get_msg() + hand_msg = hand_subscriber.read() + head_msg = head_subscriber.read() + missing = [ + name + for name, value in ( + ("camera", camera_msg), + ("SONIC state", state_msg), + ("Dex1 state", hand_msg), + ("head state", head_msg), + ) + if value is None + ] + if missing: + if diagnostics: + diagnostics.log("observation", f"waiting for {missing}") + return None + if not hand_subscriber.is_fresh(state_timeout): + if diagnostics: + diagnostics.log("observation", "Dex1 feedback is stale") + return None + if not head_subscriber.is_fresh(state_timeout): + if diagnostics: + diagnostics.log("observation", "head feedback is stale") + return None + + images = camera_msg.get("images", {}) + missing_images = set(EXPECTED_VIDEO_KEYS).difference(images) + if missing_images: + raise KeyError( + f"Camera server is missing trained views {sorted(missing_images)}; " + f"available views: {sorted(images)}" + ) + + body_q = np.asarray(state_msg.get("body_q"), dtype=np.float32) + if body_q.shape != (29,) or not np.isfinite(body_q).all(): + raise ValueError(f"body_q must be finite [29], got {body_q.shape}") + base_quat = np.asarray(state_msg.get("base_quat"), dtype=np.float64) + if base_quat.shape != (4,) or not np.isfinite(base_quat).all(): + raise ValueError(f"base_quat must be finite [4], got {base_quat}") + + left_gripper, right_gripper = parse_dex1_state(hand_msg) + measured_head = parse_head_state(head_msg) + + # Hand joints are not checkpoint state modalities. Zeros are supplied only + # to satisfy RobotModel while splitting the 29 body joints into named groups. + unused_hand = np.zeros(7, dtype=np.float32) + qpos = robot_model.get_configuration_from_actuated_joints( + body_actuated_joint_values=body_q, + left_hand_actuated_joint_values=unused_hand, + right_hand_actuated_joint_values=unused_hand, + ) + grouped = {"state": {}, "q": np.asarray(qpos, dtype=np.float32)[None, None]} + prepare_observation_for_eval(robot_model, grouped) + + state = { + key: np.ascontiguousarray(grouped["state"][key], dtype=np.float32) + for key in BODY_STATE_KEYS + } + state["left_gripper"] = np.asarray(left_gripper, dtype=np.float32).reshape(1, 1, 1) + state["right_gripper"] = np.asarray(right_gripper, dtype=np.float32).reshape(1, 1, 1) + state["head_joints"] = measured_head.reshape(1, 1, 2) + state["projected_gravity"] = np.asarray( + compute_projected_gravity(base_quat), dtype=np.float32 + ).reshape(1, 1, 3) + + observation = { + "video": { + key: _validate_image(key, images[key])[None, None] for key in EXPECTED_VIDEO_KEYS + }, + "state": state, + "language": {"annotation.human.task_description": [[language_prompt]]}, + } + validate_observation_schema(observation) + if diagnostics: + diagnostics.log( + "observation", + "ready: " + f"video={{{', '.join(f'{k}: {v.shape}' for k, v in observation['video'].items())}}}, " + f"state={{{', '.join(f'{k}: {v.shape}' for k, v in state.items())}}}", + ) + return observation + + +def _inference_worker_loop( + requests: queue.Queue, + results: queue.Queue, + stop_event: threading.Event, + busy_event: threading.Event, + prepare_observation_fn, + infer_fn, + action_horizon: int, + diagnostics: DiagnosticLogger, +) -> None: + while not stop_event.is_set(): + try: + try: + requests.get(timeout=0.1) + except queue.Empty: + continue + busy_event.set() + stage = "observation" + try: + observation = prepare_observation_fn() + if observation is None: + continue + stage = "policy" + started_at = time.monotonic() + # Standard GR00T call: deliberately no options and no previous chunk. + action, _info = infer_fn(observation) + stage = "action" + chunk = validate_action_chunk(action, action_horizon) + try: + results.put_nowait((chunk, started_at)) + except queue.Full: + results.get_nowait() + results.put_nowait((chunk, started_at)) + except Exception as error: + print(f"[inference][{stage}] rejected request: {error}", flush=True) + diagnostics.exception(stage, error) + finally: + busy_event.clear() + except Exception as error: + print(f"[inference worker] {error}", flush=True) + diagnostics.exception("worker", error) + + +def _validate_config(config: InferenceConfig) -> None: + if config.embodiment_tag.lower() != "unitree_g1_sonic": + raise ValueError("This checkpoint must use embodiment_tag=unitree_g1_sonic") + if config.action_horizon != 40: + raise ValueError("This checkpoint requires action_horizon=40") + if config.action_publish_rate <= 0 or config.rate <= 0: + raise ValueError("action_publish_rate and rate must be positive") + if config.state_timeout <= 0: + raise ValueError("state_timeout must be positive") + + +def _sleep_remaining(started_at: float, period: float) -> None: + remaining = period - (time.monotonic() - started_at) + if remaining > 0: + time.sleep(remaining) + + +def main(config: InferenceConfig) -> None: + _validate_config(config) + from gr00t.policy.server_client import PolicyClient + + diagnostics = DiagnosticLogger(config.debug) + robot_model = instantiate_g1_robot_model(waist_location="lower_and_upper_body") + policy = PolicyClient( + host=config.host, + port=config.port, + timeout_ms=config.policy_timeout_ms, + ) + print(f"Connecting to standard GR00T PolicyServer at {config.host}:{config.port}") + if not policy.ping(): + print("WARNING: PolicyServer is not reachable yet; inference requests will fail") + else: + _green("PolicyServer is reachable") + + camera = ComposedCameraClientSensor(server_ip=config.camera_host, port=config.camera_port) + robot_state = ZMQStateSubscriber(host=config.state_zmq_host, port=config.state_zmq_port) + context = zmq.Context() + hand_state = LatestMsgpackSubscriber( + context, + _endpoint(config.hand_state_host, config.hand_state_port), + parse_dex1_state, + ) + head_state = LatestMsgpackSubscriber( + context, + _endpoint(config.head_state_host, config.head_state_port), + parse_head_state, + ) + keyboard = ( + ZMQKeyboardSubscriber( + host=config.keyboard_zmq_host, + port=config.keyboard_zmq_port, + ) + if config.use_zmq_keyboard + else StdinCommandReader() + ) + + output_sockets: dict[str, zmq.Socket] = {} + if config.dry_run: + print("DRY RUN: hardware command publication is disabled") + else: + for name, port in ( + ("sonic", config.action_zmq_port), + ("hand", config.hand_command_port), + ("head", config.head_command_port), + ): + socket = context.socket(zmq.PUB) + socket.setsockopt(zmq.SNDHWM, 2) + socket.bind(_endpoint(config.output_bind_host, port)) + output_sockets[name] = socket + _green(f"{name} command socket bound on {config.output_bind_host}:{port}") + time.sleep(0.2) + + paused = True + initialized = False + cpp_running = False + cpp_mode = "OFF" + prompt = [config.prompt] + cached_chunk: dict[str, np.ndarray] | None = None + chunk_index = 0 + last_inference_time = 0.0 + frame_counter = 0 + gripper_filter = Dex1CommandFilter(max_step=config.max_gripper_step) + head_filter = HeadCommandFilter( + yaw_limits=config.head_yaw_limits, + pitch_limits=config.head_pitch_limits, + max_yaw_step=config.max_yaw_step, + max_pitch_step=config.max_pitch_step, + ) + + requests: queue.Queue = queue.Queue(maxsize=1) + results: queue.Queue = queue.Queue(maxsize=1) + stop_event = threading.Event() + busy_event = threading.Event() + worker = threading.Thread( + target=_inference_worker_loop, + args=( + requests, + results, + stop_event, + busy_event, + lambda: prepare_observation( + camera, + robot_state, + hand_state, + head_state, + robot_model, + prompt[0], + config.state_timeout, + diagnostics, + ), + policy.get_action, + config.action_horizon, + diagnostics, + ), + daemon=True, + ) + worker.start() + + def publish_targets( + motion_token: Any, + left_target: Any, + right_target: Any, + head_target: Any, + ) -> None: + nonlocal frame_counter + token = np.asarray(motion_token, dtype=np.float32).reshape(-1) + if token.shape != (64,) or not np.isfinite(token).all(): + raise ValueError(f"motion token must be finite [64], got {token.shape}") + left, right = gripper_filter.update(left_target, right_target) + head = head_filter.update(head_target) + if config.dry_run: + if frame_counter % config.action_publish_rate == 0: + print( + f"[dry-run] frame={frame_counter} grippers=({left:.3f}, {right:.3f}) " + f"head=({head[0]:.3f}, {head[1]:.3f}) token_peak={np.abs(token).max():.3f}" + ) + else: + timestamp_ns = time.time_ns() + output_sockets["sonic"].send( + pack_latent_action_message( + token, + np.asarray([frame_counter], dtype=np.int64), + ) + ) + output_sockets["hand"].send(pack_gripper_command(left, right, timestamp_ns)) + output_sockets["head"].send( + pack_head_command(float(head[0]), float(head[1]), timestamp_ns) + ) + frame_counter += 1 + + def send_cpp_command(start: bool, planner: bool) -> None: + nonlocal cpp_running, cpp_mode + if config.dry_run: + print(f"[dry-run] SONIC {'start' if start else 'stop'} command suppressed") + else: + output_sockets["sonic"].send( + build_command_message(start=start, stop=not start, planner=planner) + ) + cpp_running = start + cpp_mode = "PLANNER" if start and planner else "POSE" if start else "OFF" + print(f"SONIC control: running={cpp_running}, mode={cpp_mode}", flush=True) + + def initialize_pose() -> bool: + nonlocal initialized, cached_chunk, chunk_index, paused + if not config.dry_run and not cpp_running: + print("Cannot initialize: press k first to start SONIC in PLANNER mode") + return False + hand = hand_state.read() + head = head_state.read() + if hand is None or not hand_state.is_fresh(config.state_timeout): + print("Cannot initialize: no fresh Dex1 feedback") + return False + if head is None or not head_state.is_fresh(config.state_timeout): + print("Cannot initialize: no fresh head feedback") + return False + left, right = parse_dex1_state(hand) + measured_head = parse_head_state(head) + gripper_filter.reset(left, right) + head_filter.reset(measured_head) + publish_targets(LATENT_INITIAL_MOTION_TOKEN, left, right, measured_head) + send_cpp_command(True, planner=False) + paused = True + initialized = True + cached_chunk = None + chunk_index = 0 + print("Initial SONIC token sent; measured gripper/head pose held; policy remains paused") + return True + + def handle_command(command: str) -> None: + nonlocal paused, initialized, cached_chunk, chunk_index + normalized = command.strip() + if normalized.startswith("prompt:"): + normalized = "t " + normalized[len("prompt:") :] + if normalized.startswith("t ") and normalized[2:].strip(): + prompt[0] = normalized[2:].strip() + paused = True + cached_chunk = None + chunk_index = 0 + print(f'Prompt changed to "{prompt[0]}"; policy paused') + elif normalized == "k": + if cpp_running: + paused = True + initialized = False + send_cpp_command(False, planner=(cpp_mode == "PLANNER")) + else: + initialized = False + send_cpp_command(True, planner=True) + print("Press i after the robot is stable in PLANNER mode") + elif normalized == "i": + initialize_pose() + elif normalized == "p": + if not paused: + paused = True + print("Policy paused") + elif not config.dry_run and (not initialized or cpp_mode != "POSE"): + print("Cannot resume: complete k -> i first") + else: + cached_chunk = None + chunk_index = 0 + paused = False + print("Policy resumed; requesting a fresh standard GR00T chunk") + elif normalized in {"help", "h", "?"}: + print( + "Commands: k=start/stop SONIC, i=initialize pose, " + "p=pause/resume, t =prompt" + ) + elif normalized: + print(f"Unknown command {normalized!r}; enter help") + + print(f'Standard GR00T client ready with prompt: "{prompt[0]}"') + print("No RTC options are used. Camera views remain separate.") + print("Commands: k -> i -> p; enter 'help' for details. Ctrl-C stops publication.") + + loop_period = 1.0 / config.action_publish_rate + inference_interval = 1.0 / config.rate + try: + while True: + tick_started = time.monotonic() + command = keyboard.read_msg() + if command: + handle_command(command) + + try: + new_chunk, inference_started = results.get_nowait() + delay = time.monotonic() - inference_started + cached_chunk = new_chunk + chunk_index = calculate_latency_compensated_index( + delay, + config.action_publish_rate, + config.action_horizon, + ) + last_inference_time = time.monotonic() + _green( + f"New standard action chunk: latency={delay:.3f}s, start_index={chunk_index}" + ) + except queue.Empty: + pass + + if not paused: + should_start = should_trigger_new_inference( + cached_chunk_exists=cached_chunk is not None, + inference_thread_running=busy_event.is_set(), + time_since_last_inference=time.monotonic() - last_inference_time, + inference_interval=inference_interval, + ) + if should_start: + try: + requests.put_nowait(None) + except queue.Full: + pass + + if cached_chunk is not None: + current = min(chunk_index, config.action_horizon - 1) + publish_targets( + cached_chunk["motion_token"][current], + cached_chunk["left_hand_joints"][current], + cached_chunk["right_hand_joints"][current], + cached_chunk["head_joints"][current], + ) + chunk_index = min(chunk_index + 1, config.action_horizon - 1) + + elapsed = time.monotonic() - tick_started + if config.verbose_timing and elapsed > loop_period: + print(f"Control tick overrun: {(elapsed - loop_period) * 1000:.1f} ms") + _sleep_remaining(tick_started, loop_period) + except KeyboardInterrupt: + print("Stopping real-robot GR00T client") + finally: + stop_event.set() + worker.join(timeout=1.0) + if not config.dry_run and cpp_running: + try: + output_sockets["sonic"].send( + build_command_message( + start=False, + stop=True, + planner=(cpp_mode == "PLANNER"), + ) + ) + time.sleep(0.05) + print("Sent SONIC stop command") + except Exception as error: + print(f"WARNING: failed to send SONIC stop command: {error}") + for socket in output_sockets.values(): + socket.close(linger=0) + hand_state.close() + head_state.close() + robot_state.close() + camera.close() + keyboard.close() + policy.close() + context.term() + print("Shutdown complete") + + +if __name__ == "__main__": + main(tyro.cli(InferenceConfig)) diff --git a/GR00T-WholeBodyControl/gear_sonic/tests/test_dex1_head_inference.py b/GR00T-WholeBodyControl/gear_sonic/tests/test_dex1_head_inference.py new file mode 100644 index 0000000000000000000000000000000000000000..100401e24ff5410bf8fb49ae38652796fdd3589e --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/tests/test_dex1_head_inference.py @@ -0,0 +1,230 @@ +"""Offline tests for the custom checkpoint/real-hardware boundary.""" + +from __future__ import annotations + +import importlib +import queue +import sys +import threading +import types +import unittest + +import numpy as np + +from gear_sonic.utils.inference.dex1_head import ( + Dex1CommandFilter, + HeadCommandFilter, + parse_dex1_state, + validate_action_chunk, +) + + +def action_chunk(horizon: int = 40) -> dict: + return { + "motion_token": np.zeros((1, horizon, 64), dtype=np.float32), + "left_hand_joints": np.full((1, horizon, 1), 0.25, dtype=np.float32), + "right_hand_joints": np.full((1, horizon, 1), 0.75, dtype=np.float32), + "head_joints": np.zeros((1, horizon, 2), dtype=np.float32), + } + + +class ActionAndFilterTests(unittest.TestCase): + def test_checkpoint_action_contract_is_64_plus_1_plus_1_plus_2(self): + normalized = validate_action_chunk(action_chunk(), 40) + self.assertEqual(normalized["motion_token"].shape, (40, 64)) + self.assertEqual(normalized["left_hand_joints"].shape, (40, 1)) + self.assertEqual(normalized["right_hand_joints"].shape, (40, 1)) + self.assertEqual(normalized["head_joints"].shape, (40, 2)) + + wrong = action_chunk() + wrong["left_hand_joints"] = np.zeros((1, 40, 7), dtype=np.float32) + with self.assertRaisesRegex(ValueError, r"\[T, 1\]"): + validate_action_chunk(wrong, 40) + + def test_scalar_feedback_and_rate_limits(self): + self.assertEqual( + parse_dex1_state( + {"left_hand_joints": [0.2], "right_hand_joints": [0.8]} + ), + (0.2, 0.8), + ) + gripper = Dex1CommandFilter(max_step=0.1, left_safe=0.2, right_safe=0.8) + left, right = gripper.update([1.0], [0.0]) + self.assertAlmostEqual(left, 0.3) + self.assertAlmostEqual(right, 0.7) + + head = HeadCommandFilter() + head.reset([0.0, 0.0]) + np.testing.assert_allclose(head.update([2.0, -2.0]), [0.08, -0.06]) + + +class ObservationAndStandardPolicyTests(unittest.TestCase): + @staticmethod + def install_stubs() -> None: + def register(name: str, **members) -> None: + module = types.ModuleType(name) + module.__dict__.update(members) + sys.modules[name] = module + + register("tyro", cli=lambda config: config()) + register("zmq", PUB=1, SNDHWM=2, Socket=object, Context=object) + register("gear_sonic.camera.composed_camera", ComposedCameraClientSensor=object) + register( + "gear_sonic.data.robot_model.instantiation.g1", + instantiate_g1_robot_model=lambda **_kwargs: None, + ) + register( + "gear_sonic.scripts.run_vla_inference", + pack_latent_action_message=lambda *_args, **_kwargs: b"sonic", + ) + register( + "gear_sonic.utils.data_collection.keyboard_subscriber", + DEFAULT_ZMQ_KEYBOARD_PORT=5580, + ZMQKeyboardSubscriber=object, + ) + register( + "gear_sonic.utils.data_collection.transforms", + compute_projected_gravity=lambda _quat: np.array([0.0, 0.0, -1.0]), + ) + register( + "gear_sonic.utils.data_collection.zmq_state_subscriber", + ZMQStateSubscriber=object, + ) + register( + "gear_sonic.utils.inference.initial_poses", + LATENT_INITIAL_MOTION_TOKEN=np.zeros(64, dtype=np.float32), + ) + + def split_groups(robot_model, observation): + for name, indices in robot_model.groups.items(): + observation["state"][name] = observation["q"][..., indices] + return observation + + register( + "gear_sonic.utils.inference.vla_utils", + calculate_latency_compensated_index=lambda *_args: 0, + prepare_observation_for_eval=split_groups, + should_trigger_new_inference=lambda **_kwargs: True, + ) + register( + "gear_sonic.utils.teleop.zmq.zmq_planner_sender", + build_command_message=lambda **_kwargs: b"command", + ) + + @classmethod + def load_client(cls): + cls.install_stubs() + name = "gear_sonic.scripts.run_vla_inference_dex1_head" + sys.modules.pop(name, None) + return importlib.import_module(name) + + def test_observation_matches_checkpoint_exactly(self): + client = self.load_client() + + class Latest: + def __init__(self, value): + self.value = value + + def read(self): + return self.value + + def get_msg(self): + return self.value + + def is_fresh(self, _timeout): + return True + + class RobotModel: + groups = { + "left_leg": list(range(0, 6)), + "right_leg": list(range(6, 12)), + "waist": list(range(12, 15)), + "left_arm": list(range(15, 22)), + "right_arm": list(range(22, 29)), + "left_hand": list(range(29, 36)), + "right_hand": list(range(36, 43)), + } + + def get_configuration_from_actuated_joints(self, **kwargs): + return np.concatenate( + [ + kwargs["body_actuated_joint_values"], + kwargs["left_hand_actuated_joint_values"], + kwargs["right_hand_actuated_joint_values"], + ] + ) + + camera = Latest( + { + "images": { + "ego_view": np.zeros((4, 5, 3), dtype=np.uint8), + "left_wrist": np.zeros((2, 3, 3), dtype=np.uint8), + "right_wrist": np.zeros((2, 3, 3), dtype=np.uint8), + } + } + ) + body = Latest( + { + "body_q": np.zeros(29), + "base_quat": np.array([1.0, 0.0, 0.0, 0.0]), + } + ) + hand = Latest({"left_hand_joints": [0.2], "right_hand_joints": [0.8]}) + head = Latest({"yaw_position": 0.1, "pitch_position": -0.2}) + observation = client.prepare_observation( + camera, + body, + hand, + head, + RobotModel(), + "pick up the cup", + 0.5, + ) + self.assertEqual( + set(observation["video"]), + {"ego_view", "left_wrist", "right_wrist"}, + ) + self.assertEqual(set(observation["state"]), set(client.EXPECTED_STATE_DIMS)) + self.assertEqual(observation["state"]["left_gripper"].shape, (1, 1, 1)) + self.assertEqual(observation["state"]["right_gripper"].shape, (1, 1, 1)) + self.assertEqual(observation["state"]["head_joints"].shape, (1, 1, 2)) + self.assertNotIn("left_hand", observation["state"]) + self.assertNotIn("right_hand", observation["state"]) + + def test_worker_calls_standard_policy_without_rtc_options(self): + client = self.load_client() + requests: queue.Queue = queue.Queue(maxsize=1) + results: queue.Queue = queue.Queue(maxsize=1) + stop = threading.Event() + busy = threading.Event() + called = [] + + def standard_policy(observation): + called.append(observation) + return action_chunk(), {} + + worker = threading.Thread( + target=client._inference_worker_loop, + args=( + requests, + results, + stop, + busy, + lambda: {"standard": True}, + standard_policy, + 40, + client.DiagnosticLogger(False), + ), + daemon=True, + ) + worker.start() + requests.put(None) + chunk, _started_at = results.get(timeout=2.0) + stop.set() + worker.join(timeout=1.0) + self.assertEqual(len(called), 1) + self.assertEqual(chunk["head_joints"].shape, (40, 2)) + + +if __name__ == "__main__": + unittest.main() diff --git a/GR00T-WholeBodyControl/gear_sonic/tests/test_input_readers.py b/GR00T-WholeBodyControl/gear_sonic/tests/test_input_readers.py new file mode 100644 index 0000000000000000000000000000000000000000..2fc4e6153fd38d381af8416df79367a2d460a7c9 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/tests/test_input_readers.py @@ -0,0 +1,50 @@ +import msgpack +import msgpack_numpy as msgpack_numpy +import numpy as np + +from gear_sonic.utils.teleop.input_readers import ( + build_body_pose_sample, + decode_msgpack_byte_multi_array, +) + + +def test_decode_msgpack_byte_multi_array_from_byte_chunks(): + payload = { + "timestamp": 123456789, + "joint_positions": [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], + "joint_orientations": [[0.0, 0.0, 0.0, 1.0], [0.5, 0.5, 0.5, 0.5]], + } + packed = msgpack.packb(payload, default=msgpack_numpy.encode, use_bin_type=True) + byte_chunks = [bytes([value]) for value in packed] + + decoded = decode_msgpack_byte_multi_array( + byte_chunks, + msgpack_module=msgpack, + msgpack_numpy_module=msgpack_numpy, + ) + + assert decoded["timestamp"] == payload["timestamp"] + assert decoded["joint_positions"] == payload["joint_positions"] + assert decoded["joint_orientations"] == payload["joint_orientations"] + + +def test_build_body_pose_sample_uses_existing_teleop_shape(): + payload = { + "timestamp": 1_000_000_100, + "joint_positions": [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]], + "joint_orientations": [[0.0, 0.0, 0.0, 1.0], [0.0, 0.0, 1.0, 0.0]], + } + + sample, stamp_ns, fps_ema = build_body_pose_sample( + payload, + prev_stamp_ns=1_000_000_000, + fps_ema=0.0, + ) + + assert sample is not None + assert stamp_ns == payload["timestamp"] + assert sample["body_poses_np"].shape == (24, 7) + np.testing.assert_allclose(sample["body_poses_np"][0], np.array([0.1, 0.2, 0.3, 0.0, 0.0, 0.0, 1.0])) + np.testing.assert_allclose(sample["body_poses_np"][1], np.array([0.4, 0.5, 0.6, 0.0, 0.0, 1.0, 0.0])) + assert sample["dt"] == 1e-7 + assert fps_ema == 1.0 / 1e-7 diff --git a/GR00T-WholeBodyControl/gear_sonic/train_agent_trl.py b/GR00T-WholeBodyControl/gear_sonic/train_agent_trl.py new file mode 100644 index 0000000000000000000000000000000000000000..be4b1265334a2973f5d2e94d50a089663ccce529 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/train_agent_trl.py @@ -0,0 +1,487 @@ +#!/usr/bin/env python3 +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Fix sys.path: when running as `python gear_sonic/train_agent_trl.py`, Python adds +# gear_sonic/ to sys.path[0], causing `from trl import ...` to resolve to our local +# gear_sonic/trl/ instead of the HuggingFace trl package. Replace with repo root. +import sys +import os +_script_dir = os.path.dirname(os.path.abspath(__file__)) +_repo_root = os.path.dirname(_script_dir) +if _script_dir in sys.path: + sys.path.remove(_script_dir) +if _repo_root not in sys.path: + sys.path.insert(0, _repo_root) + +try: + import isaaclab # noqa: F401 +except ImportError: + print( + "\n" + "ERROR: Isaac Lab is required for training but not installed.\n" + "\n" + "Isaac Lab is not a pip dependency — it must be installed separately.\n" + "Follow the official guide:\n" + " https://isaac-sim.github.io/IsaacLab/main/source/setup/installation/index.html\n" + "\n" + "After installing, activate the Isaac Lab conda/venv environment\n" + "before running this script.\n" + ) + sys.exit(1) + +import glob +import logging +import os +from pathlib import Path +import re +import sys + +from filelock import FileLock +import hydra +from hydra.core.hydra_config import HydraConfig +from hydra.utils import instantiate +from loguru import logger +from omegaconf import DictConfig, OmegaConf +import wandb +import yaml + +from gear_sonic.trl.utils.common import ( + custom_instantiate, + get_filtered_state_dict, + materialize_lazy_params, + wandb_run_exists, +) +from gear_sonic.utils.common import seeding +from gear_sonic.utils.config_utils import register_rl_resolvers +from gear_sonic.utils.obs_utils import get_group_term_obs_shape + +register_rl_resolvers() + + +def resume_training(config): + if config.get("checkpoint", None) is not None: + last_existing_checkpoint = config.checkpoint + elif config.get("experiment_dir", None) is not None: + last_existing_checkpoint = os.path.join(config.experiment_dir, "last.pt") + else: + # Use experiment_dir to find the checkpoint, rather than reconstructing + # from config.project_name which can differ from the actual filesystem path. + experiment_dir_base = re.sub(r"-\d{8}_\d{6}$", "", config.experiment_dir) + checkpoints = sorted(glob.glob(os.path.join(f"{experiment_dir_base}-*", "last.pt"))) + if not checkpoints: + print(f"No checkpoint found matching {experiment_dir_base}-*/last.pt, starting fresh") + return + last_existing_checkpoint = checkpoints[-1] + experiment_dir = os.path.dirname(last_existing_checkpoint) + config.experiment_dir = experiment_dir + config.checkpoint = last_existing_checkpoint + print(f"Resuming training from {last_existing_checkpoint}") + + +def resume_checkpoint(config): + config.checkpoint = config.checkpoint + + +def create_manager_env(config, device, args_cli): + + # import wandb + + from isaaclab.envs import ( + ManagerBasedRLEnv, + ) + + from gear_sonic.envs.wrapper.manager_env_wrapper import ManagerEnvWrapper + + env_instance_cfg = custom_instantiate(config.manager_env) + + # Iteratively check the difference in attribute of env_instance_cfg1 and env_instance_cfg, print out the difference + def compare_attrs(obj1, obj2, prefix=""): + # Only compare attributes that do not start with '__' and are not methods + attrs1 = set(dir(obj1)) + attrs2 = set(dir(obj2)) + common_attrs = attrs1 & attrs2 + for attr in sorted(common_attrs): + if ( + attr.startswith("__") + or callable(getattr(obj1, attr)) + or callable(getattr(obj2, attr)) + ): + continue + try: + val1 = getattr(obj1, attr) + val2 = getattr(obj2, attr) + except Exception: + continue + # Recursively compare if both are objects with __dict__ or are dicts + if isinstance(val1, dict | DictConfig) and isinstance(val2, dict | DictConfig): + compare_attrs(val1, val2, prefix + attr + ".") + elif hasattr(val1, "__dict__") and hasattr(val2, "__dict__"): + compare_attrs(val1, val2, prefix + attr + ".") + else: + if isinstance(val1, list): + val1 = tuple(val1) + if isinstance(val2, list): + val2 = tuple(val2) + if val1 != val2: + print( + f"\nDifference found at '{prefix}{attr}':\n" + f" - env_instance_cfg1: {val1!r}\n" + f" - env_instance_cfg : {val2!r}\n" + ) + + env_instance_cfg.seed = config.seed + env_instance_cfg.sim.device = device + env_instance_cfg.config["headless"] = args_cli.headless + env = ManagerBasedRLEnv( + cfg=env_instance_cfg, render_mode="rgb_array" if not args_cli.headless else None + ) + + env = ManagerEnvWrapper(env, env_instance_cfg.config) + return env + + +@hydra.main(config_path="config", config_name="base", version_base="1.1") +def main(config: OmegaConf): + simulator_type = "IsaacSim" + env_config = config.manager_env + from transformers import HfArgumentParser + from trl import ModelConfig, PPOConfig, ScriptArguments + + # Setup model components + parser = HfArgumentParser((ScriptArguments, PPOConfig, ModelConfig)) + + if config.get("resume", False): + resume_training(config) + elif config.get("checkpoint", None) is not None: + resume_checkpoint(config) + + config.algo.trl.output_dir = str(Path(config.experiment_dir)) + + script_args, training_args, model_args = parser.parse_dict(config.algo.trl) + + # Add exp_name from main config to training_args + training_args.exp_name = config.experiment_name + + from datetime import timedelta + + from accelerate import Accelerator, DistributedDataParallelKwargs, InitProcessGroupKwargs + import torch # noqa: E402 + + ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=False) + kwargs = InitProcessGroupKwargs(timeout=timedelta(seconds=6000)) + accelerator = Accelerator( + gradient_accumulation_steps=training_args.gradient_accumulation_steps, + kwargs_handlers=[ddp_kwargs, kwargs], + ) + + device = str(accelerator.device) + if device == "cuda": + device = "cuda:0" + config.multi_gpu = accelerator.num_processes > 1 + if config.multi_gpu: + config.global_rank = accelerator.process_index + config.seed += accelerator.process_index + config.algo.config.global_rank = accelerator.process_index + config.algo.config.world_size = accelerator.num_processes + seeding(config.seed) + + meta_path = Path(config.experiment_dir) / "meta.yaml" + if meta_path.exists(): + meta = yaml.safe_load(open(meta_path)) + config.wandb.wandb_id = meta["wandb_run"] + print(f"resume wandb from run: {config.wandb.wandb_id}") + + unresolved_conf = OmegaConf.to_container(config, resolve=False) + if config.use_wandb and accelerator.is_main_process: + project_name = f"{config.project_name}" + run_name = config.experiment_dir.replace(f"{config.base_dir}/{project_name}/", "") + wandb_dir = Path(config.wandb.wandb_dir) + wandb_dir.mkdir(exist_ok=True, parents=True) + wandb_group = None if config.wandb.wandb_id is not None else config.wandb.wandb_group + logger.info(f"Saving wandb logs to {wandb_dir}") + wandb.init( + project=project_name, + entity=config.wandb.wandb_entity, + name=run_name, + sync_tensorboard=True, + config=unresolved_conf, + dir=wandb_dir, + id=config.wandb.wandb_id, + group=wandb_group, + resume="allow", + ) + + # Setup simulator similar to train_agent.py + + if simulator_type == "IsaacSim": + try: + with open("./rl/simulator/isaacsim/.isaacsim_version", encoding="utf-8") as f: + DEFAULT_ISAACSIM_VERSION = f.read().strip() + except FileNotFoundError: + DEFAULT_ISAACSIM_VERSION = "4.5" + + if DEFAULT_ISAACSIM_VERSION == "4.5": + from isaaclab.app import AppLauncher + elif DEFAULT_ISAACSIM_VERSION == "4.2": + logger.warning("Using IsaacSim 4.2, replacing isaaclab with omni.isaac.lab") + from omni.isaac.lab.app import AppLauncher # 4.2 + + # from isaaclab.app import AppLauncher # not working + # from omni.isaac.lab.app import AppLauncher + + import argparse + + parser = argparse.ArgumentParser(description="Train an RL agent with TRL.") + AppLauncher.add_app_launcher_args(parser) + + ######################################################### ZL: fix isaacsim 4.5 rendering ######################################################### + args_cli, hydra_args = parser.parse_known_args() + sys.argv = [sys.argv[0]] + hydra_args + args_cli.num_envs = config.num_envs + args_cli.seed = config.seed + args_cli.env_spacing = env_config.config.env_spacing # config.env_spacing + args_cli.output_dir = config.output_dir + # Enable cameras if enable_cameras, render_results, render_ego, or overview_camera is True + args_cli.enable_cameras = ( + env_config.config.get("enable_cameras", False) + or env_config.config.get("render_results", False) + or env_config.config.get("render_ego", False) + or env_config.config.get("overview_camera", False) + ) + args_cli.headless = config.headless + args_cli.multi_gpu = config.multi_gpu + args_cli.distributed = config.multi_gpu + args_cli.device = device + + # Base kit args (quiet logs) + args_cli.kit_args = ( + "--/log/level=error --/log/fileLogLevel=error --/log/outputStreamLevel=error" + ) + + # Allow air-gapped machines to use an experience file with online + # extension registries disabled, while preserving the default behavior. + offline_experience = os.environ.get("ISAACLAB_EXPERIENCE") + if offline_experience: + args_cli.experience = offline_experience + + # AppLauncher can't handle multiple processes creating it at the same time so we need a lock + _lock_path = "/tmp/isaaclab_app_launcher.lock" + _local_rank = int(os.environ.get("LOCAL_RANK", 0)) + with FileLock(_lock_path): + app_launcher = AppLauncher(args_cli) + + simulation_app = app_launcher.app + + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + torch.backends.cudnn.deterministic = False + torch.backends.cudnn.benchmark = False + + from gear_sonic.utils.logging import HydraLoggerBridge + + # resolve=False is important otherwise overrides + # at inference time won't work properly + # also, I believe this must be done before instantiation + + # logging to hydra log file + hydra_log_path = os.path.join(HydraConfig.get().runtime.output_dir, "train.log") + logger.remove() + logger.add(hydra_log_path, level="DEBUG") + console_log_level = os.environ.get("LOGURU_LEVEL", "INFO").upper() + logger.add(sys.stdout, level=console_log_level, colorize=True) + logging.basicConfig(level=logging.DEBUG) + logging.getLogger().addHandler(HydraLoggerBridge()) + + # Setup wandb if enabled + os.chdir(hydra.utils.get_original_cwd()) + + # Save config and meta BEFORE env creation so eval jobs can postprocess + # checkpoint configs even if training crashes during env init. + experiment_save_dir = Path(config.experiment_dir) + if accelerator.is_main_process: + experiment_save_dir.mkdir(exist_ok=True, parents=True) + logger.info(f"Saving config file to {experiment_save_dir}") + with open(experiment_save_dir / "config.yaml", "w") as file: + OmegaConf.save(unresolved_conf, file) + meta = {"wandb_run": wandb.run.id if wandb_run_exists() else None} + meta["max_train_steps"] = config.algo.config.num_learning_iterations + yaml.safe_dump(meta, open(meta_path, "w")) + print("saved meta:", meta) + + # Initialize environment + env_config.config.save_rendering_dir = str(Path(config.experiment_dir) / "renderings_training") + env_config.config.experiment_dir = str(Path(config.experiment_dir)) + + env = create_manager_env(config, device, args_cli) + if config.get("replay", False): + _save_video_path = config.get("replay_save_video", None) + env.run_replay( + start_time_step=-1, + loop=config.get("replay_loop_num", True), + save_video_path=_save_video_path, + grid_spacing=config.get("replay_grid_spacing", 2.0), + ) + os._exit(0) + if config.get("vplanner_replay", False): + vplanner_checkpoint = config.get("vplanner_checkpoint", None) + if vplanner_checkpoint is None: + raise ValueError("vplanner_checkpoint must be specified for vplanner_replay") + env.run_vplanner_replay( + checkpoint_path=vplanner_checkpoint, + max_frames=config.get("vplanner_max_frames", 500), + replan_interval=config.get("vplanner_replan_interval", 0), + speed=config.get("vplanner_speed", 1.0), + loop=config.get("vplanner_loop", True), + save_images=config.get("vplanner_save_images", False), + output_dir=config.get("vplanner_output_dir", None), + dof_noise=config.get("vplanner_dof_noise", 0.0), + dof_vel_noise=config.get("vplanner_dof_vel_noise", 0.0), + quat_noise=config.get("vplanner_quat_noise", 0.0), + ) + os._exit(0) + + ref_model = None + value_model = None + disc_model = None + # import ipdb; ipdb.set_trace() + + if config.algo.config.get("use_new_actor_critic", False): + module_dim_dict = getattr(config.algo.config, "module_dim", {}) + policy_backbone_kwargs = {} + critic_backbone_kwargs = {} + env.config["obs"]["obs_dims"]["actor_obs"] = env.env.observation_space["policy"].shape[-1] + env.config["obs"]["obs_dims"]["critic_obs"] = env.env.observation_space["critic"].shape[-1] + env.config["robot"]["algo_obs_dim_dict"]["actor_obs"] = env.env.observation_space[ + "policy" + ].shape[-1] + env.config["robot"]["algo_obs_dim_dict"]["critic_obs"] = env.env.observation_space[ + "critic" + ].shape[-1] + example_obs = env.reset(flatten_dict_obs=False) + for key in env.env.observation_space: + if key not in ["policy", "critic"]: + group_obs_dims, group_obs_names, group_obs_total_dim = get_group_term_obs_shape( + example_obs, key + ) + env.config["obs"]["group_obs_dims"][key] = group_obs_dims + env.config["obs"]["group_obs_names"][key] = group_obs_names + env.config["obs"]["obs_dims"][key] = group_obs_total_dim + env.config["robot"]["algo_obs_dim_dict"][key] = group_obs_total_dim + if config.manager_env.config.get("meta_action_dim", None) is not None: + env.config["robot"]["actions_dim"] = config.manager_env.config.meta_action_dim + else: + env.config["robot"]["actions_dim"] = env.env.action_space.shape[-1] + + policy = custom_instantiate( + config.algo.config.actor, + env_config=env.config, + algo_config=config.algo.config, + module_dim_dict=module_dim_dict, + backbone_kwargs=policy_backbone_kwargs, + _resolve=False, + ).to(device) + + if getattr(config.algo.config, "use_dagger", False): + # Get teacher input key from config or default to "teacher" + teacher_input_key = config.algo.config.get("teacher_input_key", "teacher") + ref_model = custom_instantiate( + config.algo.config.teacher_actor, + env_config=env.config, + algo_config=config.algo.config, + module_dim_dict=module_dim_dict, + _resolve=False, + input_key=teacher_input_key, + ).to(device) + if not getattr(config.algo.config, "distill_only", False): + value_model = custom_instantiate( + config.algo.config.critic, + env_config=env.config, + algo_config=config.algo.config, + module_dim_dict=module_dim_dict, + backbone_kwargs=critic_backbone_kwargs, + _resolve=False, + ).to(device) + if config.algo.config.get("use_amp", False): + disc_model = custom_instantiate( + config.algo.config.disc, + env_config=env.config, + algo_config=config.algo.config, + module_dim_dict=module_dim_dict, + _resolve=False, + ).to(device) + else: + raise ValueError("No longer supported") + + materialize_lazy_params(policy, env) + + if config.algo.config.get("pretrained_model", None) is not None: + pretrained_cfg = config.algo.config.pretrained_model + sd_key = pretrained_cfg.get("state_dict_key", "state_dict") + strict = pretrained_cfg.get("strict", True) + state_dict = torch.load(pretrained_cfg.path, map_location=device, weights_only=False)[ + sd_key + ] + for ( + module_name, + state_dict_key, + ) in pretrained_cfg.module_mapping.items(): + module = eval(module_name) + filtered_state_dict = get_filtered_state_dict(state_dict, state_dict_key) + missing, unexpected = module.load_state_dict(filtered_state_dict, strict=strict) + if missing: + logger.info(f"Pretrained loading '{module_name}': missing keys: {missing}") + if unexpected: + logger.info(f"Pretrained loading '{module_name}': unexpected keys: {unexpected}") + + accelerator.wait_for_everyone() + + callbacks = [] + for callback in config.callbacks.values(): + callbacks.append(instantiate(callback)) + + ################ + # Training + ################ + trainer = custom_instantiate( + config.trainer, + args=training_args, + config=config.algo.config, + env=env, + model=policy, + disc_model=disc_model, + value_model=value_model, + ref_model=ref_model, + use_ref_model=getattr(config.algo.config, "use_dagger", False), + train_dataset=None, + eval_dataset=None, + callbacks=callbacks, + checkpoint=config.checkpoint, + resume=config.get("resume", False), + local_seed=config.seed, + log_dir=experiment_save_dir, + accelerator=accelerator, + _resolve=False, + ) + + # Training loop + trainer.train() + + if simulator_type == "IsaacSim": + os._exit(0) + + +if __name__ == "__main__": + + main() diff --git a/GR00T-WholeBodyControl/gear_sonic/trl/__init__.py b/GR00T-WholeBodyControl/gear_sonic/trl/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/__pycache__/__init__.cpython-310.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7eb4198147641f2d266195cfb058c4fd09fd7726 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/__pycache__/__init__.cpython-310.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/__pycache__/keyboard_subscriber.cpython-310.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/__pycache__/keyboard_subscriber.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b509d3e31332c19a3855b34e9a1574c6200efe29 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/__pycache__/keyboard_subscriber.cpython-310.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/__pycache__/telemetry.cpython-310.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/__pycache__/telemetry.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..db2ceba877effea2ec4c9bdbdca9274df71b84f1 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/__pycache__/telemetry.cpython-310.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/__pycache__/transforms.cpython-310.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/__pycache__/transforms.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6f8dbd1b47afbc2e3cc1c0acce40985b434cae49 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/__pycache__/transforms.cpython-310.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/__pycache__/zmq_state_subscriber.cpython-310.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/__pycache__/zmq_state_subscriber.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5f07dd53edfcfb79ff850946d906fea3efc2e67f Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/data_collection/__pycache__/zmq_state_subscriber.cpython-310.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/__pycache__/metric_utils.cpython-310.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/__pycache__/metric_utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0f9991d999b0e123507429926652a009c8e0c59a Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/__pycache__/metric_utils.cpython-310.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/__pycache__/sensor_server.cpython-310.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/__pycache__/sensor_server.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..21ef806213327dfaa724db8e86cab19c2eacf037 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/mujoco_sim/__pycache__/sensor_server.cpython-310.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/network/__pycache__/network_utils.cpython-310.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/network/__pycache__/network_utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..293fd099bc72a631fd39cb98d101ddcfd444e1ea Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/network/__pycache__/network_utils.cpython-310.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/teleop/solver/__pycache__/solver.cpython-310.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/teleop/solver/__pycache__/solver.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ffd3ed640a3ede4f388ae89e87bcec4745c1baa6 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/teleop/solver/__pycache__/solver.cpython-310.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/teleop/solver/hand/__pycache__/g1_gripper_ik_solver.cpython-310.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/teleop/solver/hand/__pycache__/g1_gripper_ik_solver.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..902bf055a75160cd364dbb40b010d30c319b7010 Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/teleop/solver/hand/__pycache__/g1_gripper_ik_solver.cpython-310.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/teleop/solver/hand/g1_gripper_ik_solver.py b/GR00T-WholeBodyControl/gear_sonic/utils/teleop/solver/hand/g1_gripper_ik_solver.py new file mode 100644 index 0000000000000000000000000000000000000000..3fe1dde3a5b8d62b25e24f2f57b3bb20518ad3ca --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/teleop/solver/hand/g1_gripper_ik_solver.py @@ -0,0 +1,176 @@ +"""IK solver that maps fingertip distances to G1 gripper joint targets. + +Computes a 7-DOF hand joint vector from thumb-to-finger distances, +selecting a gesture (index / middle / ring / pinky) based on which +finger has the largest grip value and interpolating toward a preset +closed pose. +""" + +import numpy as np + +from gear_sonic.utils.teleop.solver.solver import Solver +class G1GripperInverseKinematicsSolver(Solver): + def __init__(self, side) -> None: + self.side = "L" if side.lower() == "left" else "R" + + def register_robot(self, robot): + pass + + def __call__(self, finger_data): + # manus data + fingertips = finger_data["position"] + + # Extract X, Y, Z positions of fingertips from the transformation matrices + positions = np.array([finger[:3, 3] for finger in fingertips]) + + # Ensure the positions are 2D arrays (N, 3) + positions = np.reshape(positions, (-1, 3)) # Ensure 2D array with shape (N, 3) + + # Fingertip positions: each finger has 5 joints, tip is at base_index + 4 + # thumb=4, index=9, middle=14, ring=19, pinky=24 + thumb_pos = positions[4, :] + index_pos = positions[4 + 5, :] + middle_pos = positions[4 + 10, :] + ring_pos = positions[4 + 15, :] + pinky_pos = positions[4 + 20, :] + + # Calculate distances for continuous grip control + # When thumb at (1,0,0) and finger at (grip_value,0,0): dist = 1.0 - grip_value + index_dist = np.linalg.norm(thumb_pos - index_pos) + middle_dist = np.linalg.norm(thumb_pos - middle_pos) + ring_dist = np.linalg.norm(thumb_pos - ring_pos) + pinky_dist = np.linalg.norm(thumb_pos - pinky_pos) + + # Dead zone threshold - ignore very small grip values, snap to full at high values + dist_threshold = 0.05 + + # Convert distance to grip amount (0.0 = open, 1.0 = closed) + index_grip = np.clip(1.0 - index_dist, 0.0, 1.0) + middle_grip = np.clip(1.0 - middle_dist, 0.0, 1.0) + ring_grip = np.clip(1.0 - ring_dist, 0.0, 1.0) + pinky_grip = np.clip(1.0 - pinky_dist, 0.0, 1.0) + + # Apply dead zone: ignore small values, snap to full near 1.0 + def apply_dead_zone(grip, threshold): + if grip < threshold: + return 0.0 + return grip + + index_grip = apply_dead_zone(index_grip, dist_threshold) + middle_grip = apply_dead_zone(middle_grip, dist_threshold) + ring_grip = apply_dead_zone(ring_grip, dist_threshold) + pinky_grip = apply_dead_zone(pinky_grip, dist_threshold) + + # Choose the active gesture based on which finger has the highest grip value + # Each gesture type maps to a different close pose + q_open = np.zeros(7) + + # Find the finger with the highest grip value + grips = [index_grip, middle_grip, ring_grip, pinky_grip] + max_grip = max(grips) + + if max_grip == 0: + # No grip - fully open + q_desired = q_open + elif index_grip == max_grip: + # Index gesture (trigger only): interpolate to index close pose + q_closed = self._get_index_close_q_desired() + q_desired = q_open + index_grip * (q_closed - q_open) + elif middle_grip == max_grip: + # Middle gesture (both pressed): interpolate to middle close pose + q_closed = self._get_middle_close_q_desired() + q_desired = q_open + middle_grip * (q_closed - q_open) + elif ring_grip == max_grip: + # Ring gesture (grip only): interpolate to ring close pose + q_closed = self._get_ring_close_q_desired() + q_desired = q_open + ring_grip * (q_closed - q_open) + else: + # Pinky gesture: interpolate to pinky close pose + q_closed = self._get_pinky_close_q_desired() + q_desired = q_open + pinky_grip * (q_closed - q_open) + + return q_desired + + def _get_index_close_q_desired(self): + q_desired = np.zeros(7) + + amp0 = 0.5 + if self.side == "L": + q_desired[0] -= amp0 + else: + q_desired[0] += amp0 + + amp = 0.7 + + q_desired[1] += amp + q_desired[2] += amp + + ampA1 = 1.5 + ampB1 = 1.5 + ampA2 = 0.6 + ampB2 = 1.5 + + q_desired[3] -= ampA1 + q_desired[4] -= ampB1 + q_desired[5] -= ampA2 + q_desired[6] -= ampB2 + + # Right hand has mirrored joint convention, so negate all targets + return q_desired if self.side == "L" else -q_desired + + def _get_middle_close_q_desired(self): + q_desired = np.zeros(7) + + amp0 = 0.0 + if self.side == "L": + q_desired[0] -= amp0 + else: + q_desired[0] += amp0 + + amp = 0.7 + + q_desired[1] += amp + q_desired[2] += amp + + ampA1 = 1.0 + ampB1 = 1.5 + ampA2 = 1.0 + ampB2 = 1.5 + + q_desired[3] -= ampA1 + q_desired[4] -= ampB1 + q_desired[5] -= ampA2 + q_desired[6] -= ampB2 + + return q_desired if self.side == "L" else -q_desired + + def _get_ring_close_q_desired(self): + q_desired = np.zeros(7) + + amp0 = -0.5 + if self.side == "L": + q_desired[0] -= amp0 + else: + q_desired[0] += amp0 + + amp = 0.7 + + q_desired[1] += amp + q_desired[2] += amp + + ampA1 = 0.6 + ampB1 = 1.5 + ampA2 = 1.5 + ampB2 = 1.5 + + q_desired[3] -= ampA1 + q_desired[4] -= ampB1 + q_desired[5] -= ampA2 + q_desired[6] -= ampB2 + + return q_desired if self.side == "L" else -q_desired + + def _get_pinky_close_q_desired(self): + q_desired = np.zeros(7) + + return q_desired if self.side == "L" else -q_desired diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/teleop/solver/solver.py b/GR00T-WholeBodyControl/gear_sonic/utils/teleop/solver/solver.py new file mode 100644 index 0000000000000000000000000000000000000000..4ff90a81d03ec8a8eb8e6d447457fd6e95e1854e --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/teleop/solver/solver.py @@ -0,0 +1,19 @@ +"""Abstract base class for teleoperation solvers (e.g., IK for hands).""" + +from abc import ABC, abstractmethod +from typing import Any + + +class Solver(ABC): + def __init__(self): + pass + + def register_robot(self, robot): + pass + + def calibrate(self, data): + pass + + @abstractmethod + def __call__(self, target) -> Any: + pass diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/teleop/vis/vr3pt_pose_visualizer.py b/GR00T-WholeBodyControl/gear_sonic/utils/teleop/vis/vr3pt_pose_visualizer.py new file mode 100644 index 0000000000000000000000000000000000000000..205f03836c9e0bea9a822a05462a3bb25a83a78e --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/teleop/vis/vr3pt_pose_visualizer.py @@ -0,0 +1,2066 @@ +""" +VR 3-Point Pose Visualizer + +A standalone PyVista-based visualizer for VR 3-point pose data (Head, Left Wrist, Right Wrist). +Can be used by any process that provides pose data as numpy arrays. + +Coordinate convention: +- X: forward (RED axis) +- Y: left (GREEN axis) +- Z: up (BLUE axis) + +Quaternion format: [qw, qx, qy, qz] (scalar-first) + +Usage: + # Basic static visualization with reference frames + visualizer = VR3PtPoseVisualizer() + visualizer.show_static() + + # Visualize with pose data + vr_3pt_pose = np.array([...]) # Shape (3, 7): [x, y, z, qw, qx, qy, qz] for each point + visualizer.show_with_vr_pose(vr_3pt_pose) + + # Real-time visualization (requires update callback) + visualizer.create_realtime_plotter() + # In your loop: + visualizer.update_vr_poses(vr_3pt_pose) + visualizer.render() + + # Visualization with G1 robot model + visualizer = VR3PtPoseVisualizer(with_g1_robot=True) + visualizer.show_with_vr_pose(vr_3pt_pose) # G1 robot will be shown at origin +""" + +import os +import time +from collections import deque +from pathlib import Path +from typing import Dict, List, Optional, Tuple + +import numpy as np +from scipy.spatial.transform import Rotation as sRot + +try: + import pyvista as pv + + PYVISTA_AVAILABLE = True +except ImportError: + pv = None + PYVISTA_AVAILABLE = False + +try: + import vtk + + VTK_AVAILABLE = True +except ImportError: + vtk = None + VTK_AVAILABLE = False + +try: + import pinocchio as pin + + PINOCCHIO_AVAILABLE = True +except ImportError: + pin = None + PINOCCHIO_AVAILABLE = False + + +# ============================================================================= +# Shared FK constants and function (display-independent, only needs Pinocchio) +# ============================================================================= + +# Key frame names for FK pose extraction +G1_LEFT_WRIST_FRAME = "left_wrist_yaw_link" +G1_RIGHT_WRIST_FRAME = "right_wrist_yaw_link" +G1_TORSO_FRAME = "torso_link" + +# Key frame offsets applied in the local frame of each link +# (from gear_sonic/config/manager_env/commands/terms/force.yaml) +G1_KEY_FRAME_OFFSETS = { + "left_wrist": np.array([0.18, -0.025, 0.0]), + "right_wrist": np.array([0.18, 0.025, 0.0]), + "torso": np.array([0.0, 0.0, 0.35]), +} + +G1_FRAME_MAPPING = { + "left_wrist": G1_LEFT_WRIST_FRAME, + "right_wrist": G1_RIGHT_WRIST_FRAME, + "torso": G1_TORSO_FRAME, +} + + +def get_g1_key_frame_poses( + robot_model, + q: np.ndarray = None, + root_position: np.ndarray = None, + apply_offset: bool = True, +) -> Dict[str, Dict[str, np.ndarray]]: + """ + Get poses (position + orientation) of G1 key frames using Pinocchio FK. + + This is a **display-independent** function — it only needs a Pinocchio robot + model, no PyVista/VTK/display. It can be used by both the visualizer and + headless calibration code. + + Args: + robot_model: Pinocchio-based robot model with cache_forward_kinematics() + and frame_placement() methods. + q: Joint configuration. If None, uses robot_model.default_body_pose. + root_position: Position of robot root. Default is origin [0, 0, 0]. + apply_offset: Whether to apply the local frame offsets. Default True. + + Returns: + Dict with keys 'left_wrist', 'right_wrist', 'torso', each containing: + - 'position': np.ndarray [x, y, z] (with offset applied in local frame) + - 'orientation_xyzw': np.ndarray [qx, qy, qz, qw] (scipy/ROS convention) + - 'orientation_wxyz': np.ndarray [qw, qx, qy, qz] (scalar-first convention) + """ + if q is None: + q = robot_model.default_body_pose + if root_position is None: + root_position = np.array([0.0, 0.0, 0.0]) + + # Update forward kinematics + robot_model.cache_forward_kinematics(q, auto_clip=False) + + result = {} + for key, frame_name in G1_FRAME_MAPPING.items(): + # Get frame placement from Pinocchio — if the frame doesn't exist, + # this is a fatal configuration error (wrong URDF or frame name). + try: + placement = robot_model.frame_placement(frame_name) + except ValueError as e: + raise RuntimeError( + f"Cannot find frame '{frame_name}' (key='{key}') in robot model. " + f"Ensure the URDF contains this frame. Original error: {e}" + ) from e + + rotation_matrix = placement.rotation + + # Apply offset in local frame, then transform to world frame + if apply_offset and key in G1_KEY_FRAME_OFFSETS: + local_offset = G1_KEY_FRAME_OFFSETS[key] + world_offset = rotation_matrix @ local_offset + position = placement.translation + world_offset + root_position + else: + position = placement.translation + root_position + + # Convert rotation matrix to quaternion using scipy + rot = sRot.from_matrix(rotation_matrix) + quat_xyzw = rot.as_quat() # scipy returns [qx, qy, qz, qw] + quat_wxyz = np.array([quat_xyzw[3], quat_xyzw[0], quat_xyzw[1], quat_xyzw[2]]) + + result[key] = { + "position": position.copy(), + "orientation_xyzw": quat_xyzw.copy(), + "orientation_wxyz": quat_wxyz.copy(), + } + + return result + + +class G1RobotVisualizer: + """ + PyVista-based G1 robot visualizer that loads STL meshes and transforms them + based on joint configurations using Pinocchio forward kinematics. + + The robot is placed at the origin (pelvis at [0, 0, 0]). + """ + + # Default robot color (dark gray for main body) + ROBOT_COLOR = "#404040" + ROBOT_OPACITY = 0.1 # Semi-transparent to see key points better + + # Key frame names and offsets — reference the shared module-level constants + LEFT_WRIST_FRAME = G1_LEFT_WRIST_FRAME + RIGHT_WRIST_FRAME = G1_RIGHT_WRIST_FRAME + TORSO_FRAME = G1_TORSO_FRAME + KEY_FRAME_OFFSETS = G1_KEY_FRAME_OFFSETS + + # Key point visualization colors + KEY_POINT_COLORS = { + "left_wrist": "lightgreen", + "right_wrist": "lightblue", + "torso": "yellow", + } + + # Key point labels (with offset indicator) + KEY_POINT_LABELS = { + "left_wrist": "L-Wrist (with offset)", + "right_wrist": "R-Wrist (with offset)", + "torso": "Torso (with offset)", + } + + # Waist joint names in the G1 robot (order: yaw, roll, pitch) + WAIST_JOINT_NAMES = ["waist_yaw_joint", "waist_roll_joint", "waist_pitch_joint"] + + def __init__(self, robot_model=None): + """ + Initialize the G1 robot visualizer. + + Args: + robot_model: Optional pre-instantiated RobotModel. If None, will create one. + """ + if not PYVISTA_AVAILABLE: + raise ImportError("PyVista is required. Install with: pip install pyvista") + if not PINOCCHIO_AVAILABLE: + raise ImportError("Pinocchio is required. Install with: pip install pin") + + # Import here to avoid circular imports and make the dependency optional + from gear_sonic.data.robot_model.instantiation.g1 import instantiate_g1_robot_model + + self.robot_model = robot_model if robot_model is not None else instantiate_g1_robot_model() + + # Get paths for mesh files + gear_sonic_root = Path(__file__).resolve().parent.parent.parent.parent + self.mesh_dir = gear_sonic_root / "data" / "robot_model" / "model_data" / "g1" / "meshes" + + # Load mesh data from Pinocchio visual model + self._load_visual_geometries() + + # Store actors for real-time updates + self.mesh_actors: Dict[str, any] = {} + self.key_point_actors: Dict[str, any] = {} + self._initialized = False + self._key_points_initialized = False + + # Cache waist joint indices for efficient updates + self._waist_joint_indices: Optional[List[int]] = None + try: + self._waist_joint_indices = self.robot_model.get_joint_group_indices("waist") + except (ValueError, AttributeError) as e: + raise RuntimeError( + f"Could not get waist joint indices from robot model. " + f"Ensure the robot model supplemental info defines a 'waist' joint group. " + f"Original error: {e}" + ) from e + + def compute_waist_joints_from_orientation( + self, + neck_quat_wxyz: np.ndarray, + scale_factor: float = 1.0, + ) -> Optional[np.ndarray]: + """ + Compute waist joint angles from VR neck orientation. + + The neck orientation from VR represents the upper body tilt. We decompose it + into yaw, roll, pitch Euler angles and map them to the waist joints. + + Args: + neck_quat_wxyz: Quaternion [qw, qx, qy, qz] representing neck orientation + (relative to root, already calibrated) + scale_factor: Scale factor for joint angles (0.0-1.0), useful for limiting range + + Returns: + np.ndarray of shape (3,) containing [waist_yaw, waist_roll, waist_pitch] + or None if waist control is not available + """ + if self._waist_joint_indices is None: + return None + + # Convert quaternion to Euler angles (ZYX = yaw, pitch, roll in extrinsic) + # G1 waist joints order: yaw, roll, pitch + quat_xyzw = np.array( + [neck_quat_wxyz[1], neck_quat_wxyz[2], neck_quat_wxyz[3], neck_quat_wxyz[0]] + ) + rot = sRot.from_quat(quat_xyzw) + + # Use ZYX Euler convention: rotation about Z (yaw), then Y (pitch), then X (roll) + # Output order is [z_angle, y_angle, x_angle] = [yaw, pitch, roll] + euler_zyx = rot.as_euler("ZYX", degrees=False) + + # Map to waist joints: [waist_yaw, waist_roll, waist_pitch] + # euler_zyx = [yaw, pitch, roll] + waist_yaw = euler_zyx[0] * scale_factor + waist_roll = euler_zyx[2] * scale_factor # X rotation + waist_pitch = euler_zyx[1] * scale_factor # Y rotation + + return np.array([waist_yaw, waist_roll, waist_pitch]) + + def apply_waist_joints_to_config( + self, + q: np.ndarray, + waist_joints: np.ndarray, + ) -> np.ndarray: + """ + Apply waist joint angles to a robot configuration. + + Args: + q: Full robot joint configuration array + waist_joints: Array of shape (3,) containing [waist_yaw, waist_roll, waist_pitch] + + Returns: + Updated robot configuration with waist joints set + """ + if self._waist_joint_indices is None or waist_joints is None: + return q + + q_new = q.copy() + for i, idx in enumerate(self._waist_joint_indices): + if i < len(waist_joints): + q_new[idx] = waist_joints[i] + return q_new + + def _load_visual_geometries(self): + """Load visual geometry info from Pinocchio's visual model.""" + self.visual_geometries: List[Dict] = [] + + visual_model = self.robot_model.pinocchio_wrapper.visual_model + model = self.robot_model.pinocchio_wrapper.model + + if len(visual_model.geometryObjects) == 0: + raise RuntimeError( + "No visual geometries found in Pinocchio visual model. " + "Check that the URDF file contains visual geometry elements." + ) + + for geom_id, geom in enumerate(visual_model.geometryObjects): + # Get the mesh file path + mesh_path = geom.meshPath + if not mesh_path: + print( + f"Warning: Visual geometry {geom_id} ({geom.name}) has no mesh path, skipping." + ) + continue + + # Get the frame this geometry is attached to + frame_id = geom.parentFrame + frame_name = model.frames[frame_id].name if frame_id < len(model.frames) else None + + # Get the local placement (geometry relative to parent frame) + local_placement = geom.placement + + self.visual_geometries.append( + { + "geom_id": geom_id, + "mesh_path": str(mesh_path), + "frame_id": frame_id, + "frame_name": frame_name, + "local_placement": local_placement, + "mesh": None, # Will be loaded later + } + ) + + def _load_mesh(self, mesh_path: str) -> "pv.PolyData": + """Load a mesh file and return PyVista PolyData. + + Raises: + FileNotFoundError: If the mesh file cannot be found. + RuntimeError: If the mesh file exists but cannot be loaded. + """ + original_path = mesh_path + if not os.path.exists(mesh_path): + # Try relative to mesh_dir + mesh_name = os.path.basename(mesh_path) + mesh_path = str(self.mesh_dir / mesh_name) + + if not os.path.exists(mesh_path): + raise FileNotFoundError( + f"Robot mesh file not found: '{original_path}'\n" + f" Also tried: '{mesh_path}'\n" + f" Mesh directory: '{self.mesh_dir}'\n" + f" Please ensure robot mesh files (STL) are present." + ) + + try: + return pv.read(mesh_path) + except Exception as e: + raise RuntimeError( + f"Failed to load robot mesh file '{mesh_path}': {e}" + ) from e + + def _get_geometry_world_transform( + self, geom_info: Dict, q: np.ndarray = None + ) -> Tuple[np.ndarray, np.ndarray]: + """ + Get the world transform (position, rotation matrix) for a visual geometry. + + Args: + geom_info: Geometry info dict from visual_geometries + q: Joint configuration. If None, uses q_zero. + + Returns: + Tuple of (position [3], rotation_matrix [3x3]) + """ + if q is None: + q = self.robot_model.q_zero + + # Update forward kinematics + self.robot_model.cache_forward_kinematics(q, auto_clip=False) + + # Get the frame's world placement + frame_placement = self.robot_model.pinocchio_wrapper.data.oMf[geom_info["frame_id"]] + + # Compose with local placement + world_placement = frame_placement * geom_info["local_placement"] + + # Extract position and rotation + position = world_placement.translation + rotation = world_placement.rotation + + return position, rotation + + def add_to_plotter( + self, + plotter: "pv.Plotter", + q: np.ndarray = None, + color: str = None, + opacity: float = None, + root_position: np.ndarray = None, + ) -> Dict[str, any]: + """ + Add G1 robot meshes to a PyVista plotter. + + Args: + plotter: PyVista plotter to add meshes to + q: Joint configuration. If None, uses default body pose. + color: Override mesh color + opacity: Override mesh opacity + root_position: Position of the robot root (pelvis). Default is origin [0, 0, 0]. + + Returns: + Dict mapping geometry names to actors for later updates + """ + if q is None: + q = self.robot_model.default_body_pose + if color is None: + color = self.ROBOT_COLOR + if opacity is None: + opacity = self.ROBOT_OPACITY + if root_position is None: + root_position = np.array([0.0, 0.0, 0.0]) + + actors = {} + + for geom_info in self.visual_geometries: + # Load mesh if not already loaded + if geom_info["mesh"] is None: + geom_info["mesh"] = self._load_mesh(geom_info["mesh_path"]) + + # Get world transform + position, rotation = self._get_geometry_world_transform(geom_info, q) + + # Apply root offset + position = position + root_position + + # Create a copy of the mesh and transform it + mesh = geom_info["mesh"].copy() + + # Build 4x4 transformation matrix + transform = np.eye(4) + transform[:3, :3] = rotation + transform[:3, 3] = position + + mesh.transform(transform) + + # Add to plotter + actor = plotter.add_mesh( + mesh, + color=color, + opacity=opacity, + smooth_shading=True, + name=geom_info["frame_name"], + ) + actors[geom_info["frame_name"]] = { + "actor": actor, + "geom_info": geom_info, + } + + self.mesh_actors = actors + self._initialized = True + return actors + + def add_to_plotter_realtime( + self, + plotter: "pv.Plotter", + q: np.ndarray = None, + color: str = None, + opacity: float = None, + root_position: np.ndarray = None, + ) -> Dict[str, any]: + """ + Add G1 robot meshes to a PyVista plotter for real-time updates. + Uses VTK transforms for efficient updates without recreating meshes. + + Args: + plotter: PyVista plotter to add meshes to + q: Initial joint configuration. If None, uses default body pose. + color: Override mesh color + opacity: Override mesh opacity + root_position: Position of the robot root (pelvis). Default is origin [0, 0, 0]. + + Returns: + Dict mapping geometry names to actors for later updates + """ + if not VTK_AVAILABLE: + raise ImportError("VTK is required for real-time mode") + + if q is None: + q = self.robot_model.default_body_pose + if color is None: + color = self.ROBOT_COLOR + if opacity is None: + opacity = self.ROBOT_OPACITY + if root_position is None: + root_position = np.array([0.0, 0.0, 0.0]) + + actors = {} + + for geom_info in self.visual_geometries: + # Load mesh if not already loaded + if geom_info["mesh"] is None: + geom_info["mesh"] = self._load_mesh(geom_info["mesh_path"]) + + # Add mesh without initial transform (we'll set it via VTK transform) + mesh = geom_info["mesh"].copy() + actor = plotter.add_mesh( + mesh, + color=color, + opacity=opacity, + smooth_shading=True, + name=geom_info["frame_name"], + ) + + actors[geom_info["frame_name"]] = { + "actor": actor, + "geom_info": geom_info, + } + + self.mesh_actors = actors + self._root_position = root_position + self._initialized = True + + # Set initial pose + self.update_pose(q, root_position) + + return actors + + def add_key_points_to_plotter( + self, + plotter: "pv.Plotter", + q: np.ndarray = None, + root_position: np.ndarray = None, + axis_length: float = 0.08, + ball_radius: float = 0.02, + show_axes: bool = True, + ) -> Dict[str, Dict[str, np.ndarray]]: + """ + Add key point visualizations (left wrist, right wrist, torso) to the plotter. + + Args: + plotter: PyVista plotter + q: Joint configuration. If None, uses default body pose. + root_position: Position of robot root. Default is origin [0, 0, 0]. + axis_length: Length of coordinate frame axes + ball_radius: Radius of position marker balls + show_axes: Whether to show coordinate frame axes + + Returns: + Dict with poses of each key point (same format as get_g1_key_frame_poses) + """ + if q is None: + q = self.robot_model.default_body_pose + if root_position is None: + root_position = np.array([0.0, 0.0, 0.0]) + + # Get key frame poses + poses = get_g1_key_frame_poses(self.robot_model, q=q, root_position=root_position) + + # Axis colors (RGB for XYZ) + axis_colors = ["red", "green", "blue"] + axis_dirs = [ + np.array([1, 0, 0]), # X + np.array([0, 1, 0]), # Y + np.array([0, 0, 1]), # Z + ] + + for key, pose in poses.items(): + position = pose["position"] + quat_xyzw = pose["orientation_xyzw"] + + # Convert quaternion to rotation matrix + rot = sRot.from_quat(quat_xyzw) + rot_matrix = rot.as_matrix() + + if show_axes: + # Add coordinate frame arrows + for i, (color, local_dir) in enumerate(zip(axis_colors, axis_dirs)): + world_dir = rot_matrix @ local_dir + arrow = pv.Arrow( + start=position, + direction=world_dir, + scale=axis_length, + tip_length=0.3, + tip_radius=0.15, + shaft_radius=0.05, + ) + plotter.add_mesh(arrow, color=color, smooth_shading=True) + + # Add colored ball at position + ball = pv.Sphere(radius=ball_radius, center=position) + plotter.add_mesh( + ball, color=self.KEY_POINT_COLORS[key], smooth_shading=True, name=f"keypoint_{key}" + ) + + # Add label (with offset indicator) + plotter.add_point_labels( + [position + np.array([0, 0, ball_radius * 2])], + [self.KEY_POINT_LABELS[key]], + font_size=10, + point_color=self.KEY_POINT_COLORS[key], + text_color="white", + always_visible=True, + shape_opacity=0.7, + ) + + return poses + + def add_key_points_realtime( + self, + plotter: "pv.Plotter", + q: np.ndarray = None, + root_position: np.ndarray = None, + axis_length: float = 0.08, + ball_radius: float = 0.02, + ) -> Dict[str, Dict[str, np.ndarray]]: + """ + Add key point visualizations for real-time updates. + + Args: + plotter: PyVista plotter + q: Initial joint configuration. If None, uses default body pose. + root_position: Position of robot root. Default is origin [0, 0, 0]. + axis_length: Length of coordinate frame axes + ball_radius: Radius of position marker balls + + Returns: + Dict with initial poses of each key point + """ + if not VTK_AVAILABLE: + raise ImportError("VTK is required for real-time mode") + + if q is None: + q = self.robot_model.default_body_pose + if root_position is None: + root_position = np.array([0.0, 0.0, 0.0]) + + self.key_point_actors = {} + + # Axis colors (RGB for XYZ) + axis_colors = ["red", "green", "blue"] + + for key in ["left_wrist", "right_wrist", "torso"]: + actors = {"arrows": [], "ball": None} + + # Create arrows for each axis + for color in axis_colors: + arrow = pv.Arrow( + start=(0, 0, 0), + direction=(1, 0, 0), + scale=axis_length, + tip_length=0.3, + tip_radius=0.15, + shaft_radius=0.05, + ) + actor = plotter.add_mesh(arrow, color=color, smooth_shading=True) + actors["arrows"].append(actor) + + # Create ball + ball = pv.Sphere(radius=ball_radius, center=(0, 0, 0)) + actors["ball"] = plotter.add_mesh( + ball, color=self.KEY_POINT_COLORS[key], smooth_shading=True + ) + + self.key_point_actors[key] = actors + + self._key_points_initialized = True + self._key_point_axis_length = axis_length + + # Set initial poses + poses = self.update_key_points(q, root_position) + return poses + + def update_key_points( + self, q: np.ndarray, root_position: np.ndarray = None + ) -> Dict[str, Dict[str, np.ndarray]]: + """ + Update key point visualizations for real-time mode. + + Args: + q: Joint configuration + root_position: Optional new root position + + Returns: + Dict with updated poses of each key point + """ + if not self._key_points_initialized or not VTK_AVAILABLE: + return {} + + if root_position is None: + root_position = getattr(self, "_root_position", np.array([0.0, 0.0, 0.0])) + + # Get key frame poses + poses = get_g1_key_frame_poses(self.robot_model, q=q, root_position=root_position) + + axis_dirs = [ + np.array([1, 0, 0]), # X + np.array([0, 1, 0]), # Y + np.array([0, 0, 1]), # Z + ] + + for key, pose in poses.items(): + if key not in self.key_point_actors: + continue + + position = pose["position"] + quat_xyzw = pose["orientation_xyzw"] + + # Convert quaternion to rotation matrix + rot = sRot.from_quat(quat_xyzw) + rot_matrix = rot.as_matrix() + + actors = self.key_point_actors[key] + + # Update each arrow's transform + for j, local_dir in enumerate(axis_dirs): + world_dir = rot_matrix @ local_dir + + # Compute rotation to align arrow (which points along X) to world_dir + x_axis = np.array([1.0, 0.0, 0.0]) + v = np.cross(x_axis, world_dir) + c = np.dot(x_axis, world_dir) + + if np.linalg.norm(v) > 1e-6: + s = np.linalg.norm(v) + vx = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]]) + arrow_rot = np.eye(3) + vx + vx @ vx * ((1 - c) / (s * s + 1e-9)) + elif c < 0: + arrow_rot = np.diag([-1.0, 1.0, -1.0]) + else: + arrow_rot = np.eye(3) + + # Create VTK transform + transform = vtk.vtkTransform() + mat = vtk.vtkMatrix4x4() + + for ri in range(3): + for ci in range(3): + mat.SetElement(ri, ci, arrow_rot[ri, ci]) + mat.SetElement(0, 3, position[0]) + mat.SetElement(1, 3, position[1]) + mat.SetElement(2, 3, position[2]) + + transform.SetMatrix(mat) + actors["arrows"][j].SetUserTransform(transform) + + # Update ball position + ball_transform = vtk.vtkTransform() + ball_transform.Translate(position[0], position[1], position[2]) + actors["ball"].SetUserTransform(ball_transform) + + return poses + + def update_pose( + self, q: np.ndarray, root_position: np.ndarray = None + ) -> Optional[Dict[str, Dict[str, np.ndarray]]]: + """ + Update robot pose for real-time visualization. + + Args: + q: Joint configuration + root_position: Optional new root position + + Returns: + Dict with updated key frame poses if key points are initialized, else None + """ + if not self._initialized or not VTK_AVAILABLE: + return None + + if root_position is None: + root_position = getattr(self, "_root_position", np.array([0.0, 0.0, 0.0])) + else: + self._root_position = root_position + + # Update forward kinematics once + self.robot_model.cache_forward_kinematics(q, auto_clip=False) + + for name, actor_info in self.mesh_actors.items(): + geom_info = actor_info["geom_info"] + actor = actor_info["actor"] + + # Get the frame's world placement + frame_placement = self.robot_model.pinocchio_wrapper.data.oMf[geom_info["frame_id"]] + world_placement = frame_placement * geom_info["local_placement"] + + position = world_placement.translation + root_position + rotation = world_placement.rotation + + # Create VTK transform + transform = vtk.vtkTransform() + mat = vtk.vtkMatrix4x4() + + for i in range(3): + for j in range(3): + mat.SetElement(i, j, rotation[i, j]) + mat.SetElement(0, 3, position[0]) + mat.SetElement(1, 3, position[1]) + mat.SetElement(2, 3, position[2]) + + transform.SetMatrix(mat) + actor.SetUserTransform(transform) + + # Also update key points if they are initialized + if self._key_points_initialized: + return self.update_key_points(q, root_position) + return None + + +class VR3PtPoseVisualizer: + """ + PyVista-based visualizer for VR 3-point pose debugging. + + Coordinate convention: + - X: forward (RED axis) + - Y: left (GREEN axis) + - Z: up (BLUE axis) + + Quaternion format: [qw, qx, qy, qz] (scalar-first) + + Reference frame: + - World frame at origin (0, 0, 0) - WHITE ball + + Head kinematic chain visualization: + - Origin (root) → torso_link (+0.05m along Z) + - torso_link → head (+0.35m along head's local Z axis) + + G1 Robot visualization (optional): + - Loads G1 robot meshes using Pinocchio + - Robot root (pelvis) placed at origin + - Key points at left wrist, right wrist, and torso + """ + + # Ball color for world reference frame + WORLD_BALL_COLOR = "white" + + # VR pose ball colors - order: [0]=L-Wrist, [1]=R-Wrist, [2]=Head + VR_BALL_COLORS = ["lightgreen", "lightblue", "orange"] + VR_POSE_LABELS = ["L-Wrist", "R-Wrist", "Head"] + + # Axis colors (RGB for XYZ) + AXIS_COLORS = ["red", "green", "blue"] + + # Head kinematic chain constants (must match pico_manager_thread_server.py) + TORSO_LINK_OFFSET_Z = 0.05 # meters from root to torso_link + HEAD_LINK_LENGTH = 0.35 # meters from torso_link to head along head's local Z + TORSO_LINK_COLOR = "purple" + HEAD_LINK_COLOR = "orange" + + # SMPL body joint visualization constants + SMPL_NUM_JOINTS = 24 + SMPL_LOWER_BODY_INDICES = [0, 1, 2, 4, 5, 7, 8, 10, 11] + SMPL_JOINT_RADIUS = 0.02 + + # SMPL kinematic tree parent indices (standard 24-joint model) + # Joint i connects to SMPL_PARENT_INDICES[i]; root (joint 0) has parent -1 + SMPL_PARENT_INDICES = [ + -1, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, # 0-11: pelvis, hips, spine, knees, ankles, feet + 9, 9, 9, 12, 13, 14, 16, 17, 18, 19, 20, 21, # 12-23: neck, collars, head, shoulders..hands + ] + SMPL_BONE_COLOR = "white" + SMPL_BONE_WIDTH = 2.0 + SMPL_BONE_OPACITY = 0.5 + + # Timing report interval (seconds) + TIMING_REPORT_INTERVAL = 5.0 + + def __init__( + self, + axis_length: float = 0.1, + axis_radius: float = 0.005, + ball_radius: float = 0.02, + with_g1_robot: bool = False, + robot_model=None, + robot_opacity: float = 0.4, + enable_waist_tracking: bool = False, + enable_smpl_vis: bool = False, + smpl_root_position: Optional[np.ndarray] = None, + ): + """ + Initialize the VR 3-point pose visualizer. + + Args: + axis_length: Length of each axis arrow + axis_radius: Radius of the axis cylinders (for arrows) + ball_radius: Radius of the position marker balls + with_g1_robot: If True, load and display G1 robot at origin + robot_model: Optional pre-instantiated RobotModel for G1 visualization + robot_opacity: Opacity of G1 robot (0.0 = invisible, 1.0 = opaque, default 0.4) + enable_waist_tracking: If True, G1 robot waist follows VR head orientation + enable_smpl_vis: If True, show SMPL body joint spheres (24 joints) + smpl_root_position: Where to anchor the SMPL skeleton root (joint 0) on + the first frame. Default [-0.3, 0.0, 0.0] (0.3m behind + the G1 robot at origin, for side-by-side comparison). + """ + if not PYVISTA_AVAILABLE: + raise ImportError("PyVista is required. Install with: pip install pyvista") + + self.axis_length = axis_length + self.axis_radius = axis_radius + self.ball_radius = ball_radius + self.plotter = None + self.vr_actors = [] # For real-time mode: list of dicts with 'arrows' and 'ball' + self._initialized = False + + # SMPL body joint visualization + self.enable_smpl_vis = enable_smpl_vis + self.smpl_joint_actors: List = [] + + # SMPL root anchoring: on the first frame, capture the root (joint 0) position + # and offset all subsequent frames so the skeleton starts at smpl_root_position. + self._smpl_initial_root: Optional[np.ndarray] = None # captured on first frame + self.smpl_root_position: np.ndarray = ( + np.array(smpl_root_position, dtype=np.float64) + if smpl_root_position is not None + else np.array([-0.3, 0.0, 0.0]) + ) + + # Pre-allocated VTK caches (populated in create_realtime_plotter) + self._vr_arrow_transforms: List = [] # 3×3 vtkTransform for arrows + self._vr_arrow_matrices: List = [] # 3×3 vtkMatrix4x4 for arrows + self._vr_ball_transforms: List = [] # 3 vtkTransform for balls + self._smpl_transforms: List = [] # 24 vtkTransform for joints + self._smpl_bone_actor = None # Single actor for all 23 bones + self._smpl_bone_cells: Optional[np.ndarray] = None # Line connectivity + + # Pre-allocated numpy arrays for arrow rotation computation + self._x_axis = np.array([1.0, 0.0, 0.0]) + self._axis_dirs = [ + np.array([1, 0, 0], dtype=np.float64), + np.array([0, 1, 0], dtype=np.float64), + np.array([0, 0, 1], dtype=np.float64), + ] + self._diag_flip = np.diag([-1.0, 1.0, -1.0]) + + # Timing instrumentation (deques of per-frame durations in seconds) + self._vis_times_vr3pt: deque = deque(maxlen=200) + self._vis_times_smpl: deque = deque(maxlen=200) + self._vis_times_render: deque = deque(maxlen=200) + self._last_timing_report: float = 0.0 + + # G1 robot visualization + self.with_g1_robot = with_g1_robot + self.robot_opacity = robot_opacity + self.enable_waist_tracking = enable_waist_tracking + self.g1_visualizer: Optional[G1RobotVisualizer] = None + self._robot_q: Optional[np.ndarray] = None # Current robot joint configuration + self._last_key_frame_poses: Optional[Dict[str, Dict[str, np.ndarray]]] = None + + if with_g1_robot: + if not PINOCCHIO_AVAILABLE: + raise ImportError( + "Pinocchio is required for G1 robot visualization. " + "Install with: pip install pin" + ) + self.g1_visualizer = G1RobotVisualizer(robot_model=robot_model) + self._robot_q = self.g1_visualizer.robot_model.default_body_pose.copy() + + def _add_coordinate_frame( + self, + plotter, + position: np.ndarray, + quat_wxyz: np.ndarray, + ball_color: str, + axis_length: float = None, + ball_radius: float = None, + label: str = "", + ): + """ + Add a coordinate frame (3 RGB arrows + colored ball) to the plotter. + + Args: + plotter: PyVista plotter + position: [x, y, z] position + quat_wxyz: [qw, qx, qy, qz] quaternion (scalar-first) + ball_color: Color of the position marker ball + axis_length: Override axis length (uses default if None) + ball_radius: Override ball radius (uses default if None) + label: Optional label for the frame + """ + if axis_length is None: + axis_length = self.axis_length + if ball_radius is None: + ball_radius = self.ball_radius + + # Convert quaternion to rotation matrix + # quat_wxyz is [qw, qx, qy, qz], scipy uses [qx, qy, qz, qw] + quat_xyzw = np.array([quat_wxyz[1], quat_wxyz[2], quat_wxyz[3], quat_wxyz[0]]) + rot = sRot.from_quat(quat_xyzw) + rot_matrix = rot.as_matrix() + + # RGB colors for XYZ axes + axis_dirs = [ + np.array([1, 0, 0]), # X + np.array([0, 1, 0]), # Y + np.array([0, 0, 1]), # Z + ] + + # Add arrows for each axis + for i, (color, local_dir) in enumerate(zip(self.AXIS_COLORS, axis_dirs)): + # Rotate local direction by the frame's orientation + world_dir = rot_matrix @ local_dir + + # Create arrow from position in world_dir direction + arrow = pv.Arrow( + start=position, + direction=world_dir, + scale=axis_length, + tip_length=0.3, + tip_radius=0.15, + shaft_radius=0.05, + ) + plotter.add_mesh(arrow, color=color, smooth_shading=True) + + # Add colored ball at position + ball = pv.Sphere(radius=ball_radius, center=position) + plotter.add_mesh(ball, color=ball_color, smooth_shading=True) + + # Add label if provided + if label: + plotter.add_point_labels( + [position + np.array([0, 0, ball_radius * 2])], + [label], + font_size=12, + point_color=ball_color, + text_color="white", + always_visible=True, + shape_opacity=0.7, + ) + + def _add_head_kinematic_chain( + self, + plotter, + head_position: np.ndarray, + origin: np.ndarray = None, + line_width: float = 3.0, + torso_ball_radius: float = None, + ): + """ + Add visualization of head kinematic chain: origin → torso_link → head. + + Args: + plotter: PyVista plotter + head_position: [x, y, z] head position (already computed via kinematic chain) + origin: [x, y, z] origin/root position (default: [0, 0, 0]) + line_width: Width of the link lines + torso_ball_radius: Radius of torso_link marker ball (default: half of self.ball_radius) + """ + if origin is None: + origin = np.array([0.0, 0.0, 0.0]) + if torso_ball_radius is None: + torso_ball_radius = self.ball_radius * 0.5 + + # Torso link position (fixed offset from origin along Z) + torso_link_pos = origin + np.array([0.0, 0.0, self.TORSO_LINK_OFFSET_Z]) + + # Link 1: Origin → torso_link + line1 = pv.Line(origin, torso_link_pos) + plotter.add_mesh(line1, color=self.TORSO_LINK_COLOR, line_width=line_width) + + # Link 2: torso_link → head + line2 = pv.Line(torso_link_pos, head_position) + plotter.add_mesh(line2, color=self.HEAD_LINK_COLOR, line_width=line_width) + + # Small ball at torso_link + torso_ball = pv.Sphere(radius=torso_ball_radius, center=torso_link_pos) + plotter.add_mesh(torso_ball, color=self.TORSO_LINK_COLOR, smooth_shading=True) + + def _add_reference_frames(self, plotter): + """Add the world reference frame to the plotter.""" + + # World frame at origin (0, 0, 0) - identity rotation, WHITE ball + identity_quat = np.array([1.0, 0.0, 0.0, 0.0]) # [qw, qx, qy, qz] + self._add_coordinate_frame( + plotter, + position=np.array([0.0, 0.0, 0.0]), + quat_wxyz=identity_quat, + ball_color=self.WORLD_BALL_COLOR, + axis_length=self.axis_length * 1.5, # Larger for world frame + ball_radius=self.ball_radius * 1.5, + label="World (origin)", + ) + + def _add_ground_and_grid(self, plotter): + """Add ground plane and grid for spatial reference.""" + # Ground plane + ground = pv.Plane(center=(0, 0, -0.005), direction=(0, 0, 1), i_size=1.5, j_size=1.5) + plotter.add_mesh(ground, color="gray", opacity=0.3) + + # Grid lines + for i in range(-7, 8): + val = i * 0.1 + # X-direction lines + line_x = pv.Line((-0.7, val, 0.001), (0.7, val, 0.001)) + plotter.add_mesh(line_x, color="darkgray", line_width=1) + # Y-direction lines + line_y = pv.Line((val, -0.7, 0.001), (val, 0.7, 0.001)) + plotter.add_mesh(line_y, color="darkgray", line_width=1) + + def _add_legend( + self, + plotter, + include_vr_poses: bool = True, + live: bool = False, + include_g1: bool = False, + include_smpl: bool = False, + ): + """Add legend text to the plotter.""" + legend_text = ( + "VR 3-Point Pose Debugger\n" + "─────────────────────────\n" + "Axes: RED=X GREEN=Y BLUE=Z\n" + "─────────────────────────\n" + "WHITE: World origin" + ) + + if include_vr_poses: + live_str = " (live)" if live else "" + legend_text += ( + "\n─────────────────────────\n" + f"VR Pose{live_str}:\n" + " LIGHTGREEN: L-Wrist\n" + " LIGHTBLUE: R-Wrist\n" + " ORANGE: Head\n" + "─────────────────────────\n" + "Head Kinematic Chain:\n" + f" PURPLE: torso_link (+{self.TORSO_LINK_OFFSET_Z}m Z)\n" + f" ORANGE line: head link ({self.HEAD_LINK_LENGTH}m)" + ) + + if include_smpl: + legend_text += ( + "\n─────────────────────────\n" + "SMPL Body (24 joints):\n" + " BLUE gradient: Lower body\n" + " RED: Upper body\n" + " WHITE lines: Skeleton" + ) + + if include_g1: + legend_text += ( + "\n─────────────────────────\n" + "G1 Robot: Root at origin\n" + "Key Points (with offset):\n" + " LIGHTGREEN: L-Wrist\n" + " LIGHTBLUE: R-Wrist\n" + " YELLOW: Torso" + ) + + plotter.add_text(legend_text, position="upper_left", font_size=9, color="white") + + # ========================================================================= + # SMPL body joint visualization + # ========================================================================= + + @staticmethod + def _get_smpl_joint_color(joint_idx: int) -> List[float]: + """Get color for SMPL joint (blue gradient for lower body, red for upper body). + + Lower body joints are colored with a blue gradient from light to dark. + Upper body joints are red. + """ + lower_indices = VR3PtPoseVisualizer.SMPL_LOWER_BODY_INDICES + if joint_idx in lower_indices: + idx_in_list = lower_indices.index(joint_idx) + t = idx_in_list / max(1, len(lower_indices) - 1) + r = 0.6 * (1 - t) + 0.0 * t + g = 0.8 * (1 - t) + 0.1 * t + b = 1.0 * (1 - t) + 0.5 * t + return [r, g, b] + return [1.0, 0.0, 0.0] # Red for upper body + + def _create_smpl_joint_actors(self): + """Pre-create 24 SMPL joint spheres + bone PolyData for real-time updates. + + Optimizations applied: + - Low-resolution spheres (8×8 instead of default 30×30) — ~14x fewer triangles + - Pre-allocated vtkTransform per joint — no per-frame Python object creation + - Single PolyData with 23 line segments for all bones — one draw call + """ + self.smpl_joint_actors = [] + self._smpl_transforms = [] + + for joint_idx in range(self.SMPL_NUM_JOINTS): + # Low-res sphere: 8×8 is plenty for r=0.02 joints + sphere = pv.Sphere( + radius=self.SMPL_JOINT_RADIUS, + center=(0, 0, 0), + theta_resolution=8, + phi_resolution=8, + ) + color = self._get_smpl_joint_color(joint_idx) + actor = self.plotter.add_mesh( + sphere, color=color, smooth_shading=True, opacity=0.8 + ) + + # Pre-allocate VTK transform and bind to actor once + t = vtk.vtkTransform() + actor.SetUserTransform(t) + + self.smpl_joint_actors.append(actor) + self._smpl_transforms.append(t) + + # Build bone connectivity (23 line segments: each child → parent) + cells = [] + for child_idx in range(1, self.SMPL_NUM_JOINTS): + parent_idx = self.SMPL_PARENT_INDICES[child_idx] + cells.extend([2, parent_idx, child_idx]) + self._smpl_bone_cells = np.array(cells, dtype=np.int64) + + # Create initial bone PolyData and add as single actor + bone_points = np.zeros((self.SMPL_NUM_JOINTS, 3), dtype=np.float64) + bone_poly = pv.PolyData(bone_points) + bone_poly.lines = self._smpl_bone_cells + self._smpl_bone_actor = self.plotter.add_mesh( + bone_poly, + color=self.SMPL_BONE_COLOR, + line_width=self.SMPL_BONE_WIDTH, + opacity=self.SMPL_BONE_OPACITY, + ) + + def update_smpl_joints(self, joints_np: np.ndarray): + """Update SMPL joint sphere positions and bone lines for real-time visualization. + + On the first call, captures the root (joint 0) position and anchors + all subsequent frames so the skeleton is placed at ``smpl_root_position``. + + Uses pre-allocated vtkTransform objects (no per-frame allocation) and + updates the bone PolyData via mapper input swap. + + Args: + joints_np: Shape (24, 3) array of joint positions in local (root-relative) frame. + """ + if not self.enable_smpl_vis or len(self.smpl_joint_actors) == 0: + return + + t0 = time.perf_counter() + n = min(len(self._smpl_transforms), len(joints_np)) + + # First-frame anchoring: capture the initial root and compute the offset + if self._smpl_initial_root is None: + self._smpl_initial_root = joints_np[0].copy() + print( + f"[VR3PtVis] SMPL root anchored: initial_root=" + f"[{self._smpl_initial_root[0]:.4f}, {self._smpl_initial_root[1]:.4f}, " + f"{self._smpl_initial_root[2]:.4f}] → anchor=" + f"[{self.smpl_root_position[0]:.4f}, {self.smpl_root_position[1]:.4f}, " + f"{self.smpl_root_position[2]:.4f}]" + ) + + # Shift all joints: subtract initial root, add desired anchor position + offset = self.smpl_root_position - self._smpl_initial_root + joints_shifted = joints_np[:n] + offset + + # Update joint sphere transforms (reuse pre-allocated vtkTransform) + for i in range(n): + t = self._smpl_transforms[i] + t.Identity() + t.Translate( + float(joints_shifted[i, 0]), + float(joints_shifted[i, 1]), + float(joints_shifted[i, 2]), + ) + + # Update bone PolyData (single draw call for all 23 bones) + if self._smpl_bone_actor is not None and self._smpl_bone_cells is not None: + bone_poly = pv.PolyData(joints_shifted.astype(np.float64)) + bone_poly.lines = self._smpl_bone_cells + self._smpl_bone_actor.GetMapper().SetInputData(bone_poly) + + self._vis_times_smpl.append(time.perf_counter() - t0) + + def reset_smpl_anchor(self): + """Reset the SMPL root anchor so it is re-captured on the next frame. + + Call this after recalibration or when the operator's pose has changed + significantly and the skeleton should be re-anchored. + """ + self._smpl_initial_root = None + print("[VR3PtVis] SMPL root anchor reset — will re-capture on next frame") + + # ========================================================================= + # Timing instrumentation + # ========================================================================= + + def _maybe_report_timing(self): + """Periodically log average timing breakdown for vis_vr3pt vs vis_both.""" + now = time.time() + if now - self._last_timing_report < self.TIMING_REPORT_INTERVAL: + return + self._last_timing_report = now + + def _avg_ms(dq: deque) -> float: + return (sum(dq) / len(dq) * 1000.0) if dq else 0.0 + + vr3pt_ms = _avg_ms(self._vis_times_vr3pt) + smpl_ms = _avg_ms(self._vis_times_smpl) + render_ms = _avg_ms(self._vis_times_render) + vr3pt_only_ms = vr3pt_ms + render_ms + both_ms = vr3pt_ms + smpl_ms + render_ms + + parts = [f"vr3pt: {vr3pt_ms:.2f}ms"] + if self.enable_smpl_vis: + parts.append(f"smpl: {smpl_ms:.2f}ms") + parts.append(f"render: {render_ms:.2f}ms") + parts.append(f"vr3pt_only: {vr3pt_only_ms:.2f}ms") + if self.enable_smpl_vis: + parts.append(f"both(vr3pt+smpl): {both_ms:.2f}ms") + + print(f"[Vis Timing] {' | '.join(parts)}") + + def show_static(self, robot_q: np.ndarray = None): + """ + Show static visualization with reference frames only (blocking). + + Args: + robot_q: Optional joint configuration for G1 robot (if with_g1_robot=True). + If None, uses default body pose. + """ + pv.set_plot_theme("dark") + plotter = pv.Plotter(window_size=(1400, 900)) + plotter.set_background("black") + + # Add ground and grid + self._add_ground_and_grid(plotter) + + # Add reference frames + self._add_reference_frames(plotter) + + # Add G1 robot if enabled + if self.with_g1_robot and self.g1_visualizer is not None: + q = robot_q if robot_q is not None else self._robot_q + self.g1_visualizer.add_to_plotter( + plotter, q=q, root_position=np.array([0.0, 0.0, 0.0]), opacity=self.robot_opacity + ) + # Add key point markers and get their poses + self._last_key_frame_poses = self.g1_visualizer.add_key_points_to_plotter( + plotter, q=q, root_position=np.array([0.0, 0.0, 0.0]) + ) + + # Set camera — zoomed out for global view + plotter.camera_position = [(1.5, -1.2, 1.2), (0.0, 0.0, 0.2), (0, 0, 1)] + + # Add legend + self._add_legend(plotter, include_vr_poses=False, include_g1=self.with_g1_robot) + + print("\n[VR3PtPoseVisualizer] Reference frame displayed:") + print(" - WHITE ball at (0, 0, 0): World frame, identity rotation") + if self.with_g1_robot: + print(" - G1 Robot: Displayed at origin") + if self._last_key_frame_poses: + print("\n Key Frame Poses (position + orientation_xyzw):") + for key, pose in self._last_key_frame_poses.items(): + pos = pose["position"] + quat = pose["orientation_xyzw"] + print( + f" {key}: pos=[{pos[0]:.4f}, {pos[1]:.4f}, {pos[2]:.4f}], " + f"quat_xyzw=[{quat[0]:.4f}, {quat[1]:.4f}, {quat[2]:.4f}, {quat[3]:.4f}]" + ) + print("\nClose the window to exit.") + + plotter.show() + + def show_with_vr_pose(self, vr_3pt_pose: np.ndarray, robot_q: np.ndarray = None): + """ + Show visualization with both reference frames and actual VR pose data. + + Args: + vr_3pt_pose: Shape (3, 7) array where each row is [x, y, z, qw, qx, qy, qz] + Row 0: L-Wrist, Row 1: R-Wrist, Row 2: Head + robot_q: Optional joint configuration for G1 robot (if with_g1_robot=True). + If None, uses default body pose. + """ + pv.set_plot_theme("dark") + plotter = pv.Plotter(window_size=(1400, 900)) + plotter.set_background("black") + + # Add ground and grid + self._add_ground_and_grid(plotter) + + # Add reference frames + self._add_reference_frames(plotter) + + # Add G1 robot if enabled + if self.with_g1_robot and self.g1_visualizer is not None: + q = robot_q if robot_q is not None else self._robot_q + self.g1_visualizer.add_to_plotter( + plotter, q=q, root_position=np.array([0.0, 0.0, 0.0]), opacity=self.robot_opacity + ) + # Add key point markers and get their poses + self._last_key_frame_poses = self.g1_visualizer.add_key_points_to_plotter( + plotter, q=q, root_position=np.array([0.0, 0.0, 0.0]) + ) + + # Add VR pose frames + for i in range(vr_3pt_pose.shape[0]): + position = vr_3pt_pose[i, :3] + quat_wxyz = vr_3pt_pose[i, 3:7] + + self._add_coordinate_frame( + plotter, + position=position, + quat_wxyz=quat_wxyz, + ball_color=self.VR_BALL_COLORS[i], + label=f"VR {self.VR_POSE_LABELS[i]}", + ) + + # Add head kinematic chain visualization (origin → torso_link → head) + # Head is at index 2 in vr_3pt_pose + head_position = vr_3pt_pose[2, :3] + self._add_head_kinematic_chain(plotter, head_position) + + # Set camera — zoomed out for global view + plotter.camera_position = [(1.5, -1.2, 1.2), (0.0, 0.0, 0.2), (0, 0, 1)] + + # Add legend + self._add_legend(plotter, include_vr_poses=True, include_g1=self.with_g1_robot) + + plotter.show() + + def create_realtime_plotter( + self, + interactive: bool = True, + window_size: tuple = (1400, 900), + with_reference_frames: bool = True, + robot_q: np.ndarray = None, + ): + """ + Create a plotter for real-time visualization with pre-created actors. + + This method initializes a plotter that can be updated efficiently without + recreating actors each frame. + + Args: + interactive: If True, enables interactive updates (non-blocking) + window_size: Window size as (width, height) + with_reference_frames: Whether to add static reference frames + robot_q: Optional initial joint configuration for G1 robot. + + Returns: + The PyVista plotter object + """ + if not VTK_AVAILABLE: + raise ImportError("VTK is required for real-time mode. Install with: pip install vtk") + + pv.set_plot_theme("dark") + self.plotter = pv.Plotter(window_size=window_size) + self.plotter.set_background("black") + + # Add ground and grid (static) + self._add_ground_and_grid(self.plotter) + + # Add reference frames (static) + if with_reference_frames: + self._add_reference_frames(self.plotter) + + # Add G1 robot if enabled (for real-time updates) + if self.with_g1_robot and self.g1_visualizer is not None: + q = robot_q if robot_q is not None else self._robot_q + self.g1_visualizer.add_to_plotter_realtime( + self.plotter, + q=q, + root_position=np.array([0.0, 0.0, 0.0]), + opacity=self.robot_opacity, + ) + # Add key point markers for real-time updates + self._last_key_frame_poses = self.g1_visualizer.add_key_points_realtime( + self.plotter, q=q, root_position=np.array([0.0, 0.0, 0.0]) + ) + + # Set camera — zoomed out for global view of SMPL + G1/VR3pt + self.plotter.camera_position = [(1.5, -1.2, 1.2), (0.0, 0.0, 0.2), (0, 0, 1)] + + # Add legend + self._add_legend( + self.plotter, + include_vr_poses=True, + live=True, + include_g1=self.with_g1_robot, + include_smpl=self.enable_smpl_vis, + ) + + # Pre-create VR pose actors with pre-allocated VTK transforms + # (3 poses × (3 arrows + 1 ball) = 12 actors, 9 transforms + 9 matrices + 3 ball transforms) + self.vr_actors = [] + self._vr_arrow_transforms = [] + self._vr_arrow_matrices = [] + self._vr_ball_transforms = [] + + for i in range(3): + pose_actors = {"arrows": [], "ball": None} + arrow_transforms_i = [] + arrow_matrices_i = [] + + # Create low-overhead arrows for each axis + for j, color in enumerate(self.AXIS_COLORS): + arrow = pv.Arrow( + start=(0, 0, 0), + direction=(1, 0, 0), + scale=0.08, + tip_length=0.3, + tip_radius=0.15, + shaft_radius=0.05, + tip_resolution=6, + shaft_resolution=6, + ) + actor = self.plotter.add_mesh(arrow, color=color, smooth_shading=True) + + # Pre-allocate transform + matrix and bind once + t = vtk.vtkTransform() + m = vtk.vtkMatrix4x4() + actor.SetUserTransform(t) + + pose_actors["arrows"].append(actor) + arrow_transforms_i.append(t) + arrow_matrices_i.append(m) + + # Low-res ball + ball = pv.Sphere( + radius=0.015, center=(0, 0, 0), theta_resolution=8, phi_resolution=8 + ) + ball_actor = self.plotter.add_mesh( + ball, color=self.VR_BALL_COLORS[i], smooth_shading=True + ) + bt = vtk.vtkTransform() + ball_actor.SetUserTransform(bt) + pose_actors["ball"] = ball_actor + + self.vr_actors.append(pose_actors) + self._vr_arrow_transforms.append(arrow_transforms_i) + self._vr_arrow_matrices.append(arrow_matrices_i) + self._vr_ball_transforms.append(bt) + + # Pre-create SMPL body joint spheres + bone PolyData (if enabled) + if self.enable_smpl_vis: + self._create_smpl_joint_actors() + + # Pre-create head kinematic chain actors (origin → torso_link → head) + # These will be updated dynamically as head position changes + origin = np.array([0.0, 0.0, 0.0]) + torso_link_pos = np.array([0.0, 0.0, self.TORSO_LINK_OFFSET_Z]) + initial_head_pos = np.array([0.0, 0.0, self.TORSO_LINK_OFFSET_Z + self.HEAD_LINK_LENGTH]) + + # Link 1: origin → torso_link (static, doesn't change) + line1 = pv.Line(origin, torso_link_pos) + self.plotter.add_mesh(line1, color=self.TORSO_LINK_COLOR, line_width=3.0) + + # Torso link ball (static) + torso_ball = pv.Sphere(radius=self.ball_radius * 0.5, center=torso_link_pos) + self.plotter.add_mesh(torso_ball, color=self.TORSO_LINK_COLOR, smooth_shading=True) + + # Link 2: torso_link → head (dynamic, needs updating) + line2 = pv.Line(torso_link_pos, initial_head_pos) + self.head_link_actor = self.plotter.add_mesh( + line2, color=self.HEAD_LINK_COLOR, line_width=3.0 + ) + # Store torso_link position for updating head link + self._torso_link_pos = torso_link_pos + + self._initialized = True + + if interactive: + self.plotter.show(interactive_update=True) + + return self.plotter + + def update_vr_poses(self, vr_3pt_pose: np.ndarray): + """ + Update VR pose actors with new pose data (for real-time mode). + + Uses pre-allocated vtkTransform and vtkMatrix4x4 objects — no per-frame + Python object creation. Transforms are bound to actors once at init time; + mutating them in-place triggers VTK re-render via MTime. + + Args: + vr_3pt_pose: Shape (3, 7) array where each row is [x, y, z, qw, qx, qy, qz] + Row 0: L-Wrist, Row 1: R-Wrist, Row 2: Head + """ + if not self._initialized or len(self.vr_actors) != 3: + return + + for i in range(min(vr_3pt_pose.shape[0], 3)): + position = vr_3pt_pose[i, :3] + quat_wxyz = vr_3pt_pose[i, 3:7] + + # Convert quaternion to rotation matrix + quat_xyzw = np.array([quat_wxyz[1], quat_wxyz[2], quat_wxyz[3], quat_wxyz[0]]) + rot_matrix = sRot.from_quat(quat_xyzw).as_matrix() + + # Update each arrow's transform (reuse pre-allocated objects) + for j, local_dir in enumerate(self._axis_dirs): + world_dir = rot_matrix @ local_dir + + # Rodrigues' rotation: align X-axis arrow to world_dir + v = np.cross(self._x_axis, world_dir) + c = float(np.dot(self._x_axis, world_dir)) + v_norm = float(np.linalg.norm(v)) + + if v_norm > 1e-6: + s = v_norm + vx = np.array( + [[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]] + ) + arrow_rot = np.eye(3) + vx + vx @ vx * ((1 - c) / (s * s + 1e-9)) + elif c < 0: + arrow_rot = self._diag_flip + else: + arrow_rot = np.eye(3) + + # Reuse pre-allocated matrix and transform + mat = self._vr_arrow_matrices[i][j] + mat.Identity() + for ri in range(3): + for ci in range(3): + mat.SetElement(ri, ci, arrow_rot[ri, ci]) + mat.SetElement(0, 3, float(position[0])) + mat.SetElement(1, 3, float(position[1])) + mat.SetElement(2, 3, float(position[2])) + self._vr_arrow_transforms[i][j].SetMatrix(mat) + + # Reuse pre-allocated ball transform + bt = self._vr_ball_transforms[i] + bt.Identity() + bt.Translate(float(position[0]), float(position[1]), float(position[2])) + + # Update head kinematic chain link (torso_link → head) + if hasattr(self, "head_link_actor") and self.head_link_actor is not None: + head_position = vr_3pt_pose[2, :3] + new_line = pv.Line(self._torso_link_pos, head_position) + self.head_link_actor.GetMapper().SetInputData(new_line) + + def update_robot_pose(self, robot_q: np.ndarray) -> Optional[Dict[str, Dict[str, np.ndarray]]]: + """ + Update G1 robot pose for real-time visualization. + + Args: + robot_q: Joint configuration for the robot + + Returns: + Dict with updated key frame poses (left_wrist, right_wrist, torso), + each containing 'position' and 'orientation_xyzw'. Returns None if + G1 robot is not enabled. + """ + if self.with_g1_robot and self.g1_visualizer is not None: + self._robot_q = robot_q.copy() + self._last_key_frame_poses = self.g1_visualizer.update_pose(robot_q) + return self._last_key_frame_poses + return None + + def update_from_vr_pose( + self, + vr_3pt_pose: np.ndarray, + waist_scale: float = 1.0, + ) -> Optional[Dict[str, Dict[str, np.ndarray]]]: + """ + Update both VR pose visualization and optionally G1 robot waist from VR pose data. + + This method: + 1. Updates the VR pose markers (L-Wrist, R-Wrist, Neck) + 2. If enable_waist_tracking is True: computes waist joint angles from VR neck + orientation and updates the G1 robot visualization + + Timing: The entire method is tracked as "vr3pt" time for the delay comparison + between vis_vr3pt vs vis_both(vr3pt+smpl). + + Args: + vr_3pt_pose: Shape (3, 7) array where each row is [x, y, z, qw, qx, qy, qz] + Row 0: L-Wrist, Row 1: R-Wrist, Row 2: Neck + waist_scale: Scale factor for waist joint angles (0.0-1.0) + + Returns: + Dict with updated key frame poses, or None if G1 robot is not enabled + or waist tracking is disabled + """ + t0 = time.perf_counter() + + # Update VR pose markers + self.update_vr_poses(vr_3pt_pose) + + result = None + + # Update G1 robot waist from VR neck orientation (only if waist tracking enabled) + if self.enable_waist_tracking and self.with_g1_robot and self.g1_visualizer is not None: + # Extract neck orientation (row 2, columns 3-7 are qw,qx,qy,qz) + neck_quat_wxyz = vr_3pt_pose[2, 3:] + + # Compute waist joints from neck orientation + waist_joints = self.g1_visualizer.compute_waist_joints_from_orientation( + neck_quat_wxyz, scale_factor=waist_scale + ) + + if waist_joints is not None and self._robot_q is not None: + # Apply waist joints to current robot configuration + new_q = self.g1_visualizer.apply_waist_joints_to_config(self._robot_q, waist_joints) + # Update robot visualization + self._robot_q = new_q + self._last_key_frame_poses = self.g1_visualizer.update_pose(new_q) + result = self._last_key_frame_poses + + self._vis_times_vr3pt.append(time.perf_counter() - t0) + return result + + @property + def last_key_frame_poses(self) -> Optional[Dict[str, Dict[str, np.ndarray]]]: + """Get the last computed key frame poses.""" + return self._last_key_frame_poses + + @property + def robot_model(self): + """Get the robot model (if G1 visualization is enabled).""" + if self.g1_visualizer is not None: + return self.g1_visualizer.robot_model + return None + + def render(self): + """Render the current frame (for real-time mode). Tracks render time and reports timing.""" + if self.plotter is not None: + t0 = time.perf_counter() + self.plotter.update() + self._vis_times_render.append(time.perf_counter() - t0) + self._maybe_report_timing() + + def close(self): + """Close the plotter window.""" + if self.plotter is not None: + self.plotter.close() + self.plotter = None + self._initialized = False + + @property + def is_open(self) -> bool: + """Check if the plotter window is still open.""" + if self.plotter is None: + return False + try: + # Check if plotter is still active + return self.plotter.ren_win is not None and not self.plotter._closed + except (AttributeError, RuntimeError): + return False + + +def run_vr3pt_visualizer_test(): + """ + Standalone test for VR 3-point pose visualizer using PyVista. + Run this to verify the reference frames are displayed correctly. + """ + print("=" * 60) + print("VR 3-Point Pose Visualizer Test (PyVista)") + print("=" * 60) + print("\nWorld reference frame (RGB axes for XYZ):") + print(" WHITE ball at origin (0, 0, 0) - World frame") + print("\nClose the window to exit.") + print("=" * 60) + + visualizer = VR3PtPoseVisualizer(axis_length=0.08, ball_radius=0.015) + visualizer.show_static() + + +def run_vr3pt_demo_with_fake_data(): + """ + Demo visualization with fake VR pose data for testing without hardware. + """ + print("=" * 60) + print("VR 3-Point Pose Demo (Fake Data)") + print("=" * 60) + + # Create fake VR 3-point pose data + # Format: [x, y, z, qw, qx, qy, qz] for each of [L-Wrist, R-Wrist, Head] + fake_pose = np.array( + [ + [0.2, 0.3, 0.4, 1.0, 0.0, 0.0, 0.0], # L-Wrist at (0.2, 0.3, 0.4), identity + [0.2, -0.3, 0.4, 1.0, 0.0, 0.0, 0.0], # R-Wrist at (0.2, -0.3, 0.4), identity + [0.0, 0.0, 0.5, 1.0, 0.0, 0.0, 0.0], # Head at (0, 0, 0.5), identity + ], + dtype=np.float32, + ) + + print(f"\nFake pose data shape: {fake_pose.shape}") + print(f" L-Wrist: pos={fake_pose[0, :3]}, quat_wxyz={fake_pose[0, 3:]}") + print(f" R-Wrist: pos={fake_pose[1, :3]}, quat_wxyz={fake_pose[1, 3:]}") + print(f" Head: pos={fake_pose[2, :3]}, quat_wxyz={fake_pose[2, 3:]}") + print("\nClose the window to exit.") + print("=" * 60) + + visualizer = VR3PtPoseVisualizer(axis_length=0.08, ball_radius=0.015) + visualizer.show_with_vr_pose(fake_pose) + + +def run_realtime_demo_with_fake_data(duration: float = 10.0, update_hz: int = 30): + """ + Real-time demo with animated fake VR pose data. + + Args: + duration: How long to run the demo in seconds + update_hz: Update rate in Hz + """ + import time + + print("=" * 60) + print("VR 3-Point Pose Real-time Demo (Animated Fake Data)") + print("=" * 60) + print(f"Duration: {duration}s, Update rate: {update_hz} Hz") + print("Close the window to exit early.") + print("=" * 60) + + visualizer = VR3PtPoseVisualizer(axis_length=0.08, ball_radius=0.015) + visualizer.create_realtime_plotter() + + start_time = time.time() + + while time.time() - start_time < duration: + if not visualizer.is_open: + break + + t = time.time() - start_time + + # Animate the fake pose + fake_pose = np.array( + [ + # L-Wrist: circular motion + [ + 0.2 + 0.1 * np.sin(t * 2), + 0.3 + 0.1 * np.cos(t * 2), + 0.4, + 1.0, + 0.0, + 0.0, + 0.0, + ], + # R-Wrist: circular motion (opposite phase) + [ + 0.2 + 0.1 * np.sin(t * 2 + np.pi), + -0.3 + 0.1 * np.cos(t * 2 + np.pi), + 0.4, + 1.0, + 0.0, + 0.0, + 0.0, + ], + # Head: slight bobbing + [0.0, 0.0, 0.5 + 0.05 * np.sin(t * 3), 1.0, 0.0, 0.0, 0.0], + ], + dtype=np.float32, + ) + + visualizer.update_vr_poses(fake_pose) + visualizer.render() + + time.sleep(1.0 / update_hz) + + visualizer.close() + print("Demo finished.") + + +def run_g1_robot_demo(): + """ + Demo visualization showing G1 robot at origin with default pose. + """ + print("=" * 60) + print("G1 Robot Visualization Demo") + print("=" * 60) + print("Loading G1 robot model...") + + visualizer = VR3PtPoseVisualizer(axis_length=0.08, ball_radius=0.015, with_g1_robot=True) + + print(f"Robot DOFs: {visualizer.robot_model.num_dofs}") + + # Get and print key frame poses + key_poses = get_g1_key_frame_poses(visualizer.robot_model) + if key_poses: + print("\nKey Frame Poses (at default body pose):") + print("-" * 50) + for key, pose in key_poses.items(): + pos = pose["position"] + quat_xyzw = pose["orientation_xyzw"] + print(f" {key}:") + print(f" position: [{pos[0]:.4f}, {pos[1]:.4f}, {pos[2]:.4f}]") + print( + f" orientation_xyzw: [{quat_xyzw[0]:.4f}, {quat_xyzw[1]:.4f}, " + f"{quat_xyzw[2]:.4f}, {quat_xyzw[3]:.4f}]" + ) + print("-" * 50) + + print("\nG1 robot displayed at origin with key point markers.") + print("Close the window to exit.") + print("=" * 60) + + visualizer.show_static() + + +def run_g1_robot_with_vr_demo(): + """ + Demo visualization showing G1 robot with fake VR pose data. + """ + print("=" * 60) + print("G1 Robot + VR Pose Demo") + print("=" * 60) + print("Loading G1 robot model...") + + visualizer = VR3PtPoseVisualizer(axis_length=0.08, ball_radius=0.015, with_g1_robot=True) + + # Create fake VR 3-point pose data + fake_pose = np.array( + [ + [0.2, 0.3, 0.4, 1.0, 0.0, 0.0, 0.0], # L-Wrist + [0.2, -0.3, 0.4, 1.0, 0.0, 0.0, 0.0], # R-Wrist + [0.0, 0.0, 0.5, 1.0, 0.0, 0.0, 0.0], # Head + ], + dtype=np.float32, + ) + + print(f"Robot DOFs: {visualizer.robot_model.num_dofs}") + print("G1 robot displayed at origin with VR pose overlay.") + print("Close the window to exit.") + print("=" * 60) + + visualizer.show_with_vr_pose(fake_pose) + + +def run_g1_realtime_demo(duration: float = 10.0, update_hz: int = 30): + """ + Real-time demo with G1 robot and animated joint movements. + + Args: + duration: How long to run the demo in seconds + update_hz: Update rate in Hz + """ + import time + + print("=" * 60) + print("G1 Robot Real-time Demo (Animated Joints)") + print("=" * 60) + print("Loading G1 robot model...") + + visualizer = VR3PtPoseVisualizer(axis_length=0.08, ball_radius=0.015, with_g1_robot=True) + + print(f"Robot DOFs: {visualizer.robot_model.num_dofs}") + print(f"Duration: {duration}s, Update rate: {update_hz} Hz") + print("Key frame poses will be printed every second.") + print("Close the window to exit early.") + print("=" * 60) + + visualizer.create_realtime_plotter() + + start_time = time.time() + last_print_time = 0 + base_q = visualizer.robot_model.default_body_pose.copy() + + # Get joint indices for arm joints + try: + left_arm_indices = visualizer.robot_model.get_joint_group_indices("left_arm") + right_arm_indices = visualizer.robot_model.get_joint_group_indices("right_arm") + except (ValueError, AttributeError) as e: + raise RuntimeError( + f"Could not get arm joint indices from robot model for the demo. " + f"Ensure the robot model defines 'left_arm' and 'right_arm' joint groups. " + f"Original error: {e}" + ) from e + + while time.time() - start_time < duration: + if not visualizer.is_open: + break + + t = time.time() - start_time + + # Animate joint positions + q = base_q.copy() + + # Animate arm joints if available + if left_arm_indices: + for i, idx in enumerate(left_arm_indices[:3]): # First 3 joints + q[idx] = base_q[idx] + 0.3 * np.sin(t * 2 + i * 0.5) + + if right_arm_indices: + for i, idx in enumerate(right_arm_indices[:3]): # First 3 joints + q[idx] = base_q[idx] + 0.3 * np.sin(t * 2 + i * 0.5 + np.pi) + + # Update robot pose and get key frame poses + key_poses = visualizer.update_robot_pose(q) + + # Print key frame poses every second + if t - last_print_time >= 1.0 and key_poses: + last_print_time = t + print(f"\n[t={t:.1f}s] Key Frame Poses:") + for key, pose in key_poses.items(): + pos = pose["position"] + quat = pose["orientation_xyzw"] + print( + f" {key}: pos=[{pos[0]:.3f}, {pos[1]:.3f}, {pos[2]:.3f}], " + f"quat_xyzw=[{quat[0]:.3f}, {quat[1]:.3f}, {quat[2]:.3f}, {quat[3]:.3f}]" + ) + + # Also animate VR poses + fake_pose = np.array( + [ + [0.2 + 0.1 * np.sin(t * 2), 0.3 + 0.1 * np.cos(t * 2), 0.4, 1.0, 0.0, 0.0, 0.0], + [ + 0.2 + 0.1 * np.sin(t * 2 + np.pi), + -0.3 + 0.1 * np.cos(t * 2 + np.pi), + 0.4, + 1.0, + 0.0, + 0.0, + 0.0, + ], + [0.0, 0.0, 0.5 + 0.05 * np.sin(t * 3), 1.0, 0.0, 0.0, 0.0], + ], + dtype=np.float32, + ) + visualizer.update_vr_poses(fake_pose) + + visualizer.render() + time.sleep(1.0 / update_hz) + + visualizer.close() + print("\nDemo finished.") + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description="VR 3-Point Pose Visualizer with G1 Robot") + parser.add_argument( + "--mode", + choices=["static", "demo", "realtime", "g1", "g1_vr", "g1_realtime"], + default="g1_vr", + help=( + "Visualization mode: " + "static (reference frames only), " + "demo (fake VR pose), " + "realtime (animated VR demo), " + "g1 (G1 robot at origin), " + "g1_vr (G1 + VR pose), " + "g1_realtime (G1 with animated joints)" + ), + ) + parser.add_argument( + "--duration", + type=float, + default=10.0, + help="Duration for realtime demo in seconds", + ) + parser.add_argument("--hz", type=int, default=30, help="Update rate for realtime demo") + + args = parser.parse_args() + + if args.mode == "static": + run_vr3pt_visualizer_test() + elif args.mode == "demo": + run_vr3pt_demo_with_fake_data() + elif args.mode == "realtime": + run_realtime_demo_with_fake_data(duration=args.duration, update_hz=args.hz) + elif args.mode == "g1": + run_g1_robot_demo() + elif args.mode == "g1_vr": + run_g1_robot_with_vr_demo() + elif args.mode == "g1_realtime": + run_g1_realtime_demo(duration=args.duration, update_hz=args.hz) diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/teleop/zmq/__pycache__/zmq_planner_sender.cpython-310.pyc b/GR00T-WholeBodyControl/gear_sonic/utils/teleop/zmq/__pycache__/zmq_planner_sender.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9743c0988fd4b014395d708f8fbfec36d33b28ec Binary files /dev/null and b/GR00T-WholeBodyControl/gear_sonic/utils/teleop/zmq/__pycache__/zmq_planner_sender.cpython-310.pyc differ diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/teleop/zmq/zmq_planner_sender.py b/GR00T-WholeBodyControl/gear_sonic/utils/teleop/zmq/zmq_planner_sender.py new file mode 100644 index 0000000000000000000000000000000000000000..caa33609c86a9d4066c25956994bd1e9c67ba1ea --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/teleop/zmq/zmq_planner_sender.py @@ -0,0 +1,224 @@ +"""Builders for ZMQ wire-format messages on the 'command', 'planner', and 'pose' topics. + +Message layout: [topic_bytes][1024-byte JSON header][packed binary payload]. +The header describes field names, dtypes, and shapes so the receiver can +deserialize without out-of-band schema knowledge. +""" + +import json +import struct +from typing import Sequence + +import numpy as np + +HEADER_SIZE = 1280 + + +def _build_header(fields: list, version: int = 1, count: int = 1) -> bytes: + header = { + "v": version, + "endian": "le", + "count": count, + "fields": fields, + } + header_json = json.dumps(header, separators=(",", ":")).encode("utf-8") + if len(header_json) > HEADER_SIZE: + raise ValueError(f"Header too large: {len(header_json)} > {HEADER_SIZE}") + return header_json.ljust(HEADER_SIZE, b"\x00") + + +def build_command_message( + start: bool, stop: bool, planner: bool, delta_heading: float | None = None +) -> bytes: + """ + Assemble a 'command' topic message: + - start: u8 (1=start control) + - stop: u8 (1=stop control) + - planner: u8 (1=planner mode, 0=streamed motion) + - delta_heading: f32 (optional, yaw relative to heading command in radians) + Returns: bytes ready to send via socket.send() + """ + fields = [ + {"name": "start", "dtype": "u8", "shape": [1]}, + {"name": "stop", "dtype": "u8", "shape": [1]}, + {"name": "planner", "dtype": "u8", "shape": [1]}, + ] + payload = b"".join( + ( + struct.pack("B", 1 if start else 0), + struct.pack("B", 1 if stop else 0), + struct.pack("B", 1 if planner else 0), + ) + ) + + if delta_heading is not None: + # Append delta_heading field to header and payload + fields.append({"name": "delta_heading", "dtype": "f32", "shape": [1]}) + payload += struct.pack(" bytes: + """ + Assemble a 'planner' topic message: + - mode: i32 (LocomotionMode enum) + - movement: f32[3] (x,y,z) + - facing: f32[3] (x,y,z) + - speed: f32 (optional, -1 for default) + - height: f32 (optional, -1 for default) + Returns: bytes ready to send via socket.send() + """ + if len(movement) != 3: + raise ValueError("movement must have length 3") + if len(facing) != 3: + raise ValueError("facing must have length 3") + + fields = [ + {"name": "mode", "dtype": "i32", "shape": [1]}, + {"name": "movement", "dtype": "f32", "shape": [3]}, + {"name": "facing", "dtype": "f32", "shape": [3]}, + {"name": "speed", "dtype": "f32", "shape": [1]}, + {"name": "height", "dtype": "f32", "shape": [1]}, + ] + + payload = b"".join( + ( + struct.pack(" bytes: + """ + Pack pose/action data into ZMQ message format: + [topic_prefix][1024-byte JSON header][concatenated binary fields] + + This is a general-purpose function for packing numpy arrays into ZMQ messages. + Supports protocol versions 3 and 4. + + Args: + pose_data: Dictionary containing numpy arrays to send + topic: Topic prefix string (default: "pose") + version: Protocol version (default: 3). Version 4 includes "count" field. + + Returns: + Packed message as bytes + + Example: + >>> data = { + ... "token_state": np.array([1.0, 2.0], dtype=np.float32), + ... "frame_index": np.array([0], dtype=np.int64) + ... } + >>> msg = pack_pose_message(data, topic="pose", version=4) + """ + # Build fields list from pose_data + fields = [] + binary_data = [] + + for key, value in pose_data.items(): + if isinstance(value, np.ndarray): + # Determine dtype string + if value.dtype == np.float32: + dtype_str = "f32" + elif value.dtype == np.float64: + dtype_str = "f64" + elif value.dtype == np.int32: + dtype_str = "i32" + elif value.dtype == np.int64: + dtype_str = "i64" + elif value.dtype == bool: + dtype_str = "bool" + else: + # Default to f32, cast if needed + dtype_str = "f32" + value = value.astype(np.float32) + + fields.append({"name": key, "dtype": dtype_str, "shape": list(value.shape)}) + + # Ensure contiguous and little-endian + if not value.flags["C_CONTIGUOUS"]: + value = np.ascontiguousarray(value) + if value.dtype.byteorder == ">": + value = value.astype(value.dtype.newbyteorder("<")) + + binary_data.append(value.tobytes()) + + # Build header using common utility + header_bytes = _build_header(fields, version=version, count=1) + + # Pack message: [topic][1024-byte header][binary data] + topic_bytes = topic.encode("utf-8") + data_bytes = b"".join(binary_data) + + packed_message = topic_bytes + header_bytes + data_bytes + return packed_message diff --git a/GR00T-WholeBodyControl/gear_sonic/utils/teleop/zmq/zmq_poller.py b/GR00T-WholeBodyControl/gear_sonic/utils/teleop/zmq/zmq_poller.py new file mode 100644 index 0000000000000000000000000000000000000000..692fe6c3c3a5c3b042093bd806d4dd0b52728b4f --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/utils/teleop/zmq/zmq_poller.py @@ -0,0 +1,36 @@ +"""Non-blocking ZMQ subscriber that keeps only the latest message (CONFLATE mode).""" + +import zmq + + +class ZMQPoller: + """Simple ZMQ subscriber for sporadic non-blocking reads.""" + + def __init__(self, host: str = "localhost", port: int = 5555, topic: str = ""): + self._context = zmq.Context() + self._socket = self._context.socket(zmq.SUB) + self._socket.setsockopt_string(zmq.SUBSCRIBE, topic) + self._socket.setsockopt(zmq.CONFLATE, 1) + self._socket.connect(f"tcp://{host}:{port}") + self._topic = topic + + def __del__(self): + self.close() + + def get_data(self): + """Get latest data or None if no data available.""" + if self._socket.poll(timeout=0): + data = self._socket.recv(zmq.NOBLOCK) + if data is None: + print("ZMQPoller: no data received") + return None + + # Strip topic prefix + return data[len(self._topic) :] + + print("ZMQPoller: no data available") + return None + + def close(self): + self._socket.close() + self._context.term() diff --git a/GR00T-WholeBodyControl/gear_sonic/version.py b/GR00T-WholeBodyControl/gear_sonic/version.py new file mode 100644 index 0000000000000000000000000000000000000000..5b5d3295fa1b94852b2efde912bb782733830b02 --- /dev/null +++ b/GR00T-WholeBodyControl/gear_sonic/version.py @@ -0,0 +1,9 @@ +"""GEAR-SONIC package version definition.""" + +_MAJOR = "0" +_MINOR = "1" +_PATCH = "0" +_SUFFIX = "" + +VERSION_SHORT = "{0}.{1}".format(_MAJOR, _MINOR) +VERSION = "{0}.{1}.{2}".format(_MAJOR, _MINOR, _PATCH) diff --git a/GR00T-WholeBodyControl/legal/Apache License - GEAR-SONIC.txt b/GR00T-WholeBodyControl/legal/Apache License - GEAR-SONIC.txt new file mode 100644 index 0000000000000000000000000000000000000000..df6a4545492b07aae64dda378640fdbbe32e5e6b --- /dev/null +++ b/GR00T-WholeBodyControl/legal/Apache License - GEAR-SONIC.txt @@ -0,0 +1,93 @@ +https://www.apache.org/licenses/LICENSE-2.0.txt +Apache License +                           Version 2.0, January 2004 +                        http://www.apache.org/licenses/ + +   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + +   1. Definitions. +      "License" shall mean the terms and conditions for use, reproduction, +      and distribution as defined by Sections 1 through 9 of this document. +      "Licensor" shall mean the copyright owner or entity authorized by +      the copyright owner that is granting the License. +      "Legal Entity" shall mean the union of the acting entity and all +      other entities that control, are controlled by, or are under common +      control with that entity. +      "You" (or "Your") shall mean an individual or Legal Entity +      exercising permissions granted by this License. +      "Source" form shall mean the preferred form for making modifications, +      including but not limited to software source code, documentation +      source, and configuration files. +      "Object" form shall mean any form resulting from mechanical +      transformation or translation of a Source form, including but +      not limited to compiled object code, generated documentation, +      and conversions to other media types. +      "Work" shall mean the work of authorship, whether in Source or +      Object form, made available under the License, as indicated by a +      copyright notice that is included in or attached to the work. +      "Derivative Works" shall mean any work, whether in Source or Object +      form, that is based on (or derived from) the Work and for which the +      editorial revisions, annotations, elaborations, or other modifications +      represent, as a whole, an original work of authorship. +      "Contribution" shall mean any work of authorship, including +      the original version of the Work and any modifications or additions +      to that Work or Derivative Works thereof, that is intentionally +      submitted to Licensor for inclusion in the Work by the copyright owner +      or by an individual or Legal Entity authorized to submit on behalf of +      the copyright owner. +      "Contributor" shall mean Licensor and any individual or Legal Entity +      on behalf of whom a Contribution has been received by Licensor and +      subsequently incorporated within the Work. + +   2. Grant of Copyright License. Subject to the terms and conditions of +      this License, each Contributor hereby grants to You a perpetual, +      worldwide, non-exclusive, no-charge, royalty-free, irrevocable +      copyright license to reproduce, prepare Derivative Works of, +      publicly display, publicly perform, sublicense, and distribute the +      Work and such Derivative Works in Source or Object form. + +   3. Grant of Patent License. Subject to the terms and conditions of +      this License, each Contributor hereby grants to You a perpetual, +      worldwide, non-exclusive, no-charge, royalty-free, irrevocable +      (except as stated in this section) patent license to make, have made, +      use, offer to sell, sell, import, and otherwise transfer the Work, +      where such license applies only to those patent claims licensable +      by such Contributor that are necessarily infringed by their +      Contribution(s) alone or by combination of their Contribution(s) +      with the Work to which such Contribution(s) was submitted. + +   4. Redistribution. You may reproduce and distribute copies of the +      Work or Derivative Works thereof in any medium, with or without +      modifications, and in Source or Object form, provided that You +      meet the following conditions: +      (a) You must give any other recipients of the Work or +          Derivative Works a copy of this License; and +      (b) You must cause any modified files to carry prominent notices +          stating that You changed the files; and +      (c) You must retain, in the Source form of any Derivative Works +          that You distribute, all copyright, patent, trademark, and +          attribution notices from the Source form of the Work. +      (d) If the Work includes a "NOTICE" text file, you must include a +          readable copy of the attribution notices. + +   5. Submission of Contributions. Unless You explicitly state otherwise, +      any Contribution intentionally submitted for inclusion in the Work +      by You to the Licensor shall be under the terms and conditions of +      this License, without any additional terms or conditions. + +   6. Trademarks. This License does not grant permission to use the trade +      names, trademarks, service marks, or product names of the Licensor. + +   7. Disclaimer of Warranty. Unless required by applicable law or +      agreed to in writing, Licensor provides the Work (and each +      Contributor provides its Contributions) on an "AS IS" BASIS, +      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND. + +   8. Limitation of Liability. In no event and under no legal theory, +      whether in tort (including negligence), contract, or otherwise, +      shall any Contributor be liable to You for damages. + +   9. Accepting Warranty or Additional Liability. While redistributing +      the Work or Derivative Works thereof, You may choose to offer, +      and charge a fee for, acceptance of support, warranty, indemnity, +      or other liability obligations. \ No newline at end of file diff --git a/GR00T-WholeBodyControl/legal/DCO - GEAR-SONIC.txt b/GR00T-WholeBodyControl/legal/DCO - GEAR-SONIC.txt new file mode 100644 index 0000000000000000000000000000000000000000..9e03879207edb2b382285240a6de45eb3ea620df --- /dev/null +++ b/GR00T-WholeBodyControl/legal/DCO - GEAR-SONIC.txt @@ -0,0 +1,11 @@ +# Contributing to this Project + +By contributing to this repository, you agree to the Developer Certificate of Origin (DCO) 1.1: + +"I certify that: +(a) The contribution was created in whole or in part by me and I have the right to submit it under the open source license indicated in the file; or +(b) The contribution is based upon previous work that, to the best of my knowledge, is covered under an appropriate open source license and I have the right under that license to submit that work with modifications... +(c) The contribution was provided directly to me by some other person who certified (a), (b) or (c) and I have not modified it. +(d) I understand and agree that this project and the contribution are public and that a record of the contribution is maintained indefinitely." + +Full DCO text: https://developercertificate.org/ \ No newline at end of file diff --git a/GR00T-WholeBodyControl/legal/THIRD-PARTY SOFTWARE NOTICES - GEAR-SONIC.txt b/GR00T-WholeBodyControl/legal/THIRD-PARTY SOFTWARE NOTICES - GEAR-SONIC.txt new file mode 100644 index 0000000000000000000000000000000000000000..3d4d065c72c93dd62ae6d136a5ca9c84b4249ec3 --- /dev/null +++ b/GR00T-WholeBodyControl/legal/THIRD-PARTY SOFTWARE NOTICES - GEAR-SONIC.txt @@ -0,0 +1,20 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +#     http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# This file contains modifications of/dependencies on: +# - unitree_sdk2_python (BSD 3-Clause) +# - XRoboToolkit-PC-Service-Pybind (MIT) +# - robosuite (Apache 2.0) +# - tmux-plugins/tpm (MIT) +# - BeyondMimic (MIT) \ No newline at end of file diff --git a/GR00T-WholeBodyControl/legal/THIRD-PARTY SOFTWARE NOTICES AND ASSET LICENSES - GEAR-SONIC.txt b/GR00T-WholeBodyControl/legal/THIRD-PARTY SOFTWARE NOTICES AND ASSET LICENSES - GEAR-SONIC.txt new file mode 100644 index 0000000000000000000000000000000000000000..347c97a331b41762815488e76d88252cef21d702 --- /dev/null +++ b/GR00T-WholeBodyControl/legal/THIRD-PARTY SOFTWARE NOTICES AND ASSET LICENSES - GEAR-SONIC.txt @@ -0,0 +1,51 @@ +THIRD-PARTY SOFTWARE NOTICES AND ASSET LICENSES +----------------------------------------------- + +This project incorporates the following third-party components: + +1. unitree_sdk2_python: Copyright (c) Unitree Robotics. (BSD 3-Clause) +2. XRoboToolkit-PC-Service-Pybind: Copyright (c) XR-Robotics. (MIT) +3. robosuite: Copyright (c) ARISE Initiative. (Apache 2.0) +4. tmux-plugins/tpm: Copyright (c) tmux-plugins. (MIT) +5. BeyondMimic: Copyright (c) HybridRobotics. (MIT) + +-------------------------------------------------------------------------------- +FULL LICENSE TEXTS +-------------------------------------------------------------------------------- + +[BSD 3-Clause License] - Applied to: unitree_sdk2_python +Redistribution and use in source and binary forms, with or without modification, +are permitted provided that the following conditions are met: +1. Redistributions of source code must retain the above copyright notice, this +   list of conditions and the following disclaimer. +2. Redistributions in binary form must reproduce the above copyright notice, +   this list of conditions and the following disclaimer in the documentation +   and/or other materials provided with the distribution. +3. Neither the name of the copyright holder nor the names of its contributors +   may be used to endorse or promote products derived from this software without +   specific prior written permission. +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. + +--- + +[MIT License] - Applied to: XRoboToolkit, tpm, BeyondMimic +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. + +--- + +[Apache License 2.0] - Applied to: robosuite +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at: http://www.apache.org/licenses/LICENSE-2.0 \ No newline at end of file diff --git a/GR00T-WholeBodyControl/media/Pipeline.jpg b/GR00T-WholeBodyControl/media/Pipeline.jpg new file mode 100644 index 0000000000000000000000000000000000000000..51989f0cf1f0ff5137993ad03eb890f84fe2323b Binary files /dev/null and b/GR00T-WholeBodyControl/media/Pipeline.jpg differ diff --git a/GR00T-WholeBodyControl/media/gear_sonic_header.png b/GR00T-WholeBodyControl/media/gear_sonic_header.png new file mode 100644 index 0000000000000000000000000000000000000000..f147051b9b13fd60c9e412ba1f6b3a9cdc8d78cf Binary files /dev/null and b/GR00T-WholeBodyControl/media/gear_sonic_header.png differ diff --git a/GR00T-WholeBodyControl/motionbricks/.gitignore b/GR00T-WholeBodyControl/motionbricks/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..505918761be6be89eb21c66d231ccfde0f166aa0 --- /dev/null +++ b/GR00T-WholeBodyControl/motionbricks/.gitignore @@ -0,0 +1,36 @@ +# Python +__pycache__/ +*.py[cod] +*.egg-info/ +dist/ +build/ +*.egg + +# Recordings (local test artifacts) +recordings/ + +# Unused intermediate checkpoints +out/motionbricks_pose/version_1/checkpoints/model-step=0850000.ckpt +out/motionbricks_root/version_1/checkpoints/model-step=1350000.ckpt + +# Test logs +out/testing/ + +# IDE +.idea/ +.vscode/ +*.swp + +# OS +.DS_Store +Thumbs.db + +# Override monorepo root .gitignore: the motionbricks Python package uses +# `data/` and `models/` directory names for source code, not runtime +# artifacts. Re-include only the specific source directories. +!motionbricks/data/ +!motionbricks/data/** +!motionbricks/motion_backbone/models/ +!motionbricks/motion_backbone/models/** +!motionbricks/vqvae/models/ +!motionbricks/vqvae/models/**