scripts/data_generation/process_rl_trajectories.py

import copy
import json
import os
import shutil
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

import gymnasium as gym
import h5py
import numpy as np
import tyro
from tqdm import tqdm

import mani_skill.envs
from mani_skill.trajectory import utils
from mani_skill.utils.visualization.misc import images_to_video


@dataclass
class Args:
    runs_path: str
    out_dir: str
    dry_run: Optional[bool] = True


def main():
    args = tyro.cli(Args)

    if args.dry_run:
        print("Dry run, skipping actual processing")

    # Dictionary to store paths for each environment experiment
    env_paths = {}

    # List all subfolders in runs_path
    for env_name in os.listdir(args.runs_path):
        env_dir = Path(args.runs_path) / env_name

        if not env_dir.is_dir():
            continue

        # Look for checkpoint and trajectory files
        ckpt_path = env_dir / "final_ckpt.pt"
        traj_path = env_dir / "test_videos" / "trajectory.h5"
        traj_metadata_path = env_dir / "test_videos" / "trajectory.json"

        # Only store if both files exist
        if not (ckpt_path.exists() and traj_path.exists()):
            print(
                f"Skipping {env_name} because checkpoint or trajectory file does not exist"
            )
            continue

        env_paths[env_name] = {
            "checkpoint": str(ckpt_path),
            "trajectory": str(traj_path),
            "metadata": str(traj_metadata_path),
        }
    high_fail_rate_envs = []
    for env_name, env_path in env_paths.items():
        print(f"Processing {env_name}")
        traj_path = env_path["trajectory"]
        file = h5py.File(traj_path, "r")
        metadata_path = env_path["metadata"]
        with open(metadata_path, "r") as f:
            metadata = json.load(f)
        env_id = metadata["env_info"]["env_id"]
        new_metadata = copy.deepcopy(metadata)
        new_metadata["episodes"] = []
        control_mode = new_metadata["env_info"]["env_kwargs"]["control_mode"]
        sim_backend = new_metadata["env_info"]["env_kwargs"]["sim_backend"]
        traj_filename = f"{env_id}/rl/trajectory.none.{control_mode}.{sim_backend}"
        out_trajectory_path = os.path.join(args.out_dir, f"{traj_filename}.h5")

        if not args.dry_run:
            os.makedirs(os.path.dirname(out_trajectory_path), exist_ok=True)
            out_file = h5py.File(out_trajectory_path, "w")

        failed_count = 0
        truncated_count = 0
        avg_episode_length = 0
        original_episode_count = len(metadata["episodes"])
        first_success_indexes = []
        recorded_sample_video = False
        for episode in tqdm(metadata["episodes"]):
            traj_id = f"traj_{episode['episode_id']}"
            traj = file[traj_id]
            success = np.array(traj["success"])
            if not success.any():
                # this failed
                failed_count += 1
                continue
            # truncate until last success
            success_indexes = success.nonzero()[0]
            last_success_index = int(success_indexes[-1])
            first_success_index = int(success_indexes[0])
            first_success_indexes.append(first_success_index)
            if last_success_index != len(success) - 1:
                truncated_count += 1
            avg_episode_length += last_success_index + 1

            def recursive_copy_and_slice(
                key, source_group, target_group, add_last_frame=False
            ):
                if key == "obs" or key == "rewards":
                    return
                if isinstance(target_group, h5py.Dataset):
                    if not add_last_frame and ("obs" == key or "env_states" == key):
                        add_last_frame = True
                    source_group.create_dataset(
                        key,
                        data=target_group[: last_success_index + 1 + add_last_frame],
                    )
                elif isinstance(target_group, h5py.Group):
                    if not add_last_frame and ("obs" == key or "env_states" == key):
                        add_last_frame = True
                    source_group.create_group(key, track_order=True)
                    for k in target_group.keys():
                        recursive_copy_and_slice(
                            k,
                            source_group[key],
                            target_group[k],
                            add_last_frame=add_last_frame,
                        )

            if not args.dry_run:
                recursive_copy_and_slice(traj_id, out_file, traj)
            new_episode = copy.deepcopy(episode)
            new_episode["success"] = True
            new_episode["elapsed_steps"] = last_success_index + 1
            new_metadata["episodes"].append(new_episode)

            if not args.dry_run:
                if not recorded_sample_video:
                    recorded_sample_video = True
                    env_kwargs = copy.deepcopy(new_metadata["env_info"]["env_kwargs"])
                    env_kwargs["num_envs"] = 1
                    env_kwargs["sim_backend"] = "physx_cpu"
                    env_kwargs["human_render_camera_configs"] = {
                        "shader_pack": "rt-med"
                    }
                    env = gym.make(env_id, **env_kwargs)
                    env.reset(
                        seed=episode["episode_seed"], **new_episode["reset_kwargs"]
                    )
                    imgs = []
                    env_states = utils.dict_to_list_of_dicts(
                        out_file[traj_id]["env_states"]
                    )
                    for step in range(new_episode["elapsed_steps"]):
                        env.set_state_dict(env_states[step])
                        imgs.append(env.render_rgb_array().cpu().numpy()[0])
                    env.close()
                    images_to_video(
                        imgs,
                        output_dir=os.path.join(args.out_dir, env_id, "rl"),
                        video_name=f"sample_{control_mode}",
                        fps=30,
                    )
        final_episode_count = len(new_metadata["episodes"])
        avg_episode_length /= final_episode_count
        avg_steps_to_first_success = np.mean(first_success_indexes)
        print(
            f"{env_id}: Failed: {failed_count}/{original_episode_count}, Truncated: {truncated_count}/{original_episode_count}, Final Episodes: {final_episode_count}, Avg Episode Length: {avg_episode_length}, Avg Steps to First Success: {avg_steps_to_first_success}"
        )
        if failed_count / original_episode_count >= 0.05:
            high_fail_rate_envs.append(
                (env_name, failed_count / original_episode_count)
            )

        new_metadata["source_type"] = "rl"
        new_metadata[
            "source_desc"
        ] = "Demonstrations generated by rolling out a PPO dense reward trained policy"
        if not args.dry_run:
            with open(
                os.path.join(
                    args.out_dir,
                    f"{traj_filename}.json",
                ),
                "w",
            ) as f:
                json.dump(new_metadata, f, indent=2)
            print(f"Saved to {os.path.join(args.out_dir, f'{traj_filename}.json')}")
            out_file.close()

            # Copy checkpoint to output dir
            checkpoint_path = env_path["checkpoint"]
            checkpoint_out_path = os.path.join(
                args.out_dir, f"{env_id}/rl/ppo_{control_mode}_ckpt.pt"
            )
            os.makedirs(os.path.dirname(checkpoint_out_path), exist_ok=True)
            shutil.copy(checkpoint_path, checkpoint_out_path)

    for env_name, fail_rate in high_fail_rate_envs:
        print(
            f"Warning: {env_name} has {fail_rate*100:0.1f} >= 5% failed episodes. Need a better policy."
        )


if __name__ == "__main__":
    main()