utils.py

import glob
import os
import pickle

import torch
import numpy as np
import gymnasium as gym
from huggingface_hub.utils import EntryNotFoundError
from huggingface_sb3 import load_from_hub
from moviepy.video.compositing.concatenate import concatenate_videoclips
from moviepy.video.io.VideoFileClip import VideoFileClip
from rl_zoo3 import ALGOS
from gymnasium.wrappers import RecordVideo
from stable_baselines3.common.running_mean_std import RunningMeanStd

import os
import tarfile
import urllib.request


def install_mujoco():
    mujoco_url = "https://mujoco.org/download/mujoco210-linux-x86_64.tar.gz"
    mujoco_file = "mujoco210-linux-x86_64.tar.gz"
    mujoco_dir = "mujoco210"

    # Check if the directory already exists
    if not os.path.exists("mujoco210"):
        # Download Mujoco if not exists
        print("Downloading Mujoco...")
        urllib.request.urlretrieve(mujoco_url, mujoco_file)

        # Extract Mujoco
        print("Extracting Mujoco...")
        with tarfile.open(mujoco_file, "r:gz") as tar:
            tar.extractall()

        # Clean up the downloaded tar file
        os.remove(mujoco_file)

        print("Mujoco installed successfully!")
    else:
        print("Mujoco already installed.")

    # Set environment variable MUJOCO_PY_MUJOCO_PATH
    os.environ["MUJOCO_PY_MUJOCO_PATH"] = os.path.abspath(mujoco_dir)

    ld_library_path = os.environ.get("LD_LIBRARY_PATH", "")
    mujoco_bin_path = os.path.join(os.path.abspath(mujoco_dir), "bin")
    if mujoco_bin_path not in ld_library_path:
        os.environ["LD_LIBRARY_PATH"] = ld_library_path + ":" + mujoco_bin_path



class NormalizeObservation(gym.Wrapper):
    def __init__(self, env: gym.Env, clip_obs: float, obs_rms: RunningMeanStd, epsilon: float):
        gym.Wrapper.__init__(self, env)
        self.clip_obs = clip_obs
        self.obs_rms = obs_rms
        self.epsilon = epsilon

    def step(self, action):
        observation, reward, terminated, truncated, info = self.env.step(action)
        observation = self.normalize(np.array([observation]))[0]
        return observation, reward, terminated, truncated, info

    def reset(self, **kwargs):
        observation, info = self.env.reset(**kwargs)
        return self.normalize(np.array([observation]))[0], info

    def normalize(self, obs):
        return np.clip((obs - self.obs_rms.mean) / np.sqrt(self.obs_rms.var + self.epsilon), -self.clip_obs, self.clip_obs)


class CreateDataset(gym.Wrapper):
    def __init__(self, env: gym.Env):
        gym.Wrapper.__init__(self, env)
        self.observations = []
        self.actions = []
        self.last_observation = None

    def step(self, action):
        self.observations.append(self.last_observation)
        self.actions.append(action)
        observation, reward, terminated, truncated, info = self.env.step(action)
        self.last_observation = observation
        return observation, reward, terminated, truncated, info

    def reset(self, **kwargs):
        observation, info = self.env.reset(**kwargs)
        self.last_observation = observation
        return observation, info

    def get_dataset(self):
        if isinstance(self.env.action_space, gym.spaces.Box) and self.env.action_space.shape != (1,):
            actions = np.vstack(self.actions)
        else:
            actions = np.hstack(self.actions)
        return np.vstack(self.observations), actions


def rollouts(env, policy, num_episodes=1):
    for episode in range(num_episodes):
        done = False
        observation, _ = env.reset()
        while not done:
            action = policy(observation)
            observation, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated
    env.close()


def generate_dataset_from_expert(algo, env_name, num_train_episodes=5, num_test_episodes=2, force=False):
    if env_name.startswith("Swimmer") or env_name.startswith("Hopper"):
        install_mujoco()
    dataset_path = os.path.join("datasets", f"{algo}-{env_name}.pt")
    video_path = os.path.join("videos", f"{algo}-{env_name}.mp4")
    if os.path.exists(dataset_path) and os.path.exists(video_path) and not force:
        return dataset_path, video_path
    repo_id = f"sb3/{algo}-{env_name}"
    policy_file = f"{algo}-{env_name}.zip"

    expert_path = load_from_hub(repo_id, policy_file)
    try:
        vec_normalize_path = load_from_hub(repo_id, "vec_normalize.pkl")
        with open(vec_normalize_path, "rb") as f:
            vec_normalize = pickle.load(f)
            if vec_normalize.norm_obs:
                vec_normalize_params = {"clip_obs": vec_normalize.clip_obs, "obs_rms": vec_normalize.obs_rms, "epsilon": vec_normalize.epsilon}
            else:
                vec_normalize_params = None
    except EntryNotFoundError:
        vec_normalize_params = None

    expert = ALGOS[algo].load(expert_path)
    train_env = gym.make(env_name)
    train_env = CreateDataset(train_env)
    if vec_normalize_params is not None:
        train_env = NormalizeObservation(train_env, **vec_normalize_params)
    test_env = gym.make(env_name, render_mode="rgb_array")
    test_env = CreateDataset(test_env)
    if vec_normalize_params is not None:
        test_env = NormalizeObservation(test_env, **vec_normalize_params)
    test_env = RecordVideo(test_env, video_folder="videos", episode_trigger=lambda x: True, name_prefix=f"{algo}-{env_name}")

    def policy(obs):
        return expert.predict(obs, deterministic=True)[0]

    os.makedirs("videos", exist_ok=True)
    rollouts(train_env, policy, num_train_episodes)
    rollouts(test_env, policy, num_test_episodes)

    train_observations, train_actions = train_env.get_dataset()
    test_observations, test_actions = test_env.get_dataset()

    dataset = {
        "train_input": torch.from_numpy(train_observations),
        "test_input": torch.from_numpy(test_observations),
        "train_label": torch.from_numpy(train_actions),
        "test_label": torch.from_numpy(test_actions)
    }

    os.makedirs("datasets", exist_ok=True)
    torch.save(dataset, dataset_path)

    video_files = glob.glob(os.path.join("videos", f"{algo}-{env_name}-episode*.mp4"))
    clips = [VideoFileClip(file) for file in video_files]
    final_clip = concatenate_videoclips(clips)
    final_clip.write_videofile(video_path, codec="libx264", fps=24)

    return dataset_path, video_path


if __name__ == "__main__":
    generate_dataset_from_expert("ppo", "CartPole-v1", force=True)