prompt: You will be provided with a partial code base and an issue statement explaining a problem to resolve.
  <issue>
  PPO example does not terminate properly
  ### Configuration
  
  Running the PPO example for a short number of frames  in order to reproduce as fast as possible on a cloud VM with a V100 GPU. Config python3.7, flax 0.2.2, jax 0.2.1, jaxlib 0.1.55 .
  
  Command run:
  `python ppo_main.py --config.game=Qbert --config.total_frames=4000`
  
  ### Problem you have encountered:
  
  Program does not exit. One can `print('Done')` after `ppo_lib.train` in `ppo_main` but there is an open thread and program can't exit (even after adding `raise SystemExit`).
  
  ### Extra comments
  
  Added extra line in `main`  ` tf.config.experimental.set_visible_devices([],'GPU')` in order for the program to run properly with `tensorflow-gpu`, this is common in other `flax/examples`. 
  
  </issue>
  <code>
  [start of examples/ppo/ppo_main.py]
  1 # Copyright 2020 The Flax Authors.
  2 #
  3 # Licensed under the Apache License, Version 2.0 (the "License");
  4 # you may not use this file except in compliance with the License.
  5 # You may obtain a copy of the License at
  6 #
  7 #     http://www.apache.org/licenses/LICENSE-2.0
  8 #
  9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS,
  11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 # See the License for the specific language governing permissions and
  13 # limitations under the License.
  14 
  15 import os
  16 from absl import flags
  17 from absl import app
  18 import jax
  19 import jax.random
  20 from ml_collections import config_flags
  21 
  22 import ppo_lib
  23 import models
  24 import env_utils
  25 
  26 FLAGS = flags.FLAGS
  27 
  28 flags.DEFINE_string(
  29     'logdir', default='/tmp/ppo_training',
  30     help=('Directory to save checkpoints and logging info.'))
  31 
  32 config_flags.DEFINE_config_file(
  33     'config', os.path.join(os.path.dirname(__file__), 'default_config.py'),
  34     'File path to the default configuration file.')
  35 
  36 def main(argv):
  37   config = FLAGS.config
  38   game = config.game + 'NoFrameskip-v4'
  39   num_actions = env_utils.get_num_actions(game)
  40   print(f'Playing {game} with {num_actions} actions')
  41   key = jax.random.PRNGKey(0)
  42   key, subkey = jax.random.split(key)
  43   model = models.create_model(subkey, num_outputs=num_actions)
  44   optimizer = models.create_optimizer(model, learning_rate=config.learning_rate)
  45   del model
  46   optimizer = ppo_lib.train(optimizer, config, FLAGS.logdir)
  47 
  48 if __name__ == '__main__':
  49   app.run(main)
  50
  [end of examples/ppo/ppo_main.py]
  [start of examples/ppo/agent.py]
  1 # Copyright 2020 The Flax Authors.
  2 #
  3 # Licensed under the Apache License, Version 2.0 (the "License");
  4 # you may not use this file except in compliance with the License.
  5 # You may obtain a copy of the License at
  6 #
  7 #     http://www.apache.org/licenses/LICENSE-2.0
  8 #
  9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS,
  11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 # See the License for the specific language governing permissions and
  13 # limitations under the License.
  14 
  15 """Agent utilities, incl. choosing the move and running in separate process."""
  16 
  17 import multiprocessing
  18 import collections
  19 import jax
  20 import numpy as onp
  21 
  22 import env_utils
  23 
  24 @jax.jit
  25 def policy_action(model, state):
  26   """Forward pass of the network."""
  27   out = model(state)
  28   return out
  29 
  30 
  31 ExpTuple = collections.namedtuple(
  32     'ExpTuple', ['state', 'action', 'reward', 'value', 'log_prob', 'done'])
  33 
  34 
  35 class RemoteSimulator:
  36   """Wrap functionality for an agent emulating Atari in a separate process.
  37 
  38   An object of this class is created for every agent.
  39   """
  40 
  41   def __init__(self, game: str):
  42     """Start the remote process and create Pipe() to communicate with it."""
  43     parent_conn, child_conn = multiprocessing.Pipe()
  44     self.proc = multiprocessing.Process(
  45         target=rcv_action_send_exp, args=(child_conn, game))
  46     self.conn = parent_conn
  47     self.proc.start()
  48 
  49 
  50 def rcv_action_send_exp(conn, game: str):
  51   """Run the remote agents.
  52 
  53   Receive action from the main learner, perform one step of simulation and
  54   send back collected experience.
  55   """
  56   env = env_utils.create_env(game, clip_rewards=True)
  57   while True:
  58     obs = env.reset()
  59     done = False
  60     # Observations fetched from Atari env need additional batch dimension.
  61     state = obs[None, ...]
  62     while not done:
  63       conn.send(state)
  64       action = conn.recv()
  65       obs, reward, done, _ = env.step(action)
  66       next_state = obs[None, ...] if not done else None
  67       experience = (state, action, reward, done)
  68       conn.send(experience)
  69       if done:
  70         break
  71       state = next_state
  72
  [end of examples/ppo/agent.py]
  </code>
  
  I need you to solve this issue by generating a single patch file that I can apply directly to this repository using git apply. Please respond with a single patch file in the following format.
  <patch>
  diff --git a/file.py b/file.py
  --- a/file.py
  +++ b/file.py
  @@ -1,27 +1,35 @@
   def euclidean(a, b):
  -    while b:
  -        a, b = b, a % b
  -    return a
  +    if b == 0:
  +        return a
  +    return euclidean(b, a % b)
   
   
   def bresenham(x0, y0, x1, y1):
       points = []
       dx = abs(x1 - x0)
       dy = abs(y1 - y0)
  -    sx = 1 if x0 < x1 else -1
  -    sy = 1 if y0 < y1 else -1
  -    err = dx - dy
  +    x, y = x0, y0
  +    sx = -1 if x0 > x1 else 1
  +    sy = -1 if y0 > y1 else 1
   
  -    while True:
  -        points.append((x0, y0))
  -        if x0 == x1 and y0 == y1:
  -            break
  -        e2 = 2 * err
  -        if e2 > -dy:
  -            err -= dy
  -            x0 += sx
  -        if e2 < dx:
  -            err += dx
  -            y0 += sy
  +    if dx > dy:
  +        err = dx / 2.0
  +        while x != x1:
  +            points.append((x, y))
  +            err -= dy
  +            if err < 0:
  +                y += sy
  +                err += dx
  +            x += sx
  +    else:
  +        err = dy / 2.0
  +        while y != y1:
  +            points.append((x, y))
  +            err -= dx
  +            if err < 0:
  +                x += sx
  +                err += dy
  +            y += sy
  +
  +    points.append((x, y))
       return points
  
  </patch>
