def main(_: Any) -> None:
# Environment.
environment_factory = functools.partial(
debugging_utils.make_environment,
env_name=FLAGS.env_name,
action_space=FLAGS.action_space,
)
# Networks.
network_factory = lp_utils.partial_kwargs(
madqn.make_default_networks,
archecture_type=ArchitectureType.recurrent)
# Checkpointer appends "Checkpoints" to checkpoint_dir
checkpoint_dir = f"{FLAGS.base_dir}/{FLAGS.mava_id}"
# Log every [log_every] seconds.
log_every = 10
logger_factory = functools.partial(
logger_utils.make_logger,
directory=FLAGS.base_dir,
to_terminal=True,
to_tensorboard=True,
time_stamp=FLAGS.mava_id,
time_delta=log_every,
)
# Distributed program.
program = madqn.MADQN(
environment_factory=environment_factory,
network_factory=network_factory,
logger_factory=logger_factory,
num_executors=1,
exploration_scheduler_fn=LinearExplorationScheduler,
epsilon_min=0.05,
epsilon_decay=5e-4,
optimizer=snt.optimizers.Adam(learning_rate=1e-4),
checkpoint_subpath=checkpoint_dir,
trainer_fn=madqn.training.MADQNRecurrentTrainer,
executor_fn=madqn.execution.MADQNRecurrentExecutor,
batch_size=32,
).build()
# Ensure only trainer runs on gpu, while other processes run on cpu.
gpu_id = -1
env_vars = {"CUDA_VISIBLE_DEVICES": str(gpu_id)}
local_resources = {
"trainer": [],
"evaluator": PythonProcess(env=env_vars),
"executor": PythonProcess(env=env_vars),
}
# Launch.
lp.launch(
program,
lp.LaunchType.LOCAL_MULTI_PROCESSING,
terminal="current_terminal",
local_resources=local_resources,
)
def main(_: Any) -> None:
# environment
environment_factory = functools.partial(
smac_utils.make_environment, map_name=FLAGS.map_name
)
# Networks.
network_factory = lp_utils.partial_kwargs(
custom_recurrent_network,
q_networks_layer_sizes=[128, 128],
)
# Checkpointer appends "Checkpoints" to checkpoint_dir
checkpoint_dir = f"{FLAGS.base_dir}/{FLAGS.mava_id}"
# Log every [log_every] seconds.
log_every = 10
logger_factory = functools.partial(
logger_utils.make_logger,
directory=FLAGS.base_dir,
to_terminal=True,
to_tensorboard=True,
time_stamp=FLAGS.mava_id,
time_delta=log_every,
)
# distributed program
program = madqn.MADQN(
environment_factory=environment_factory,
network_factory=network_factory,
logger_factory=logger_factory,
num_executors=1,
exploration_scheduler_fn=LinearExplorationScheduler,
epsilon_min=0.05,
epsilon_decay=1e-5,
optimizer=snt.optimizers.RMSProp(learning_rate=1e-5),
checkpoint_subpath=checkpoint_dir,
batch_size=32,
executor_variable_update_period=100,
target_update_period=200,
max_gradient_norm=10.0,
trainer_fn=madqn.training.MADQNRecurrentTrainer,
executor_fn=madqn.execution.MADQNRecurrentExecutor,
).build()
# launch
gpu_id = -1
env_vars = {"CUDA_VISIBLE_DEVICES": str(gpu_id)}
local_resources = {
"trainer": [],
"evaluator": PythonProcess(env=env_vars),
"executor": PythonProcess(env=env_vars),
}
lp.launch(
program,
lp.LaunchType.LOCAL_MULTI_PROCESSING,
terminal="current_terminal",
local_resources=local_resources,
)
def test_madqn_on_debugging_env(self) -> None:
"""Tests that the system can run on the simple spread
debugging environment without crashing."""
# environment
environment_factory = functools.partial(
debugging_utils.make_environment,
env_name="simple_spread",
action_space="discrete",
)
# networks
network_factory = lp_utils.partial_kwargs(make_networks)
# system
system = madqn.MADQN(
environment_factory=environment_factory,
network_factory=network_factory,
num_executors=2,
batch_size=32,
min_replay_size=32,
max_replay_size=1000,
optimizer=snt.optimizers.Adam(learning_rate=1e-3),
checkpoint=False,
)
program = system.build()
(trainer_node, ) = program.groups["trainer"]
trainer_node.disable_run()
# Launch gpu config - don't use gpu
gpu_id = -1
env_vars = {"CUDA_VISIBLE_DEVICES": str(gpu_id)}
local_resources = {
"trainer": PythonProcess(env=env_vars),
"evaluator": PythonProcess(env=env_vars),
"executor": PythonProcess(env=env_vars),
}
lp.launch(
program,
launch_type="test_mt",
local_resources=local_resources,
)
trainer: mava.Trainer = trainer_node.create_handle().dereference()
for _ in range(5):
trainer.step()