Fixes a minor bug in handling of random numbers in CQL.
@@ -390,7 +390,8 @@ class CQLLearner(acme.Learner):
self._update_step = jax.jit(self._update_step)
# Create initial state.
- key_policy, key_q, _ = jax.random.split(random_key, 3)
+ key_policy, key_q, training_state_key = jax.random.split(random_key, 3)
+ del random_key
policy_params = networks.policy_network.init(key_policy)
policy_optimizer_state = policy_optimizer.init(policy_params)
critic_params = networks.critic_network.init(key_q)
@@ -402,7 +403,7 @@ class CQLLearner(acme.Learner):
policy_params=policy_params,
critic_params=critic_params,
target_critic_params=critic_params,
- key=random_key,
+ key=training_state_key,
steps=0)
if adaptive_entropy_coefficient:
