prompt: 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()
  
  ```
  
golden_diff: 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,
  
verification_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"}]}