Turn DistributedD4PG into a function.
@@ -27,19 +27,29 @@ from acme.jax.layouts import local_layout
from acme.utils import counting
from acme.utils import loggers
import dm_env
+import launchpad as lp
NetworkFactory = Callable[[specs.EnvironmentSpec], d4pg_networks.D4PGNetworks]
+# TODO(stanczyk): Remove DistributedD4PG once not used.
class DistributedD4PG(distributed_layout.DistributedLayout):
- """Program definition for D4PG.
+ """Distributed program definition for SAC.
- We are in the process of migrating towards a more modular agent configuration.
- This is maintained now for compatibility.
+ DEPRECATED: Use distributed_sac function instead.
"""
- def __init__(
- self,
+ 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_d4pg(
+ *self.args, name=name, program=program, **self.kwargs)
+
+
+def make_distributed_d4pg(
environment_factory: Callable[[bool], dm_env.Environment],
network_factory: NetworkFactory,
random_seed: int,
@@ -60,7 +70,9 @@ class DistributedD4PG(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 D4PG program."""
config = d4pg_config.D4PGConfig(
discount=discount,
learning_rate=1e-4,
@@ -99,13 +111,12 @@ class DistributedD4PG(distributed_layout.DistributedLayout):
policy_factory=_eval_policy_network,
log_to_bigtable=log_to_bigtable)
]
-
- super().__init__(
+ return distributed_layout.make_distributed_program(
seed=random_seed,
environment_factory=lambda seed: environment_factory(False),
network_factory=network_factory,
builder=builder,
- policy_network=_policy_network,
+ policy_network_factory=_policy_network,
evaluator_factories=evaluator_factories,
num_actors=num_actors,
environment_spec=environment_spec,
@@ -114,7 +125,8 @@ class DistributedD4PG(distributed_layout.DistributedLayout):
actor_logger_fn=distributed_layout.get_default_logger_fn(
log_to_bigtable, log_every),
prefetch_size=config.prefetch_size,
- )
+ name=name,
+ program=program)
class D4PG(local_layout.LocalLayout):
