def test_simple_distribution(self): simple_action_space = gym.spaces.Discrete(3) simple_num_logits = calc_num_logits(simple_action_space) self.assertEqual(simple_num_logits, simple_action_space.n) simple_logits = torch.rand(self.batch_size, simple_num_logits) simple_action_distribution = get_action_distribution(simple_action_space, simple_logits) simple_actions = simple_action_distribution.sample() self.assertEqual(list(simple_actions.shape), [self.batch_size]) self.assertTrue(all(0 <= a < simple_action_space.n for a in simple_actions))
def test_tuple_sanity_check(self): num_spaces, num_actions = 3, 2 simple_space = gym.spaces.Discrete(num_actions) spaces = [simple_space for _ in range(num_spaces)] tuple_space = gym.spaces.Tuple(spaces) self.assertTrue(calc_num_logits(tuple_space), num_spaces * num_actions) simple_logits = torch.zeros(1, num_actions) tuple_logits = torch.zeros(1, calc_num_logits(tuple_space)) simple_distr = get_action_distribution(simple_space, simple_logits) tuple_distr = get_action_distribution(tuple_space, tuple_logits) tuple_entropy = tuple_distr.entropy() self.assertEqual(tuple_entropy, simple_distr.entropy() * num_spaces) simple_logprob = simple_distr.log_prob(torch.ones(1)) tuple_logprob = tuple_distr.log_prob(torch.ones(1, num_spaces)) self.assertEqual(tuple_logprob, simple_logprob * num_spaces)
def test_gumbel_trick(self): """ We use a Gumbel noise which seems to be faster compared to using pytorch multinomial. Here we test that those are actually equivalent. """ timing = Timing() torch.backends.cudnn.enabled = True torch.backends.cudnn.benchmark = True with torch.no_grad(): action_space = gym.spaces.Discrete(8) num_logits = calc_num_logits(action_space) device_type = 'cpu' device = torch.device(device_type) logits = torch.rand(self.batch_size, num_logits, device=device) * 10.0 - 5.0 if device_type == 'cuda': torch.cuda.synchronize(device) count_gumbel, count_multinomial = np.zeros([action_space.n]), np.zeros([action_space.n]) # estimate probability mass by actually sampling both ways num_samples = 20000 action_distribution = get_action_distribution(action_space, logits) sample_actions_log_probs(action_distribution) action_distribution.sample_gumbel() with timing.add_time('gumbel'): for i in range(num_samples): action_distribution = get_action_distribution(action_space, logits) samples_gumbel = action_distribution.sample_gumbel() count_gumbel[samples_gumbel[0]] += 1 action_distribution = get_action_distribution(action_space, logits) action_distribution.sample() with timing.add_time('multinomial'): for i in range(num_samples): action_distribution = get_action_distribution(action_space, logits) samples_multinomial = action_distribution.sample() count_multinomial[samples_multinomial[0]] += 1 estimated_probs_gumbel = count_gumbel / float(num_samples) estimated_probs_multinomial = count_multinomial / float(num_samples) log.debug('Gumbel estimated probs: %r', estimated_probs_gumbel) log.debug('Multinomial estimated probs: %r', estimated_probs_multinomial) log.debug('Sampling timing: %s', timing) time.sleep(0.1) # to finish logging
def __init__(self, cfg, core_out_size, action_space): super().__init__(cfg, action_space) assert not cfg.adaptive_stddev assert is_continuous_action_space(self.action_space), \ 'Non-adaptive stddev makes sense only for continuous action spaces' num_action_outputs = calc_num_logits(action_space) # calculate only action means using the policy neural network self.distribution_linear = nn.Linear(core_out_size, num_action_outputs // 2) # stddev is a single learned parameter initial_stddev = torch.empty([num_action_outputs // 2]) initial_stddev.fill_(math.log(self.cfg.initial_stddev)) self.learned_stddev = nn.Parameter(initial_stddev, requires_grad=True)
def test_tuple_distribution(self): num_spaces = random.randint(1, 4) spaces = [gym.spaces.Discrete(random.randint(2, 5)) for _ in range(num_spaces)] action_space = gym.spaces.Tuple(spaces) num_logits = calc_num_logits(action_space) logits = torch.rand(self.batch_size, num_logits) self.assertEqual(num_logits, sum(s.n for s in action_space.spaces)) action_distribution = get_action_distribution(action_space, logits) tuple_actions = action_distribution.sample() self.assertEqual(list(tuple_actions.shape), [self.batch_size, num_spaces]) log_probs = action_distribution.log_prob(tuple_actions) self.assertEqual(list(log_probs.shape), [self.batch_size]) entropy = action_distribution.entropy() self.assertEqual(list(entropy.shape), [self.batch_size])
def __init__(self, cfg, action_space): super().__init__() self.k: int = cfg.cpc_forward_steps self.time_subsample: int = cfg.cpc_time_subsample self.forward_subsample: int = cfg.cpc_forward_subsample self.hidden_size: int = cfg.hidden_size self.num_actions: int = calc_num_actions(action_space) if isinstance(action_space, gym.spaces.Discrete): self.action_sizes = [action_space.n] else: self.action_sizes = [space.n for space in action_space.spaces] self.rnn = nn.GRU(32 * self.num_actions, cfg.hidden_size) self.action_embed = nn.Embedding(calc_num_logits(action_space), 32) self.predictor = nn.Sequential( nn.Linear(2 * self.hidden_size, self.hidden_size), nn.ReLU(True), nn.Linear(self.hidden_size, self.hidden_size), nn.ReLU(True), nn.Linear(self.hidden_size, 1), )
def __init__(self, cfg, core_out_size, action_space): super().__init__(cfg, action_space) num_action_outputs = calc_num_logits(action_space) self.distribution_linear = nn.Linear(core_out_size, num_action_outputs)
def __init__(self, cfg, num_agents, obs_space, action_space): self.cfg = cfg self.num_agents = num_agents self.envs_per_split = cfg.num_envs_per_worker // cfg.worker_num_splits self.num_traj_buffers = self.calc_num_trajectory_buffers() num_actions = calc_num_actions(action_space) num_action_logits = calc_num_logits(action_space) hidden_size = get_hidden_size(self.cfg) log.debug('Allocating shared memory for trajectories') self.tensors = TensorDict() # policy inputs obs_dict = TensorDict() self.tensors['obs'] = obs_dict if isinstance(obs_space, spaces.Dict): for space_name, space in obs_space.spaces.items(): obs_dict[space_name] = self.init_tensor(space.dtype, space.shape) else: raise Exception('Only Dict observations spaces are supported') # env outputs self.tensors['rewards'] = self.init_tensor(torch.float32, [1]) self.tensors['dones'] = self.init_tensor(torch.bool, [1]) # policy outputs policy_outputs = [ ('actions', num_actions), ('action_logits', num_action_logits), ('log_prob_actions', 1), ('values', 1), ('policy_version', 1), ('rnn_states', hidden_size) ] policy_outputs = [PolicyOutput(*po) for po in policy_outputs] policy_outputs = sorted(policy_outputs, key=lambda policy_output: policy_output.name) for po in policy_outputs: self.tensors[po.name] = self.init_tensor(torch.float32, [po.size]) ensure_memory_shared(self.tensors) # this is for performance optimization # indexing in numpy arrays is faster than in PyTorch tensors self.tensors_individual_transitions = self.tensor_dict_to_numpy(len(self.tensor_dimensions())) self.tensor_trajectories = self.tensor_dict_to_numpy(len(self.tensor_dimensions()) - 1) # create a shared tensor to indicate when the learner is done with the trajectory buffer and # it can be used to store the next trajectory traj_buffer_available_shape = [ self.cfg.num_workers, self.cfg.worker_num_splits, self.envs_per_split, self.num_agents, self.num_traj_buffers, ] self.is_traj_tensor_available = torch.ones(traj_buffer_available_shape, dtype=torch.uint8) self.is_traj_tensor_available.share_memory_() self.is_traj_tensor_available = to_numpy(self.is_traj_tensor_available, 2) # copying small policy outputs (e.g. individual value predictions & action logits) to shared memory is a # bottleneck on the policy worker. For optimization purposes we create additional tensors to hold # just concatenated policy outputs. Rollout workers parse the data and add it to the trajectory buffers # in a proper format policy_outputs_combined_size = sum(po.size for po in policy_outputs) policy_outputs_shape = [ self.cfg.num_workers, self.cfg.worker_num_splits, self.envs_per_split, self.num_agents, policy_outputs_combined_size, ] self.policy_outputs = policy_outputs self.policy_output_tensors = torch.zeros(policy_outputs_shape, dtype=torch.float32) self.policy_output_tensors.share_memory_() self.policy_output_tensors = to_numpy(self.policy_output_tensors, 4) self.policy_versions = torch.zeros([self.cfg.num_policies], dtype=torch.int32) self.policy_versions.share_memory_() # a list of boolean flags to be shared among components that indicate that experience collection should be # temporarily stopped (e.g. due to too much experience accumulated on the learner) self.stop_experience_collection = torch.ones([self.cfg.num_policies], dtype=torch.bool) self.stop_experience_collection.share_memory_()
def __init__(self, cfg, num_agents, obs_space, action_space): self.cfg = cfg self.num_agents = num_agents self.envs_per_split = cfg.num_envs_per_worker // cfg.worker_num_splits self.num_traj_buffers = self.calc_num_trajectory_buffers() num_actions = calc_num_actions(action_space) num_action_logits = calc_num_logits(action_space) hidden_size = get_hidden_size(self.cfg) log.debug('Allocating shared memory for trajectories') self._tensors = TensorDict() # policy inputs obs_dict = TensorDict() self._tensors['obs'] = obs_dict if isinstance(obs_space, spaces.Dict): for space_name, space in obs_space.spaces.items(): obs_dict[space_name] = self.init_tensor( space.dtype, space.shape) else: raise Exception('Only Dict observations spaces are supported') # env outputs self._tensors['rewards'] = self.init_tensor(torch.float32, [1]) self._tensors['rewards'].fill_( -42.42) # if we're using uninitialized values it will be obvious self._tensors['dones'] = self.init_tensor(torch.bool, [1]) self._tensors['dones'].fill_(True) self._tensors['policy_id'] = self.init_tensor(torch.int, [1]) self._tensors['policy_id'].fill_( -1 ) # -1 is an invalid policy index, experience from policy "-1" is always ignored # policy outputs policy_outputs = [('actions', num_actions), ('action_logits', num_action_logits), ('log_prob_actions', 1), ('values', 1), ('policy_version', 1), ('rnn_states', hidden_size)] policy_outputs = [PolicyOutput(*po) for po in policy_outputs] policy_outputs = sorted(policy_outputs, key=lambda policy_output: policy_output.name) for po in policy_outputs: self._tensors[po.name] = self.init_tensor(torch.float32, [po.size]) ensure_memory_shared(self._tensors) # this is for performance optimization # indexing in numpy arrays is faster than in PyTorch tensors self.tensors = self.tensor_dict_to_numpy() # copying small policy outputs (e.g. individual value predictions & action logits) to shared memory is a # bottleneck on the policy worker. For optimization purposes we create additional tensors to hold # just concatenated policy outputs. Rollout workers parse the data and add it to the trajectory buffers # in a proper format policy_outputs_combined_size = sum(po.size for po in policy_outputs) policy_outputs_shape = [ self.cfg.num_workers, self.cfg.worker_num_splits, self.envs_per_split, self.num_agents, policy_outputs_combined_size, ] self.policy_outputs = policy_outputs self._policy_output_tensors = torch.zeros(policy_outputs_shape, dtype=torch.float32) self._policy_output_tensors.share_memory_() self.policy_output_tensors = self._policy_output_tensors.numpy() self._policy_versions = torch.zeros([self.cfg.num_policies], dtype=torch.int32) self._policy_versions.share_memory_() self.policy_versions = self._policy_versions.numpy() # a list of boolean flags to be shared among components that indicate that experience collection should be # temporarily stopped (e.g. due to too much experience accumulated on the learner) self._stop_experience_collection = torch.ones([self.cfg.num_policies], dtype=torch.bool) self._stop_experience_collection.share_memory_() self.stop_experience_collection = self._stop_experience_collection.numpy( ) queue_max_size_bytes = self.num_traj_buffers * 40 # 40 bytes to encode an int should be enough? self.free_buffers_queue = faster_fifo.Queue( max_size_bytes=queue_max_size_bytes) # since all buffers are initially free, we add all buffer indices to the queue self.free_buffers_queue.put_many_nowait( [int(i) for i in np.arange(self.num_traj_buffers)])