We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
On policy algos stop learning midway
Avg return either drops dramatically or the run stops completely due to NaN errors. Could affect off policy as well.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `src/garage/envs/normalized_env.py`
Content:
```
1 """An environment wrapper that normalizes action, observation and reward."""
2 import akro
3 import numpy as np
4 
5 from garage import EnvStep, Wrapper
6 
7 
8 class NormalizedEnv(Wrapper):
9     """An environment wrapper for normalization.
10 
11     This wrapper normalizes action, and optionally observation and reward.
12 
13     Args:
14         env (Environment): An environment instance.
15         scale_reward (float): Scale of environment reward.
16         normalize_obs (bool): If True, normalize observation.
17         normalize_reward (bool): If True, normalize reward. scale_reward is
18             applied after normalization.
19         expected_action_scale (float): Assuming action falls in the range of
20             [-expected_action_scale, expected_action_scale] when normalize it.
21         flatten_obs (bool): Flatten observation if True.
22         obs_alpha (float): Update rate of moving average when estimating the
23             mean and variance of observations.
24         reward_alpha (float): Update rate of moving average when estimating the
25             mean and variance of rewards.
26 
27     """
28 
29     def __init__(
30         self,
31         env,
32         scale_reward=1.,
33         normalize_obs=False,
34         normalize_reward=False,
35         expected_action_scale=1.,
36         flatten_obs=True,
37         obs_alpha=0.001,
38         reward_alpha=0.001,
39     ):
40         super().__init__(env)
41 
42         self._scale_reward = scale_reward
43         self._normalize_obs = normalize_obs
44         self._normalize_reward = normalize_reward
45         self._expected_action_scale = expected_action_scale
46         self._flatten_obs = flatten_obs
47 
48         self._obs_alpha = obs_alpha
49         flat_obs_dim = self._env.observation_space.flat_dim
50         self._obs_mean = np.zeros(flat_obs_dim)
51         self._obs_var = np.ones(flat_obs_dim)
52 
53         self._reward_alpha = reward_alpha
54         self._reward_mean = 0.
55         self._reward_var = 1.
56 
57     def reset(self):
58         """Call reset on wrapped env.
59 
60         Returns:
61             numpy.ndarray: The first observation conforming to
62                 `observation_space`.
63             dict: The episode-level information.
64                 Note that this is not part of `env_info` provided in `step()`.
65                 It contains information of he entire episode， which could be
66                 needed to determine the first action (e.g. in the case of
67                 goal-conditioned or MTRL.)
68 
69         """
70         first_obs, episode_info = self._env.reset()
71         if self._normalize_obs:
72             return self._apply_normalize_obs(first_obs), episode_info
73         else:
74             return first_obs, episode_info
75 
76     def step(self, action):
77         """Call step on wrapped env.
78 
79         Args:
80             action (np.ndarray): An action provided by the agent.
81 
82         Returns:
83             EnvStep: The environment step resulting from the action.
84 
85         Raises:
86             RuntimeError: if `step()` is called after the environment has been
87                 constructed and `reset()` has not been called.
88 
89         """
90         if isinstance(self.action_space, akro.Box):
91             # rescale the action when the bounds are not inf
92             lb, ub = self.action_space.low, self.action_space.high
93             if np.all(lb != -np.inf) and np.all(ub != -np.inf):
94                 scaled_action = lb + (action + self._expected_action_scale) * (
95                     0.5 * (ub - lb) / self._expected_action_scale)
96                 scaled_action = np.clip(scaled_action, lb, ub)
97             else:
98                 scaled_action = action
99         else:
100             scaled_action = action
101 
102         es = self._env.step(scaled_action)
103         next_obs = es.observation
104         reward = es.reward
105 
106         if self._normalize_obs:
107             next_obs = self._apply_normalize_obs(next_obs)
108         if self._normalize_reward:
109             reward = self._apply_normalize_reward(reward)
110 
111         return EnvStep(env_spec=es.env_spec,
112                        action=es.action,
113                        reward=reward * self._scale_reward,
114                        observation=next_obs,
115                        env_info=es.env_info,
116                        step_type=es.step_type)
117 
118     def _update_obs_estimate(self, obs):
119         flat_obs = self._env.observation_space.flatten(obs)
120         self._obs_mean = (
121             1 - self._obs_alpha) * self._obs_mean + self._obs_alpha * flat_obs
122         self._obs_var = (
123             1 - self._obs_alpha) * self._obs_var + self._obs_alpha * np.square(
124                 flat_obs - self._obs_mean)
125 
126     def _update_reward_estimate(self, reward):
127         self._reward_mean = (1 - self._reward_alpha) * \
128             self._reward_mean + self._reward_alpha * reward
129         self._reward_var = (
130             1 - self._reward_alpha
131         ) * self._reward_var + self._reward_alpha * np.square(
132             reward - self._reward_mean)
133 
134     def _apply_normalize_obs(self, obs):
135         """Compute normalized observation.
136 
137         Args:
138             obs (np.ndarray): Observation.
139 
140         Returns:
141             np.ndarray: Normalized observation.
142 
143         """
144         self._update_obs_estimate(obs)
145         flat_obs = self._env.observation_space.flatten(obs)
146         normalized_obs = (flat_obs -
147                           self._obs_mean) / (np.sqrt(self._obs_var) + 1e-8)
148         if not self._flatten_obs:
149             normalized_obs = self._env.observation_space.unflatten(
150                 self._env.observation_space, normalized_obs)
151         return normalized_obs
152 
153     def _apply_normalize_reward(self, reward):
154         """Compute normalized reward.
155 
156         Args:
157             reward (float): Reward.
158 
159         Returns:
160             float: Normalized reward.
161 
162         """
163         self._update_reward_estimate(reward)
164         return reward / (np.sqrt(self._reward_var) + 1e-8)
165 
166 
167 normalize = NormalizedEnv
168
```
--- 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/src/garage/envs/normalized_env.py b/src/garage/envs/normalized_env.py
--- a/src/garage/envs/normalized_env.py
+++ b/src/garage/envs/normalized_env.py
@@ -109,7 +109,7 @@
             reward = self._apply_normalize_reward(reward)
 
         return EnvStep(env_spec=es.env_spec,
-                       action=es.action,
+                       action=action,
                        reward=reward * self._scale_reward,
                        observation=next_obs,
                        env_info=es.env_info,

{"golden_diff": "diff --git a/src/garage/envs/normalized_env.py b/src/garage/envs/normalized_env.py\n--- a/src/garage/envs/normalized_env.py\n+++ b/src/garage/envs/normalized_env.py\n@@ -109,7 +109,7 @@\n             reward = self._apply_normalize_reward(reward)\n \n         return EnvStep(env_spec=es.env_spec,\n-                       action=es.action,\n+                       action=action,\n                        reward=reward * self._scale_reward,\n                        observation=next_obs,\n                        env_info=es.env_info,\n", "issue": "On policy algos stop learning midway\nAvg return either drops dramatically or the run stops completely due to NaN errors. Could affect off policy as well.\n", "before_files": [{"content": "\"\"\"An environment wrapper that normalizes action, observation and reward.\"\"\"\nimport akro\nimport numpy as np\n\nfrom garage import EnvStep, Wrapper\n\n\nclass NormalizedEnv(Wrapper):\n    \"\"\"An environment wrapper for normalization.\n\n    This wrapper normalizes action, and optionally observation and reward.\n\n    Args:\n        env (Environment): An environment instance.\n        scale_reward (float): Scale of environment reward.\n        normalize_obs (bool): If True, normalize observation.\n        normalize_reward (bool): If True, normalize reward. scale_reward is\n            applied after normalization.\n        expected_action_scale (float): Assuming action falls in the range of\n            [-expected_action_scale, expected_action_scale] when normalize it.\n        flatten_obs (bool): Flatten observation if True.\n        obs_alpha (float): Update rate of moving average when estimating the\n            mean and variance of observations.\n        reward_alpha (float): Update rate of moving average when estimating the\n            mean and variance of rewards.\n\n    \"\"\"\n\n    def __init__(\n        self,\n        env,\n        scale_reward=1.,\n        normalize_obs=False,\n        normalize_reward=False,\n        expected_action_scale=1.,\n        flatten_obs=True,\n        obs_alpha=0.001,\n        reward_alpha=0.001,\n    ):\n        super().__init__(env)\n\n        self._scale_reward = scale_reward\n        self._normalize_obs = normalize_obs\n        self._normalize_reward = normalize_reward\n        self._expected_action_scale = expected_action_scale\n        self._flatten_obs = flatten_obs\n\n        self._obs_alpha = obs_alpha\n        flat_obs_dim = self._env.observation_space.flat_dim\n        self._obs_mean = np.zeros(flat_obs_dim)\n        self._obs_var = np.ones(flat_obs_dim)\n\n        self._reward_alpha = reward_alpha\n        self._reward_mean = 0.\n        self._reward_var = 1.\n\n    def reset(self):\n        \"\"\"Call reset on wrapped env.\n\n        Returns:\n            numpy.ndarray: The first observation conforming to\n                `observation_space`.\n            dict: The episode-level information.\n                Note that this is not part of `env_info` provided in `step()`.\n                It contains information of he entire episode\uff0c which could be\n                needed to determine the first action (e.g. in the case of\n                goal-conditioned or MTRL.)\n\n        \"\"\"\n        first_obs, episode_info = self._env.reset()\n        if self._normalize_obs:\n            return self._apply_normalize_obs(first_obs), episode_info\n        else:\n            return first_obs, episode_info\n\n    def step(self, action):\n        \"\"\"Call step on wrapped env.\n\n        Args:\n            action (np.ndarray): An action provided by the agent.\n\n        Returns:\n            EnvStep: The environment step resulting from the action.\n\n        Raises:\n            RuntimeError: if `step()` is called after the environment has been\n                constructed and `reset()` has not been called.\n\n        \"\"\"\n        if isinstance(self.action_space, akro.Box):\n            # rescale the action when the bounds are not inf\n            lb, ub = self.action_space.low, self.action_space.high\n            if np.all(lb != -np.inf) and np.all(ub != -np.inf):\n                scaled_action = lb + (action + self._expected_action_scale) * (\n                    0.5 * (ub - lb) / self._expected_action_scale)\n                scaled_action = np.clip(scaled_action, lb, ub)\n            else:\n                scaled_action = action\n        else:\n            scaled_action = action\n\n        es = self._env.step(scaled_action)\n        next_obs = es.observation\n        reward = es.reward\n\n        if self._normalize_obs:\n            next_obs = self._apply_normalize_obs(next_obs)\n        if self._normalize_reward:\n            reward = self._apply_normalize_reward(reward)\n\n        return EnvStep(env_spec=es.env_spec,\n                       action=es.action,\n                       reward=reward * self._scale_reward,\n                       observation=next_obs,\n                       env_info=es.env_info,\n                       step_type=es.step_type)\n\n    def _update_obs_estimate(self, obs):\n        flat_obs = self._env.observation_space.flatten(obs)\n        self._obs_mean = (\n            1 - self._obs_alpha) * self._obs_mean + self._obs_alpha * flat_obs\n        self._obs_var = (\n            1 - self._obs_alpha) * self._obs_var + self._obs_alpha * np.square(\n                flat_obs - self._obs_mean)\n\n    def _update_reward_estimate(self, reward):\n        self._reward_mean = (1 - self._reward_alpha) * \\\n            self._reward_mean + self._reward_alpha * reward\n        self._reward_var = (\n            1 - self._reward_alpha\n        ) * self._reward_var + self._reward_alpha * np.square(\n            reward - self._reward_mean)\n\n    def _apply_normalize_obs(self, obs):\n        \"\"\"Compute normalized observation.\n\n        Args:\n            obs (np.ndarray): Observation.\n\n        Returns:\n            np.ndarray: Normalized observation.\n\n        \"\"\"\n        self._update_obs_estimate(obs)\n        flat_obs = self._env.observation_space.flatten(obs)\n        normalized_obs = (flat_obs -\n                          self._obs_mean) / (np.sqrt(self._obs_var) + 1e-8)\n        if not self._flatten_obs:\n            normalized_obs = self._env.observation_space.unflatten(\n                self._env.observation_space, normalized_obs)\n        return normalized_obs\n\n    def _apply_normalize_reward(self, reward):\n        \"\"\"Compute normalized reward.\n\n        Args:\n            reward (float): Reward.\n\n        Returns:\n            float: Normalized reward.\n\n        \"\"\"\n        self._update_reward_estimate(reward)\n        return reward / (np.sqrt(self._reward_var) + 1e-8)\n\n\nnormalize = NormalizedEnv\n", "path": "src/garage/envs/normalized_env.py"}], "after_files": [{"content": "\"\"\"An environment wrapper that normalizes action, observation and reward.\"\"\"\nimport akro\nimport numpy as np\n\nfrom garage import EnvStep, Wrapper\n\n\nclass NormalizedEnv(Wrapper):\n    \"\"\"An environment wrapper for normalization.\n\n    This wrapper normalizes action, and optionally observation and reward.\n\n    Args:\n        env (Environment): An environment instance.\n        scale_reward (float): Scale of environment reward.\n        normalize_obs (bool): If True, normalize observation.\n        normalize_reward (bool): If True, normalize reward. scale_reward is\n            applied after normalization.\n        expected_action_scale (float): Assuming action falls in the range of\n            [-expected_action_scale, expected_action_scale] when normalize it.\n        flatten_obs (bool): Flatten observation if True.\n        obs_alpha (float): Update rate of moving average when estimating the\n            mean and variance of observations.\n        reward_alpha (float): Update rate of moving average when estimating the\n            mean and variance of rewards.\n\n    \"\"\"\n\n    def __init__(\n        self,\n        env,\n        scale_reward=1.,\n        normalize_obs=False,\n        normalize_reward=False,\n        expected_action_scale=1.,\n        flatten_obs=True,\n        obs_alpha=0.001,\n        reward_alpha=0.001,\n    ):\n        super().__init__(env)\n\n        self._scale_reward = scale_reward\n        self._normalize_obs = normalize_obs\n        self._normalize_reward = normalize_reward\n        self._expected_action_scale = expected_action_scale\n        self._flatten_obs = flatten_obs\n\n        self._obs_alpha = obs_alpha\n        flat_obs_dim = self._env.observation_space.flat_dim\n        self._obs_mean = np.zeros(flat_obs_dim)\n        self._obs_var = np.ones(flat_obs_dim)\n\n        self._reward_alpha = reward_alpha\n        self._reward_mean = 0.\n        self._reward_var = 1.\n\n    def reset(self):\n        \"\"\"Call reset on wrapped env.\n\n        Returns:\n            numpy.ndarray: The first observation conforming to\n                `observation_space`.\n            dict: The episode-level information.\n                Note that this is not part of `env_info` provided in `step()`.\n                It contains information of he entire episode\uff0c which could be\n                needed to determine the first action (e.g. in the case of\n                goal-conditioned or MTRL.)\n\n        \"\"\"\n        first_obs, episode_info = self._env.reset()\n        if self._normalize_obs:\n            return self._apply_normalize_obs(first_obs), episode_info\n        else:\n            return first_obs, episode_info\n\n    def step(self, action):\n        \"\"\"Call step on wrapped env.\n\n        Args:\n            action (np.ndarray): An action provided by the agent.\n\n        Returns:\n            EnvStep: The environment step resulting from the action.\n\n        Raises:\n            RuntimeError: if `step()` is called after the environment has been\n                constructed and `reset()` has not been called.\n\n        \"\"\"\n        if isinstance(self.action_space, akro.Box):\n            # rescale the action when the bounds are not inf\n            lb, ub = self.action_space.low, self.action_space.high\n            if np.all(lb != -np.inf) and np.all(ub != -np.inf):\n                scaled_action = lb + (action + self._expected_action_scale) * (\n                    0.5 * (ub - lb) / self._expected_action_scale)\n                scaled_action = np.clip(scaled_action, lb, ub)\n            else:\n                scaled_action = action\n        else:\n            scaled_action = action\n\n        es = self._env.step(scaled_action)\n        next_obs = es.observation\n        reward = es.reward\n\n        if self._normalize_obs:\n            next_obs = self._apply_normalize_obs(next_obs)\n        if self._normalize_reward:\n            reward = self._apply_normalize_reward(reward)\n\n        return EnvStep(env_spec=es.env_spec,\n                       action=action,\n                       reward=reward * self._scale_reward,\n                       observation=next_obs,\n                       env_info=es.env_info,\n                       step_type=es.step_type)\n\n    def _update_obs_estimate(self, obs):\n        flat_obs = self._env.observation_space.flatten(obs)\n        self._obs_mean = (\n            1 - self._obs_alpha) * self._obs_mean + self._obs_alpha * flat_obs\n        self._obs_var = (\n            1 - self._obs_alpha) * self._obs_var + self._obs_alpha * np.square(\n                flat_obs - self._obs_mean)\n\n    def _update_reward_estimate(self, reward):\n        self._reward_mean = (1 - self._reward_alpha) * \\\n            self._reward_mean + self._reward_alpha * reward\n        self._reward_var = (\n            1 - self._reward_alpha\n        ) * self._reward_var + self._reward_alpha * np.square(\n            reward - self._reward_mean)\n\n    def _apply_normalize_obs(self, obs):\n        \"\"\"Compute normalized observation.\n\n        Args:\n            obs (np.ndarray): Observation.\n\n        Returns:\n            np.ndarray: Normalized observation.\n\n        \"\"\"\n        self._update_obs_estimate(obs)\n        flat_obs = self._env.observation_space.flatten(obs)\n        normalized_obs = (flat_obs -\n                          self._obs_mean) / (np.sqrt(self._obs_var) + 1e-8)\n        if not self._flatten_obs:\n            normalized_obs = self._env.observation_space.unflatten(\n                self._env.observation_space, normalized_obs)\n        return normalized_obs\n\n    def _apply_normalize_reward(self, reward):\n        \"\"\"Compute normalized reward.\n\n        Args:\n            reward (float): Reward.\n\n        Returns:\n            float: Normalized reward.\n\n        \"\"\"\n        self._update_reward_estimate(reward)\n        return reward / (np.sqrt(self._reward_var) + 1e-8)\n\n\nnormalize = NormalizedEnv\n", "path": "src/garage/envs/normalized_env.py"}]}