Merge pull request #435 from muupan/iqn-tuple-obs

Make IQN work with tuple observations
chainer · Apr 5, 2019 · e6e47d3 · e6e47d3
2 parents db8fa5c + 4d1f98f
commit e6e47d3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/chainerrl/agents/iqn.py b/chainerrl/agents/iqn.py
@@ -202,7 +202,7 @@ def _compute_target_values(self, exp_batch):
             chainer.Variable: (batch_size, N_prime).
         """
         batch_next_state = exp_batch['next_state']
-        batch_size = batch_next_state.shape[0]
+        batch_size = len(exp_batch['reward'])
         taus_tilde = self.xp.random.uniform(
             0, 1, size=(batch_size, self.quantile_thresholds_K)).astype('f')