Pass separate policy optimizer and critic optimizer to the D4PG agent.
@@ -52,6 +52,8 @@ class D4PG(agent.Agent):
batch_size: int = 256,
prefetch_size: int = 4,
target_update_period: int = 100,
+ policy_optimizer: snt.Optimizer = None,
+ critic_optimizer: snt.Optimizer = None,
min_replay_size: int = 1000,
max_replay_size: int = 1000000,
samples_per_insert: float = 32.0,
@@ -75,6 +77,8 @@ class D4PG(agent.Agent):
prefetch_size: size to prefetch from replay.
target_update_period: number of learner steps to perform before updating
the target networks.
+ policy_optimizer: optimizer for the policy network updates.
+ critic_optimizer: optimizer for the critic network updates.
min_replay_size: minimum replay size before updating.
max_replay_size: maximum replay size.
samples_per_insert: number of samples to take from replay for every insert
@@ -146,8 +150,10 @@ class D4PG(agent.Agent):
actor = actors.FeedForwardActor(behavior_network, adder=adder)
# Create optimizers.
- policy_optimizer = snt.optimizers.Adam(learning_rate=1e-4)
- critic_optimizer = snt.optimizers.Adam(learning_rate=1e-4)
+ policy_optimizer = policy_optimizer or snt.optimizers.Adam(
+ learning_rate=1e-4)
+ critic_optimizer = critic_optimizer or snt.optimizers.Adam(
+ learning_rate=1e-4)
# The learner updates the parameters (and initializes them).
learner = learning.D4PGLearner(
