Repository: ray-project/ray

Issue:
Qmix Bug with Truncated Episodes or when max_seq_len is set
<!--
General questions should be asked on the mailing list ray-dev@googlegroups.com.
Questions about how to use Ray should be asked on
[StackOverflow](https://stackoverflow.com/questions/tagged/ray).

Before submitting an issue, please fill out the following form.
-->

### System information
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**:  Ubuntu 18.04
- **Ray installed from (source or binary)**: source
- **Ray version**:  0.6.2
- **Python version**:  3.5.5
- **Exact command to reproduce**: (running on custom env)

<!--
You can obtain the Ray version with

python -c "import ray; print(ray.__version__)"
-->

### Describe the problem
<!-- Describe the problem clearly here. -->

Qmix is set up to add padded q-values of -9999999 to the end of every sequence. These q-values are not being correctly masked out. They are ignored in terminal states, due to line 108 in qmix_policy_graph.py, however they are being used in non-terminal states that happen to be at the end of a sequence. This is causing astronomical TD-errors on my custom env with 1 agent.

The issue is also a bit more complicated than just masking them out, as the code tries to do: There does need to be a next state (observation) to complete the bootstrapped return. If it is simply masked out, then certain states will never receive a back propagated loss.

I would suggest also passing in next_states to the loss function (new_obs in samples dictionary I believe). If terminal states are not recorded with a next observation, then one can be added as padding, since it will be 0-ed out by (1 - terminated) in line 108 anyway.

This issue can be verified by truncating episodes and noting that reward + self.gamma*-9999999 is contained in the variable masked_td_error.

### Source code / logs
<!-- Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached. Try to provide a reproducible test case that is the bare minimum necessary to generate the problem. -->



Discussion:
To clarify, this only happens if batch mode is set to truncate episodes?
I assume it would also be an issue with complete episodes but with a max_seq_len less than the length of the episode. It seems to be an issue if there is a non-terminal state at the end of a sequence.
I have attached code that I believe reproduces the same bug. The commands are:

python3 twostep_game_no_bug.py --run QMIX

python3 twostep_game_bug.py --run QMIX

the variable masked_td_error seems to still contain the dummy constant in the bug version.

[Qmix bug.zip](https://github.com/ray-project/ray/files/3116258/Qmix.bug.zip)

Hm I see. The proposed fix makes sense to me, and it also seems ok to potentially ignore the last state when truncating. Do you want to make a patch?
Sure. I can't do it this second, but should have some time next week.

Resolving patch:
diff --git a/python/ray/rllib/agents/qmix/qmix_policy_graph.py b/python/ray/rllib/agents/qmix/qmix_policy_graph.py
--- a/python/ray/rllib/agents/qmix/qmix_policy_graph.py
+++ b/python/ray/rllib/agents/qmix/qmix_policy_graph.py
@@ -46,16 +46,19 @@ def __init__(self,
         self.double_q = double_q
         self.gamma = gamma
 
-    def forward(self, rewards, actions, terminated, mask, obs, action_mask):
+    def forward(self, rewards, actions, terminated, mask, obs, next_obs,
+                action_mask, next_action_mask):
         """Forward pass of the loss.
 
         Arguments:
-            rewards: Tensor of shape [B, T-1, n_agents]
-            actions: Tensor of shape [B, T-1, n_agents]
-            terminated: Tensor of shape [B, T-1, n_agents]
-            mask: Tensor of shape [B, T-1, n_agents]
+            rewards: Tensor of shape [B, T, n_agents]
+            actions: Tensor of shape [B, T, n_agents]
+            terminated: Tensor of shape [B, T, n_agents]
+            mask: Tensor of shape [B, T, n_agents]
             obs: Tensor of shape [B, T, n_agents, obs_size]
+            next_obs: Tensor of shape [B, T, n_agents, obs_size]
             action_mask: Tensor of shape [B, T, n_agents, n_actions]
+            next_action_mask: Tensor of shape [B, T, n_agents, n_actions]
         """
 
         B, T = obs.size(0), obs.size(1)
@@ -68,9 +71,9 @@ def forward(self, rewards, actions, terminated, mask, obs, action_mask):
             mac_out.append(q)
         mac_out = th.stack(mac_out, dim=1)  # Concat over time
 
-        # Pick the Q-Values for the actions taken -> [B * n_agents, T-1]
+        # Pick the Q-Values for the actions taken -> [B * n_agents, T]
         chosen_action_qvals = th.gather(
-            mac_out[:, :-1], dim=3, index=actions.unsqueeze(3)).squeeze(3)
+            mac_out, dim=3, index=actions.unsqueeze(3)).squeeze(3)
 
         # Calculate the Q-Values necessary for the target
         target_mac_out = []
@@ -79,32 +82,37 @@ def forward(self, rewards, actions, terminated, mask, obs, action_mask):
             for s in self.target_model.state_init()
         ]
         for t in range(T):
-            target_q, target_h = _mac(self.target_model, obs[:, t], target_h)
+            target_q, target_h = _mac(self.target_model, next_obs[:, t],
+                                      target_h)
             target_mac_out.append(target_q)
-
-        # We don't need the first timesteps Q-Value estimate for targets
-        target_mac_out = th.stack(
-            target_mac_out[1:], dim=1)  # Concat across time
+        target_mac_out = th.stack(target_mac_out, dim=1)  # Concat across time
 
         # Mask out unavailable actions
-        target_mac_out[action_mask[:, 1:] == 0] = -9999999
+        ignore_action = (next_action_mask == 0) & (mask == 1).unsqueeze(-1)
+        target_mac_out[ignore_action] = -np.inf
 
         # Max over target Q-Values
         if self.double_q:
             # Get actions that maximise live Q (for double q-learning)
-            mac_out[action_mask == 0] = -9999999
-            cur_max_actions = mac_out[:, 1:].max(dim=3, keepdim=True)[1]
+            ignore_action = (action_mask == 0) & (mask == 1).unsqueeze(-1)
+            mac_out = mac_out.clone()  # issue 4742
+            mac_out[ignore_action] = -np.inf
+            cur_max_actions = mac_out.max(dim=3, keepdim=True)[1]
             target_max_qvals = th.gather(target_mac_out, 3,
                                          cur_max_actions).squeeze(3)
         else:
             target_max_qvals = target_mac_out.max(dim=3)[0]
 
+        assert target_max_qvals.min().item() != -np.inf, \
+            "target_max_qvals contains a masked action; \
+            there may be a state with no valid actions."
+
         # Mix
         if self.mixer is not None:
             # TODO(ekl) add support for handling global state? This is just
             # treating the stacked agent obs as the state.
-            chosen_action_qvals = self.mixer(chosen_action_qvals, obs[:, :-1])
-            target_max_qvals = self.target_mixer(target_max_qvals, obs[:, 1:])
+            chosen_action_qvals = self.mixer(chosen_action_qvals, obs)
+            target_max_qvals = self.target_mixer(target_max_qvals, next_obs)
 
         # Calculate 1-step Q-Learning targets
         targets = rewards + self.gamma * (1 - terminated) * target_max_qvals
@@ -239,48 +247,53 @@ def compute_actions(self,
     def learn_on_batch(self, samples):
         obs_batch, action_mask = self._unpack_observation(
             samples[SampleBatch.CUR_OBS])
+        next_obs_batch, next_action_mask = self._unpack_observation(
+            samples[SampleBatch.NEXT_OBS])
         group_rewards = self._get_group_rewards(samples[SampleBatch.INFOS])
 
         # These will be padded to shape [B * T, ...]
-        [rew, action_mask, act, dones, obs], initial_states, seq_lens = \
+        [rew, action_mask, next_action_mask, act, dones, obs, next_obs], \
+            initial_states, seq_lens = \
             chop_into_sequences(
                 samples[SampleBatch.EPS_ID],
                 samples[SampleBatch.UNROLL_ID],
                 samples[SampleBatch.AGENT_INDEX], [
-                    group_rewards, action_mask, samples[SampleBatch.ACTIONS],
-                    samples[SampleBatch.DONES], obs_batch
+                    group_rewards, action_mask, next_action_mask,
+                    samples[SampleBatch.ACTIONS], samples[SampleBatch.DONES],
+                    obs_batch, next_obs_batch
                 ],
                 [samples["state_in_{}".format(k)]
                  for k in range(len(self.get_initial_state()))],
                 max_seq_len=self.config["model"]["max_seq_len"],
-                dynamic_max=True,
-                _extra_padding=1)
-        # TODO(ekl) adding 1 extra unit of padding here, since otherwise we
-        # lose the terminating reward and the Q-values will be unanchored!
-        B, T = len(seq_lens), max(seq_lens) + 1
+                dynamic_max=True)
+        B, T = len(seq_lens), max(seq_lens)
 
         def to_batches(arr):
             new_shape = [B, T] + list(arr.shape[1:])
             return th.from_numpy(np.reshape(arr, new_shape))
 
-        rewards = to_batches(rew)[:, :-1].float()
-        actions = to_batches(act)[:, :-1].long()
+        rewards = to_batches(rew).float()
+        actions = to_batches(act).long()
         obs = to_batches(obs).reshape([B, T, self.n_agents,
                                        self.obs_size]).float()
         action_mask = to_batches(action_mask)
+        next_obs = to_batches(next_obs).reshape(
+            [B, T, self.n_agents, self.obs_size]).float()
+        next_action_mask = to_batches(next_action_mask)
 
         # TODO(ekl) this treats group termination as individual termination
         terminated = to_batches(dones.astype(np.float32)).unsqueeze(2).expand(
-            B, T, self.n_agents)[:, :-1]
+            B, T, self.n_agents)
+
+        # Create mask for where index is < unpadded sequence length
         filled = (np.reshape(np.tile(np.arange(T), B), [B, T]) <
                   np.expand_dims(seq_lens, 1)).astype(np.float32)
-        mask = th.from_numpy(filled).unsqueeze(2).expand(B, T,
-                                                         self.n_agents)[:, :-1]
-        mask[:, 1:] = mask[:, 1:] * (1 - terminated[:, :-1])
+        mask = th.from_numpy(filled).unsqueeze(2).expand(B, T, self.n_agents)
 
         # Compute loss
         loss_out, mask, masked_td_error, chosen_action_qvals, targets = \
-            self.loss(rewards, actions, terminated, mask, obs, action_mask)
+            self.loss(rewards, actions, terminated, mask, obs,
+                      next_obs, action_mask, next_action_mask)
 
         # Optimise
         self.optimiser.zero_grad()


Test patch: