Migrate the ACME example agent to JAX-based ACME APIs.

examples/run_acme_agent.py

@@ -13,21 +13,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Acme DQN agent interacting with AndroidEnv."""
+"""Acme JAX DQN agent interacting with AndroidEnv."""
 
 from absl import app
 from absl import flags
 from absl import logging
-import acme
-from acme import specs
 from acme import wrappers as acme_wrappers
-from acme.agents.tf import dqn
-from acme.tf import networks
+from acme.agents.jax import dqn
+from acme.jax import experiments
+from acme.jax import networks as acme_jax_networks
+from acme.jax import utils as acme_jax_utils
 from android_env import loader
 from android_env.components import config_classes
 from android_env.wrappers import discrete_action_wrapper
+from android_env.wrappers import flat_interface_wrapper
 from android_env.wrappers import float_pixels_wrapper
 from android_env.wrappers import image_rescale_wrapper
+import haiku as hk
 
 # Simulator args
 flags.DEFINE_string('avd_name', None, 'Name of AVD to use.')
@@ -42,7 +44,7 @@
 flags.DEFINE_string('task_path', None, 'Path to task textproto file.')
 
 # Experiment args
-flags.DEFINE_integer('num_episodes', 100, 'Number of episodes.')
+flags.DEFINE_integer('num_steps', 1000, 'Number of steps to train.')
 
 FLAGS = flags.FLAGS
 
@@ -53,42 +55,80 @@ def apply_wrappers(env):
   env = image_rescale_wrapper.ImageRescaleWrapper(
       env, zoom_factors=(0.25, 0.25))
   env = float_pixels_wrapper.FloatPixelsWrapper(env)
+  env = flat_interface_wrapper.FlatInterfaceWrapper(
+      env, flat_actions=True, flat_observations=True
+  )
   env = acme_wrappers.SinglePrecisionWrapper(env)
   return env
 
 
+def make_network(environment_spec) -> dqn.DQNNetworks:
+  """Creates networks for training DQN."""
+  num_actions = environment_spec.actions.num_values
+  network_fn = acme_jax_networks.dqn_atari_network(num_actions)
+  network_hk = hk.without_apply_rng(hk.transform(network_fn))
+  obs = acme_jax_utils.add_batch_dim(
+      acme_jax_utils.zeros_like(environment_spec.observations)
+  )
+  network = acme_jax_networks.FeedForwardNetwork(
+      init=lambda rng: network_hk.init(rng, obs), apply=network_hk.apply
+  )
+  typed_network = acme_jax_networks.non_stochastic_network_to_typed(network)
+  return dqn.DQNNetworks(policy_network=typed_network)
+
+
 def main(_):
 
-  config = config_classes.AndroidEnvConfig(
-      task=config_classes.FilesystemTaskConfig(path=FLAGS.task_path),
-      simulator=config_classes.EmulatorConfig(
-          emulator_launcher=config_classes.EmulatorLauncherConfig(
-              emulator_path=FLAGS.emulator_path,
-              android_sdk_root=FLAGS.android_sdk_root,
-              android_avd_home=FLAGS.android_avd_home,
-              avd_name=FLAGS.avd_name,
-              run_headless=FLAGS.run_headless,
-          ),
-          adb_controller=config_classes.AdbControllerConfig(
-              adb_path=FLAGS.adb_path
-          ),
-      ),
+  def env_factory(seed):
+    del seed
+    config = config_classes.AndroidEnvConfig(
+        task=config_classes.FilesystemTaskConfig(path=FLAGS.task_path),
+        simulator=config_classes.EmulatorConfig(
+            emulator_launcher=config_classes.EmulatorLauncherConfig(
+                emulator_path=FLAGS.emulator_path,
+                android_sdk_root=FLAGS.android_sdk_root,
+                android_avd_home=FLAGS.android_avd_home,
+                avd_name=FLAGS.avd_name,
+                run_headless=FLAGS.run_headless,
+            ),
+            adb_controller=config_classes.AdbControllerConfig(
+                adb_path=FLAGS.adb_path
+            ),
+        ),
+    )
+    env = loader.load(config)
+    env = apply_wrappers(env)
+    return env
+
+  # Construct the agent config.
+  agent_config = dqn.DQNConfig(
+      discount=0.99,
+      eval_epsilon=0.0,
+      learning_rate=5e-5,
+      n_step=1,
+      epsilon=0.01,
+      target_update_period=2000,
+      min_replay_size=10,
+      max_replay_size=1000,
+      samples_per_insert=2.0,
+      batch_size=10,
+  )
+
+  loss_fn = dqn.PrioritizedDoubleQLearning(
+      discount=agent_config.discount, max_abs_reward=1.0
   )
-  with loader.load(config) as env:
 
-    env = apply_wrappers(env)
-    env_spec = specs.make_environment_spec(env)
-
-    agent = dqn.DQN(
-        environment_spec=env_spec,
-        network=networks.DQNAtariNetwork(
-            num_actions=env_spec.actions.num_values),
-        batch_size=10,
-        samples_per_insert=2,
-        min_replay_size=10)
-
-    loop = acme.EnvironmentLoop(env, agent)
-    loop.run(num_episodes=FLAGS.num_episodes)
+  dqn_builder = dqn.DQNBuilder(agent_config, loss_fn=loss_fn)
+
+  experiment_config = experiments.ExperimentConfig(
+      builder=dqn_builder,
+      environment_factory=env_factory,
+      network_factory=make_network,
+      seed=1,
+      max_num_actor_steps=FLAGS.num_steps,
+  )
+
+  experiments.run_experiment(experiment_config)
 
 
 if __name__ == '__main__':