When I use tianshou.utils.net.common.Recurrent as policy model, how to control the sequence length?

[GF-Huang]

We know the torch.nn.LSTM input shape is (batch_size, sequence_length, hidden_size), how to control the sequence_length without modify the obs (returned by custom env's env.step()) shape?

[Trinkle23897]

stack_num in replay buffer

[GF-Huang]

Hello, if I add stack_num=15, it will make the shape not match: RuntimeError: mat1 and mat2 shapes cannot be multiplied (390x30 and 2x3) (2 * 15 = 30). If no stack_num specified it works well.

Could you help me fix it?

Policy

PPOPolicy(
  (actor): Actor(
    (preprocess): Recurrent(
      (nn): LSTM(128, 128, batch_first=True)
      (fc1): Linear(in_features=2, out_features=128, bias=True)
      (fc2): Linear(in_features=128, out_features=3, bias=True)
    )
    (last): MLP(
      (model): Sequential(
        (0): Linear(in_features=3, out_features=3, bias=True)
      )
    )
  )
  (critic): Critic(
    (preprocess): Net(
      (model): MLP(
        (model): Sequential(
          (0): Linear(in_features=2, out_features=3, bias=True)
        )
      )
    )
    (last): MLP(
      (model): Sequential(
        (0): Linear(in_features=3, out_features=1, bias=True)
      )
    )
  )
  (_actor_critic): ActorCritic(
    (actor): Actor(
      (preprocess): Recurrent(
        (nn): LSTM(128, 128, batch_first=True)
        (fc1): Linear(in_features=2, out_features=128, bias=True)
        (fc2): Linear(in_features=128, out_features=3, bias=True)
      )
      (last): MLP(
        (model): Sequential(
          (0): Linear(in_features=3, out_features=3, bias=True)
        )
      )
    )
    (critic): Critic(
      (preprocess): Net(
        (model): MLP(
          (model): Sequential(
            (0): Linear(in_features=2, out_features=3, bias=True)
          )
        )
      )
      (last): MLP(
        (model): Sequential(
          (0): Linear(in_features=3, out_features=1, bias=True)
        )
      )
    )
  )
)

Exception

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
...
---> trainer.run()

File d:\Miniconda3\lib\site-packages\tianshou\trainer\base.py:439, in BaseTrainer.run(self)
    437 try:
    438     self.is_run = True
--> 439     deque(self, maxlen=0)  # feed the entire iterator into a zero-length deque
    440     info = gather_info(
    441         self.start_time, self.train_collector, self.test_collector,
    442         self.best_reward, self.best_reward_std
    443     )
    444 finally:

File d:\Miniconda3\lib\site-packages\tianshou\trainer\base.py:298, in BaseTrainer.__next__(self)
    295         result["n/st"] = int(self.gradient_step)
    296         t.update()
--> 298     self.policy_update_fn(data, result)
    299     t.set_postfix(**data)
    301 if t.n <= t.total and not self.stop_fn_flag:

File d:\Miniconda3\lib\site-packages\tianshou\trainer\onpolicy.py:131, in OnpolicyTrainer.policy_update_fn(self, data, result)
    129 """Perform one on-policy update."""
    130 assert self.train_collector is not None
--> 131 losses = self.policy.update(
    132     0,
    133     self.train_collector.buffer,
    134     batch_size=self.batch_size,
    135     repeat=self.repeat_per_collect,
    136 )
    137 self.train_collector.reset_buffer(keep_statistics=True)
    138 step = max([1] + [len(v) for v in losses.values() if isinstance(v, list)])

File d:\Miniconda3\lib\site-packages\tianshou\policy\base.py:276, in BasePolicy.update(self, sample_size, buffer, **kwargs)
    274 batch, indices = buffer.sample(sample_size)
    275 self.updating = True
--> 276 batch = self.process_fn(batch, buffer, indices)
    277 result = self.learn(batch, **kwargs)
    278 self.post_process_fn(batch, buffer, indices)

File d:\Miniconda3\lib\site-packages\tianshou\policy\modelfree\ppo.py:95, in PPOPolicy.process_fn(self, batch, buffer, indices)
     92 if self._recompute_adv:
     93     # buffer input `buffer` and `indices` to be used in `learn()`.
     94     self._buffer, self._indices = buffer, indices
---> 95 batch = self._compute_returns(batch, buffer, indices)
     96 batch.act = to_torch_as(batch.act, batch.v_s)
     97 old_log_prob = []

File d:\Miniconda3\lib\site-packages\tianshou\policy\modelfree\a2c.py:89, in A2CPolicy._compute_returns(self, batch, buffer, indices)
     87 with torch.no_grad():
     88     for minibatch in batch.split(self._batch, shuffle=False, merge_last=True):
---> 89         v_s.append(self.critic(minibatch.obs))
     90         v_s_.append(self.critic(minibatch.obs_next))
     91 batch.v_s = torch.cat(v_s, dim=0).flatten()  # old value

File d:\Miniconda3\lib\site-packages\torch\nn\modules\module.py:1130, in Module._call_impl(self, *input, **kwargs)
   1126 # If we don't have any hooks, we want to skip the rest of the logic in
   1127 # this function, and just call forward.
   1128 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1129         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1130     return forward_call(*input, **kwargs)
   1131 # Do not call functions when jit is used
   1132 full_backward_hooks, non_full_backward_hooks = [], []

File d:\Miniconda3\lib\site-packages\tianshou\utils\net\discrete.py:120, in Critic.forward(self, obs, **kwargs)
    116 def forward(
    117     self, obs: Union[np.ndarray, torch.Tensor], **kwargs: Any
    118 ) -> torch.Tensor:
    119     """Mapping: s -> V(s)."""
--> 120     logits, _ = self.preprocess(obs, state=kwargs.get("state", None))
    121     return self.last(logits)

File d:\Miniconda3\lib\site-packages\torch\nn\modules\module.py:1130, in Module._call_impl(self, *input, **kwargs)
   1126 # If we don't have any hooks, we want to skip the rest of the logic in
   1127 # this function, and just call forward.
   1128 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1129         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1130     return forward_call(*input, **kwargs)
   1131 # Do not call functions when jit is used
   1132 full_backward_hooks, non_full_backward_hooks = [], []

File d:\Miniconda3\lib\site-packages\tianshou\utils\net\common.py:210, in Net.forward(self, obs, state, info)
    203 def forward(
    204     self,
    205     obs: Union[np.ndarray, torch.Tensor],
    206     state: Any = None,
    207     info: Dict[str, Any] = {},
    208 ) -> Tuple[torch.Tensor, Any]:
    209     """Mapping: obs -> flatten (inside MLP)-> logits."""
--> 210     logits = self.model(obs)
    211     bsz = logits.shape[0]
    212     if self.use_dueling:  # Dueling DQN

File d:\Miniconda3\lib\site-packages\torch\nn\modules\module.py:1130, in Module._call_impl(self, *input, **kwargs)
   1126 # If we don't have any hooks, we want to skip the rest of the logic in
   1127 # this function, and just call forward.
   1128 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1129         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1130     return forward_call(*input, **kwargs)
   1131 # Do not call functions when jit is used
   1132 full_backward_hooks, non_full_backward_hooks = [], []

File d:\Miniconda3\lib\site-packages\tianshou\utils\net\common.py:109, in MLP.forward(self, obs)
    107 if self.flatten_input:
    108     obs = obs.flatten(1)
--> 109 return self.model(obs)

File d:\Miniconda3\lib\site-packages\torch\nn\modules\module.py:1130, in Module._call_impl(self, *input, **kwargs)
   1126 # If we don't have any hooks, we want to skip the rest of the logic in
   1127 # this function, and just call forward.
   1128 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1129         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1130     return forward_call(*input, **kwargs)
   1131 # Do not call functions when jit is used
   1132 full_backward_hooks, non_full_backward_hooks = [], []

File d:\Miniconda3\lib\site-packages\torch\nn\modules\container.py:139, in Sequential.forward(self, input)
    137 def forward(self, input):
    138     for module in self:
--> 139         input = module(input)
    140     return input

File d:\Miniconda3\lib\site-packages\torch\nn\modules\module.py:1130, in Module._call_impl(self, *input, **kwargs)
   1126 # If we don't have any hooks, we want to skip the rest of the logic in
   1127 # this function, and just call forward.
   1128 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1129         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1130     return forward_call(*input, **kwargs)
   1131 # Do not call functions when jit is used
   1132 full_backward_hooks, non_full_backward_hooks = [], []

File d:\Miniconda3\lib\site-packages\torch\nn\modules\linear.py:114, in Linear.forward(self, input)
    113 def forward(self, input: Tensor) -> Tensor:
--> 114     return F.linear(input, self.weight, self.bias)

RuntimeError: mat1 and mat2 shapes cannot be multiplied (390x30 and 2x3)

Additional

If I add stack_num=15 to test_ppo.py collector's replay buffer, it also raise similar exception. mat1 and mat2 shapes cannot be multiplied (256x60 and 4x64)

tianshou/test/discrete/test_ppo.py

Lines 115 to 118 in 99c99bb

	# collector
	train_collector = Collector(
	policy, train_envs, VectorReplayBuffer(args.buffer_size, len(train_envs))
	)

[GF-Huang]

When I use same preprocess_net for both actor and critic, issue solved, thanks.