Sim Library Improvements (#116)

* add cmd arg for h5 output path with sensible default

* save current actions to h5

* add h5out as arg to produce_sim_data as well

* compute actions before saving them

* typo fix: args.log_h5 -> log_h5

* fix: rename actions to prev_actions, and set prev_action to curr_action at the end of loop

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,

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

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,
     )