def actor(self, random_key, replay, variable_source, counter, actor_id): """The actor process.""" adder = self._builder.make_adder(replay) environment_key, actor_key = jax.random.split(random_key) # Create environment and policy core. # Environments normally require uint32 as a seed. environment = self._environment_factory( utils.sample_uint32(environment_key)) networks = self._network_factory(specs.make_environment_spec(environment)) policy_network = self._policy_network(networks) actor = self._builder.make_actor(actor_key, policy_network, adder, variable_source) # Create logger and counter. counter = counting.Counter(counter, 'actor') # Only actor #0 will write to bigtable in order not to spam it too much. logger = self._actor_logger_fn(actor_id) # Create the loop to connect environment and agent. return environment_loop.EnvironmentLoop(environment, actor, counter, logger, observers=self._observers)
def evaluator( random_key: networks_lib.PRNGKey, variable_source: core.VariableSource, counter: counting.Counter, make_actor: MakeActorFn, ): """The evaluation process.""" # Create environment and evaluator networks environment_key, actor_key = jax.random.split(random_key) # Environments normally require uint32 as a seed. environment = environment_factory(utils.sample_uint32(environment_key)) networks = network_factory(specs.make_environment_spec(environment)) actor = make_actor(actor_key, policy_factory(networks), variable_source) # Create logger and counter. counter = counting.Counter(counter, 'evaluator') if logger_fn is not None: logger = logger_fn('evaluator', 'actor_steps') else: logger = loggers.make_default_logger('evaluator', log_to_bigtable, steps_key='actor_steps') # Create the run loop and return it. return environment_loop.EnvironmentLoop(environment, actor, counter, logger, observers=observers)
def evaluator( random_key: types.PRNGKey, variable_source: core.VariableSource, counter: counting.Counter, make_actor: MakeActorFn, ): """The evaluation process.""" # Create environment and evaluator networks environment_key, actor_key = jax.random.split(random_key) # Environments normally require uint32 as a seed. environment = environment_factory(utils.sample_uint32(environment_key)) environment_spec = specs.make_environment_spec(environment) networks = network_factory(environment_spec) policy = policy_factory(networks, environment_spec, True) actor = make_actor(actor_key, policy, environment_spec, variable_source) # Create logger and counter. counter = counting.Counter(counter, 'evaluator') logger = logger_factory('evaluator', 'actor_steps', 0) # Create the run loop and return it. return environment_loop.EnvironmentLoop(environment, actor, counter, logger, observers=observers)
def build_actor( random_key: networks_lib.PRNGKey, replay: reverb.Client, variable_source: core.VariableSource, counter: counting.Counter, actor_id: ActorId, ) -> environment_loop.EnvironmentLoop: """The actor process.""" environment_key, actor_key = jax.random.split(random_key) # Create environment and policy core. # Environments normally require uint32 as a seed. environment = experiment.environment_factory( utils.sample_uint32(environment_key)) environment_spec = specs.make_environment_spec(environment) networks = experiment.network_factory(environment_spec) policy_network = config.make_policy(experiment=experiment, networks=networks, environment_spec=environment_spec, evaluation=False) adder = experiment.builder.make_adder(replay, environment_spec, policy_network) actor = experiment.builder.make_actor(actor_key, policy_network, environment_spec, variable_source, adder) # Create logger and counter. counter = counting.Counter(counter, 'actor') logger = experiment.logger_factory('actor', counter.get_steps_key(), actor_id) # Create the loop to connect environment and agent. return environment_loop.EnvironmentLoop(environment, actor, counter, logger, observers=experiment.observers)