Support composition of Launchpad-based agents with other Launchpad nodes / agents in a single Launchpad program.
@@ -211,8 +211,9 @@ class DistributedLayout:
def coordinator(self, counter: counting.Counter, max_actor_steps: int):
return lp_utils.StepsLimiter(counter, max_actor_steps)
- def build(self, name='agent'):
+ def build(self, name='agent', program: Optional[lp.Program] = None):
"""Build the distributed agent topology."""
+ if not program:
program = lp.Program(name=name)
key = jax.random.PRNGKey(self._seed)
