Nuke FeedForwardActor from JAX based agents
@@ -103,81 +103,6 @@ class GenericActor(core.Actor, Generic[actor_core.State, actor_core.Extras]):
self._variable_client.update(wait)
-# TODO(raveman): Migrate all users of FeedForwardActor to GenericActor and
-# remove this class.
-class FeedForwardActor(core.Actor):
- """A simple feed-forward actor implemented in JAX.
-
- An actor based on a policy which takes observations and outputs actions. It
- also adds experiences to replay and updates the actor weights from the policy
- on the learner.
- """
-
- def __init__(
- self,
- policy: FeedForwardPolicy,
- random_key: network_lib.PRNGKey,
- variable_client: variable_utils.VariableClient,
- adder: Optional[adders.Adder] = None,
- has_extras: bool = False,
- backend: Optional[str] = 'cpu',
- ):
- """Initializes a feed forward actor.
-
- Args:
- policy: A policy network taking observation and returning an action, if
- `has_extras=False`, and returning an (action, extras) tuple if
- `has_extras=True`.
- random_key: Random key.
- variable_client: The variable client to get policy parameters from.
- 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
- self._extras: types.NestedArray = ()
-
- # Adding batch dimension inside jit is much more efficient than outside.
- def batched_policy(
- params: network_lib.Params, key: network_lib.PRNGKey,
- observation: network_lib.Observation
- ) -> Tuple[Union[network_lib.Action, Tuple[
- network_lib.Action, types.NestedArray]], network_lib.PRNGKey]:
- # TODO(b/161332815): Make JAX Actor work with batched or unbatched inputs.
- key, key2 = jax.random.split(key)
- 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=backend)
-
- self._adder = adder
- self._variable_client = variable_client
-
- def select_action(self,
- observation: network_lib.Observation) -> types.NestedArray:
- result, self._random_key = self._policy(self._variable_client.params,
- self._random_key, observation)
- if self._has_extras:
- action, self._extras = result
- else:
- action = result
- return utils.to_numpy(action)
-
- def observe_first(self, timestep: dm_env.TimeStep):
- if self._adder:
- self._adder.add_first(timestep)
-
- def observe(self, action: network_lib.Action, next_timestep: dm_env.TimeStep):
- if self._adder:
- self._adder.add(action, next_timestep, extras=self._extras)
-
- def update(self, wait: bool = False):
- self._variable_client.update(wait)
-
-
# TODO(raveman): Migrate all users of RecurrentActor to GenericActor and
# remove this class.
class RecurrentActor(core.Actor):
