Turn DistributedSAC into a function.
@@ -28,16 +28,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.SACNetworks]
+# TODO(stanczyk): Remove DistributedTD3 once not used.
class DistributedSAC(distributed_layout.DistributedLayout):
- """Distributed program definition for SAC."""
+ """Distributed program definition for SAC.
- def __init__(
- self,
- environment_factory: jax_types.EnvironmentFactory,
+ DEPRECATED: Use distributed_sac 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_sac(
+ *self.args, name=name, program=program, **self.kwargs)
+
+
+def make_distributed_sac(environment_factory: jax_types.EnvironmentFactory,
network_factory: NetworkFactory,
config: sac_config.SACConfig,
seed: int,
@@ -48,7 +61,9 @@ class DistributedSAC(distributed_layout.DistributedLayout):
normalize_input: bool = True,
evaluator_factories: Optional[Sequence[
distributed_layout.EvaluatorFactory]] = None,
- ):
+ name: str = 'agent',
+ program: Optional[lp.Program] = None):
+ """Builds distributed SAC program."""
logger_fn = functools.partial(
loggers.make_default_logger,
'learner',
@@ -79,12 +94,12 @@ class DistributedSAC(distributed_layout.DistributedLayout):
policy_factory=eval_policy_factory,
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=sac_builder,
- policy_network=networks.apply_policy_and_sample,
+ policy_network_factory=networks.apply_policy_and_sample,
evaluator_factories=evaluator_factories,
num_actors=num_actors,
max_number_of_steps=max_number_of_steps,
@@ -94,7 +109,8 @@ class DistributedSAC(distributed_layout.DistributedLayout):
log_to_bigtable, log_every),
checkpointing_config=distributed_layout.CheckpointingConfig(),
make_snapshot_models=networks.default_models_to_snapshot,
- )
+ name=name,
+ program=program)
class SAC(local_layout.LocalLayout):
