diff --git a/sim/produce_sim_data.py b/sim/produce_sim_data.py new file mode 100644 index 00000000..f4bae38a --- /dev/null +++ b/sim/produce_sim_data.py @@ -0,0 +1,60 @@ +import argparse +import multiprocessing as mp +import subprocess +import time + +from tqdm import tqdm + + +def run_simulation(sim_idx: int, args: argparse.Namespace) -> None: + """Run a single simulation instance with the given parameters.""" + cmd = [ + "python", + "sim/sim2sim.py", + "--embodiment", + args.embodiment, + "--load_model", + args.load_model, + "--log_h5", + "--no_render", + ] + + try: + subprocess.run(cmd, check=True) + except subprocess.CalledProcessError as e: + print(f"Simulation {sim_idx} failed with error: {e}") + + +def run_parallel_sims(num_threads: int, args: argparse.Namespace) -> None: + """Run multiple simulation instances in parallel with a delay between each start.""" + processes = [] + delay_between_starts = 0.1 # Adjust the delay (in seconds) as needed + + for idx in range(num_threads): + p = mp.Process(target=run_simulation, args=(idx, args)) + p.start() + processes.append(p) + time.sleep(delay_between_starts) # Introduce a delay before starting the next process + + # Wait for all processes to finish with a progress bar + for p in tqdm(processes, desc="Running parallel simulations"): + p.join() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Run parallel simulations.") + parser.add_argument("--num_threads", type=int, default=10, help="Number of parallel simulations to run") + parser.add_argument("--embodiment", default="stompypro", type=str, help="Embodiment name") + parser.add_argument("--load_model", default="examples/walking_pro.onnx", type=str, help="Path to model to load") + + args = parser.parse_args() + + # Run 100 examples total, in parallel batches + num_examples = 2000 + num_batches = (num_examples + args.num_threads - 1) // args.num_threads + + for batch in range(num_batches): + examples_remaining = num_examples - (batch * args.num_threads) + threads_this_batch = min(args.num_threads, examples_remaining) + print(f"\nRunning batch {batch+1}/{num_batches} ({threads_this_batch} simulations)") + run_parallel_sims(threads_this_batch, args) diff --git a/sim/sim2sim.py b/sim/sim2sim.py index 7c034bac..d7b9339f 100755 --- a/sim/sim2sim.py +++ b/sim/sim2sim.py @@ -1,44 +1,17 @@ -# SPDX-License-Identifier: BSD-3-Clause -# -# 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. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -# -# Copyright (c) 2024 Beijing RobotEra TECHNOLOGY CO.,LTD. All rights reserved. """ Difference setup python sim/play.py --task mini_ppo --sim_device cpu python sim/sim2sim.py --load_model examples/standing_pro.pt --embodiment stompypro python sim/sim2sim.py --load_model examples/standing_micro.pt --embodiment stompymicro """ + import argparse import math import os +import uuid from copy import deepcopy from dataclasses import dataclass -from datetime import datetime -from typing import Any, Dict, List, Tuple, Union +from typing import Dict, List, Tuple, Union import h5py import mujoco @@ -52,10 +25,11 @@ from sim.model_export import ActorCfg, convert_model_to_onnx -def log_hdf5(data_name, num_actions, now, stop_state_log): +def log_hdf5(data_name, num_actions, stop_state_log): # Create data directory if it doesn't exist - os.makedirs(data_name, exist_ok=True) - h5_file = h5py.File(f"{data_name}/{now}.h5", "w") + # os.makedirs(data_name, exist_ok=True) + idd = str(uuid.uuid4()) + h5_file = h5py.File(f"{data_name}/{idd}.h5", "w") # Create dataset for actions max_timesteps = stop_state_log @@ -68,18 +42,14 @@ def log_hdf5(data_name, num_actions, now, stop_state_log): dset_3D_command = h5_file.create_dataset( "observations/3D_command", (max_timesteps, 3), dtype=np.float32 ) # x, y, yaw commands - dset_q = h5_file.create_dataset( - "observations/q", (max_timesteps, num_actions), dtype=np.float32 - ) # joint positions + dset_q = h5_file.create_dataset("observations/q", (max_timesteps, num_actions), dtype=np.float32) # joint positions dset_dq = h5_file.create_dataset( "observations/dq", (max_timesteps, num_actions), dtype=np.float32 ) # joint velocities dset_ang_vel = h5_file.create_dataset( "observations/ang_vel", (max_timesteps, 3), dtype=np.float32 ) # root angular velocity - dset_euler = h5_file.create_dataset( - "observations/euler", (max_timesteps, 3), dtype=np.float32 - ) # root orientation + dset_euler = h5_file.create_dataset("observations/euler", (max_timesteps, 3), dtype=np.float32) # root orientation h5_dict = { "actions": dset_actions, @@ -92,6 +62,7 @@ def log_hdf5(data_name, num_actions, now, stop_state_log): } return h5_file, h5_dict + @dataclass class Sim2simCfg: sim_duration: float = 60.0 @@ -176,6 +147,7 @@ def run_mujoco( model_info: Dict[str, Union[float, List[float], str]], keyboard_use: bool = False, log_h5: bool = False, + render: bool = True, ) -> None: """ Run the Mujoco simulation using the provided policy and configuration. @@ -218,7 +190,9 @@ def run_mujoco( data.qvel = np.zeros_like(data.qvel) data.qacc = np.zeros_like(data.qacc) - viewer = mujoco_viewer.MujocoViewer(model, data) + + if render: + viewer = mujoco_viewer.MujocoViewer(model, data) target_q = np.zeros((model_info["num_actions"]), dtype=np.double) actions = np.zeros((model_info["num_actions"]), dtype=np.double) @@ -241,8 +215,7 @@ def run_mujoco( if log_h5: stop_state_log = int(cfg.sim_duration / cfg.dt) - now = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") - h5_file, h5_dict = log_hdf5(embodiment, model_info["num_actions"], now, stop_state_log) + h5_file, h5_dict = log_hdf5(embodiment, model_info["num_actions"], stop_state_log) # Initialize variables for tracking upright steps and average speed upright_steps = 0 @@ -317,10 +290,12 @@ def run_mujoco( data.ctrl = tau mujoco.mj_step(model, data) - viewer.render() + if render: + viewer.render() count_lowlevel += 1 - viewer.close() + if render: + viewer.close() # Calculate average speed if step_count > 0: @@ -333,9 +308,9 @@ def run_mujoco( print(f"Average speed: {average_speed:.4f} m/s") if args.log_h5: - print(f"Saving data to {os.path.abspath(f'{embodiment}/{now}.h5')}") h5_file.close() + def parse_modelmeta( modelmeta: List[Tuple[str, str]], verbose: bool = False, @@ -367,6 +342,8 @@ def parse_modelmeta( parser.add_argument("--load_model", type=str, required=True, help="Path to run to load from.") parser.add_argument("--keyboard_use", action="store_true", help="keyboard_use") parser.add_argument("--log_h5", action="store_true", help="log_h5") + parser.add_argument("--no_render", action="store_false", dest="render", help="Disable rendering") + parser.set_defaults(render=True) args = parser.parse_args() if args.keyboard_use: @@ -374,7 +351,7 @@ def parse_modelmeta( pygame.init() pygame.display.set_caption("Simulation Control") else: - x_vel_cmd, y_vel_cmd, yaw_vel_cmd = 0.2, 0.0, 0.0 + x_vel_cmd, y_vel_cmd, yaw_vel_cmd = np.random.uniform(0.1, 0.5), 0.0, 0.0 policy_cfg = ActorCfg(embodiment=args.embodiment) if args.embodiment == "stompypro": @@ -404,4 +381,4 @@ def parse_modelmeta( verbose=True, ) - run_mujoco(args.embodiment, policy, cfg, model_info, args.keyboard_use, args.log_h5) + run_mujoco(args.embodiment, policy, cfg, model_info, args.keyboard_use, args.log_h5, args.render)