217 lines
7.0 KiB
Python
217 lines
7.0 KiB
Python
# %%
|
|
#!/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
|
|
# %%
|