(sry if bad english i’m not native)
Hello There,
I am currently trying to make a Chess AI using A3C methods.
So i followed a tutorial on how to implement such methods with the gym cartpole environment.
After finishing the tutorial i’m now trying to change the environment from cartpole to a chess environment.
I’m using this chess environment.
So, my problem is, when changing the output layer to match my needs with the environment, the model is stuck doing the forward, doing nothing, looks like it sleeps.
Here is the entire code:
import gym
import torch as T
import torch.multiprocessing as mp
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from torch.distributions import Categorical
from gym_chess import ChessEnvV1
class SharedAdam(T.optim.Adam):
def __init__(self, params, lr=1e-3, betas=(0.9, 0.99), eps=1e-8,
weight_decay=0):
super(SharedAdam, self).__init__(params, lr=lr, betas=betas, eps=eps,
weight_decay=weight_decay)
for group in self.param_groups:
for p in group['params']:
state = self.state[p]
state['step'] = 0
state['exp_avg'] = T.zeros_like(p.data)
state['exp_avg_sq'] = T.zeros_like(p.data)
state['exp_avg'].share_memory_()
state['exp_avg_sq'].share_memory_()
class ActorCritic(nn.Module):
def __init__(self, input_dims, n_actions, gamma=0.99):
super(ActorCritic, self).__init__()
self.gamma = gamma
self.pi1 = nn.Linear(*input_dims, 128)
self.v1 = nn.Linear(*input_dims, 128)
self.pi = nn.Linear(128, n_actions)
self.v = nn.Linear(128, 1)
self.rewards = []
self.actions = []
self.states = []
def remember(self, state, action, reward):
self.states.append(state)
self.actions.append(action)
self.rewards.append(reward)
def clear_memory(self):
self.states = []
self.actions = []
self.rewards = []
def forward(self, state):
pi1 = F.relu(self.pi1(state))
v1 = F.relu(self.v1(state))
pi = self.pi(pi1)
v = self.v(v1)
return pi, v
def calc_R(self, done):
states = T.tensor(self.states, dtype=T.float)
_, v = self.forward(states)
R = v[-1]*(1-int(done))
batch_return = []
for reward in self.rewards[::-1]:
R = reward + self.gamma*R
batch_return.append(R)
batch_return.reverse()
batch_return = T.tensor(batch_return, dtype=T.float)
return batch_return
def calc_loss(self, done):
states = T.tensor(self.states, dtype=T.float)
actions = T.tensor(self.actions, dtype=T.float)
returns = self.calc_R(done)
pi, values = self.forward(states)
values = values.squeeze()
critic_loss = (returns-values)**2
probs = T.softmax(pi, dim=1)
dist = Categorical(probs)
log_probs = dist.log_prob(actions)
actor_loss = -log_probs*(returns-values)
total_loss = (critic_loss + actor_loss).mean()
return total_loss
def choose_action(self, observation):
state = T.tensor([observation], dtype=T.float)
pi, v = self.forward(state)
probs = T.softmax(pi, dim=1)
dist = Categorical(probs)
action = dist.sample().numpy()[0]
return action
class Agent(mp.Process):
def __init__(self, global_actor_critic, optimizer, input_dims, n_actions,
gamma, lr, name, global_ep_idx, env_id):
super(Agent, self).__init__()
self.local_actor_critic = ActorCritic(input_dims, n_actions, gamma)
self.global_actor_critic = global_actor_critic
self.name = 'w%02i' % name
self.episode_idx = global_ep_idx
self.env = gym.make(env_id)
self.optimizer = optimizer
def run(self):
t_step = 1
while self.episode_idx.value < N_GAMES:
done = False
observation = self.env.reset()
score = 0
self.local_actor_critic.clear_memory()
while not done:
action = self.local_actor_critic.choose_action(np.array(observation).flatten())
print(action)
observation_, reward, done, info = self.env.step(action)
print('played: ', reward, done, info)
score += reward
self.local_actor_critic.remember(observation, action, reward)
if t_step % T_MAX == 0 or done:
loss = self.local_actor_critic.calc_loss(done)
self.optimizer.zero_grad()
loss.backward()
for local_param, global_param in zip(
self.local_actor_critic.parameters(),
self.global_actor_critic.parameters()):
global_param._grad = local_param.grad
self.optimizer.step()
self.local_actor_critic.load_state_dict(
self.global_actor_critic.state_dict())
self.local_actor_critic.clear_memory()
t_step += 1
observation = observation_
with self.episode_idx.get_lock():
self.episode_idx.value += 1
print(self.name, 'episode ', self.episode_idx.value, 'reward %.1f' % score)
if __name__ == '__main__':
lr = 1e-4
env_id = 'ChessVsSelf-v1'
n_actions = 4096
input_dims = [64]
N_GAMES = 3000
T_MAX = 5
global_actor_critic = ActorCritic(input_dims, n_actions)
global_actor_critic.share_memory()
optim = SharedAdam(global_actor_critic.parameters(), lr=lr,
betas=(0.92, 0.999))
global_ep = mp.Value('i', 0)
workers = [Agent(global_actor_critic,
optim,
input_dims,
n_actions,
gamma=0.99,
lr=lr,
name=i,
global_ep_idx=global_ep,
env_id=env_id) for i in range(mp.cpu_count())]
[w.start() for w in workers]
[w.join() for w in workers]
So in the main i define my n_actions, which is the output dimension for the model.
When doing it with cartpole, n_actions = 2, it works perfectly well, 0 problem.
BUT when changing my n_actions to 4096, which is the output i need for the environment, the code is stuck during the forward, more precisely doing the pi = self.pi(pi1) line (i’ve tried putting a print before and after this line and the second print never happens).
What’s even more intriguing is that no computation is done when it’s stuck like this, my CPUs are all at 10% which is is just their idle state. And i have no error logs from pytorch as well.
Also worth noting that i tried changing the n_actions to something like 200, the forward works, but increasing it up to 300 it doesn’t work anymore.
If that matters, i’m on Fedora 29, python 3.7.7, and pytorch 1.8.1+cu102.
(I tried on different computers and the same behaviour happens.)
So yeah i’m here stuck with a problem that i really don’t understand, i can’t imagine why the forward is blocking here when increasing the output dimension.
Any help greatly appreciated.