Project import generated by Copybara.

PiperOrigin-RevId: 621552931

dopamine/discrete_domains/run_experiment.py

@@ -35,6 +35,9 @@
 from dopamine.jax.agents.implicit_quantile import implicit_quantile_agent as jax_implicit_quantile_agent
 from dopamine.jax.agents.quantile import quantile_agent as jax_quantile_agent
 from dopamine.jax.agents.rainbow import rainbow_agent as jax_rainbow_agent
+from dopamine.labs.moes.agents import dqn_moe_agent
+from dopamine.labs.moes.agents import full_rainbow_moe_agent
+from dopamine.labs.moes.agents import rainbow_100k_moe_agent
 from dopamine.metrics import collector_dispatcher
 from dopamine.metrics import statistics_instance
 import gin.tf
@@ -120,6 +123,18 @@ def create_agent(
     return jax_implicit_quantile_agent.JaxImplicitQuantileAgent(
         num_actions=environment.action_space.n, summary_writer=summary_writer
     )
+  elif agent_name == 'moe_dqn':
+    return dqn_moe_agent.DQNMoEAgent(
+        num_actions=environment.action_space.n, summary_writer=summary_writer
+    )
+  elif agent_name == 'moe_full_rainbow':
+    return full_rainbow_moe_agent.JaxFullRainbowMoEAgent(
+        num_actions=environment.action_space.n, summary_writer=summary_writer
+    )
+  elif agent_name == 'moe_der':
+    return rainbow_100k_moe_agent.Atari100kRainbowMoEAgent(
+        num_actions=environment.action_space.n, summary_writer=summary_writer
+    )
   else:
     raise ValueError('Unknown agent: {}'.format(agent_name))
 

dopamine/labs/moes/README.md

@@ -0,0 +1,4 @@
+# Mixture of Experts (MoEs)
+
+This is the code accompanying the paper "Mixtures of Experts Unlock Parameter
+Scaling for Deep RL" (https://arxiv.org/abs/2402.08609).