We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
DD-PPO does not all reduce gradients
## 🐛 Bug

DD-PPO does not all reduce gradients during the backward call, because `reducer.prepare_for_backward` is not being called during training process.

The problem is in this line: https://github.com/facebookresearch/habitat-api/blob/v0.1.4/habitat_baselines/rl/ddppo/algo/ddppo.py#L96

```
class DecentralizedDistributedMixin:

    ...
    def before_backward(self, loss):
         # ...
         self.reducer.prepare_for_backward(..)


# Mixin goes second that way the PPO __init__ will still be called
class DDPPO(PPO, DecentralizedDistributedMixin):  
    # Here PPO and Mixin both have "before_backward" method, 
    # DDPPO will call PPO's not the Mixin's.
    pass
```

And here is a quick fix:
```
class DecentralizedDistributedMixin:
    ...


# Mixin goes second that way the PPO __init__ will still be called
class DDPPO(PPO, DecentralizedDistributedMixin):  

      # Move before_backward to DDPPO
      def before_backward(self, loss):
         # ...
         self.reducer.prepare_for_backward(..)
```


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `habitat_baselines/rl/ddppo/algo/ddppo.py`
Content:
```
1 #!/usr/bin/env python3
2 
3 # Copyright (c) Facebook, Inc. and its affiliates.
4 # This source code is licensed under the MIT license found in the
5 # LICENSE file in the root directory of this source tree.
6 
7 from typing import Tuple
8 
9 import torch
10 import torch.distributed as distrib
11 
12 from habitat_baselines.common.rollout_storage import RolloutStorage
13 from habitat_baselines.rl.ppo import PPO
14 
15 EPS_PPO = 1e-5
16 
17 
18 def distributed_mean_and_var(
19     values: torch.Tensor,
20 ) -> Tuple[torch.Tensor, torch.Tensor]:
21     r"""Computes the mean and variances of a tensor over multiple workers.
22 
23     This method is equivalent to first collecting all versions of values and
24     then computing the mean and variance locally over that
25 
26     :param values: (*,) shaped tensors to compute mean and variance over.  Assumed
27                         to be solely the workers local copy of this tensor,
28                         the resultant mean and variance will be computed
29                         over _all_ workers version of this tensor.
30     """
31     assert distrib.is_initialized(), "Distributed must be initialized"
32 
33     world_size = distrib.get_world_size()
34     mean = values.mean()
35     distrib.all_reduce(mean)
36     mean /= world_size
37 
38     sq_diff = (values - mean).pow(2).mean()
39     distrib.all_reduce(sq_diff)
40     var = sq_diff / world_size
41 
42     return mean, var
43 
44 
45 class DecentralizedDistributedMixin:
46     def _get_advantages_distributed(
47         self, rollouts: RolloutStorage
48     ) -> torch.Tensor:
49         advantages = rollouts.returns[:-1] - rollouts.value_preds[:-1]
50         if not self.use_normalized_advantage:
51             return advantages
52 
53         mean, var = distributed_mean_and_var(advantages)
54 
55         return (advantages - mean) / (var.sqrt() + EPS_PPO)
56 
57     def init_distributed(self, find_unused_params: bool = True) -> None:
58         r"""Initializes distributed training for the model
59 
60         1. Broadcasts the model weights from world_rank 0 to all other workers
61         2. Adds gradient hooks to the model
62 
63         :param find_unused_params: Whether or not to filter out unused parameters
64                                    before gradient reduction.  This *must* be True if
65                                    there are any parameters in the model that where unused in the
66                                    forward pass, otherwise the gradient reduction
67                                    will not work correctly.
68         """
69         # NB: Used to hide the hooks from the nn.Module,
70         # so they don't show up in the state_dict
71         class Guard:
72             def __init__(self, model, device):
73                 if torch.cuda.is_available():
74                     self.ddp = torch.nn.parallel.DistributedDataParallel(
75                         model, device_ids=[device], output_device=device
76                     )
77                 else:
78                     self.ddp = torch.nn.parallel.DistributedDataParallel(model)
79 
80         self._ddp_hooks = Guard(self.actor_critic, self.device)
81         self.get_advantages = self._get_advantages_distributed
82 
83         self.reducer = self._ddp_hooks.ddp.reducer
84         self.find_unused_params = find_unused_params
85 
86     def before_backward(self, loss):
87         super().before_backward(loss)
88 
89         if self.find_unused_params:
90             self.reducer.prepare_for_backward([loss])
91         else:
92             self.reducer.prepare_for_backward([])
93 
94 
95 # Mixin goes second that way the PPO __init__ will still be called
96 class DDPPO(PPO, DecentralizedDistributedMixin):
97     pass
98
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/habitat_baselines/rl/ddppo/algo/ddppo.py b/habitat_baselines/rl/ddppo/algo/ddppo.py
--- a/habitat_baselines/rl/ddppo/algo/ddppo.py
+++ b/habitat_baselines/rl/ddppo/algo/ddppo.py
@@ -92,6 +92,5 @@
             self.reducer.prepare_for_backward([])
 
 
-# Mixin goes second that way the PPO __init__ will still be called
-class DDPPO(PPO, DecentralizedDistributedMixin):
+class DDPPO(DecentralizedDistributedMixin, PPO):
     pass

{"golden_diff": "diff --git a/habitat_baselines/rl/ddppo/algo/ddppo.py b/habitat_baselines/rl/ddppo/algo/ddppo.py\n--- a/habitat_baselines/rl/ddppo/algo/ddppo.py\n+++ b/habitat_baselines/rl/ddppo/algo/ddppo.py\n@@ -92,6 +92,5 @@\n             self.reducer.prepare_for_backward([])\n \n \n-# Mixin goes second that way the PPO __init__ will still be called\n-class DDPPO(PPO, DecentralizedDistributedMixin):\n+class DDPPO(DecentralizedDistributedMixin, PPO):\n     pass\n", "issue": "DD-PPO does not all reduce gradients\n## \ud83d\udc1b Bug\r\n\r\nDD-PPO does not all reduce gradients during the backward call, because `reducer.prepare_for_backward` is not being called during training process.\r\n\r\nThe problem is in this line: https://github.com/facebookresearch/habitat-api/blob/v0.1.4/habitat_baselines/rl/ddppo/algo/ddppo.py#L96\r\n\r\n```\r\nclass DecentralizedDistributedMixin:\r\n\r\n    ...\r\n    def before_backward(self, loss):\r\n         # ...\r\n         self.reducer.prepare_for_backward(..)\r\n\r\n\r\n# Mixin goes second that way the PPO __init__ will still be called\r\nclass DDPPO(PPO, DecentralizedDistributedMixin):  \r\n    # Here PPO and Mixin both have \"before_backward\" method, \r\n    # DDPPO will call PPO's not the Mixin's.\r\n    pass\r\n```\r\n\r\nAnd here is a quick fix:\r\n```\r\nclass DecentralizedDistributedMixin:\r\n    ...\r\n\r\n\r\n# Mixin goes second that way the PPO __init__ will still be called\r\nclass DDPPO(PPO, DecentralizedDistributedMixin):  \r\n\r\n      # Move before_backward to DDPPO\r\n      def before_backward(self, loss):\r\n         # ...\r\n         self.reducer.prepare_for_backward(..)\r\n```\r\n\n", "before_files": [{"content": "#!/usr/bin/env python3\n\n# Copyright (c) Facebook, Inc. and its affiliates.\n# This source code is licensed under the MIT license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom typing import Tuple\n\nimport torch\nimport torch.distributed as distrib\n\nfrom habitat_baselines.common.rollout_storage import RolloutStorage\nfrom habitat_baselines.rl.ppo import PPO\n\nEPS_PPO = 1e-5\n\n\ndef distributed_mean_and_var(\n    values: torch.Tensor,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n    r\"\"\"Computes the mean and variances of a tensor over multiple workers.\n\n    This method is equivalent to first collecting all versions of values and\n    then computing the mean and variance locally over that\n\n    :param values: (*,) shaped tensors to compute mean and variance over.  Assumed\n                        to be solely the workers local copy of this tensor,\n                        the resultant mean and variance will be computed\n                        over _all_ workers version of this tensor.\n    \"\"\"\n    assert distrib.is_initialized(), \"Distributed must be initialized\"\n\n    world_size = distrib.get_world_size()\n    mean = values.mean()\n    distrib.all_reduce(mean)\n    mean /= world_size\n\n    sq_diff = (values - mean).pow(2).mean()\n    distrib.all_reduce(sq_diff)\n    var = sq_diff / world_size\n\n    return mean, var\n\n\nclass DecentralizedDistributedMixin:\n    def _get_advantages_distributed(\n        self, rollouts: RolloutStorage\n    ) -> torch.Tensor:\n        advantages = rollouts.returns[:-1] - rollouts.value_preds[:-1]\n        if not self.use_normalized_advantage:\n            return advantages\n\n        mean, var = distributed_mean_and_var(advantages)\n\n        return (advantages - mean) / (var.sqrt() + EPS_PPO)\n\n    def init_distributed(self, find_unused_params: bool = True) -> None:\n        r\"\"\"Initializes distributed training for the model\n\n        1. Broadcasts the model weights from world_rank 0 to all other workers\n        2. Adds gradient hooks to the model\n\n        :param find_unused_params: Whether or not to filter out unused parameters\n                                   before gradient reduction.  This *must* be True if\n                                   there are any parameters in the model that where unused in the\n                                   forward pass, otherwise the gradient reduction\n                                   will not work correctly.\n        \"\"\"\n        # NB: Used to hide the hooks from the nn.Module,\n        # so they don't show up in the state_dict\n        class Guard:\n            def __init__(self, model, device):\n                if torch.cuda.is_available():\n                    self.ddp = torch.nn.parallel.DistributedDataParallel(\n                        model, device_ids=[device], output_device=device\n                    )\n                else:\n                    self.ddp = torch.nn.parallel.DistributedDataParallel(model)\n\n        self._ddp_hooks = Guard(self.actor_critic, self.device)\n        self.get_advantages = self._get_advantages_distributed\n\n        self.reducer = self._ddp_hooks.ddp.reducer\n        self.find_unused_params = find_unused_params\n\n    def before_backward(self, loss):\n        super().before_backward(loss)\n\n        if self.find_unused_params:\n            self.reducer.prepare_for_backward([loss])\n        else:\n            self.reducer.prepare_for_backward([])\n\n\n# Mixin goes second that way the PPO __init__ will still be called\nclass DDPPO(PPO, DecentralizedDistributedMixin):\n    pass\n", "path": "habitat_baselines/rl/ddppo/algo/ddppo.py"}], "after_files": [{"content": "#!/usr/bin/env python3\n\n# Copyright (c) Facebook, Inc. and its affiliates.\n# This source code is licensed under the MIT license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom typing import Tuple\n\nimport torch\nimport torch.distributed as distrib\n\nfrom habitat_baselines.common.rollout_storage import RolloutStorage\nfrom habitat_baselines.rl.ppo import PPO\n\nEPS_PPO = 1e-5\n\n\ndef distributed_mean_and_var(\n    values: torch.Tensor,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n    r\"\"\"Computes the mean and variances of a tensor over multiple workers.\n\n    This method is equivalent to first collecting all versions of values and\n    then computing the mean and variance locally over that\n\n    :param values: (*,) shaped tensors to compute mean and variance over.  Assumed\n                        to be solely the workers local copy of this tensor,\n                        the resultant mean and variance will be computed\n                        over _all_ workers version of this tensor.\n    \"\"\"\n    assert distrib.is_initialized(), \"Distributed must be initialized\"\n\n    world_size = distrib.get_world_size()\n    mean = values.mean()\n    distrib.all_reduce(mean)\n    mean /= world_size\n\n    sq_diff = (values - mean).pow(2).mean()\n    distrib.all_reduce(sq_diff)\n    var = sq_diff / world_size\n\n    return mean, var\n\n\nclass DecentralizedDistributedMixin:\n    def _get_advantages_distributed(\n        self, rollouts: RolloutStorage\n    ) -> torch.Tensor:\n        advantages = rollouts.returns[:-1] - rollouts.value_preds[:-1]\n        if not self.use_normalized_advantage:\n            return advantages\n\n        mean, var = distributed_mean_and_var(advantages)\n\n        return (advantages - mean) / (var.sqrt() + EPS_PPO)\n\n    def init_distributed(self, find_unused_params: bool = True) -> None:\n        r\"\"\"Initializes distributed training for the model\n\n        1. Broadcasts the model weights from world_rank 0 to all other workers\n        2. Adds gradient hooks to the model\n\n        :param find_unused_params: Whether or not to filter out unused parameters\n                                   before gradient reduction.  This *must* be True if\n                                   there are any parameters in the model that where unused in the\n                                   forward pass, otherwise the gradient reduction\n                                   will not work correctly.\n        \"\"\"\n        # NB: Used to hide the hooks from the nn.Module,\n        # so they don't show up in the state_dict\n        class Guard:\n            def __init__(self, model, device):\n                if torch.cuda.is_available():\n                    self.ddp = torch.nn.parallel.DistributedDataParallel(\n                        model, device_ids=[device], output_device=device\n                    )\n                else:\n                    self.ddp = torch.nn.parallel.DistributedDataParallel(model)\n\n        self._ddp_hooks = Guard(self.actor_critic, self.device)\n        self.get_advantages = self._get_advantages_distributed\n\n        self.reducer = self._ddp_hooks.ddp.reducer\n        self.find_unused_params = find_unused_params\n\n    def before_backward(self, loss):\n        super().before_backward(loss)\n\n        if self.find_unused_params:\n            self.reducer.prepare_for_backward([loss])\n        else:\n            self.reducer.prepare_for_backward([])\n\n\nclass DDPPO(DecentralizedDistributedMixin, PPO):\n    pass\n", "path": "habitat_baselines/rl/ddppo/algo/ddppo.py"}]}