Add an option to update variable source params once per episode
@@ -50,6 +50,7 @@ class GenericActor(core.Actor, Generic[actor_core.State, actor_core.Extras]):
adder: Optional[adders.Adder] = None,
jit: bool = True,
backend: Optional[str] = 'cpu',
+ per_episode_update: bool = False
):
"""Initializes a feed forward actor.
@@ -60,6 +61,8 @@ class GenericActor(core.Actor, Generic[actor_core.State, actor_core.Extras]):
adder: An adder to add experiences to.
jit: Whether or not to jit the passed ActorCore's pure functions.
backend: Which backend to use when jitting the policy.
+ per_episode_update: if True, updates variable client params once at the
+ beginning of each episode
"""
self._random_key = random_key
self._variable_client = variable_client
@@ -74,6 +77,7 @@ class GenericActor(core.Actor, Generic[actor_core.State, actor_core.Extras]):
self._init = actor.init
self._policy = actor.select_action
self._get_extras = actor.get_extras
+ self._per_episode_update = per_episode_update
def select_action(self,
observation: network_lib.Observation) -> types.NestedArray:
@@ -86,6 +90,8 @@ class GenericActor(core.Actor, Generic[actor_core.State, actor_core.Extras]):
self._state = self._init(key)
if self._adder:
self._adder.add_first(timestep)
+ if self._per_episode_update:
+ self._variable_client.update_and_wait()
def observe(self, action: network_lib.Action, next_timestep: dm_env.TimeStep):
if self._adder:
@@ -93,4 +99,5 @@ class GenericActor(core.Actor, Generic[actor_core.State, actor_core.Extras]):
action, next_timestep, extras=self._get_extras(self._state))
def update(self, wait: bool = False):
+ if not self._per_episode_update:
self._variable_client.update(wait)
