Turn DistributedTD3 into a function.
@@ -27,16 +27,29 @@ from acme.jax.layouts import distributed_layout
from acme.jax.layouts import local_layout
from acme.utils import counting
from acme.utils import loggers
+import launchpad as lp
NetworkFactory = Callable[[specs.EnvironmentSpec], networks.TD3Networks]
+# TODO(stanczyk): Remove DistributedTD3 once not used.
class DistributedTD3(distributed_layout.DistributedLayout):
- """Distributed program definition for TD3."""
+ """Distributed program definition for TD3.
- def __init__(
- self,
- environment_factory: jax_types.EnvironmentFactory,
+ DEPRECATED: Use distributed_td3 function instead.
+ """
+
+ def __init__(self, *args, **kwargs):
+ self.args = args
+ self.kwargs = kwargs
+
+ def build(self, name='agent', program: Optional[lp.Program] = None):
+ """Build the distributed agent topology."""
+ return make_distributed_td3(
+ *self.args, name=name, program=program, **self.kwargs)
+
+
+def make_distributed_td3(environment_factory: jax_types.EnvironmentFactory,
environment_spec: specs.EnvironmentSpec,
network_factory: NetworkFactory,
config: td3_config.TD3Config,
@@ -47,7 +60,9 @@ class DistributedTD3(distributed_layout.DistributedLayout):
log_every: float = 10.0,
evaluator_factories: Optional[Sequence[
distributed_layout.EvaluatorFactory]] = None,
- ):
+ name: str = 'agent',
+ program: Optional[lp.Program] = None):
+ """Builds distributed TD3 program."""
logger_fn = functools.partial(
loggers.make_default_logger,
'learner',
@@ -76,12 +91,12 @@ class DistributedTD3(distributed_layout.DistributedLayout):
policy_factory=eval_network_fn,
log_to_bigtable=log_to_bigtable)
]
- super().__init__(
+ return distributed_layout.make_distributed_program(
seed=seed,
environment_factory=environment_factory,
network_factory=network_factory,
builder=td3_builder,
- policy_network=policy_network_fn,
+ policy_network_factory=policy_network_fn,
evaluator_factories=evaluator_factories,
num_actors=num_actors,
max_number_of_steps=max_number_of_steps,
@@ -89,7 +104,8 @@ class DistributedTD3(distributed_layout.DistributedLayout):
log_to_bigtable=log_to_bigtable,
actor_logger_fn=distributed_layout.get_default_logger_fn(
log_to_bigtable, log_every),
- )
+ name=name,
+ program=program)
class TD3(local_layout.LocalLayout):
