File size: 2,105 Bytes
1f3a93e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 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) |