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

Nellousan: <p>(sry if bad english i’m not native)<br>
Hello There,</p>
<p>I am currently trying to make a Chess AI using A3C methods.<br>
So i followed <a href="https://www.youtube.com/watch?v=OcIx_TBu90Q" rel="noopener nofollow ugc">a tutorial</a> on how to implement such methods with the gym cartpole environment.</p>
<p>After finishing the tutorial i’m now trying to change the environment from cartpole to a chess environment.<br>
I’m using <a href="https://github.com/genyrosk/gym-chess" rel="noopener nofollow ugc">this chess environment</a>.</p>
<p>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.</p>
<p>Here is the entire code:</p>
<pre><code class="lang-auto">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 &lt; 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]
</code></pre>
<p>So in the main i define my n_actions, which is the output dimension for the model.<br>
When doing it with cartpole, n_actions = 2, it works perfectly well, 0 problem.<br>
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 <code>pi = self.pi(pi1)</code> line (i’ve tried putting a print before and after this line and the second print never happens).</p>
<p>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.</p>
<p>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.</p>
<p>If that matters, i’m on Fedora 29, python 3.7.7, and pytorch 1.8.1+cu102.<br>
(I tried on different computers and the same behaviour happens.)</p>
<p>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.</p>
<p>Any help greatly appreciated.</p>

AndrewSpano: <p>Did you manage to solve this? I am having the same issue.</p>