(self)
Trains RL for one epoch.
Trains RL for one epoch.
Trains RL for one epoch .
def train_epoch(self):
    """Trains RL for one epoch."""
    # Copy policy state accumulated during data collection to the trainer.
    self._policy_trainer.model_state = self._policy_collect_model.state

    # Copy policy weights and state to value trainer.
    if self._n_shared_layers > 0:
      _copy_model_weights_and_state(
          0, self._n_shared_layers, self._policy_trainer, self._value_trainer
      )

    # Update the target value network.
    self._value_eval_model.weights = self._value_trainer.model_weights
    self._value_eval_model.state = self._value_trainer.model_state

    n_value_evals = rl_training.remaining_evals(
        self._value_trainer.step,
        self._epoch,
        self._value_train_steps_per_epoch,
        self._value_evals_per_epoch)
    for _ in range(n_value_evals):
      self._value_trainer.train_epoch(
          self._value_train_steps_per_epoch // self._value_evals_per_epoch,
          self._value_eval_steps,
      )
      # Update the target value network.
      self._value_eval_model.weights = self._value_trainer.model_weights
      self._value_eval_model.state = self._value_trainer.model_state

    # Copy value weights and state to policy trainer.
    if self._n_shared_layers > 0:
      _copy_model_weights_and_state(
          0, self._n_shared_layers, self._value_trainer, self._policy_trainer
      )
    n_policy_evals = rl_training.remaining_evals(
        self._policy_trainer.step,
        self._epoch,
        self._policy_train_steps_per_epoch,
        self._policy_evals_per_epoch)
    # Check if there was a restart after value training finishes and policy not.
    stopped_after_value = (n_value_evals == 0 and
                           n_policy_evals < self._policy_evals_per_epoch)
    should_copy_weights = self._n_shared_layers > 0 and not stopped_after_value
    if should_copy_weights:
      _copy_model_weights_and_state(
          0, self._n_shared_layers, self._value_trainer, self._policy_trainer
      )

    # Update the target value network.
    self._value_eval_model.weights = self._value_trainer.model_weights
    self._value_eval_model.state = self._value_trainer.model_state

    for _ in range(n_policy_evals):
      self._policy_trainer.train_epoch(
          self._policy_train_steps_per_epoch // self._policy_evals_per_epoch,
          self._policy_eval_steps,
      )
def train_epoch ( self ) : # Copy policy state accumulated during data collection to the trainer. self . _policy_trainer . model_state = self . _policy_collect_model . state # Copy policy weights and state to value trainer. if self . _n_shared_layers > 0 : _copy_model_weights_and_state ( 0 , self . _n_shared_layers , self . _policy_trainer , self . _value_trainer ) # Update the target value network. self . _value_eval_model . weights = self . _value_trainer . model_weights self . _value_eval_model . state = self . _value_trainer . model_state n_value_evals = rl_training . remaining_evals ( self . _value_trainer . step , self . _epoch , self . _value_train_steps_per_epoch , self . _value_evals_per_epoch ) for _ in range ( n_value_evals ) : self . _value_trainer . train_epoch ( self . _value_train_steps_per_epoch // self . _value_evals_per_epoch , self . _value_eval_steps , ) # Update the target value network. self . _value_eval_model . weights = self . _value_trainer . model_weights self . _value_eval_model . state = self . _value_trainer . model_state # Copy value weights and state to policy trainer. if self . _n_shared_layers > 0 : _copy_model_weights_and_state ( 0 , self . _n_shared_layers , self . _value_trainer , self . _policy_trainer ) n_policy_evals = rl_training . remaining_evals ( self . _policy_trainer . step , self . _epoch , self . _policy_train_steps_per_epoch , self . _policy_evals_per_epoch ) # Check if there was a restart after value training finishes and policy not. stopped_after_value = ( n_value_evals == 0 and n_policy_evals < self . _policy_evals_per_epoch ) should_copy_weights = self . _n_shared_layers > 0 and not stopped_after_value if should_copy_weights : _copy_model_weights_and_state ( 0 , self . _n_shared_layers , self . _value_trainer , self . _policy_trainer ) # Update the target value network. self . _value_eval_model . weights = self . _value_trainer . model_weights self . _value_eval_model . state = self . _value_trainer . model_state for _ in range ( n_policy_evals ) : self . _policy_trainer . train_epoch ( self . _policy_train_steps_per_epoch // self . _policy_evals_per_epoch , self . _policy_eval_steps , )