golden_diff: diff --git a/examples/ppo/agent.py b/examples/ppo/agent.py
  --- a/examples/ppo/agent.py
  +++ b/examples/ppo/agent.py
  @@ -43,6 +43,7 @@
       parent_conn, child_conn = multiprocessing.Pipe()
       self.proc = multiprocessing.Process(
           target=rcv_action_send_exp, args=(child_conn, game))
  +    self.proc.daemon = True
       self.conn = parent_conn
       self.proc.start()
   
  diff --git a/examples/ppo/ppo_main.py b/examples/ppo/ppo_main.py
  --- a/examples/ppo/ppo_main.py
  +++ b/examples/ppo/ppo_main.py
  @@ -19,6 +19,8 @@
   import jax.random
   from ml_collections import config_flags
   
  +import tensorflow as tf
  +
   import ppo_lib
   import models
   import env_utils
  @@ -34,6 +36,9 @@
       'File path to the default configuration file.')
   
   def main(argv):
  +  # Make sure tf does not allocate gpu memory.
  +  tf.config.experimental.set_visible_devices([], 'GPU')
  +
     config = FLAGS.config
     game = config.game + 'NoFrameskip-v4'
     num_actions = env_utils.get_num_actions(game)
  
verification_info: {"golden_diff": "diff --git a/examples/ppo/agent.py b/examples/ppo/agent.py\n--- a/examples/ppo/agent.py\n+++ b/examples/ppo/agent.py\n@@ -43,6 +43,7 @@\n     parent_conn, child_conn = multiprocessing.Pipe()\n     self.proc = multiprocessing.Process(\n         target=rcv_action_send_exp, args=(child_conn, game))\n+    self.proc.daemon = True\n     self.conn = parent_conn\n     self.proc.start()\n \ndiff --git a/examples/ppo/ppo_main.py b/examples/ppo/ppo_main.py\n--- a/examples/ppo/ppo_main.py\n+++ b/examples/ppo/ppo_main.py\n@@ -19,6 +19,8 @@\n import jax.random\n from ml_collections import config_flags\n \n+import tensorflow as tf\n+\n import ppo_lib\n import models\n import env_utils\n@@ -34,6 +36,9 @@\n     'File path to the default configuration file.')\n \n def main(argv):\n+  # Make sure tf does not allocate gpu memory.\n+  tf.config.experimental.set_visible_devices([], 'GPU')\n+\n   config = FLAGS.config\n   game = config.game + 'NoFrameskip-v4'\n   num_actions = env_utils.get_num_actions(game)\n", "issue": "PPO example does not terminate properly\n### Configuration\r\n\r\nRunning the PPO example for a short number of frames  in order to reproduce as fast as possible on a cloud VM with a V100 GPU. Config python3.7, flax 0.2.2, jax 0.2.1, jaxlib 0.1.55 .\r\n\r\nCommand run:\r\n`python ppo_main.py --config.game=Qbert --config.total_frames=4000`\r\n\r\n### Problem you have encountered:\r\n\r\nProgram does not exit. One can `print('Done')` after `ppo_lib.train` in `ppo_main` but there is an open thread and program can't exit (even after adding `raise SystemExit`).\r\n\r\n### Extra comments\r\n\r\nAdded extra line in `main`  ` tf.config.experimental.set_visible_devices([],'GPU')` in order for the program to run properly with `tensorflow-gpu`, this is common in other `flax/examples`. \n", "before_files": [{"content": "# Copyright 2020 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nfrom absl import flags\nfrom absl import app\nimport jax\nimport jax.random\nfrom ml_collections import config_flags\n\nimport ppo_lib\nimport models\nimport env_utils\n\nFLAGS = flags.FLAGS\n\nflags.DEFINE_string(\n    'logdir', default='/tmp/ppo_training',\n    help=('Directory to save checkpoints and logging info.'))\n\nconfig_flags.DEFINE_config_file(\n    'config', os.path.join(os.path.dirname(__file__), 'default_config.py'),\n    'File path to the default configuration file.')\n\ndef main(argv):\n  config = FLAGS.config\n  game = config.game + 'NoFrameskip-v4'\n  num_actions = env_utils.get_num_actions(game)\n  print(f'Playing {game} with {num_actions} actions')\n  key = jax.random.PRNGKey(0)\n  key, subkey = jax.random.split(key)\n  model = models.create_model(subkey, num_outputs=num_actions)\n  optimizer = models.create_optimizer(model, learning_rate=config.learning_rate)\n  del model\n  optimizer = ppo_lib.train(optimizer, config, FLAGS.logdir)\n\nif __name__ == '__main__':\n  app.run(main)\n", "path": "examples/ppo/ppo_main.py"}, {"content": "# Copyright 2020 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Agent utilities, incl. choosing the move and running in separate process.\"\"\"\n\nimport multiprocessing\nimport collections\nimport jax\nimport numpy as onp\n\nimport env_utils\n\n@jax.jit\ndef policy_action(model, state):\n  \"\"\"Forward pass of the network.\"\"\"\n  out = model(state)\n  return out\n\n\nExpTuple = collections.namedtuple(\n    'ExpTuple', ['state', 'action', 'reward', 'value', 'log_prob', 'done'])\n\n\nclass RemoteSimulator:\n  \"\"\"Wrap functionality for an agent emulating Atari in a separate process.\n\n  An object of this class is created for every agent.\n  \"\"\"\n\n  def __init__(self, game: str):\n    \"\"\"Start the remote process and create Pipe() to communicate with it.\"\"\"\n    parent_conn, child_conn = multiprocessing.Pipe()\n    self.proc = multiprocessing.Process(\n        target=rcv_action_send_exp, args=(child_conn, game))\n    self.conn = parent_conn\n    self.proc.start()\n\n\ndef rcv_action_send_exp(conn, game: str):\n  \"\"\"Run the remote agents.\n\n  Receive action from the main learner, perform one step of simulation and\n  send back collected experience.\n  \"\"\"\n  env = env_utils.create_env(game, clip_rewards=True)\n  while True:\n    obs = env.reset()\n    done = False\n    # Observations fetched from Atari env need additional batch dimension.\n    state = obs[None, ...]\n    while not done:\n      conn.send(state)\n      action = conn.recv()\n      obs, reward, done, _ = env.step(action)\n      next_state = obs[None, ...] if not done else None\n      experience = (state, action, reward, done)\n      conn.send(experience)\n      if done:\n        break\n      state = next_state\n", "path": "examples/ppo/agent.py"}]}