We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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`. 

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `examples/ppo/ppo_main.py`
Content:
```
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
```

Path: `examples/ppo/agent.py`
Content:
```
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 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/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)

{"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"}], "after_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 tensorflow as tf\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  # 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  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.proc.daemon = True\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"}]}