Spaces:
Paused
Paused
| import argparse | |
| import h5py | |
| import sys | |
| import numpy as np | |
| DEFAULT_PATH = "/data/hongzefu/data_0225/record_dataset_VideoUnmaskSwap.h5" | |
| def _is_equal(a, b): | |
| if a is None and b is None: | |
| return True | |
| if a is None or b is None: | |
| return False | |
| if isinstance(a, np.ndarray) and isinstance(b, np.ndarray): | |
| return np.array_equal(a, b) | |
| try: | |
| return bool(np.all(a == b)) | |
| except: | |
| return False | |
| def inspect_actions(filepath, target_timestep=None, window=10): | |
| print(f"Inspecting HDF5 file: {filepath}") | |
| try: | |
| with h5py.File(filepath, 'r') as f: | |
| # Sort episodes naturally (e.g., episode_0, episode_1, etc.) | |
| episodes = [k for k in f.keys() if k.startswith('episode_')] | |
| episodes.sort(key=lambda x: int(x.split('_')[1])) | |
| for ep_name in episodes: | |
| ep_group = f[ep_name] | |
| print(f"--- {ep_name} ---") | |
| # Sort timesteps naturally | |
| timesteps = [k for k in ep_group.keys() if k.startswith('timestep_')] | |
| timesteps.sort(key=lambda x: int(x.split('_')[1])) | |
| last_choice_action = None | |
| last_is_subgoal_boundary = None | |
| last_waypoint_action = None | |
| last_is_video_demo = None | |
| skip_count = 0 | |
| for ts_name in timesteps: | |
| ts_group = ep_group[ts_name] | |
| # Try to find the requested fields | |
| choice_action = None | |
| is_subgoal_boundary = None | |
| is_video_demo = None | |
| waypoint_action = None | |
| # Action group typically contains choice_action and waypoint_action | |
| if 'action' in ts_group: | |
| action_group = ts_group['action'] | |
| if 'choice_action' in action_group: | |
| choice_action = np.array(action_group['choice_action']) | |
| if 'waypoint_action' in action_group: | |
| waypoint_action = np.array(action_group['waypoint_action']) | |
| # info group typically contains is_subgoal_boundary and is_video_demo | |
| if 'info' in ts_group: | |
| info_group = ts_group['info'] | |
| if 'is_subgoal_boundary' in info_group: | |
| is_subgoal_boundary = np.array(info_group['is_subgoal_boundary']) | |
| if 'is_video_demo' in info_group: | |
| is_video_demo = np.array(info_group['is_video_demo']) | |
| # For safety, check if they are directly under timestep or elsewhere | |
| if choice_action is None and 'choice_action' in ts_group: | |
| choice_action = np.array(ts_group['choice_action']) | |
| if is_subgoal_boundary is None and 'is_subgoal_boundary' in ts_group: | |
| is_subgoal_boundary = np.array(ts_group['is_subgoal_boundary']) | |
| if is_video_demo is None and 'is_video_demo' in ts_group: | |
| is_video_demo = np.array(ts_group['is_video_demo']) | |
| if waypoint_action is None and 'waypoint_action' in ts_group: | |
| waypoint_action = np.array(ts_group['waypoint_action']) | |
| should_skip = False | |
| if last_choice_action is not None: | |
| same_choice = _is_equal(choice_action, last_choice_action) | |
| same_boundary = _is_equal(is_subgoal_boundary, last_is_subgoal_boundary) | |
| same_vd = _is_equal(is_video_demo, last_is_video_demo) | |
| same_wp = _is_equal(waypoint_action, last_waypoint_action) | |
| should_skip = same_choice and same_boundary and same_vd and same_wp | |
| if target_timestep is not None: | |
| ts_idx = int(ts_name.split('_')[1]) | |
| if abs(ts_idx - target_timestep) <= window: | |
| should_skip = False | |
| if should_skip: | |
| skip_count += 1 | |
| continue | |
| if skip_count > 0: | |
| print(f" ... ({skip_count} identical timesteps skipped) ...") | |
| skip_count = 0 | |
| print(f" {ts_name}:") | |
| print(f" choice_action: {choice_action}") | |
| print(f" is_subgoal_boundary: {is_subgoal_boundary}") | |
| print(f" is_video_demo: {is_video_demo}") | |
| print(f" waypoint_action: {waypoint_action}") | |
| last_choice_action = choice_action | |
| last_is_subgoal_boundary = is_subgoal_boundary | |
| last_is_video_demo = is_video_demo | |
| last_waypoint_action = waypoint_action | |
| if skip_count > 0: | |
| print(f" ... ({skip_count} identical timesteps skipped) ...") | |
| skip_count = 0 | |
| except Exception as e: | |
| print(f"Error reading HDF5 file: {e}") | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser(description="Inspect HDF5 dataset actions.") | |
| parser.add_argument("filepath", type=str, nargs="?", default=DEFAULT_PATH, | |
| help="Path to the HDF5 file.") | |
| parser.add_argument("-t", "--timestep", type=int, default=None, | |
| help="Specific timestep to not omit even if identical.") | |
| parser.add_argument("-w", "--window", type=int, default=10, | |
| help="Window around specified timestep to not omit (default: 10).") | |
| args = parser.parse_args() | |
| inspect_actions(args.filepath, target_timestep=args.timestep, window=args.window) | |