from stable_baselines3 import SAC
from training_env import TrainingEnv
from stable_baselines3.common.vec_env import VecFrameStack,DummyVecEnv
from stable_baselines3.common.evaluation import evaluate_policy

TEST_BEGIN_TIME = ["2022-09-05"]

env = TrainingEnv(TEST_BEGIN_TIME)
env = DummyVecEnv([lambda:env])
model = SAC.load("./best_model",env)

mean_reward, std_reward = evaluate_policy(model, env,1)
print(f"{mean_reward} {std_reward}")
obs = env.reset()
while True:
    action, _ = model.predict(obs,deterministic=True)
    obs, rewards, dones, info = env.step(action)
    #env.render()
    if dones :
        env.close()
        break