# %% #!/usr/bin/env python3 from env_pinball import CustomEnv import os import pickle import random import numpy as np import torch import torch.nn as nn from torch.distributions.normal import Normal from torch.utils.tensorboard import SummaryWriter env = CustomEnv(devicenum=3) writer = SummaryWriter(log_dir='./tensorboard/DRL') # %% class Policy_Network(nn.Module): """Parametrized Policy Network.""" def __init__(self, obs_space_dims: int, action_space_dims: int): """Initializes a neural network that estimates the mean and standard deviation of a normal distribution from which an action is sampled from. Args: obs_space_dims: Dimension of the observation space action_space_dims: Dimension of the action space """ super().__init__() hidden_space1 = 256 # Nothing special with 16, feel free to change hidden_space2 = 256 # Nothing special with 32, feel free to change # Shared Network self.shared_net = nn.Sequential( nn.Linear(obs_space_dims, hidden_space1), nn.Tanh(), nn.Linear(hidden_space1, hidden_space2), nn.Tanh(), ) # Policy Mean specific Linear Layer self.policy_mean_net = nn.Sequential( nn.Linear(hidden_space2, action_space_dims) ) # Policy Std Dev specific Linear Layer self.policy_stddev_net = nn.Sequential( nn.Linear(hidden_space2, action_space_dims) ) def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Conditioned on the observation, returns the mean and standard deviation of a normal distribution from which an action is sampled from. Args: x: Observation from the environment Returns: action_means: predicted mean of the normal distribution action_stddevs: predicted standard deviation of the normal distribution """ shared_features = self.shared_net(x.float()) action_means = self.policy_mean_net(shared_features) action_means = torch.tanh(action_means) action_stddevs = torch.log( 1 + torch.exp(self.policy_stddev_net(shared_features)) ) return action_means, action_stddevs # %% class REINFORCE: """REINFORCE algorithm.""" def __init__(self, obs_space_dims: int, action_space_dims: int): """Initializes an agent that learns a policy via REINFORCE algorithm [1] to solve the task at hand (Inverted Pendulum v4). Args: obs_space_dims: Dimension of the observation space action_space_dims: Dimension of the action space """ # Hyperparameters self.learning_rate = 1e-4 # Learning rate for policy optimization self.gamma = 0.99 # Discount factor self.eps = 1e-6 # small number for mathematical stability self.probs = [] # Stores probability values of the sampled action self.rewards = [] # Stores the corresponding rewards self.net = Policy_Network(obs_space_dims, action_space_dims) self.optimizer = torch.optim.AdamW(self.net.parameters(), lr=self.learning_rate) def sample_action(self, state: np.ndarray) -> float: """Returns an action, conditioned on the policy and observation. Args: state: Observation from the environment Returns: action: Action to be performed """ state = torch.tensor(np.array([state])) action_means, action_stddevs = self.net(state) # create a normal distribution from the predicted # mean and standard deviation and sample an action distrib = Normal(action_means[0] + self.eps, action_stddevs[0] + self.eps) action = distrib.sample() prob = distrib.log_prob(action) action = torch.tanh(action) action = action.numpy() self.probs.append(prob) return action def update(self): """Updates the policy network's weights.""" running_g = 0 gs = [] # Discounted return (backwards) - [::-1] will return an array in reverse for R in self.rewards[::-1]: running_g = R + self.gamma * running_g gs.insert(0, running_g) deltas = torch.tensor(gs) loss = 0 # minimize -1 * prob * reward obtained for log_prob, delta in zip(self.probs, deltas): loss += log_prob.mean() * delta * (-1) # Update the policy network self.optimizer.zero_grad() loss.backward() self.optimizer.step() # Empty / zero out all episode-centric/related variables self.probs = [] self.rewards = [] # %% total_num_episodes = int(5e3) # Total number of episodes obs_space_dims = 6 action_space_dims = 3 rewards_over_seeds = [] MAX_REWARD = 0 # Check if there is a saved state if os.path.exists('saved_state.pkl'): with open('saved_state.pkl', 'rb') as f: i_seed, episode, agent, reward_over_episodes, rewards_over_seeds, MAX_REWARD = pickle.load(f) os.remove('saved_state.pkl') # Remove the saved state else: i_seed = 0 episode = 0 agent = None reward_over_episodes = None for seed in [1][i_seed:]: # Fibonacci seeds # set seed torch.manual_seed(seed) random.seed(seed) np.random.seed(seed) # Reinitialize agent every seed if agent is None or reward_over_episodes is None: agent = REINFORCE(obs_space_dims, action_space_dims) reward_over_episodes = [] while episode < total_num_episodes+1: obs, info = env.reset(Ccost=0.2+episode/total_num_episodes*0.6) steps = 0 done = False terminated = False truncated = False reward_over_steps = [] while not done: action = agent.sample_action(obs) obs, reward, terminated, truncated, info = env.step(action) agent.rewards.append(reward) reward_over_steps.append(reward) steps += 1 done = terminated or truncated avg_reward = np.mean(reward_over_steps[-64:]) reward_over_episodes.append(np.array([avg_reward], dtype=np.float32)) agent.update() if episode % 10 == 0: # print("Episode:", episode, "Average Reward:", int(avg_reward)) writer.add_scalar('Average Reward', int(avg_reward), episode) if avg_reward > MAX_REWARD: MAX_REWARD = avg_reward with open('saved_model_'+str(seed)+'.pkl', 'wb') as f: pickle.dump((episode + 1, agent, reward_over_episodes, MAX_REWARD), f) # Save the current state at the end of each episode with open('saved_state.pkl', 'wb') as f: pickle.dump((i_seed, episode + 1, agent, reward_over_episodes, rewards_over_seeds, MAX_REWARD), f) episode += 1 episode = 0 MAX_REWARD = 0 i_seed += 1 rewards_over_seeds.append(reward_over_episodes) agent = None # Reset the agent reward_over_episodes = None # Reset the reward_over_episodes # %%