Add logging for h5

sim/compare_h5.py

@@ -0,0 +1,117 @@
+"""
+Compare two HDF5 files and analyze their differences.
+
+Usage:
+    python compare_h5.py --isaac-file path/to/isaac.h5 --mujoco-file path/to/mujoco.h5
+
+Example:
+    python sim/compare_h5.py \
+        --isaac-file runs/h5_out/stompypro/2024-12-02_20-04-51/env_0/stompypro_env_0/h5_out/2024-12-02_20-04-51__053b2b5b-21c9-497c-b637-b66935dfe475.h5 \
+        --mujoco-file sim/resources/stompypro/h5_out/2024-12-02_21-10-41__5820ade7-9fc0-46df-8469-2f305480bcae.h5
+
+    python sim/compare_h5.py \
+        --isaac-file runs/h5_out/stompypro/2024-12-02_20-04-51/env_1/stompypro_env_1/h5_out/2024-12-02_20-04-51__d2016a6f-a486-4e86-8c5e-c9addd2cc13e.h5 \
+        --mujoco-file sim/resources/stompypro/h5_out/2024-12-02_21-10-41__5820ade7-9fc0-46df-8469-2f305480bcae.h5
+
+"""
+
+import h5py
+import numpy as np
+from pathlib import Path
+import argparse
+
+def load_h5_file(file_path):
+    """Load H5 file and return a dictionary of datasets"""
+    data = {}
+    with h5py.File(file_path, 'r') as f:
+        # Recursively load all datasets
+        def load_group(group, prefix=''):
+            for key in group.keys():
+                path = f"{prefix}/{key}" if prefix else key
+                if isinstance(group[key], h5py.Group):
+                    load_group(group[key], path)
+                else:
+                    data[path] = group[key][:]
+        load_group(f)
+    return data
+
+def compare_h5_files(issac_path, mujoco_path):
+    """Compare two H5 files and print differences"""
+    print(f"\nLoading files:")
+    print(f"Isaac: {issac_path}")
+    print(f"Mujoco: {mujoco_path}")
+
+    # Load both files
+    issac_data = load_h5_file(issac_path)
+    mujoco_data = load_h5_file(mujoco_path)
+
+    print("\nFile lengths:")
+    print(f"Isaac datasets: {len(issac_data)}")
+    print(f"Mujoco datasets: {len(mujoco_data)}")
+
+    print("\nDataset shapes:")
+    print("\nIsaac shapes:")
+    for key, value in issac_data.items():
+        print(f"{key}: {value.shape}")
+
+    print("\nMujoco shapes:")
+    for key, value in mujoco_data.items():
+        print(f"{key}: {value.shape}")
+
+    # Find common keys
+    common_keys = set(issac_data.keys()) & set(mujoco_data.keys())
+    print(f"\nCommon datasets: {len(common_keys)}")
+
+    # Find uncommon keys
+    issac_only_keys = set(issac_data.keys()) - common_keys
+    mujoco_only_keys = set(mujoco_data.keys()) - common_keys
+    print(f"\nIsaac only datasets: {len(issac_only_keys)}")
+    for key in issac_only_keys:
+        print(f"Isaac only: {key}")
+    print(f"\nMujoco only datasets: {len(mujoco_only_keys)}")
+    for key in mujoco_only_keys:
+        print(f"Mujoco only: {key}")
+
+    # Compare data for common keys
+    for key in common_keys:
+        issac_arr = issac_data[key]
+        mujoco_arr = mujoco_data[key]
+
+        print(f"\n========== For {key} ===============")
+
+        if issac_arr.shape != mujoco_arr.shape:
+            print(f"\n{key} - Shape mismatch: Isaac {issac_arr.shape} vs Mujoco {mujoco_arr.shape}")
+
+        # Calculate differences
+        min_shape = min(issac_arr.shape[0], mujoco_arr.shape[0])
+        if issac_arr.shape != mujoco_arr.shape:
+            raise ValueError(f"Shapes do not match for {key}. Cannot compare datasets with different shapes.")
+        diff = np.abs(issac_arr[:min_shape] - mujoco_arr[:min_shape])
+        max_diff = np.max(diff)
+        mean_diff = np.mean(diff)
+
+        print(f"Max difference: {max_diff:.6f}")
+        print(f"Mean difference: {mean_diff:.6f}\n")
+
+        start_idx = 0
+        display_timesteps = 10
+        end_idx = start_idx + display_timesteps
+
+        np.set_printoptions(formatter={'float': '{:0.6f}'.format}, suppress=True)
+        print("Isaac:\n", issac_arr[start_idx:end_idx])
+        print("Mujoco:\n", mujoco_arr[start_idx:end_idx])
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Compare two H5 files from Isaac and Mujoco simulations')
+    parser.add_argument('--isaac-file', required=True, help='Path to Isaac simulation H5 file')
+    parser.add_argument('--mujoco-file', required=True, help='Path to Mujoco simulation H5 file')
+
+    args = parser.parse_args()
+
+    print(f"Isaac path: {args.isaac_file}")
+    print(f"Mujoco path: {args.mujoco_file}")
+
+    compare_h5_files(args.isaac_file, args.mujoco_file)
+
+

