Add an option to choose a backend for running a policy.
@@ -56,6 +56,7 @@ class FeedForwardActor(core.Actor):
variable_client: variable_utils.VariableClient,
adder: Optional[adders.Adder] = None,
has_extras: bool = False,
+ backend: Optional[str] = 'cpu',
):
"""Initializes a feed forward actor.
@@ -68,6 +69,7 @@ class FeedForwardActor(core.Actor):
adder: An adder to add experiences to.
has_extras: Flag indicating whether the policy returns extra
information (e.g. q-values) in addition to an action.
+ backend: Which backend to use for running the policy.
"""
self._random_key = random_key
self._has_extras = has_extras
@@ -84,7 +86,7 @@ class FeedForwardActor(core.Actor):
observation = utils.add_batch_dim(observation)
output = policy(params, key2, observation)
return utils.squeeze_batch_dim(output), key
- self._policy = jax.jit(batched_policy, backend='cpu')
+ self._policy = jax.jit(batched_policy, backend=backend)
self._adder = adder
self._client = variable_client
@@ -128,6 +130,7 @@ class RecurrentActor(core.Actor):
variable_client: variable_utils.VariableClient,
adder: Optional[adders.Adder] = None,
has_extras: bool = False,
+ backend: Optional[str] = 'cpu',
):
"""Initializes a recurrent actor.
@@ -145,6 +148,7 @@ class RecurrentActor(core.Actor):
added to the adder.
has_extras: Flag indicating whether the recurrent policy returns extra
information (e.g. q-values) in addition to an action.
+ backend: Which backend to use for running the policy.
"""
self._random_key = random_key
self._has_extras = has_extras
@@ -163,7 +167,7 @@ class RecurrentActor(core.Actor):
output, new_state = recurrent_policy(params, key2, observation,
core_state)
return output, new_state, key
- self._recurrent_policy = jax.jit(batched_recurrent_policy, backend='cpu')
+ self._recurrent_policy = jax.jit(batched_recurrent_policy, backend=backend)
self._initial_state = self._prev_state = self._state = initial_core_state
self._adder = adder
