Machine Learning Solver ¶

!pip install mathy_envs>=0.12.1 torch

import gymnasium as gym
import numpy as np
import torch
from mathy_envs import MathyEnv
from mathy_envs.gym import MathyGymEnv

# Environment difficulty level (options: 'easy', 'normal', 'hard')
env_difficulty = "easy"

# Environment names to train on, based on environment types and difficulty
env_types = [
    "poly",
    # "poly-blockers",
    # "poly-combine",
    # "poly-commute",
    # "poly-grouping",
    # "poly-like-terms-haystack",
    # "binomial",
    # "complex",
]
env_names = [f"mathy-{t}-{env_difficulty}-v0" for t in env_types]

env: MathyGymEnv = gym.make(env_names[0])
base_env: MathyEnv = env.unwrapped.mathy
print(f"Environment: {base_env.get_env_namespace()}")
print(f"Num Actions: {base_env.action_size}")
print(f"Rules      : {[e.name for e in base_env.rules]}")

# Learning rate for the actor network
lr_actor = 0.0003

# Learning rate for the critic network
lr_critic = 0.001

# Discount factor for future rewards
gamma = 0.99

# Number of epochs to update the policy
K_epochs = 80

# Clip parameter for PPO, used in policy update
eps_clip = 0.2

# Random seed setting (0 = no random seed)
random_seed = 1337

# Device to run the training on (CPU or CUDA)
device = torch.device("cpu" if not torch.cuda.is_available() else "cuda:0")

# Dimension of the hidden layer in the critic network
critic_hidden_dim = 64

# Where to save the model
checkpoint_path = "ppo.pth"

# Whether or not to use masked action selection. This makes the problems significantly easier when
# true because the action space is sparse with most possible actions being invalid. When false, the
# agent must learn to avoid invalid actions itself, making the problems much more challenging given
# the action space on the order of hundreds or thousands of possible actions for each state.
use_masked_actions = True

class RolloutBuffer:
    def __init__(self):
        self.actions = []
        self.states = []
        self.logprobs = []
        self.rewards = []
        self.state_values = []
        self.is_terminals = []

    def clear(self):
        del self.actions[:]
        del self.states[:]
        del self.logprobs[:]
        del self.rewards[:]
        del self.state_values[:]
        del self.is_terminals[:]

import torch.nn as nn
from torch.distributions import Categorical


class ActorCritic(nn.Module):
    def __init__(self, state_dim, action_dim):
        super(ActorCritic, self).__init__()
        self.device = device
        self.action_dim = action_dim

        # Actor network
        self.actor = nn.Sequential(
            nn.Linear(state_dim, critic_hidden_dim),
            nn.Tanh(),
            nn.Linear(critic_hidden_dim, critic_hidden_dim),
            nn.Tanh(),
            nn.Linear(critic_hidden_dim, action_dim),
            nn.Softmax(dim=-1),
        )

        # Critic network
        self.critic = nn.Sequential(
            nn.Linear(state_dim, critic_hidden_dim),
            nn.Tanh(),
            nn.Linear(critic_hidden_dim, critic_hidden_dim),
            nn.Tanh(),
            nn.Linear(critic_hidden_dim, 1),
        )

    def act(self, state):
        action_probs = self.actor(state)
        if use_masked_actions:
            mask = state[-action_probs.shape[0] :]
            action_probs = action_probs * mask
            action_probs = action_probs / torch.sum(action_probs)
        dist = Categorical(action_probs)
        action = dist.sample()
        action_logprob = dist.log_prob(action)
        state_val = self.critic(state)

        return action.detach(), action_logprob.detach(), state_val.detach()

    def evaluate(self, state, action):
        action_probs = self.actor(state)
        dist = Categorical(action_probs)
        action_logprobs = dist.log_prob(action)
        dist_entropy = dist.entropy()
        state_values = self.critic(state)

        return action_logprobs, state_values, dist_entropy

import torch
import torch.nn as nn


class PPO:
    def __init__(self, state_dim: int, action_dim: int):
        self.device = device
        self.gamma = gamma
        self.eps_clip = eps_clip
        self.K_epochs = K_epochs

        self.buffer = RolloutBuffer()
        self.policy = ActorCritic(state_dim, action_dim).to(device)
        self.optimizer = torch.optim.Adam(
            [
                {"params": self.policy.actor.parameters(), "lr": lr_actor},
                {"params": self.policy.critic.parameters(), "lr": lr_critic},
            ]
        )

        self.policy_old = ActorCritic(state_dim, action_dim).to(device)
        self.policy_old.load_state_dict(self.policy.state_dict())

        self.MseLoss = nn.MSELoss()

    def select_action(self, state):
        state = torch.FloatTensor(state).to(self.device)
        with torch.no_grad():
            action, action_logprob, state_val = self.policy_old.act(state)

        self.buffer.states.append(state)
        self.buffer.actions.append(action)
        self.buffer.logprobs.append(action_logprob)
        self.buffer.state_values.append(state_val)

        return action.item()

    def update(self):
        # Monte Carlo estimate of returns
        rewards = []
        discounted_reward = 0
        for reward, is_terminal in zip(
            reversed(self.buffer.rewards), reversed(self.buffer.is_terminals)
        ):
            if is_terminal:
                discounted_reward = 0
            discounted_reward = reward + (self.gamma * discounted_reward)
            rewards.insert(0, discounted_reward)

        # Normalizing the rewards
        rewards = torch.tensor(rewards, dtype=torch.float32).to(self.device)
        rewards = (rewards - rewards.mean()) / (rewards.std() + 1e-7)

        # convert list to tensor
        old_states = (
            torch.squeeze(torch.stack(self.buffer.states, dim=0))
            .detach()
            .to(self.device)
        )
        old_actions = (
            torch.squeeze(torch.stack(self.buffer.actions, dim=0))
            .detach()
            .to(self.device)
        )
        old_logprobs = (
            torch.squeeze(torch.stack(self.buffer.logprobs, dim=0))
            .detach()
            .to(self.device)
        )
        old_state_values = (
            torch.squeeze(torch.stack(self.buffer.state_values, dim=0))
            .detach()
            .to(self.device)
        )

        # calculate advantages
        advantages = rewards.detach() - old_state_values.detach()

        # Optimize policy for K epochs
        for _ in range(self.K_epochs):
            # Evaluating old actions and values
            logprobs, state_values, dist_entropy = self.policy.evaluate(
                old_states, old_actions
            )

            # match state_values tensor dimensions with rewards tensor
            state_values = torch.squeeze(state_values)

            # Finding the ratio (pi_theta / pi_theta__old)
            ratios = torch.exp(logprobs - old_logprobs.detach())

            # Finding Surrogate Loss
            surr1 = ratios * advantages
            surr2 = (
                torch.clamp(ratios, 1 - self.eps_clip, 1 + self.eps_clip) * advantages
            )

            # final loss of clipped objective PPO
            loss = (
                -torch.min(surr1, surr2)
                + 0.5 * self.MseLoss(state_values, rewards)
                - 0.01 * dist_entropy
            )

            # take gradient step
            self.optimizer.zero_grad()
            loss.mean().backward()
            self.optimizer.step()

        # Copy new weights into old policy
        self.policy_old.load_state_dict(self.policy.state_dict())

        # clear buffer
        self.buffer.clear()

    def save(self, checkpoint_path):
        torch.save(self.policy_old.state_dict(), checkpoint_path)

    def load(self, checkpoint_path):
        checkpoint = torch.load(checkpoint_path, map_location=self.device)
        self.policy_old.load_state_dict(checkpoint)
        self.policy.load_state_dict(checkpoint)

def train(checkpoint_path: str, max_steps: int = 1_000_000):
    print(
        f"Device set to: {torch.cuda.get_device_name(device) if torch.cuda.is_available() else 'CPU'}"
    )
    print(f"Training environments: {', '.join(env_names)}")

    max_ep_len = 50  # Max timesteps in one episode
    print_freq = 20_000  # Frequency for printing average reward
    save_model_freq = int(1e5)  # Model saving frequency
    update_timestep = max_ep_len * 4  # update policy every n timesteps

    envs = [
        gym.make(name, invalid_action_response="raise", verbose=False)
        for name in env_names
    ]
    env = envs[0]  # Select an environment

    # Initialize the PPO agent
    ppo_agent = PPO(env.observation_space.shape[0], env.action_space.n)

    # Training variables
    time_step = 0
    i_episode = 0
    total_reward = 0

    # Training loop
    while time_step <= max_steps:
        state, _ = env.reset()
        episode_reward = 0

        for t in range(1, max_ep_len + 1):
            action = ppo_agent.select_action(state)
            state, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated

            ppo_agent.buffer.rewards.append(reward)
            ppo_agent.buffer.is_terminals.append(done)

            time_step += 1
            episode_reward += reward

            if time_step % update_timestep == 0:
                ppo_agent.update()

            # Print average reward
            if time_step % print_freq == 0:
                avg_reward = total_reward / i_episode if i_episode > 0 else 0
                print(
                    f"Episode: {i_episode} \t Timestep: {time_step} \t Average Reward: {avg_reward:.2f}"
                )

            # Save model
            if time_step % save_model_freq == 0:
                print(f"Saving model at timestep {time_step}")
                ppo_agent.save(checkpoint_path)

            if done:
                break

        total_reward += episode_reward
        i_episode += 1

    print("Training completed.")

train(checkpoint_path, 300_000)

def test(checkpoint_path: str):
    envs = [
        gym.make(name, invalid_action_response="raise", verbose=True)
        for name in env_names
    ]
    assert len(envs) > 0, "No environments found"
    env = envs[0]
    ppo_agent = PPO(env.observation_space.shape[0], env.action_space.n)

    print(f"\nloading network from : {checkpoint_path}\n", flush=True)
    ppo_agent.load(checkpoint_path)

    total_test_episodes = 10  # total num of testing episodes
    test_running_reward = 0

    for ep in range(1, total_test_episodes + 1):
        env = np.random.choice(envs)
        ep_reward = 0
        state, _ = env.reset()
        done = False

        while not done:
            action = ppo_agent.select_action(state)
            state, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated
            ep_reward += reward

        test_running_reward += ep_reward
        emoji = "✅" if ep_reward >= 0.0 else "🔴"
        print(f"[{ep}]{emoji} Reward: {round(ep_reward, 2)}")
        ep_reward = 0

    avg_test_reward = test_running_reward / total_test_episodes
    print(f"Average test reward: {round(avg_test_reward, 2)}")

test(checkpoint_path)