diff --git a/sim/h5_logger.py b/sim/h5_logger.py index a1812fe7..81f123bb 100644 --- a/sim/h5_logger.py +++ b/sim/h5_logger.py @@ -6,25 +6,34 @@ import h5py import matplotlib.pyplot as plt import numpy as np +from datetime import datetime class HDF5Logger: - def __init__(self, data_name: str, num_actions: int, max_timesteps: int, num_observations: int): + def __init__(self, data_name: str, num_actions: int, max_timesteps: int, num_observations: int, h5_out_dir: str = "sim/resources/"): self.data_name = data_name self.num_actions = num_actions self.max_timesteps = max_timesteps self.num_observations = num_observations self.max_threshold = 1e3 # Adjust this threshold as needed + self.h5_out_dir = h5_out_dir self.h5_file, self.h5_dict = self._create_h5_file() self.current_timestep = 0 def _create_h5_file(self): # Create a unique file ID idd = str(uuid.uuid4()) - h5_file = h5py.File(f"{self.data_name}/{idd}.h5", "w") + timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + + curr_h5_out_dir = f"{self.h5_out_dir}/{self.data_name}/h5_out/" + os.makedirs(curr_h5_out_dir, exist_ok=True) + + h5_file_path = f"{curr_h5_out_dir}/{timestamp}__{idd}.h5" + print(f"Saving HDF5 data to {h5_file_path}") + h5_file = h5py.File(h5_file_path, "w") # Create datasets for logging actions and observations - dset_actions = h5_file.create_dataset("prev_actions", (self.max_timesteps, self.num_actions), dtype=np.float32) + dset_prev_actions = h5_file.create_dataset("prev_actions", (self.max_timesteps, self.num_actions), dtype=np.float32) dset_2D_command = h5_file.create_dataset("observations/2D_command", (self.max_timesteps, 2), dtype=np.float32) dset_3D_command = h5_file.create_dataset("observations/3D_command", (self.max_timesteps, 3), dtype=np.float32) dset_q = h5_file.create_dataset("observations/q", (self.max_timesteps, self.num_actions), dtype=np.float32) @@ -33,10 +42,12 @@ def _create_h5_file(self): dset_euler = h5_file.create_dataset("observations/euler", (self.max_timesteps, 3), dtype=np.float32) dset_t = h5_file.create_dataset("observations/t", (self.max_timesteps, 1), dtype=np.float32) dset_buffer = h5_file.create_dataset("observations/buffer", (self.max_timesteps, self.num_observations), dtype=np.float32) + dset_curr_actions = h5_file.create_dataset("curr_actions", (self.max_timesteps, self.num_actions), dtype=np.float32) # Map datasets for easy access h5_dict = { - "prev_actions": dset_actions, + "prev_actions": dset_prev_actions, + "curr_actions": dset_curr_actions, "2D_command": dset_2D_command, "3D_command": dset_3D_command, "joint_pos": dset_q, diff --git a/sim/produce_sim_data.py b/sim/produce_sim_data.py index 69f7e19e..a85f8f33 100644 --- a/sim/produce_sim_data.py +++ b/sim/produce_sim_data.py @@ -17,6 +17,8 @@ def run_simulation(sim_idx: int, args: argparse.Namespace) -> None: "--load_model", args.load_model, "--log_h5", + "--h5_out_dir", + args.h5_out_dir, "--no_render", ] @@ -48,6 +50,7 @@ def run_parallel_sims(num_threads: int, args: argparse.Namespace) -> None: 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") parser.add_argument("--num_examples", default=10000, type=int, help="Number of examples to run") + parser.add_argument("--h5_out_dir", default="sim/resources", type=str, help="Directory to save HDF5 files") args = parser.parse_args() # Run 100 examples total, in parallel batches diff --git a/sim/sim2sim.py b/sim/sim2sim.py index 9945067e..292462b5 100755 --- a/sim/sim2sim.py +++ b/sim/sim2sim.py @@ -110,6 +110,7 @@ def run_mujoco( keyboard_use: bool = False, log_h5: bool = False, render: bool = True, + h5_out_dir: str = "sim/resources", ) -> None: """ Run the Mujoco simulation using the provided policy and configuration. @@ -158,7 +159,7 @@ def run_mujoco( 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) + prev_actions = np.zeros((model_info["num_actions"]), dtype=np.double) hist_obs = np.zeros((model_info["num_observations"]), dtype=np.double) count_lowlevel = 0 @@ -178,7 +179,13 @@ def run_mujoco( if log_h5: stop_state_log = int(cfg.sim_duration / cfg.dt) / cfg.decimation - logger = HDF5Logger(embodiment, model_info["num_actions"], stop_state_log, model_info["num_observations"]) + logger = HDF5Logger( + data_name=embodiment, + num_actions=model_info["num_actions"], + max_timesteps=stop_state_log, + num_observations=model_info["num_observations"], + h5_out_dir=h5_out_dir + ) # Initialize variables for tracking upright steps and average speed upright_steps = 0 @@ -224,14 +231,17 @@ def run_mujoco( input_data["dof_pos.1"] = cur_pos_obs.astype(np.float32) input_data["dof_vel.1"] = cur_vel_obs.astype(np.float32) - input_data["prev_actions.1"] = actions.astype(np.float32) + input_data["prev_actions.1"] = prev_actions.astype(np.float32) input_data["imu_ang_vel.1"] = omega.astype(np.float32) input_data["imu_euler_xyz.1"] = eu_ang.astype(np.float32) input_data["buffer.1"] = hist_obs.astype(np.float32) - if args.log_h5: + positions, curr_actions, hist_obs = policy.run(None, input_data) + target_q = positions + + if log_h5: logger.log_data({ "t": np.array([count_lowlevel * cfg.dt], dtype=np.float32), "2D_command": np.array( @@ -244,14 +254,14 @@ def run_mujoco( "3D_command": np.array([x_vel_cmd, y_vel_cmd, yaw_vel_cmd], dtype=np.float32), "joint_pos": cur_pos_obs.astype(np.float32), "joint_vel": cur_vel_obs.astype(np.float32), - "prev_actions": actions.astype(np.float32), + "prev_actions": prev_actions.astype(np.float32), + "curr_actions": curr_actions.astype(np.float32), "ang_vel": omega.astype(np.float32), "euler_rotation": eu_ang.astype(np.float32), "buffer": hist_obs.astype(np.float32) }) - - positions, actions, hist_obs = policy.run(None, input_data) - target_q = positions + + prev_actions = curr_actions # Generate PD control tau = pd_control(target_q, q, kps, dq, kds, default) # Calc torques @@ -277,7 +287,7 @@ def run_mujoco( print(f"Number of upright steps: {upright_steps}") print(f"Average speed: {average_speed:.4f} m/s") - if args.log_h5: + if log_h5: logger.close() @@ -312,6 +322,7 @@ 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("--h5_out_dir", type=str, default="sim/resources", help="Directory to save HDF5 files") parser.add_argument("--no_render", action="store_false", dest="render", help="Disable rendering") parser.set_defaults(render=True) args = parser.parse_args() @@ -363,4 +374,5 @@ def parse_modelmeta( args.keyboard_use, args.log_h5, args.render, + args.h5_out_dir, )