# Buggy code
def learn(self, batch, batch_size=None, repeat=1, **kwargs):
    self._batch = batch_size
    r = batch.returns
    if self._rew_norm and r.std() > self.__eps:
        batch.returns = (r - r.mean()) / r.std()
    losses, actor_losses, vf_losses, ent_losses = [], [], [], []
    for _ in range(repeat):
        for b in batch.split(batch_size):
            self.optim.zero_grad()
            dist = self(b).dist
            v = self.critic(b.obs)
            a = torch.tensor(b.act, device=v.device)
            r = torch.tensor(b.returns, device=v.device)
            a_loss = -(dist.log_prob(a) * (r - v).detach()).mean()
            vf_loss = F.mse_loss(r[:, None], v)
            ent_loss = dist.entropy().mean()
            loss = a_loss + self._w_vf * vf_loss - self._w_ent * ent_loss
            loss.backward()
            if self._grad_norm:
                nn.utils.clip_grad_norm_(
                    self.model.parameters(), max_norm=self._grad_norm
                )
            self.optim.step()
            actor_losses.append(a_loss.item())
            vf_losses.append(vf_loss.item())
            ent_losses.append(ent_loss.item())
            losses.append(loss.item())
    return {
        "loss": losses,
        "loss/actor": actor_losses,
        "loss/vf": vf_losses,
        "loss/ent": ent_losses,
    }


# Traceback (AttributeError)
python test/discrete/test_a2c_with_il.py --max-grad-norm 1
Epoch #1:   0%|                                                                                                                                                                                                      | 0/1000 [00:00<?, ?it/s]
Traceback (most recent call last):
File "test/discrete/test_a2c_with_il.py", line 137, in <module>
test_a2c()
File "test/discrete/test_a2c_with_il.py", line 98, in test_a2c
writer=writer)
File "/research/dept6/zlhe/miniconda3/envs/fp/lib/python3.7/site-packages/tianshou/trainer/onpolicy.py", line 89, in onpolicy_trainer
train_collector.sample(0), batch_size, repeat_per_collect)
File "/research/dept6/zlhe/miniconda3/envs/fp/lib/python3.7/site-packages/tianshou/policy/modelfree/a2c.py", line 101, in learn
self.model.parameters(), max_norm=self._grad_norm)
AttributeError: 'NoneType' object has no attribute 'parameters'

# Fixed code
def learn(self, batch, batch_size=None, repeat=1, **kwargs):
    self._batch = batch_size
    r = batch.returns
    if self._rew_norm and r.std() > self.__eps:
        batch.returns = (r - r.mean()) / r.std()
    losses, actor_losses, vf_losses, ent_losses = [], [], [], []
    for _ in range(repeat):
        for b in batch.split(batch_size):
            self.optim.zero_grad()
            dist = self(b).dist
            v = self.critic(b.obs)
            a = torch.tensor(b.act, device=v.device)
            r = torch.tensor(b.returns, device=v.device)
            a_loss = -(dist.log_prob(a) * (r - v).detach()).mean()
            vf_loss = F.mse_loss(r[:, None], v)
            ent_loss = dist.entropy().mean()
            loss = a_loss + self._w_vf * vf_loss - self._w_ent * ent_loss
            loss.backward()
            if self._grad_norm:
                nn.utils.clip_grad_norm_(
                    list(self.actor.parameters()) + list(self.critic.parameters()),
                    max_norm=self._grad_norm,
                )
            self.optim.step()
            actor_losses.append(a_loss.item())
            vf_losses.append(vf_loss.item())
            ent_losses.append(ent_loss.item())
            losses.append(loss.item())
    return {
        "loss": losses,
        "loss/actor": actor_losses,
        "loss/vf": vf_losses,
        "loss/ent": ent_losses,
    }