evaluate.py

import gymnasium as gym
import torch
from agent import Policy, Agent

def evaluate(model_path, mode='reinforce'):
    # Turn on human rendering so we can watch!
    env = gym.make('Hopper-v4', render_mode='human')
    
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.shape[0]
    
    # Initialize the policy and agent
    policy = Policy(state_dim, action_dim)
    agent = Agent(policy, device='cpu', mode=mode)
    
    # Load the trained brain
    agent.policy.load_state_dict(torch.load(model_path))
    agent.policy.eval() # Set network to evaluation mode
    
    print(f"Evaluating {model_path}...")
    
    for ep in range(5): # Watch 5 episodes
        state, _ = env.reset()
        done = False
        ep_reward = 0
        
        while not done:
            # Note the evaluation=True flag! This stops random flailing.
            action, _ = agent.get_action(state, evaluation=True)
            state, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated
            ep_reward += reward
            
            env.render()
            
        print(f"Episode {ep + 1} Score: {ep_reward:.2f}")
        
    env.close()

if __name__ == '__main__':
    # Change the filename below to whichever model you want to watch
    # Options will be: 'reinforce_best_model.pth', 'actor_critic_best_model.pth'
    evaluate('actor_critic_best_model.pth', mode='actor_critic')