class ExperienceReplay(object): def __init__(self, agent, enviroment, batch_size): self._replay_buffer = TFUniformReplayBuffer( data_spec=agent.collect_data_spec, batch_size=enviroment.batch_size, max_length=100000) self._random_policy = RandomTFPolicy(enviroment.time_step_spec(), enviroment.action_spec()) self._fill_buffer(enviroment, self._random_policy, steps=100) self.dataset = self._replay_buffer.as_dataset( num_parallel_calls=3, sample_batch_size=batch_size, num_steps=2, single_deterministic_pass=False).prefetch(3) self.iterator = iter(self.dataset) def _fill_buffer(self, enviroment, policy, steps): for _ in range(steps): self.timestamp_data(enviroment, policy) def timestamp_data(self, environment, policy): time_step = environment.current_time_step() action_step = policy.action(time_step) next_time_step = environment.step(action_step.action) timestamp_trajectory = trajectory.from_transition( time_step, action_step, next_time_step) self._replay_buffer.add_batch(timestamp_trajectory)
def collect_step(self, env: tf_py_environment.TFPyEnvironment, policy: tf_policy.Base, replay_buffer: TFUniformReplayBuffer): time_step = env.current_time_step() action_step = policy.action(time_step) next_time_step = env.step(action_step.action) traj = trajectory.from_transition(time_step, action_step, next_time_step) replay_buffer.add_batch(traj)
def collect_step(environment: TFEnvironment, policy: TFPyPolicy, replay_buffer: TFUniformReplayBuffer): """ Coleta uma iteração com o ambiente e devolve o resultado m :param environment: :param policy: :param replay_buffer: :return: """ time_step = environment.current_time_step() action_step = policy.action(time_step) next_time_step = environment.step(action_step.action) traj = trajectory.from_transition(time_step, action_step, next_time_step) # Add trajectory to the replay buffer replay_buffer.add_batch(traj)
class TFUniformReplayBuffer(TFReplayBufferAbstract): def _init_replay_buffer(self, batch_size, data_spec): self._batch_size = batch_size buffer_config = { "batch_size": self._batch_size, "data_spec": data_spec, "max_length": 1 } tf.compat.v2.summary.scalar(name="replay_buffer_size", data=self._batch_size) self._replay_buffer = TFReplayBuffer(**buffer_config) def add_batch(self, traj_dict): """ add a trajectory to the replay buffer Params traj (dict[dim]:numpy): a dict of tensors representing the trajectory to be added it to the replay buffer """ collect_spec_dict = self.collect_data_spec._asdict() traj_tf, traj_spec = build_tf_trajectory(traj_dict, collect_spec_dict) if not self._replay_buffer: batch_size = len(traj_dict["observation"]) self._init_replay_buffer(batch_size, traj_spec) self._replay_buffer.add_batch(traj_tf) def get_batch(self, batch_size): if batch_size is None: batch_size = self._batch_size # TODO: convert the replay buffer to a dataset and iterate over it traj, metadata = self._replay_buffer.get_next( sample_batch_size=batch_size) return traj, metadata
class SyncUniformExperienceReplayer(ExperienceReplayer): """ For synchronous off-policy training. Example algorithms: DDPG, SAC """ def __init__(self, experience_spec, batch_size): self._buffer = TFUniformReplayBuffer(experience_spec, batch_size) self._data_iter = None def observe(self, exp, env_ids=None): """ For the sync driver, `exp` has the shape (`env_batch_size`, ...) with `num_envs`==1 and `unroll_length`==1. This function always ignores `env_ids`. """ self._buffer.add_batch(exp) def replay(self, sample_batch_size, mini_batch_length): if self._data_iter is None: dataset = self._buffer.as_dataset( num_parallel_calls=3, sample_batch_size=sample_batch_size, num_steps=mini_batch_length).prefetch(3) self._data_iter = iter(dataset) return next(self._data_iter) def replay_all(self): return self._buffer.gather_all() def clear(self): self._buffer.clear() @property def batch_size(self): return self._buffer._batch_size
class SyncUniformExperienceReplayer(ExperienceReplayer): """ For synchronous off-policy training. Example algorithms: DDPG, SAC """ def __init__(self, experience_spec, batch_size): # TFUniformReplayBuffer does not support list in spec, we have to do # some conversion. self._experience_spec = experience_spec self._exp_has_list = nest_utils.nest_contains_list(experience_spec) tuple_experience_spec = nest_utils.nest_list_to_tuple(experience_spec) self._buffer = TFUniformReplayBuffer(tuple_experience_spec, batch_size) self._data_iter = None def _list_to_tuple(self, exp): if self._exp_has_list: return nest_utils.nest_list_to_tuple(exp) else: return exp def _tuple_to_list(self, exp): if self._exp_has_list: return nest_utils.nest_tuple_to_list(exp, self._experience_spec) else: return exp def observe(self, exp, env_ids=None): """ For the sync driver, `exp` has the shape (`env_batch_size`, ...) with `num_envs`==1 and `unroll_length`==1. This function always ignores `env_ids`. """ self._buffer.add_batch(self._list_to_tuple(exp)) def replay(self, sample_batch_size, mini_batch_length): """Get a random batch. Args: sample_batch_size (int): number of sequences mini_batch_length (int): the length of each sequence Returns: Experience: experience batch in batch major (B, T, ...) tf_uniform_replay_buffer.BufferInfo: information about the batch """ if self._data_iter is None: dataset = self._buffer.as_dataset( num_parallel_calls=3, sample_batch_size=sample_batch_size, num_steps=mini_batch_length).prefetch(3) self._data_iter = iter(dataset) exp, info = next(self._data_iter) return self._tuple_to_list(exp), info def replay_all(self): return self._tuple_to_list(self._buffer.gather_all()) def clear(self): self._buffer.clear() @property def batch_size(self): return self._buffer._batch_size
class TFAgentsPPOAgent(RLAgent): def __init__(self, name=None, actor_net=None, value_net=None, predictor=None, keep_models_fixed=False, featurizer=None): super().__init__(name, predictor, keep_models_fixed, featurizer) action_spec = BoundedTensorSpec((1, ), tf.int64, 0, ACTION_DIMENSIONS - 1) # we store both mask and the actual observation in the observation # given to the agent in order to get an association between these two # see also https://github.com/tensorflow/agents/issues/125#issuecomment-496583325 observation_spec = { 'state': TensorSpec((self.featurizer.state_dimension(), ), tf.float32), 'mask': TensorSpec((ACTION_DIMENSIONS, ), tf.float32) } layers = equal_spacing_fc(5, self.featurizer.state_dimension()) if actor_net is None: self.actor_net = MaskedActorNetwork(observation_spec, action_spec, layers) else: self.actor_net = actor_net if value_net is None: self.value_net = DummyMaskedValueNetwork(observation_spec, fc_layer_params=layers) else: self.value_net = value_net self.agent = tf_agents.agents.ppo.ppo_agent.PPOAgent( time_step_spec=ts.time_step_spec(observation_spec), action_spec=action_spec, actor_net=self.actor_net, value_net=self.value_net, optimizer=tf.compat.v1.train.AdamOptimizer(learning_rate=1e-5), discount_factor=1, use_gae=True, use_td_lambda_return=True, lambda_value=0.85, num_epochs=30, # the observations are dicts { 'state': ..., 'mask': ... } # normalization does not make any sense for the mask normalize_observations=False, ) if actor_net is not None or value_net is not None: self.agent.initialize() else: self._create_train_checkpointer() # All the variables are in fact successfully restored but this is # not done immediately but only once some shapes are known. # Therefore, if the shapes are never known, the variables are not # restored. This is no problem in self play, where all of the shapes # are known after the first training but it is a problem when playing # against old versions because often some of the old versions aren't # used (and also the value net is never used because the old versions # aren't trained). It isn't an error but tensorflow gives warnings at # the end which are confusing if one doesn't know this. # Therefore we silence those warnings with .expect_partial(). # For more information see # https://github.com/tensorflow/tensorflow/issues/27937#issuecomment-484683443 # https://github.com/tensorflow/tensorflow/issues/27937#issuecomment-488356053 self.train_checkpointer.initialize_or_restore().expect_partial() # it seems like there is also agent.policy. I still don't understand when # one should use which and why but this one works for now. self.policy = self.agent.collect_policy # because tf_agents wants the data as trajectories # (prev_time_step, action, new_time_step), we have to store the prev_time_step # until we have the new_time_step to build the trajectory at which point # the new prev_time_step is the new_time_step # this variable is to keep track of the prev_time_step self.last_time_step = None # even though PPO is on policy, storing the stuff for a bit seems to be ok # and the examples in the tf_agents repo also use one self.replay_buffer = TFUniformReplayBuffer( self.agent.collect_data_spec, batch_size=1, max_length=REPLAY_BUFFER_SIZE) self.replay_buffer_position = 0 self.clone_counter = 0 def _create_train_checkpointer(self): self.train_checkpointer = tf_agents.utils.common.Checkpointer( ckpt_dir=os.path.join(MODELS_PATH, self.name, 'Agent'), agent=self.agent) def _add_trajectory(self, prev_time_step, action, new_time_step): """Add a trajectory (prev_time_step, action, new_time_step) to the replay buffer Also train the agent on the whole buffer if it is full. """ traj = tf_agents.trajectories.trajectory.from_transition( prev_time_step, action, new_time_step) self.replay_buffer.add_batch(traj) self.replay_buffer_position += 1 if self.replay_buffer_position == REPLAY_BUFFER_SIZE + 1: if not self.keep_models_fixed: self.agent.train(self.replay_buffer.gather_all()) self.replay_buffer_position = 0 self.replay_buffer.clear() def act(self, observation, valid_action_mask): observation = { 'state': np.array(observation, dtype=np.float32), 'mask': valid_action_mask } if self.last_time_step is None: # a new episode started self.last_time_step = _to_tf_timestep(ts.restart(observation)) self.last_action_step = self.policy.action(self.last_time_step) return self.last_action_step.action.numpy()[0, 0] new_time_step = _to_tf_timestep( ts.transition(observation, self.prev_reward)) self._add_trajectory(self.last_time_step, self.last_action_step, new_time_step) self.last_time_step = new_time_step self.last_action_step = self.policy.action(new_time_step) self.prev_reward = None return self.last_action_step.action.numpy()[0, 0] def observe(self, reward, terminal): if not terminal: self.prev_reward = reward return # even when the episode ends, tf_agents expects some observation # additionally to the reward. Because that makes no sense for us, # we just give it an observation consisting of all-zeros new_time_step = _to_tf_timestep( ts.termination( { 'state': np.zeros(self.featurizer.state_dimension()), 'mask': np.zeros(ACTION_DIMENSIONS) }, reward)) self._add_trajectory(self.last_time_step, self.last_action_step, new_time_step) self.last_time_step = None self.last_action_step = None self.prev_reward = None def clone(self, name=None): """Return a clone of this agent with networks & predictor shared""" if name is None: self.clone_counter += 1 name = self.name + 'Clone' + str(self.clone_counter) return TFAgentsPPOAgent(name=name, actor_net=self.actor_net, value_net=self.value_net, predictor=self.predictor, keep_models_fixed=self.keep_models_fixed, featurizer=self.featurizer) def save_models(self): """Save actor, critic and predictor Args: global_step: the current game number, is appended to the filenames of the saved models """ if self.keep_models_fixed: return super().save_models(os.path.join(MODELS_PATH, self.name)) if not hasattr(self, 'train_checkpointer'): self._create_train_checkpointer() self.train_checkpointer.save(0)