def run(self): """ Run a learning step of the model :return: ([float], [float], [float], [bool], [float], [float]) observations, states, rewards, masks, actions, values """ mb_obs, mb_rewards, mb_actions, mb_values, mb_dones = [], [], [], [], [] mb_states = self.states ep_infos = [] for _ in range(self.n_steps): actions, values, states, _ = self.model.step( self.obs, self.states, self.dones) mb_obs.append(np.copy(self.obs)) mb_actions.append(actions) mb_values.append(values) mb_dones.append(self.dones) clipped_actions = actions # Clip the actions to avoid out of bound error if isinstance(self.env.action_space, gym.spaces.Box): clipped_actions = np.clip(actions, self.env.action_space.low, self.env.action_space.high) obs, rewards, dones, infos = self.env.step(clipped_actions) for info in infos: maybe_ep_info = info.get('episode') if maybe_ep_info is not None: ep_infos.append(maybe_ep_info) self.states = states self.dones = dones self.obs = obs mb_rewards.append(rewards) mb_dones.append(self.dones) # batch of steps to batch of rollouts mb_obs = np.asarray(mb_obs, dtype=self.obs.dtype).swapaxes( 1, 0).reshape(self.batch_ob_shape) mb_rewards = np.asarray(mb_rewards, dtype=np.float32).swapaxes(0, 1) mb_actions = np.asarray(mb_actions, dtype=self.env.action_space.dtype).swapaxes( 0, 1) mb_values = np.asarray(mb_values, dtype=np.float32).swapaxes(0, 1) mb_dones = np.asarray(mb_dones, dtype=np.bool).swapaxes(0, 1) mb_masks = mb_dones[:, :-1] mb_dones = mb_dones[:, 1:] true_rewards = np.copy(mb_rewards) last_values = self.model.value(self.obs, self.states, self.dones).tolist() # discount/bootstrap off value fn for n, (rewards, dones, value) in enumerate(zip(mb_rewards, mb_dones, last_values)): rewards = rewards.tolist() dones = dones.tolist() if dones[-1] == 0: rewards = discount_with_dones(rewards + [value], dones + [0], self.gamma)[:-1] else: rewards = discount_with_dones(rewards, dones, self.gamma) mb_rewards[n] = rewards # convert from [n_env, n_steps, ...] to [n_steps * n_env, ...] mb_rewards = mb_rewards.reshape(-1, *mb_rewards.shape[2:]) mb_actions = mb_actions.reshape(-1, *mb_actions.shape[2:]) mb_values = mb_values.reshape(-1, *mb_values.shape[2:]) mb_masks = mb_masks.reshape(-1, *mb_masks.shape[2:]) true_rewards = true_rewards.reshape(-1, *true_rewards.shape[2:]) return mb_obs, mb_states, mb_rewards, mb_masks, mb_actions, mb_values, ep_infos, true_rewards
def run(self): """ Run a learning step of the model :return: ([float], [float], [float], [bool], [float], [float]) observations, states, rewards, masks, actions, values """ mb_obs, mb_rewards, mb_actions, mb_values, mb_dones = [], [], [], [], [] # TODO: add mb_raw_rewards = [] mb_states = self.states for _ in range(self.n_steps): actions, values, states, _ = self.model.step( self.obs, self.states, self.dones) mb_obs.append(np.copy(self.obs)) mb_actions.append(actions) mb_values.append(values) mb_dones.append(self.dones) clipped_actions = actions # Clip the actions to avoid out of bound error if isinstance(self.env.action_space, gym.spaces.Box): clipped_actions = np.clip(actions, self.env.action_space.low, self.env.action_space.high) # obs, rewards, dones, _ = self.env.step(clipped_actions) obs, raw_rewards, dones, _ = self.env.step(clipped_actions) mb_raw_rewards.append(raw_rewards) rewards = np.sign(raw_rewards) self.states = states self.dones = dones if hasattr(self.model, 'sil'): self.model.sil.step(self.obs, actions, raw_rewards, dones) else: raise ValueError('Model doesnot have SIL module.') self.obs = obs mb_rewards.append(rewards) mb_dones.append(self.dones) # batch of steps to batch of rollouts mb_obs = np.asarray(mb_obs, dtype=self.obs.dtype).swapaxes( 1, 0).reshape(self.batch_ob_shape) mb_rewards = np.asarray(mb_rewards, dtype=np.float32).swapaxes(0, 1) mb_actions = np.asarray(mb_actions, dtype=np.int32).swapaxes(0, 1) mb_values = np.asarray(mb_values, dtype=np.float32).swapaxes(0, 1) mb_dones = np.asarray(mb_dones, dtype=np.bool).swapaxes(0, 1) mb_raw_rewards = np.asarray(mb_raw_rewards, dtype=np.float32).swapaxes(0, 1) mb_masks = mb_dones[:, :-1] mb_dones = mb_dones[:, 1:] true_rewards = np.copy(mb_rewards) last_values = self.model.value(self.obs, self.states, self.dones).tolist() # discount/bootstrap off value fn for n, (rewards, dones, value) in enumerate(zip(mb_rewards, mb_dones, last_values)): rewards = rewards.tolist() dones = dones.tolist() if dones[-1] == 0: rewards = discount_with_dones(rewards + [value], dones + [0], self.gamma)[:-1] else: rewards = discount_with_dones(rewards, dones, self.gamma) mb_rewards[n] = rewards # convert from [n_env, n_steps, ...] to [n_steps * n_env, ...] mb_rewards = mb_rewards.reshape(-1, *mb_rewards.shape[2:]) mb_actions = mb_actions.reshape(-1, *mb_actions.shape[2:]) mb_values = mb_values.reshape(-1, *mb_values.shape[2:]) mb_masks = mb_masks.reshape(-1, *mb_masks.shape[2:]) return mb_obs, mb_states, mb_rewards, mb_masks, mb_actions, mb_values, true_rewards, mb_raw_rewards