Make td3 examples use run_agent method.
"""Example running TD3 on the control suite."""
from absl import flags
-import acme
from acme import specs
from acme.agents.jax import td3
from absl import app
import helpers
-from acme.utils import counting
-from acme.utils import experiment_utils
-import jax
+from acme.jax import runners
FLAGS = flags.FLAGS
flags.DEFINE_integer('num_steps', 1000000,
@@ -41,40 +38,23 @@ def main(_):
environment_spec = specs.make_environment_spec(environment)
agent_networks = td3.make_networks(environment_spec)
- # Construct the agent.
config = td3.TD3Config(num_sgd_steps_per_step=FLAGS.num_sgd_steps_per_step)
- agent = td3.TD3(
- environment_spec, agent_networks, config=config, seed=FLAGS.seed)
-
- # Create the environment loop used for training.
- train_logger = experiment_utils.make_experiment_logger(
- label='train', steps_key='train_steps')
- train_loop = acme.EnvironmentLoop(
- environment,
- agent,
- counter=counting.Counter(prefix='train'),
- logger=train_logger)
-
- # Create the evaluation actor and loop.
- eval_logger = experiment_utils.make_experiment_logger(
- label='eval', steps_key='eval_steps')
- eval_actor = agent.builder.make_actor(
- random_key=jax.random.PRNGKey(FLAGS.seed),
+ td3_builder = td3.TD3Builder(config)
policy_network = td3.get_default_behavior_policy(
- agent_networks, environment_spec.actions, sigma=0.),
- variable_source=agent)
- eval_env = helpers.make_environment(task=FLAGS.env_name)
- eval_loop = acme.EnvironmentLoop(
- eval_env,
- eval_actor,
- counter=counting.Counter(prefix='eval'),
- logger=eval_logger)
-
- assert FLAGS.num_steps % FLAGS.eval_every == 0
- for _ in range(FLAGS.num_steps // FLAGS.eval_every):
- eval_loop.run(num_episodes=5)
- train_loop.run(num_steps=FLAGS.eval_every)
- eval_loop.run(num_episodes=5)
+ networks=agent_networks,
+ action_specs=environment_spec.actions,
+ sigma=config.sigma)
+ eval_policy_network = td3.get_default_behavior_policy(
+ agent_networks, environment_spec.actions, sigma=0.)
+ runners.run_agent(
+ builder=td3_builder,
+ environment=environment,
+ num_steps=FLAGS.num_steps,
+ eval_every=FLAGS.eval_every,
+ seed=FLAGS.seed,
+ networks=agent_networks,
+ policy_network=policy_network,
+ eval_policy_network=eval_policy_network)
if __name__ == '__main__':
