(A3C) Forward does nothing when increasing size of output layer

(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,
        super(SharedAdam, self).__init__(params, lr=lr, betas=betas, eps=eps,

        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)


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):

    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 = 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
            while not done:
                action = self.local_actor_critic.choose_action(np.array(observation).flatten())
                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)
                    for local_param, global_param in zip(
                        global_param._grad = local_param.grad
                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)
    optim = SharedAdam(global_actor_critic.parameters(), lr=lr, 
                        betas=(0.92, 0.999))
    global_ep = mp.Value('i', 0)

    workers = [Agent(global_actor_critic,
                    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.


Did you manage to solve this? I am having the same issue.