sim/h5_logger.py

@@ -4,7 +4,6 @@
 from typing import Dict
 
 import h5py
-import matplotlib.pyplot as plt
 import numpy as np
 from datetime import datetime
 
@@ -117,6 +116,7 @@ def _plot_dataset(name: str, data: np.ndarray):
             name (str): Name of the dataset.
             data (np.ndarray): Data to be plotted.
         """
+        import matplotlib.pyplot as plt # dependency issues with python 3.8
         plt.figure(figsize=(10, 5))
         if data.ndim == 2:  # Handle multi-dimensional data
             for i in range(data.shape[1]):

sim/play.py

@@ -2,31 +2,34 @@
 """Play a trained policy in the environment.
 
 Run:
-    python sim/play.py --task g1 --log_h5
-    python sim/play.py --task stompymini --log_h5
+    python sim/play.py --task stompypro --log_h5
 """
 import argparse
 import copy
 import logging
 import os
+import time
+import uuid
 from datetime import datetime
 from typing import Any, Union
+import math
 
 import cv2
 import h5py
 import numpy as np
 from isaacgym import gymapi
+import torch  # isort: skip
 from tqdm import tqdm
 
-logger = logging.getLogger(__name__)
-
 from sim.env import run_dir  # noqa: E402
 from sim.envs import task_registry  # noqa: E402
 from sim.model_export import ActorCfg, convert_model_to_onnx  # noqa: E402
 from sim.utils.helpers import get_args  # noqa: E402
 from sim.utils.logger import Logger  # noqa: E402
 
-import torch  # isort: skip
+from sim.h5_logger import HDF5Logger
+
+logger = logging.getLogger(__name__)
 
 
 def export_policy_as_jit(actor_critic: Any, path: Union[str, os.PathLike]) -> None:
@@ -41,9 +44,10 @@ def play(args: argparse.Namespace) -> None:
     logger.info("Configuring environment and training settings...")
     env_cfg, train_cfg = task_registry.get_cfgs(name=args.task)
 
-    # override some parameters for testing
-    env_cfg.env.num_envs = min(env_cfg.env.num_envs, 1)
-    env_cfg.sim.max_gpu_contact_pairs = 2**10
+    num_parallel_envs = 2
+    env_cfg.env.num_envs = num_parallel_envs
+    env_cfg.sim.max_gpu_contact_pairs = 2**10 * num_parallel_envs
+
     if args.trimesh:
         env_cfg.terrain.mesh_type = "trimesh"
     else:
@@ -78,49 +82,35 @@ def play(args: argparse.Namespace) -> None:
         export_policy_as_jit(ppo_runner.alg.actor_critic, path)
         print("Exported policy as jit script to: ", path)
 
-    # export policy as a onnx module (used to run it on web)
-    if args.export_onnx:
-        path = ppo_runner.alg.actor_critic
-        convert_model_to_onnx(path, ActorCfg(), save_path="policy.onnx")
-        print("Exported policy as onnx to: ", path)
-
     # Prepare for logging
     env_logger = Logger(env.dt)
     robot_index = 0
     joint_index = 1
-    stop_state_log = 1000
+    env_steps_to_run = 1000
+
     now = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
     if args.log_h5:
-        h5_file = h5py.File(f"data{now}.h5", "w")
-
-        # Create dataset for actions
-        max_timesteps = stop_state_log
-        num_dof = env.num_dof
-        dset_actions = h5_file.create_dataset("actions", (max_timesteps, num_dof), dtype=np.float32)
-
-        # Create dataset of observations
-        buf_len = len(env.obs_history)  # length of observation buffer
-        dset_2D_command = h5_file.create_dataset(
-            "observations/2D_command", (max_timesteps, buf_len, 2), dtype=np.float32
-        )  # sin and cos commands
-        dset_3D_command = h5_file.create_dataset(
-            "observations/3D_command", (max_timesteps, buf_len, 3), dtype=np.float32
-        )  # x, y, yaw commands
-        dset_q = h5_file.create_dataset(
-            "observations/q", (max_timesteps, buf_len, num_dof), dtype=np.float32
-        )  # joint positions
-        dset_dq = h5_file.create_dataset(
-            "observations/dq", (max_timesteps, buf_len, num_dof), dtype=np.float32
-        )  # joint velocities
-        dset_obs_actions = h5_file.create_dataset(
-            "observations/actions", (max_timesteps, buf_len, num_dof), dtype=np.float32
-        )  # actions
-        dset_ang_vel = h5_file.create_dataset(
-            "observations/ang_vel", (max_timesteps, buf_len, 3), dtype=np.float32
-        )  # root angular velocity
-        dset_euler = h5_file.create_dataset(
-            "observations/euler", (max_timesteps, buf_len, 3), dtype=np.float32
-        )  # root orientation
+        # Create directory for HDF5 files
+        h5_dir = run_dir() / "h5_out" / args.task / now
+        h5_dir.mkdir(parents=True, exist_ok=True)
+
+        # Get observation dimensions
+        num_actions = env.num_dof
+        obs_buffer = env.obs_buf.shape[1]
+        prev_actions = np.zeros((num_actions), dtype=np.double)
+
+        h5_loggers = []
+        for env_idx in range(env_cfg.env.num_envs):
+            h5_dir = run_dir() / "h5_out" / args.task / now / f"env_{env_idx}"
+            h5_dir.mkdir(parents=True, exist_ok=True)
+
+            h5_loggers.append(HDF5Logger(
+                data_name=f"{args.task}_env_{env_idx}",
+                num_actions=num_actions,
+                max_timesteps=env_steps_to_run,
+                num_observations=obs_buffer,
+                h5_out_dir=str(h5_dir)
+            ))
 
     if args.render:
         camera_properties = gymapi.CameraProperties()
@@ -151,17 +141,15 @@ def play(args: argparse.Namespace) -> None:
             os.mkdir(experiment_dir)
         video = cv2.VideoWriter(dir, fourcc, 50.0, (1920, 1080))
 
-    for t in tqdm(range(stop_state_log)):
+    for t in tqdm(range(env_steps_to_run)):
         actions = policy(obs.detach())
-        if args.log_h5:
-            dset_actions[t] = actions.detach().numpy()
+
         if args.fix_command:
             env.commands[:, 0] = 0.5
             env.commands[:, 1] = 0.0
             env.commands[:, 2] = 0.0
             env.commands[:, 3] = 0.0
         obs, critic_obs, rews, dones, infos = env.step(actions.detach())
-        print(f"IMU: {obs[0, (3 * env.num_actions + 5) + 3 : (3 * env.num_actions + 5) + 2 * 3]}")
 
         if args.render:
             env.gym.fetch_results(env.sim, True)
@@ -187,17 +175,6 @@ def play(args: argparse.Namespace) -> None:
         base_vel_yaw = env.base_ang_vel[robot_index, 2].item()
         contact_forces_z = env.contact_forces[robot_index, env.feet_indices, 2].cpu().numpy()
 
-        if args.log_h5:
-            for i in range(buf_len):
-                cur_obs = env.obs_history[i].tolist()[0]
-                dset_2D_command[t, i] = cur_obs[0:2]  # sin and cos commands
-                dset_3D_command[t, i] = cur_obs[2:5]  # x, y, yaw commands
-                dset_q[t, i] = cur_obs[5 : 5 + num_dof]  # joint positions
-                dset_dq[t, i] = cur_obs[5 + num_dof : 5 + 2 * num_dof]  # joint velocities
-                dset_obs_actions[t, i] = cur_obs[5 + 2 * num_dof : 5 + 3 * num_dof]  # actions
-                dset_ang_vel[t, i] = cur_obs[5 + 3 * num_dof : 8 + 3 * num_dof]  # root angular velocity
-                dset_euler[t, i] = cur_obs[8 + 3 * num_dof : 11 + 3 * num_dof]  # root orientation
-
         env_logger.log_states(
             {
                 "dof_pos_target": dof_pos_target,
@@ -214,20 +191,46 @@ def play(args: argparse.Namespace) -> None:
                 "contact_forces_z": contact_forces_z,
             }
         )
+        actions = actions.detach().cpu().numpy()
+        if args.log_h5:
+            # Extract the current observation
+            for env_idx in range(env_cfg.env.num_envs):
+                h5_loggers[env_idx].log_data({
+                    "t": np.array([t * env.dt], dtype=np.float32),
+                    "2D_command": np.array(
+                        [
+                            np.sin(2 * math.pi * t * env.dt / env.cfg.rewards.cycle_time),
+                            np.cos(2 * math.pi * t * env.dt / env.cfg.rewards.cycle_time),
+                        ],
+                        dtype=np.float32,
+                    ),
+                    "3D_command": np.array(env.commands[env_idx, :3].cpu().numpy(), dtype=np.float32),
+                    "joint_pos": np.array(env.dof_pos[env_idx].cpu().numpy(), dtype=np.float32), 
+                    "joint_vel": np.array(env.dof_vel[env_idx].cpu().numpy(), dtype=np.float32),
+                    "prev_actions": prev_actions[env_idx].astype(np.float32),
+                    "curr_actions": actions[env_idx].astype(np.float32),
+                    "ang_vel": env.base_ang_vel[env_idx].cpu().numpy().astype(np.float32),
+                    "euler_rotation": env.base_euler_xyz[env_idx].cpu().numpy().astype(np.float32),
+                    "buffer": env.obs_buf[env_idx].cpu().numpy().astype(np.float32)
+                })
+
+            prev_actions = actions
+
         if infos["episode"]:
             num_episodes = env.reset_buf.sum().item()
             if num_episodes > 0:
                 env_logger.log_rewards(infos["episode"], num_episodes)
 
     env_logger.print_rewards()
-    env_logger.plot_states()
 
     if args.render:
         video.release()
 
     if args.log_h5:
-        print("Saving data to " + os.path.abspath(f"data{now}.h5"))
-        h5_file.close()
+        # print(f"Saving HDF5 file to {h5_logger.h5_file_path}") # TODO use code from kdatagen
+        for h5_logger in h5_loggers:
+            h5_logger.close()
+        print(f"HDF5 file(s) saved!")
 
 
 if __name__ == "__main__":

sim/utils/helpers.py

@@ -232,6 +232,12 @@ def get_args() -> argparse.Namespace:
             "default": False,
             "help": "Enable HDF5 logging",
         },
+        {
+            "name": "--log_krec",
+            "action": "store_true",
+            "default": False,
+            "help": "Enable KRec logging",
+        },
         {
             "name": "--trimesh",
             "action": "store_true",