from stable_baselines import PPO2
import numpy as np
from env.create_vec_env import create_vec_env


def evaluate(model, env, num_steps=1000):
    """
    Evaluate a RL agent
    :param model: (BaseRLModel object) the RL Agent
    :param num_steps: (int) number of timesteps to evaluate it
    :return: (float) Mean reward
    """
    episode_rewards = [[0.0] for _ in range(env.num_envs)]
    obs = env.reset()
    for i in range(num_steps):
        actions, _states = model.predict(obs)
        obs, rewards, dones, info = env.step(actions)

        # Stats
        for i in range(env.num_envs):
            episode_rewards[i][-1] += rewards[i]
            if dones[i]:
                episode_rewards[i].append(0.0)

    mean_rewards = [0.0 for _ in range(env.num_envs)]
    n_episodes = 0
    for i in range(env.num_envs):
        mean_rewards[i] = np.mean(episode_rewards[i])
        n_episodes += len(episode_rewards[i])

    # Compute mean reward
    mean_reward = round(np.mean(mean_rewards), 1)

    return mean_reward


if __name__ == '__main__':
    # Load the trained agent
    model = PPO2.load("logs/policy9.zip")
    env = create_vec_env(visibility=12, safe_info=True, safe_states=True, supervisor='cbf', safety_distance_logical=None,
                         repl=1, worker_start=1000, lidar_num_bins=15)

    mean_reward_current = evaluate(model=model, env=env, num_steps=2500)

from stable_baselines import PPO2
import numpy as np
from env.create_vec_env import create_vec_env


def evaluate(model, env, num_steps=1000):
    """
    Evaluate a RL agent
    :param model: (BaseRLModel object) the RL Agent
    :param num_steps: (int) number of timesteps to evaluate it
    :return: (float) Mean reward
    """
    episode_rewards = [[0.0] for _ in range(env.num_envs)]
    obs = env.reset()
    for i in range(num_steps):
        actions, _states = model.predict(obs)
        obs, rewards, dones, info = env.step(actions)

        # Stats
        for i in range(env.num_envs):
            episode_rewards[i][-1] += rewards[i]
            if dones[i]:
                episode_rewards[i].append(0.0)

    mean_rewards = [0.0 for _ in range(env.num_envs)]
    n_episodes = 0
    for i in range(env.num_envs):
        mean_rewards[i] = np.mean(episode_rewards[i])
        n_episodes += len(episode_rewards[i])

    # Compute mean reward
    mean_reward = round(np.mean(mean_rewards), 1)

    return mean_reward


if __name__ == '__main__':
    # Load the trained agent
    model = PPO2.load("logs/policy9.zip")
    env = create_vec_env(visibility=12, safe_info=True, safe_states=True, supervisor='cbf', safety_distance_logical=None,
                         repl=1, worker_start=1000, lidar_num_bins=15)

    mean_reward_current = evaluate(model=model, env=env, num_steps=2500)

Evaluate a RL agent     :param model: (BaseRLModel object) the RL Agent     :param num_steps: (int) number of timesteps to evaluate it     :return: (float) Mean reward # Stats # Compute mean reward # Load the trained agent