runtime/env_manager/rbs_gym/scripts/evaluate.py

#!/usr/bin/env -S python3 -O

import argparse
import os
from typing import Dict

import numpy as np
import torch as th
import yaml
from stable_baselines3.common.utils import set_random_seed
from stable_baselines3.common.vec_env import DummyVecEnv, VecEnv, VecEnvWrapper

from rbs_gym import envs as gz_envs
from rbs_gym.utils import create_test_env, get_latest_run_id, get_saved_hyperparams
from rbs_gym.utils.utils import ALGOS, StoreDict, str2bool


def main(args: Dict):

    if args.exp_id == 0:
        args.exp_id = get_latest_run_id(
            os.path.join(args.log_folder, args.algo), args.env
        )
        print(f"Loading latest experiment, id={args.exp_id}")

    # Sanity checks
    if args.exp_id > 0:
        log_path = os.path.join(args.log_folder, args.algo, f"{args.env}_{args.exp_id}")
    else:
        log_path = os.path.join(args.log_folder, args.algo)

    assert os.path.isdir(log_path), f"The {log_path} folder was not found"

    found = False
    for ext in ["zip"]:
        model_path = os.path.join(log_path, f"{args.env}.{ext}")
        found = os.path.isfile(model_path)
        if found:
            break

    if args.load_best:
        model_path = os.path.join(log_path, "best_model.zip")
        found = os.path.isfile(model_path)

    if args.load_checkpoint is not None:
        model_path = os.path.join(
            log_path, f"rl_model_{args.load_checkpoint}_steps.zip"
        )
        found = os.path.isfile(model_path)

    if not found:
        raise ValueError(
            f"No model found for {args.algo} on {args.env}, path: {model_path}"
        )

    off_policy_algos = ["qrdqn", "dqn", "ddpg", "sac", "her", "td3", "tqc"]

    if args.algo in off_policy_algos:
        args.n_envs = 1

    set_random_seed(args.seed)

    if args.num_threads > 0:
        if args.verbose > 1:
            print(f"Setting torch.num_threads to {args.num_threads}")
        th.set_num_threads(args.num_threads)

    stats_path = os.path.join(log_path, args.env)
    hyperparams, stats_path = get_saved_hyperparams(
        stats_path, norm_reward=args.norm_reward, test_mode=True
    )

    # load env_kwargs if existing
    env_kwargs = {}
    args_path = os.path.join(log_path, args.env, "args.yml")
    if os.path.isfile(args_path):
        with open(args_path, "r") as f:
            # pytype: disable=module-attr
            loaded_args = yaml.load(f, Loader=yaml.UnsafeLoader)
            if loaded_args["env_kwargs"] is not None:
                env_kwargs = loaded_args["env_kwargs"]
    # overwrite with command line arguments
    if args.env_kwargs is not None:
        env_kwargs.update(args.env_kwargs)

    log_dir = args.reward_log if args.reward_log != "" else None

    env = create_test_env(
        args.env,
        n_envs=args.n_envs,
        stats_path=stats_path,
        seed=args.seed,
        log_dir=log_dir,
        should_render=not args.no_render,
        hyperparams=hyperparams,
        env_kwargs=env_kwargs,
    )

    kwargs = dict(seed=args.seed)
    if args.algo in off_policy_algos:
        # Dummy buffer size as we don't need memory to evaluate the trained agent
        kwargs.update(dict(buffer_size=1))


    custom_objects = {'observation_space': env.observation_space, 'action_space': env.action_space}

    model = ALGOS[args.algo].load(model_path, env=env, custom_objects=custom_objects, **kwargs)

    obs = env.reset()

    # Deterministic by default
    stochastic = args.stochastic
    deterministic = not stochastic

    print(
        f"Evaluating for {args.n_episodes} episodes with a",
        "deterministic" if deterministic else "stochastic",
        "policy.",
    )

    state = None
    episode_reward = 0.0
    episode_rewards, episode_lengths, success_episode_lengths = [], [], []
    ep_len = 0
    episode = 0
    # For HER, monitor success rate
    successes = []
    while episode < args.n_episodes:
        action, state = model.predict(obs, state=state, deterministic=deterministic)
        obs, reward, done, infos = env.step(action)
        if not args.no_render:
            env.render("human")

        episode_reward += reward[0]
        ep_len += 1

        if done and args.verbose > 0:
            episode += 1
            print(f"--- Episode {episode}/{args.n_episodes}")
            # NOTE: for env using VecNormalize, the mean reward
            # is a normalized reward when `--norm_reward` flag is passed
            print(f"Episode Reward: {episode_reward:.2f}")
            episode_rewards.append(episode_reward)
            print("Episode Length", ep_len)
            episode_lengths.append(ep_len)
            if infos[0].get("is_success") is not None:
                print("Success?:", infos[0].get("is_success", False))
                successes.append(infos[0].get("is_success", False))
                if infos[0].get("is_success"):
                    success_episode_lengths.append(ep_len)
                print(f"Current success rate: {100 * np.mean(successes):.2f}%")
            episode_reward = 0.0
            ep_len = 0
            state = None

    if args.verbose > 0 and len(successes) > 0:
        print(f"Success rate: {100 * np.mean(successes):.2f}%")

    if args.verbose > 0 and len(episode_rewards) > 0:
        print(
            f"Mean reward: {np.mean(episode_rewards):.2f} "
            f"+/- {np.std(episode_rewards):.2f}"
        )

    if args.verbose > 0 and len(episode_lengths) > 0:
        print(
            f"Mean episode length: {np.mean(episode_lengths):.2f} "
            f"+/- {np.std(episode_lengths):.2f}"
        )

    if args.verbose > 0 and len(success_episode_lengths) > 0:
        print(
            f"Mean episode length of successful episodes: {np.mean(success_episode_lengths):.2f} "
            f"+/- {np.std(success_episode_lengths):.2f}"
        )

    # Workaround for https://github.com/openai/gym/issues/893
    if not args.no_render:
        if args.n_envs == 1 and "Bullet" not in args.env and isinstance(env, VecEnv):
            # DummyVecEnv
            # Unwrap env
            while isinstance(env, VecEnvWrapper):
                env = env.venv
            if isinstance(env, DummyVecEnv):
                env.envs[0].env.close()
            else:
                env.close()
        else:
            # SubprocVecEnv
            env.close()


