def _stack_path_data(self, paths): max_path = max([len(path['observations']) for path in paths]) observations = self._stack_padding(paths, 'observations', max_path) actions = self._stack_padding(paths, 'actions', max_path) rewards = self._stack_padding(paths, 'rewards', max_path) dones = self._stack_padding(paths, 'dones', max_path) returns = self._stack_padding(paths, 'returns', max_path) advantages = self._stack_padding(paths, 'advantages', max_path) env_infos = utils.stack_tensor_dict_list( [path["env_infos"] for path in paths], max_path) agent_infos = utils.stack_tensor_dict_list( [path["agent_infos"] for path in paths], max_path) return observations, actions, rewards, dones, returns, advantages, env_infos, agent_infos
def _pad_and_stack_paths(self, paths): mask, observations, actions, rewards, baselines, adjusted_rewards, env_infos, agent_infos = [], [], [], [], [], [], [], [] for path in paths: # zero-pad paths if they don't have full length + create mask path_length = path["observations"].shape[0] assert self.max_path_length >= path_length mask.append(self._pad(np.ones(path_length), path_length)) observations.append(self._pad(path["observations"], path_length)) actions.append(self._pad(path["actions"], path_length)) rewards.append(self._pad(path["discounted_rewards"], path_length)) baselines.append(self._pad(path["baselines"], path_length)) adjusted_rewards.append( self._pad(path["adjusted_rewards"], path_length)) env_infos.append( dict([(key, self._pad(array, path_length)) for key, array in path["env_infos"].items()])) agent_infos.append( (dict([(key, self._pad(array, path_length)) for key, array in path["agent_infos"].items()]))) # stack mask = np.stack(mask, axis=0) # shape: (batch_size, max_path_length) observations = np.stack( observations, axis=0) # shape: (batch_size, max_path_length, ndim_act) actions = np.stack( actions, axis=0) # shape: (batch_size, max_path_length, ndim_obs) rewards = np.stack(rewards, axis=0) # shape: (batch_size, max_path_length) baselines = np.stack(baselines, axis=0) adjusted_rewards = np.stack( adjusted_rewards, axis=0) # shape: (batch_size, max_path_length) env_infos = utils.stack_tensor_dict_list( env_infos ) # dict of ndarrays of shape: (batch_size, max_path_length, ?) agent_infos = utils.stack_tensor_dict_list( agent_infos ) # dict of ndarrays of shape: (batch_size, max_path_length, ?) return mask, observations, actions, rewards, baselines, adjusted_rewards, env_infos, agent_infos
def obtain_samples(self, log=False, log_prefix='', random=False): """ Collect batch_size trajectories from each task Args: log (boolean): whether to log sampling times log_prefix (str) : prefix for logger random (boolean): whether the actions are random Returns: (dict) : A dict of paths of size [meta_batch_size] x (batch_size) x [5] x (max_path_length) """ # initial setup / preparation paths = OrderedDict() for i in range(self.meta_batch_size): paths[i] = [] n_samples = 0 running_paths = [ _get_empty_running_paths_dict() for _ in range(self.vec_env.num_envs) ] pbar = ProgBar(self.total_samples) policy_time, env_time = 0, 0 policy = self.policy policy.reset(dones=[True] * self.meta_batch_size) # initial reset of meta_envs obses = self.vec_env.reset() while n_samples < self.total_samples: # execute policy t = time.time() obs_per_task = np.split(np.asarray(obses), self.meta_batch_size) if random: actions = np.stack([[self.env.action_space.sample()] for _ in range(len(obses))], axis=0) agent_infos = [[{ 'mean': np.zeros_like(self.env.action_space.sample()), 'log_std': np.zeros_like(self.env.action_space.sample()) }] * self.envs_per_task] * self.meta_batch_size else: actions, agent_infos = policy.get_actions(obs_per_task) policy_time += time.time() - t # step environments t = time.time() actions = np.concatenate(actions) # stack meta batch next_obses, rewards, dones, env_infos = self.vec_env.step(actions) env_time += time.time() - t # stack agent_infos and if no infos were provided (--> None) create empty dicts agent_infos, env_infos = self._handle_info_dicts( agent_infos, env_infos) new_samples = 0 for idx, observation, action, reward, env_info, agent_info, done in zip( itertools.count(), obses, actions, rewards, env_infos, agent_infos, dones): # append new samples to running paths if isinstance(reward, np.ndarray): reward = reward[0] running_paths[idx]["observations"].append(observation) running_paths[idx]["actions"].append(action) running_paths[idx]["rewards"].append(reward) running_paths[idx]["dones"].append(done) running_paths[idx]["env_infos"].append(env_info) running_paths[idx]["agent_infos"].append(agent_info) # if running path is done, add it to paths and empty the running path if done: paths[idx // self.envs_per_task].append( dict( observations=np.asarray( running_paths[idx]["observations"]), actions=np.asarray(running_paths[idx]["actions"]), rewards=np.asarray(running_paths[idx]["rewards"]), dones=np.asarray(running_paths[idx]["dones"]), env_infos=utils.stack_tensor_dict_list( running_paths[idx]["env_infos"]), agent_infos=utils.stack_tensor_dict_list( running_paths[idx]["agent_infos"]), )) new_samples += len(running_paths[idx]["rewards"]) running_paths[idx] = _get_empty_running_paths_dict() pbar.update(new_samples) n_samples += new_samples obses = next_obses pbar.stop() self.total_timesteps_sampled += self.total_samples if log: logger.logkv(log_prefix + "PolicyExecTime", policy_time) logger.logkv(log_prefix + "EnvExecTime", env_time) return paths
def obtain_samples(self, log=False, log_prefix='', random=False): """ Collect batch_size trajectories from each task Args: log (boolean): whether to log sampling times log_prefix (str) : prefix for logger random (boolean): whether the actions are random Returns: (dict) : A dict of paths of size [meta_batch_size] x (batch_size) x [5] x (max_path_length) """ # initial setup / preparation paths = OrderedDict() for i in range(self.meta_batch_size): paths[i] = [] running_paths = _get_empty_running_paths_dict() pbar = ProgBar(self.total_samples) policy_time, env_time = 0, 0 policy = self.policy for idx in range(self.meta_batch_size): ts = 0 n_samples = 0 init_obs = np.expand_dims(self.env.reset(), 0).copy() obses = [init_obs for _ in range(self.meta_batch_size)] policy.reset(dones=[True] * self.meta_batch_size) while n_samples < self.samples_per_task: # execute policy t = time.time() if random: actions = np.stack([[self.env.action_space.sample()] for _ in range(len(obses))], axis=0) agent_infos = [[{ 'mean': np.zeros_like(self.env.action_space.sample()), 'log_std': np.zeros_like(self.env.action_space.sample()) }] * self.envs_per_task] * self.meta_batch_size else: actions, agent_infos = policy.get_actions(obses) policy_time += time.time() - t # step environments t = time.time() action, agent_info = actions[idx][0], agent_infos[idx][0] observation = obses[idx][0].copy() next_obs, reward, done, env_info = self.env.step(action) ts += 1 done = done or ts >= self.max_path_length if done: next_obs = self.env.reset() # time.sleep(1) ts = 0 env_time += time.time() - t new_samples = 0 # append new samples to running paths if isinstance(reward, np.ndarray): reward = reward[0] running_paths["observations"].append(observation) running_paths["actions"].append(action) running_paths["rewards"].append(reward) running_paths["dones"].append(done) running_paths["env_infos"].append(env_info) running_paths["agent_infos"].append(agent_info) # if running path is done, add it to paths and empty the running path if done: paths[idx].append( dict( observations=np.asarray( running_paths["observations"]), actions=np.asarray(running_paths["actions"]), rewards=np.asarray(running_paths["rewards"]), dones=np.asarray(running_paths["dones"]), env_infos=utils.stack_tensor_dict_list( running_paths["env_infos"]), agent_infos=utils.stack_tensor_dict_list( running_paths["agent_infos"]), )) new_samples += len(running_paths["rewards"]) running_paths = _get_empty_running_paths_dict() pbar.update(new_samples) n_samples += new_samples obses[idx][0] = next_obs self.total_timesteps_sampled += n_samples pbar.stop() if log: logger.logkv(log_prefix + "PolicyExecTime", policy_time) logger.logkv(log_prefix + "EnvExecTime", env_time) return paths
def obtain_samples(self, log=False, log_prefix='', random=False, deterministic=False, sinusoid=False, verbose=False): """ Collect batch_size trajectories from each task Args: log (boolean): whether to log sampling times log_prefix (str) : prefix for logger random (boolean): whether the actions are random Returns: (dict) : A dict of paths of size [meta_batch_size] x (batch_size) x [5] x (max_path_length) """ # initial setup / preparation paths = [] n_samples = 0 running_paths = [_get_empty_running_paths_dict() for _ in range(self.vec_env.num_envs)] if verbose: pbar = ProgBar(self.total_samples) policy_time, env_time = 0, 0 policy = self.policy policy.reset(dones=[True] * self.vec_env.num_envs) # initial reset of meta_envs obses = np.asarray(self.vec_env.reset()) while n_samples < self.total_samples: # execute policy t = time.time() if self.vae is not None: obses = np.array(obses) obses = self.vae.encode(obses) if random: actions = np.stack([self.env.action_space.sample() for _ in range(self.vec_env.num_envs)], axis=0) agent_infos = {} elif deterministic: actions, agent_infos = policy.get_actions(obses) actions = [a_i['mean'] for a_i in agent_infos] elif sinusoid: action_space = self.env.action_space.shape[0] num_envs = self.vec_env.num_envs actions = np.stack([policy.get_sinusoid_actions(action_space, t/policy.horizon * 2 * np.pi) for _ in range(num_envs)], axis=0) agent_infos = dict() else: obses = np.array(obses) actions, agent_infos = policy.get_actions(obses) policy_time += time.time() - t # step environments t = time.time() next_obses, rewards, dones, env_infos = self.vec_env.step(actions) env_time += time.time() - t # stack agent_infos and if no infos were provided (--> None) create empty dicts agent_infos, env_infos = self._handle_info_dicts(agent_infos, env_infos) new_samples = 0 for idx, observation, action, reward, env_info, agent_info, done in zip(itertools.count(), obses, actions, rewards, env_infos, agent_infos, dones): # append new samples to running paths if isinstance(reward, np.ndarray): reward = reward[0] running_paths[idx]["observations"].append(observation) running_paths[idx]["actions"].append(action) running_paths[idx]["rewards"].append(reward) running_paths[idx]["dones"].append(done) running_paths[idx]["env_infos"].append(env_info) running_paths[idx]["agent_infos"].append(agent_info) # if running path is done, add it to paths and empty the running path if done: paths.append(dict( observations=np.asarray(running_paths[idx]["observations"]), actions=np.asarray(running_paths[idx]["actions"]), rewards=np.asarray(running_paths[idx]["rewards"]), dones=np.asarray(running_paths[idx]["dones"]), env_infos=utils.stack_tensor_dict_list(running_paths[idx]["env_infos"]), agent_infos=utils.stack_tensor_dict_list(running_paths[idx]["agent_infos"]), )) new_samples += len(running_paths[idx]["rewards"]) running_paths[idx] = _get_empty_running_paths_dict() if verbose: pbar.update(self.vec_env.num_envs) n_samples += new_samples obses = next_obses if verbose: pbar.stop() self.total_timesteps_sampled += self.total_samples if log: logger.logkv(log_prefix + "TimeStepsCtr", self.total_timesteps_sampled) logger.logkv(log_prefix + "PolicyExecTime", policy_time) logger.logkv(log_prefix + "EnvExecTime", env_time) return paths
def obtain_samples(self, log=False, log_prefix='', random=False): """ Collect batch_size trajectories from each task Args: log (boolean): whether to log sampling times log_prefix (str) : prefix for logger random (boolean): whether the actions are random Returns: (dict) : A dict of paths of size [meta_batch_size] x (batch_size) x [5] x (max_path_length) """ # initial setup / preparation paths = [] n_samples = 0 running_paths = _get_empty_running_paths_dict() if log: pbar = ProgBar(self.total_samples) policy_time, env_time = 0, 0 policy = self.policy policy.reset(dones=[True]) # initial reset of meta_envs obs = np.asarray(self.env.reset()) ts = 0 while n_samples < self.total_samples: # execute policy t = time.time() if random: action = self.env.action_space.sample() agent_info = {} else: action, agent_info = policy.get_action(obs) if action.ndim == 2: action = action[0] policy_time += time.time() - t # step environments t = time.time() next_obs, reward, done, env_info = self.env.step(action) ts += 1 done = done or ts >= self.max_path_length if done: next_obs = self.env.reset() ts = 0 env_time += time.time() - t new_samples = 0 # append new samples to running paths if isinstance(reward, np.ndarray): reward = reward[0] running_paths["observations"].append(obs) running_paths["actions"].append(action) running_paths["rewards"].append(reward) running_paths["dones"].append(done) running_paths["env_infos"].append(env_info) running_paths["agent_infos"].append(agent_info) # if running path is done, add it to paths and empty the running path if done: paths.append( dict( observations=np.asarray(running_paths["observations"]), actions=np.asarray(running_paths["actions"]), rewards=np.asarray(running_paths["rewards"]), dones=np.asarray(running_paths["dones"]), env_infos=utils.stack_tensor_dict_list( running_paths["env_infos"]), agent_infos=utils.stack_tensor_dict_list( running_paths["agent_infos"]), )) new_samples += len(running_paths["rewards"]) running_paths = _get_empty_running_paths_dict() if log: pbar.update(new_samples) n_samples += new_samples obs = next_obs if log: pbar.stop() self.total_timesteps_sampled += self.total_samples if log: logger.logkv(log_prefix + "PolicyExecTime", policy_time) logger.logkv(log_prefix + "EnvExecTime", env_time) return paths
from asynch_mb.samplers.base import SampleProcessor