def processed_experience_spec(self):
"""Spec for processed experience.
Returns:
Spec for the experience returned by preprocess_experience().
"""
if self._processed_experience_spec is not None:
return self._processed_experience_spec
batch_size = 4
exp = common.zeros_from_spec(self.experience_spec, batch_size)
transformed_exp = self.transform_timestep(exp)
processed_exp = self.preprocess_experience(transformed_exp)
self._processed_experience_spec = self.experience_spec._replace(
observation=common.extract_spec(processed_exp.observation),
rollout_info=common.extract_spec(processed_exp.rollout_info))
if not self._use_rollout_state:
self._procesed_experience_spec = \
self._processed_experience_spec._replace(state=())
return self._processed_experience_spec
def rollout_info_spec(self):
"""The spec for the PolicyInfo.info returned from rollout()."""
if self._rollout_info_spec is not None:
return self._rollout_info_spec
batch_size = 4
time_step = common.zeros_from_spec(self.time_step_spec, batch_size)
state = common.zeros_from_spec(self.train_state_spec, batch_size)
policy_step = self.rollout(self.transform_timestep(time_step), state,
RLAlgorithm.PREPARE_SPEC)
self._rollout_info_spec = common.extract_spec(policy_step.info)
return self._rollout_info_spec
def train_step_info_spec(self):
"""The spec for the PolicyInfo.info returned from train_step()."""
if self._train_step_info_spec is not None:
return self._train_step_info_spec
batch_size = 4
processed_exp = common.zeros_from_spec(self.processed_experience_spec,
batch_size)
state = common.zeros_from_spec(self.train_state_spec, batch_size)
policy_step = self.train_step(processed_exp, state)
self._train_step_info_spec = common.extract_spec(policy_step.info)
return self._train_step_info_spec
def _prepare_specs(self, algorithm):
"""Prepare various tensor specs."""
time_step = self.get_initial_time_step()
self._time_step_spec = common.extract_spec(time_step)
self._action_spec = self._env.action_spec()
policy_step = algorithm.rollout(
algorithm.transform_timestep(time_step), self._initial_state)
info_spec = common.extract_spec(policy_step.info)
self._policy_step_spec = PolicyStep(
action=self._action_spec,
state=algorithm.train_state_spec,
info=info_spec)
self._action_distribution_spec = tf.nest.map_structure(
common.to_distribution_spec, algorithm.action_distribution_spec)
self._action_dist_param_spec = tf.nest.map_structure(
lambda spec: spec.input_params_spec,
self._action_distribution_spec)
algorithm.prepare_off_policy_specs(time_step)
def prepare_off_policy_specs(self, time_step: ActionTimeStep):
"""Prepare various tensor specs for off_policy training.
prepare_off_policy_specs is called by OffPolicyDriver._prepare_spec().
"""
self._env_batch_size = time_step.step_type.shape[0]
self._time_step_spec = common.extract_spec(time_step)
initial_state = common.get_initial_policy_state(
self._env_batch_size, self.train_state_spec)
transformed_timestep = self.transform_timestep(time_step)
policy_step = self.rollout(transformed_timestep, initial_state)
info_spec = common.extract_spec(policy_step.info)
self._action_distribution_spec = tf.nest.map_structure(
common.to_distribution_spec, self.action_distribution_spec)
self._action_dist_param_spec = tf.nest.map_structure(
lambda spec: spec.input_params_spec,
self._action_distribution_spec)
self._experience_spec = Experience(
step_type=self._time_step_spec.step_type,
reward=self._time_step_spec.reward,
discount=self._time_step_spec.discount,
observation=self._time_step_spec.observation,
prev_action=self._action_spec,
action=self._action_spec,
info=info_spec,
action_distribution=self._action_dist_param_spec,
state=self.train_state_spec if self._use_rollout_state else ())
action_dist_params = common.zero_tensor_from_nested_spec(
self._experience_spec.action_distribution, self._env_batch_size)
action_dist = nested_distributions_from_specs(
self._action_distribution_spec, action_dist_params)
exp = Experience(step_type=time_step.step_type,
reward=time_step.reward,
discount=time_step.discount,
observation=time_step.observation,
prev_action=time_step.prev_action,
action=time_step.prev_action,
info=policy_step.info,
action_distribution=action_dist,
state=initial_state if self._use_rollout_state else
())
transformed_exp = self.transform_timestep(exp)
processed_exp = self.preprocess_experience(transformed_exp)
self._processed_experience_spec = self._experience_spec._replace(
observation=common.extract_spec(processed_exp.observation),
info=common.extract_spec(processed_exp.info))
policy_step = common.algorithm_step(
algorithm_step_func=self.train_step,
time_step=processed_exp,
state=initial_state)
info_spec = common.extract_spec(policy_step.info)
self._training_info_spec = TrainingInfo(
action_distribution=self._action_dist_param_spec, info=info_spec)