if __name__ == "__main__":

    parser = argparse.ArgumentParser()

    # Environment and its parameters
    parser.add_argument(
        "--env", type=str, default="Reach-Gazebo-v0", help="Environment ID"
    )
    parser.add_argument(
        "--env-kwargs",
        type=str,
        nargs="+",
        action=StoreDict,
        help="Optional keyword argument to pass to the env constructor",
    )
    parser.add_argument("--n-envs", type=int, default=1, help="Number of environments")

    # Algorithm
    parser.add_argument(
        "--algo",
        type=str,
        choices=list(ALGOS.keys()),
        required=False,
        default="sac",
        help="RL algorithm to use during the training",
    )
    parser.add_argument(
        "--num-threads",
        type=int,
        default=-1,
        help="Number of threads for PyTorch (-1 to use default)",
    )

    # Test duration
    parser.add_argument(
        "-n",
        "--n-episodes",
        type=int,
        default=200,
        help="Number of evaluation episodes",
    )

    # Random seed
    parser.add_argument("--seed", type=int, default=0, help="Random generator seed")

    # Model to test
    parser.add_argument(
        "-f", "--log-folder", type=str, default="logs", help="Path to the log directory"
    )
    parser.add_argument(
        "--exp-id",
        type=int,
        default=0,
        help="Experiment ID (default: 0: latest, -1: no exp folder)",
    )
    parser.add_argument(
        "--load-best",
        type=str2bool,
        default=False,
        help="Load best model instead of last model if available",
    )
    parser.add_argument(
        "--load-checkpoint",
        type=int,
        help="Load checkpoint instead of last model if available, you must pass the number of timesteps corresponding to it",
    )

    # Deterministic/stochastic actions
    parser.add_argument(
        "--stochastic",
        type=str2bool,
        default=False,
        help="Use stochastic actions instead of deterministic",
    )

    # Logging
    parser.add_argument(
        "--reward-log", type=str, default="reward_logs", help="Where to log reward"
    )
    parser.add_argument(
        "--norm-reward",
        type=str2bool,
        default=False,
        help="Normalize reward if applicable (trained with VecNormalize)",
    )

    # Disable render
    parser.add_argument(
        "--no-render",
        type=str2bool,
        default=False,
        help="Do not render the environment (useful for tests)",
    )

    # Verbosity
    parser.add_argument(
        "--verbose", type=int, default=1, help="Verbose mode (0: no output, 1: INFO)"
    )

    args, unknown = parser.parse_known_args()

    main(args)