def main(_): # Create an environment and create the spec. environment, environment_spec = _build_environment( FLAGS.environment_name, max_steps=FLAGS.max_steps_per_episode) if FLAGS.model_name: loaded_network = load_wb_model(FLAGS.model_name, FLAGS.model_tag) if FLAGS.stochastic: head = networks.StochasticSamplingHead() else: head = lambda q: trfl.epsilon_greedy(q, epsilon=FLAGS.epsilon ).sample() policy_network = snt.Sequential([ loaded_network, head, ]) actor = actors.FeedForwardActor(policy_network) else: actor = RandomActor(environment_spec) recorder = DemonstrationRecorder(environment, actor) recorder.collect_n_episodes(FLAGS.n_episodes) recorder.make_tf_dataset() recorder.save(FLAGS.save_dir)
def actor( self, replay: reverb.Client, variable_source: acme.VariableSource, counter: counting.Counter, ) -> acme.EnvironmentLoop: """The actor process.""" action_spec = self._environment_spec.actions observation_spec = self._environment_spec.observations # Create environment and target networks to act with. environment = self._environment_factory(False) agent_networks = self._network_factory(action_spec, self._num_critic_heads) # Make sure observation network is defined. observation_network = agent_networks.get('observation', tf.identity) # Create a stochastic behavior policy. behavior_network = snt.Sequential([ observation_network, agent_networks['policy'], networks.StochasticSamplingHead(), ]) # Ensure network variables are created. tf2_utils.create_variables(behavior_network, [observation_spec]) policy_variables = {'policy': behavior_network.variables} # Create the variable client responsible for keeping the actor up-to-date. variable_client = tf2_variable_utils.VariableClient(variable_source, policy_variables, update_period=1000) # Make sure not to use a random policy after checkpoint restoration by # assigning variables before running the environment loop. variable_client.update_and_wait() # Component to add things into replay. adder = adders.NStepTransitionAdder( client=replay, n_step=self._n_step, max_in_flight_items=self._max_in_flight_items, discount=self._additional_discount) # Create the agent. actor = actors.FeedForwardActor(policy_network=behavior_network, adder=adder, variable_client=variable_client) # Create logger and counter; actors will not spam bigtable. counter = counting.Counter(counter, 'actor') logger = loggers.make_default_logger('actor', save_data=False, time_delta=self._log_every, steps_key='actor_steps') # Create the run loop and return it. return acme.EnvironmentLoop(environment, actor, counter, logger)
def __init__(self, policy_network: snt.RNNCore, critic_network: networks.CriticDeepRNN, target_policy_network: snt.RNNCore, target_critic_network: networks.CriticDeepRNN, dataset: tf.data.Dataset, accelerator_strategy: Optional[tf.distribute.Strategy] = None, behavior_network: Optional[snt.Module] = None, cwp_network: Optional[snt.Module] = None, policy_optimizer: Optional[snt.Optimizer] = None, critic_optimizer: Optional[snt.Optimizer] = None, discount: float = 0.99, target_update_period: int = 100, num_action_samples_td_learning: int = 1, num_action_samples_policy_weight: int = 4, baseline_reduce_function: str = 'mean', clipping: bool = True, policy_improvement_modes: str = 'exp', ratio_upper_bound: float = 20., beta: float = 1.0, counter: Optional[counting.Counter] = None, logger: Optional[loggers.Logger] = None, checkpoint: bool = False): """Initializes the learner. Args: policy_network: the online (optimized) policy. critic_network: the online critic. target_policy_network: the target policy (which lags behind the online policy). target_critic_network: the target critic. dataset: dataset to learn from, whether fixed or from a replay buffer (see `acme.datasets.reverb.make_reverb_dataset` documentation). accelerator_strategy: the strategy used to distribute computation, whether on a single, or multiple, GPU or TPU; as supported by tf.distribute. behavior_network: The network to snapshot under `policy` name. If None, snapshots `policy_network` instead. cwp_network: CWP network to snapshot: samples actions from the policy and weighs them with the critic, then returns the action by sampling from the softmax distribution using critic values as logits. Used only for snapshotting, not training. policy_optimizer: the optimizer to be applied to the policy loss. critic_optimizer: the optimizer to be applied to the distributional Bellman loss. discount: discount to use for TD updates. target_update_period: number of learner steps to perform before updating the target networks. num_action_samples_td_learning: number of action samples to use to estimate expected value of the critic loss w.r.t. stochastic policy. num_action_samples_policy_weight: number of action samples to use to estimate the advantage function for the CRR weighting of the policy loss. baseline_reduce_function: one of 'mean', 'max', 'min'. Way of aggregating values from `num_action_samples` estimates of the value function. clipping: whether to clip gradients by global norm. policy_improvement_modes: one of 'exp', 'binary', 'all'. CRR mode which determines how the advantage function is processed before being multiplied by the policy loss. ratio_upper_bound: if policy_improvement_modes is 'exp', determines the upper bound of the weight (i.e. the weight is min(exp(advantage / beta), upper_bound) ). beta: if policy_improvement_modes is 'exp', determines the beta (see above). counter: counter object used to keep track of steps. logger: logger object to be used by learner. checkpoint: boolean indicating whether to checkpoint the learner. """ if accelerator_strategy is None: accelerator_strategy = snt.distribute.Replicator() self._accelerator_strategy = accelerator_strategy self._policy_improvement_modes = policy_improvement_modes self._ratio_upper_bound = ratio_upper_bound self._num_action_samples_td_learning = num_action_samples_td_learning self._num_action_samples_policy_weight = num_action_samples_policy_weight self._baseline_reduce_function = baseline_reduce_function self._beta = beta # When running on TPUs we have to know the amount of memory required (and # thus the sequence length) at the graph compilation stage. At the moment, # the only way to get it is to sample from the dataset, since the dataset # does not have any metadata, see b/160672927 to track this upcoming # feature. sample = next(dataset.as_numpy_iterator()) self._sequence_length = sample.action.shape[1] self._counter = counter or counting.Counter() self._logger = logger or loggers.TerminalLogger('learner', time_delta=1.) self._discount = discount self._clipping = clipping self._target_update_period = target_update_period with self._accelerator_strategy.scope(): # Necessary to track when to update target networks. self._num_steps = tf.Variable(0, dtype=tf.int32) # (Maybe) distributing the dataset across multiple accelerators. distributed_dataset = self._accelerator_strategy.experimental_distribute_dataset( dataset) self._iterator = iter(distributed_dataset) # Create the optimizers. self._critic_optimizer = critic_optimizer or snt.optimizers.Adam( 1e-4) self._policy_optimizer = policy_optimizer or snt.optimizers.Adam( 1e-4) # Store online and target networks. self._policy_network = policy_network self._critic_network = critic_network self._target_policy_network = target_policy_network self._target_critic_network = target_critic_network # Expose the variables. self._variables = { 'critic': self._target_critic_network.variables, 'policy': self._target_policy_network.variables, } # Create a checkpointer object. self._checkpointer = None self._snapshotter = None if checkpoint: self._checkpointer = tf2_savers.Checkpointer( objects_to_save={ 'counter': self._counter, 'policy': self._policy_network, 'critic': self._critic_network, 'target_policy': self._target_policy_network, 'target_critic': self._target_critic_network, 'policy_optimizer': self._policy_optimizer, 'critic_optimizer': self._critic_optimizer, 'num_steps': self._num_steps, }, time_delta_minutes=30.) raw_policy = snt.DeepRNN( [policy_network, networks.StochasticSamplingHead()]) critic_mean = networks.CriticDeepRNN( [critic_network, networks.StochasticMeanHead()]) objects_to_save = { 'raw_policy': raw_policy, 'critic': critic_mean, } if behavior_network is not None: objects_to_save['policy'] = behavior_network if cwp_network is not None: objects_to_save['cwp_policy'] = cwp_network self._snapshotter = tf2_savers.Snapshotter( objects_to_save=objects_to_save, time_delta_minutes=30) # Timestamp to keep track of the wall time. self._walltime_timestamp = time.time()
def __init__( self, environment_spec: specs.EnvironmentSpec, policy_network: snt.Module, critic_network: snt.Module, observation_network: types.TensorTransformation = tf.identity, discount: float = 0.99, batch_size: int = 256, prefetch_size: int = 4, target_policy_update_period: int = 100, target_critic_update_period: int = 100, min_replay_size: int = 1000, max_replay_size: int = 1000000, samples_per_insert: float = 32.0, policy_loss_module: snt.Module = None, policy_optimizer: snt.Optimizer = None, critic_optimizer: snt.Optimizer = None, n_step: int = 5, num_samples: int = 20, clipping: bool = True, logger: loggers.Logger = None, counter: counting.Counter = None, checkpoint: bool = True, replay_table_name: str = adders.DEFAULT_PRIORITY_TABLE, ): """Initialize the agent. Args: environment_spec: description of the actions, observations, etc. policy_network: the online (optimized) policy. critic_network: the online critic. observation_network: optional network to transform the observations before they are fed into any network. discount: discount to use for TD updates. batch_size: batch size for updates. prefetch_size: size to prefetch from replay. target_policy_update_period: number of updates to perform before updating the target policy network. target_critic_update_period: number of updates to perform before updating the target critic network. min_replay_size: minimum replay size before updating. max_replay_size: maximum replay size. samples_per_insert: number of samples to take from replay for every insert that is made. policy_loss_module: configured MPO loss function for the policy optimization; defaults to sensible values on the control suite. See `acme/tf/losses/mpo.py` for more details. policy_optimizer: optimizer to be used on the policy. critic_optimizer: optimizer to be used on the critic. n_step: number of steps to squash into a single transition. num_samples: number of actions to sample when doing a Monte Carlo integration with respect to the policy. clipping: whether to clip gradients by global norm. logger: logging object used to write to logs. counter: counter object used to keep track of steps. checkpoint: boolean indicating whether to checkpoint the learner. replay_table_name: string indicating what name to give the replay table. """ # Create a replay server to add data to. replay_table = reverb.Table( name=adders.DEFAULT_PRIORITY_TABLE, sampler=reverb.selectors.Uniform(), remover=reverb.selectors.Fifo(), max_size=max_replay_size, rate_limiter=reverb.rate_limiters.MinSize(min_size_to_sample=1), signature=adders.NStepTransitionAdder.signature(environment_spec)) self._server = reverb.Server([replay_table], port=None) # The adder is used to insert observations into replay. address = f'localhost:{self._server.port}' adder = adders.NStepTransitionAdder(client=reverb.Client(address), n_step=n_step, discount=discount) # The dataset object to learn from. dataset = datasets.make_reverb_dataset( table=replay_table_name, client=reverb.TFClient(address), batch_size=batch_size, prefetch_size=prefetch_size, environment_spec=environment_spec, transition_adder=True) # Make sure observation network is a Sonnet Module. observation_network = tf2_utils.to_sonnet_module(observation_network) # Create target networks before creating online/target network variables. target_policy_network = copy.deepcopy(policy_network) target_critic_network = copy.deepcopy(critic_network) target_observation_network = copy.deepcopy(observation_network) # Get observation and action specs. act_spec = environment_spec.actions obs_spec = environment_spec.observations emb_spec = tf2_utils.create_variables(observation_network, [obs_spec]) # Create the behavior policy. behavior_network = snt.Sequential([ observation_network, policy_network, networks.StochasticSamplingHead(), ]) # Create variables. tf2_utils.create_variables(policy_network, [emb_spec]) tf2_utils.create_variables(critic_network, [emb_spec, act_spec]) tf2_utils.create_variables(target_policy_network, [emb_spec]) tf2_utils.create_variables(target_critic_network, [emb_spec, act_spec]) tf2_utils.create_variables(target_observation_network, [obs_spec]) # Create the actor which defines how we take actions. actor = actors.FeedForwardActor(policy_network=behavior_network, adder=adder) # Create optimizers. policy_optimizer = policy_optimizer or snt.optimizers.Adam(1e-4) critic_optimizer = critic_optimizer or snt.optimizers.Adam(1e-4) # The learner updates the parameters (and initializes them). learner = learning.MPOLearner( policy_network=policy_network, critic_network=critic_network, observation_network=observation_network, target_policy_network=target_policy_network, target_critic_network=target_critic_network, target_observation_network=target_observation_network, policy_loss_module=policy_loss_module, policy_optimizer=policy_optimizer, critic_optimizer=critic_optimizer, clipping=clipping, discount=discount, num_samples=num_samples, target_policy_update_period=target_policy_update_period, target_critic_update_period=target_critic_update_period, dataset=dataset, logger=logger, counter=counter, checkpoint=checkpoint) super().__init__(actor=actor, learner=learner, min_observations=max(batch_size, min_replay_size), observations_per_step=float(batch_size) / samples_per_insert)
def main(_): # TODO(yutian): Create environment. # # Create an environment and grab the spec. # raw_environment = bsuite.load_and_record_to_csv( # bsuite_id=FLAGS.bsuite_id, # results_dir=FLAGS.results_dir, # overwrite=FLAGS.overwrite, # ) # environment = single_precision.SinglePrecisionWrapper(raw_environment) # environment_spec = specs.make_environment_spec(environment) # TODO(yutian): Create dataset. # Build the dataset. # if hasattr(raw_environment, 'raw_env'): # raw_environment = raw_environment.raw_env # # batch_dataset = bsuite_demonstrations.make_dataset(raw_environment) # # Combine with demonstration dataset. # transition = functools.partial( # _n_step_transition_from_episode, n_step=1, additional_discount=1.) # # dataset = batch_dataset.map(transition) # # # Batch and prefetch. # dataset = dataset.batch(FLAGS.batch_size, drop_remainder=True) # dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE) # Create the networks to optimize. networks = make_networks(environment_spec.actions) treatment_net = networks['treatment_net'] instrumental_net = networks['instrumental_net'] policy_net = networks['policy_net'] # If the agent is non-autoregressive use epsilon=0 which will be a greedy # policy. evaluator_net = snt.Sequential([ policy_net, # Sample actions. acme_nets.StochasticSamplingHead() ]) # Ensure that we create the variables before proceeding (maybe not needed). tf2_utils.create_variables(policy_net, [environment_spec.observations]) # TODO(liyuan): set the proper input spec using environment_spec.observations # and environment_spec.actions. tf2_utils.create_variables(treatment_net, [environment_spec.observations]) tf2_utils.create_variables( instrumental_net, [environment_spec.observations, environment_spec.actions]) counter = counting.Counter() learner_counter = counting.Counter(counter, prefix='learner') # Create the actor which defines how we take actions. evaluator_net = actors.FeedForwardActor(evaluator_net) eval_loop = acme.EnvironmentLoop(environment=environment, actor=evaluator_net, counter=counter, logger=loggers.TerminalLogger( 'evaluation', time_delta=1.)) # The learner updates the parameters (and initializes them). learner = learning.DFIVLearner( treatment_net=treatment_net, instrumental_net=instrumental_net, policy_net=policy_net, treatment_learning_rate=FLAGS.treatment_learning_rate, instrumental_learning_rate=FLAGS.instrumental_learning_rate, policy_learning_rate=FLAGS.policy_learning_rate, dataset=dataset, counter=learner_counter) # Run the environment loop. while True: for _ in range(FLAGS.evaluate_every): learner.step() learner_counter.increment(learner_steps=FLAGS.evaluate_every) eval_loop.run(FLAGS.evaluation_episodes)