import numpy as np
from sb3_contrib import RecurrentPPO
from stable_baselines3.common.evaluation import evaluate_policy
from helper import *


if __name__ == '__main__':

    env = get_dummy_env(1,True)

    model = RecurrentPPO.load("./model/rppo/best_model_1")


    #mean_reward, std_reward = evaluate_policy(model, env,1)
        
    #print(f"{mean_reward} {std_reward}")

    obs = env.reset()
    # cell and hidden state of the LSTM
    lstm_states = None
    # Episode start signals are used to reset the lstm states
    episode_starts = np.ones((env.num_envs,), dtype=bool)

    while True:
        action, lstm_states = model.predict(obs, state=lstm_states, episode_start=episode_starts)
  
        obs, rewards, dones, info = env.step(action)
        episode_starts = dones
        env.render()
        
        if dones :
            env.close()
            break