def evaluate_true(self, X: Tensor) -> Tensor: r"""Evaluate the GMMs.""" # This needs to be reinstantiated because MVN apparently does not # have a `to` method to make it device/dtype agnostic. mvn = MultivariateNormal(loc=self.gmm_pos, covariance_matrix=self.gmm_covar) view_shape = ( X.shape[:-1] + torch.Size([1] * (self.gmm_pos.ndim - 1)) + self.gmm_pos.shape[-1:] ) expand_shape = X.shape[:-1] + self.gmm_pos.shape pdf_X = mvn.log_prob(X.view(view_shape).expand(expand_shape)).exp() # Multiply by -1 to make this a minimization problem by default return -(self.gmm_norm * pdf_X).sum(dim=-1)
def act(self, state): if self.has_continuous_action_space: action_mean = self.actor(state) cov_mat = torch.diag(self.action_var).unsqueeze(dim=0) dist = MultivariateNormal(action_mean, cov_mat) else: action_probs = self.actor(state) dist = Categorical(action_probs) action = dist.sample() action_logprob = dist.log_prob(action) return action.detach(), action_logprob.detach()
def _get_dist(self): """ Get the :class:`~pyro.distributions.GaussianHMM` distribution that corresponds to a :class:`LinearlyCoupledMaternGP`. """ trans_matrix, process_covar = self.kernel.transition_matrix_and_covariance( dt=self.dt) trans_matrix = block_diag_embed(trans_matrix) process_covar = block_diag_embed(process_covar) loc = self.A.new_zeros(self.full_state_dim) trans_dist = MultivariateNormal(loc, process_covar) return dist.GaussianHMM(self._get_init_dist(), trans_matrix, trans_dist, self._get_obs_matrix(), self._get_obs_dist())
def test_multivariate_normal_prior_log_prob(self, cuda=False): device = torch.device("cuda") if cuda else torch.device("cpu") mean = torch.tensor([0.0, 1.0], device=device) cov = torch.eye(2, device=device) prior = MultivariateNormalPrior(mean, covariance_matrix=cov) dist = MultivariateNormal(mean, covariance_matrix=cov) self.assertFalse(prior.log_transform) t = torch.tensor([-1, 0.5], device=device) self.assertTrue(torch.equal(prior.log_prob(t), dist.log_prob(t))) t = torch.tensor([[-1, 0.5], [1.5, -2.0]], device=device) self.assertTrue(torch.equal(prior.log_prob(t), dist.log_prob(t))) with self.assertRaises(RuntimeError): prior.log_prob(torch.zeros(3, device=device))
def test_api_snpe_c_posterior_correction( sample_with_mcmc, mcmc_method, prior_str, set_seed ): """Test that leakage correction applied to sampling works, with both MCMC and rejection. Args: set_seed: fixture for manual seeding """ num_dim = 2 x_o = zeros(1, num_dim) # likelihood_mean will be likelihood_shift+theta likelihood_shift = -1.0 * ones(num_dim) likelihood_cov = 0.3 * eye(num_dim) if prior_str == "gaussian": prior_mean = zeros(num_dim) prior_cov = eye(num_dim) prior = MultivariateNormal(loc=prior_mean, covariance_matrix=prior_cov) else: prior = utils.BoxUniform(-2.0 * ones(num_dim), 2.0 * ones(num_dim)) def simulator(theta): return linear_gaussian(theta, likelihood_shift, likelihood_cov) simulator, prior = prepare_for_sbi(simulator, prior) inference = SNPE_C( prior, density_estimator="maf", simulation_batch_size=50, sample_with_mcmc=sample_with_mcmc, mcmc_method=mcmc_method, show_progress_bars=False, ) theta, x = simulate_for_sbi(simulator, prior, 1000) _ = inference.append_simulations(theta, x).train(max_num_epochs=5) posterior = inference.build_posterior() posterior = posterior.set_sample_with_mcmc(sample_with_mcmc).set_mcmc_method( mcmc_method ) # Posterior should be corrected for leakage even if num_rounds just 1. samples = posterior.sample((10,), x=x_o) # Evaluate the samples to check correction factor. posterior.log_prob(samples, x=x_o)
def test_api_snl_sampling_methods(sampling_method: str, prior_str: str, set_seed): """Runs SNL on linear Gaussian and tests sampling from posterior via mcmc. Args: mcmc_method: which mcmc method to use for sampling prior_str: use gaussian or uniform prior set_seed: fixture for manual seeding """ num_dim = 2 num_samples = 10 num_trials = 2 # HMC with uniform prior needs good likelihood. num_simulations = 10000 if sampling_method == "hmc" else 1000 x_o = zeros((num_trials, num_dim)) # Test for multiple chains is cheap when vectorized. num_chains = 3 if sampling_method == "slice_np_vectorized" else 1 if sampling_method == "rejection": sample_with = "rejection" else: sample_with = "mcmc" if prior_str == "gaussian": prior = MultivariateNormal(loc=zeros(num_dim), covariance_matrix=eye(num_dim)) else: prior = utils.BoxUniform(-1.0 * ones(num_dim), ones(num_dim)) simulator, prior = prepare_for_sbi(diagonal_linear_gaussian, prior) inference = SNL(prior, show_progress_bars=False) theta, x = simulate_for_sbi(simulator, prior, num_simulations, simulation_batch_size=50) _ = inference.append_simulations(theta, x).train(max_num_epochs=5) posterior = inference.build_posterior( sample_with=sample_with, mcmc_method=sampling_method).set_default_x(x_o) posterior.sample( sample_shape=(num_samples, ), x=x_o, mcmc_parameters={ "thin": 3, "num_chains": num_chains }, )
def _set_maybe_z_scored_prior(self) -> None: r"""Compute and store potentially standardized prior (if z-scoring was done). The proposal posterior is: $pp(\theta|x) = 1/Z * q(\theta|x) * prop(\theta) / p(\theta)$ Let's denote z-scored theta by `a`: a = (theta - mean) / std Then pp'(a|x) = 1/Z_2 * q'(a|x) * prop'(a) / p'(a)$ The ' indicates that the evaluation occurs in standardized space. The constant scaling factor has been absorbed into Z_2. From the above equation, we see that we need to evaluate the prior **in standardized space**. We build the standardized prior in this function. The standardize transform that is applied to the samples theta does not use the exact prior mean and std (due to implementation issues). Hence, the z-scored prior will not be exactly have mean=0 and std=1. """ if self.z_score_theta: scale = self._neural_net._transform._transforms[0]._scale shift = self._neural_net._transform._transforms[0]._shift # Following the definintion of the linear transform in # `standardizing_transform` in `sbiutils.py`: # shift=-mean / std # scale=1 / std # Solving these equations for mean and std: estim_prior_std = 1 / scale estim_prior_mean = -shift * estim_prior_std # Compute the discrepancy of the true prior mean and std and the mean and # std that was empirically estimated from samples. # N(theta|m,s) = N((theta-m_e)/s_e|(m-m_e)/s_e, s/s_e) # Above: m,s are true prior mean and std. m_e,s_e are estimated prior mean # and std (estimated from samples and used to build standardize transform). almost_zero_mean = (self._prior.mean - estim_prior_mean) / estim_prior_std almost_one_std = torch.sqrt(self._prior.variance) / estim_prior_std if isinstance(self._prior, MultivariateNormal): self._maybe_z_scored_prior = MultivariateNormal( almost_zero_mean, torch.diag(almost_one_std)) else: range_ = torch.sqrt(almost_one_std * 3.0) self._maybe_z_scored_prior = utils.BoxUniform( almost_zero_mean - range_, almost_zero_mean + range_) else: self._maybe_z_scored_prior = self._prior
def forward(self, input_, action=None): """ """ x = torch.relu(self.fc1(input_)) x = self.bn1(x) x = torch.relu(self.fc2(torch.cat([x, input_],dim=1))) x = self.bn2(x) x = torch.relu(self.fc3(torch.cat([x, input_],dim=1))) x = self.bn3(x) x = torch.relu(self.fc4(torch.cat([x, input_],dim=1))) action_value = torch.tanh(self.action_values(x)) entries = torch.tanh(self.matrix_entries(x)) V = self.value(x) action_value = action_value.unsqueeze(-1) # create lower-triangular matrix L = torch.zeros((input_.shape[0], self.action_size, self.action_size)).to(device) # get lower triagular indices tril_indices = torch.tril_indices(row=self.action_size, col=self.action_size, offset=0) # fill matrix with entries L[:, tril_indices[0], tril_indices[1]] = entries L.diagonal(dim1=1,dim2=2).exp_() # calculate state-dependent, positive-definite square matrix P = L*L.transpose(2, 1) Q = None if action is not None: # calculate Advantage: A = (-0.5 * torch.matmul(torch.matmul((action.unsqueeze(-1) - action_value).transpose(2, 1), P), (action.unsqueeze(-1) - action_value))).squeeze(-1) Q = A + V # add noise to action mu: dist = MultivariateNormal(action_value.squeeze(-1), torch.inverse(P)) #dist = Normal(action_value.squeeze(-1), 1) action = dist.sample() action = torch.clamp(action, min=-1, max=1) #wandb.log({"Action Noise": action.detach().cpu().numpy() - action_value.squeeze(-1).detach().cpu().numpy()}) return action, Q, V
def test_c2st_snle_external_data_on_linearGaussian(set_seed): """Test whether SNPE C infers well a simple example with available ground truth. Args: set_seed: fixture for manual seeding """ num_dim = 2 device = "cpu" configure_default_device(device) x_o = zeros(1, num_dim) num_samples = 1000 # likelihood_mean will be likelihood_shift+theta likelihood_shift = -1.0 * ones(num_dim) likelihood_cov = 0.3 * eye(num_dim) prior_mean = zeros(num_dim) prior_cov = eye(num_dim) prior = MultivariateNormal(loc=prior_mean, covariance_matrix=prior_cov) gt_posterior = true_posterior_linear_gaussian_mvn_prior( x_o[0], likelihood_shift, likelihood_cov, prior_mean, prior_cov) target_samples = gt_posterior.sample((num_samples, )) def simulator(theta): return linear_gaussian(theta, likelihood_shift, likelihood_cov) infer = SNL( *prepare_for_sbi(simulator, prior), simulation_batch_size=1000, show_progress_bars=False, device=device, ) external_theta = prior.sample((1000, )) external_x = simulator(external_theta) infer.provide_presimulated(external_theta, external_x) posterior = infer( num_rounds=1, num_simulations_per_round=1000, training_batch_size=100, ).set_default_x(x_o) samples = posterior.sample((num_samples, )) # Compute the c2st and assert it is near chance level of 0.5. check_c2st(samples, target_samples, alg="snpe_c")
def loss( self, oh, ce, mask, recon_oh, recon_ce_mean, recon_ce_log_var, gamma_d, z_mean, z_log_var, avg=True): # NL1 for oh NL1 = -(oh * (recon_oh + self.det).log() ).sum(1) # cross entropy loss # NL2 for ce dist = MultivariateNormal( loc=recon_ce_mean, covariance_matrix=torch.diag_embed( recon_ce_log_var.exp().sqrt())) NL2 = (-dist.log_prob(ce.transpose(0, 1)).transpose(0, 1) * mask).sum(1) NL = NL1 + NL2 # KLD_for pi KLD1 = -torch.sum(gamma_d * torch.log(self.pi.unsqueeze(0) / gamma_d + self.det), 1) # KLD2 for all domains logvar_division = self.log_var_d.unsqueeze(0) var_division = torch.exp( z_log_var.unsqueeze(1) - self.log_var_d.unsqueeze(0)) diff = z_mean.unsqueeze(1) - self.mean_d.unsqueeze(0) diff_term = diff.pow(2) / torch.exp(self.log_var_d.unsqueeze(0)) KLD21 = torch.sum( logvar_division + var_division + diff_term, 2) KLD21 = 0.5 * torch.sum(gamma_d * KLD21, 1) KLD22 = -0.5 * torch.sum(1 + z_log_var, 1) KLD2 = KLD21 + KLD22 KLD = KLD1 + KLD2 loss = NL + KLD # in training mode, return averaged loss. In testing mode, return # individual loss if avg: return loss.mean() else: return loss
def test_conditional_corrcoeff(corr): """ Test whether the conditional correlation coefficient is computed correctly. """ d = MultivariateNormal( torch.tensor([0.6, 5.0]), torch.tensor([[0.1, corr], [corr, 10.0]]) ) estimated_corr = conditional_corrcoeff( density=d, condition=torch.ones(1, 2), limits=torch.tensor([[-2.0, 3.0], [-70, 90]]), resolution=500, )[0, 1] assert torch.abs(corr - estimated_corr) < 1e-3
def test_conditional_pairplot(): """ This only tests whether `conditional.pairplot()` runs without errors. If does not test its correctness. See `test_conditional_density_2d` for a test on `eval_conditional_density`, which is the core building block of `conditional.pairplot()` """ d = MultivariateNormal( torch.tensor([0.6, 5.0]), torch.tensor([[0.1, 0.99], [0.99, 10.0]]) ) _ = conditional_pairplot( density=d, condition=torch.ones(1, 2), limits=torch.tensor([[-1.0, 1.0], [-30, 30]]), )
def _pred_point(self, goal_embed, im_shape, min_std=0.03): if self._2_point is None: return point_dist = self._2_point(goal_embed[:,0]) mu = point_dist[:,:2] c1, c2, c3 = F.softplus(point_dist[:,2])[:,None], point_dist[:,3][:,None], F.softplus(point_dist[:,4])[:,None] scale_tril = torch.cat((c1 + min_std, torch.zeros_like(c2), c2, c3 + min_std), dim=1).reshape((-1, 2, 2)) mu, scale_tril = [x.unsqueeze(1).unsqueeze(1) for x in (mu, scale_tril)] point_dist = MultivariateNormal(mu, scale_tril=scale_tril) h = torch.linspace(-1, 1, im_shape[0]).reshape((1, -1, 1, 1)).repeat((1, 1, im_shape[1], 1)) w = torch.linspace(-1, 1, im_shape[1]).reshape((1, 1, -1, 1)).repeat((1, im_shape[0], 1, 1)) hw = torch.cat((h, w), 3).repeat((goal_embed.shape[0], 1, 1, 1)).to(goal_embed.device) return point_dist.log_prob(hw)
def act(self, state, memory=None): # action_mean = self.actor(state) action_mean = self.actor(torch.FloatTensor(state).to(device)) cov_mat = torch.diag(self.action_var).to(device) dist = MultivariateNormal(action_mean, cov_mat) action = dist.sample() action_logprob = dist.log_prob(action) if (memory): memory.states.append(state) memory.actions.append(action) memory.logprobs.append(action_logprob) return action.detach()
def evaluate(self, states, actions): #states = torch.stack(states) states_input = states.view(-1, *states.shape[-3:]) actor_critic_input = self.conv(states_input).view(-1, self.size) action_mean = self.actor(actor_critic_input) action_var = self.action_var.expand_as(action_mean) cov_mat = torch.diag_embed(action_var).to(device) dist = MultivariateNormal(action_mean, cov_mat) actions = actions.view(-1, self.action_dim) action_logprobs = dist.log_prob(actions).view(states.shape[:-3]) dist_entropy = dist.entropy().view(states.shape[:-3]) state_value = self.critic(actor_critic_input).view(states.shape[:-3]) return action_logprobs, torch.squeeze(state_value), dist_entropy
def log_prob(self, x): log_prob = 0 for layer in self.layers[::-1]: x, log_prob_change = layer.g(x) log_prob = log_prob_change + log_prob if self.prior is None: norm_prior = MultivariateNormal(torch.zeros(self.num_vars).to(x.device), torch.eye(self.num_vars).to(x.device)) log_prob += norm_prior.log_prob(x) else: log_prob += self.prior.log_prob(x) return log_prob
def _construct_mvn(x, w): """ Constructs a multivariate normal distribution of weighted samples. :param x: The samples :type x: torch.Tensor :param w: The weights :type w: torch.Tensor :rtype: MultivariateNormal """ mean = (x * w.unsqueeze(-1)).sum(0) centralized = x - mean cov = torch.matmul(w * centralized.t(), centralized) return MultivariateNormal(mean, scale_tril=torch.cholesky(cov))
def __init__(self, data, mu_prior, alpha_prior, W_df_prior, W_prior, G_df_prior, rate_prior): d = W_prior.shape[0] self.W_prior = Wishart({ 'df': torch.tensor([W_df_prior], dtype=torch.float64), 'W': torch.from_numpy(W_prior.astype(np.float64)) }) self.nu_prior = Gamma(torch.tensor([G_df_prior], dtype=torch.float64), torch.tensor([rate_prior], dtype=torch.float64)) self.mu_prior = MultivariateNormal( loc=mu_prior * torch.ones(d, dtype=torch.float64), covariance_matrix=alpha_prior * torch.eye(d, dtype=torch.float64)) self.data = data
def get_dist(self, duration=None): """ Get the :class:`~pyro.distributions.GaussianHMM` distribution that corresponds to ``obs_dim``-many independent Matern GPs. :param int duration: Optional size of the time axis ``event_shape[0]``. This is required when sampling from homogeneous HMMs whose parameters are not expanded along the time axis. """ trans_matrix, process_covar = self.kernel.transition_matrix_and_covariance(dt=self.dt) trans_dist = MultivariateNormal(self.obs_matrix.new_zeros(self.obs_dim, 1, self.kernel.state_dim), process_covar.unsqueeze(-3)) trans_matrix = trans_matrix.unsqueeze(-3) return dist.GaussianHMM(self._get_init_dist(), trans_matrix, trans_dist, self.obs_matrix, self._get_obs_dist(), duration=duration)
def __init__(self, **kwargs): super(CPGGAN, self).__init__(**kwargs) if self.mode == 'train': # path to numpy arrays containing the calculated mean and cov matrices for calculating a multivariate gaussian path_to_lm_mean = kwargs.get('lm_mean', ARRAY_LANDMARKS_28_MEAN) path_to_lm_cov = kwargs.get('lm_cov', ARRAY_LANDMARKS_28_COV) path_to_lr_mean = kwargs.get('lr_mean', ARRAY_LOWRES_4_MEAN) path_to_lr_cov = kwargs.get('lr_cov', ARRAY_LOWRES_4_COV) # ================================================== # Currently only preparation for extensions # ================================================== # gaussian distribution of our landmarks self.landmarks_mean = np.load(path_to_lm_mean) self.landmarks_cov = np.load(path_to_lm_cov) self.landmarks_mean = torch.from_numpy(self.landmarks_mean) self.landmarks_cov = torch.from_numpy(self.landmarks_cov) self.distribution_landmarks = MultivariateNormal( loc=self.landmarks_mean.type(torch.float64), covariance_matrix=self.landmarks_cov.type(torch.float64)) # gaussian distribution of our low res pixel map self.lowres_mean = np.load(path_to_lr_mean) self.lowres_cov = np.load(path_to_lr_cov) self.lowres_mean = torch.from_numpy(self.lowres_mean) self.lowres_cov = torch.from_numpy(self.lowres_cov) self.distribution_lowres = MultivariateNormal( loc=self.lowres_mean.type(torch.float64), covariance_matrix=self.lowres_cov.type(torch.float64)) # static noise for calculating the validation self.static_landmarks = 2 * (self.distribution_landmarks.sample( (self.batch_size, )).type(torch.float32) - 0.5) self.static_lowres = 2 * (self.distribution_lowres.sample( (self.batch_size, )).type(torch.float32) - 0.5) # Static noise for anonymization self.anonymization_noise = self.noise(1)
def evaluate(self, state, action): action_mean = torch.squeeze(self.actor(state)) action_var = self.action_var.expand_as(action_mean) cov_mat = torch.diag_embed(action_var).to(device) #cov_mat = torch.diag(action_var).to(device) dist = MultivariateNormal(action_mean, cov_mat) action_logprobs = dist.log_prob(torch.squeeze(action)) dist_entropy = dist.entropy() state_value = self.critic(state) return action_logprobs, torch.squeeze(state_value), dist_entropy
def forecast(self, targets, dts): """ :param torch.Tensor targets: A 2-dimensional tensor of real-valued targets of shape `(T, obs_dim)`, where `T` is the length of the time series and `obs_dim` is the dimension of the real-valued targets at each time step. These represent the training data that are conditioned on for the purpose of making forecasts. :param torch.Tensor dts: A 1-dimensional tensor of times to forecast into the future, with zero corresponding to the time of the final target `targets[-1]`. :returns torch.distributions.MultivariateNormal: Returns a predictive MultivariateNormal distribution with batch shape `(S,)` and event shape `(obs_dim,)`, where `S` is the size of `dts`. """ filtering_state = self._filter(targets) predicted_mean, predicted_covar = self._forecast(dts, filtering_state) return MultivariateNormal(predicted_mean, predicted_covar)
def act(self, state, memory): action_mean = self.actor(state) cov_mat = torch.diag(self.action_var).to(device) dist = MultivariateNormal(action_mean, cov_mat) action = action_mean + self.action_std * torch.randn(action_mean.shape) # action = dist.sample() action_logprob = dist.log_prob(action) memory.states.append(state) memory.actions.append(action) memory.logprobs.append(action_logprob) return action.detach()
def act(self, ob, state, memory): latent = self.encoder(ob) state = torch.cat((latent, state), 1) action_mean = self.actor(state) dist = MultivariateNormal(action_mean, torch.diag(self.action_var).to(device)) action = dist.sample() action_logprob = dist.log_prob(action) memory.states.append(state) # memory.obs.append(ob) memory.actions.append(action) memory.logprobs.append(action_logprob) return action.detach()
def log_ratio(self, x, px_mean, px_var, qz_m, qz_v, z, return_full=False): if self.learn_var: log_px_z = (Normal(px_mean, torch.sqrt(px_var)).log_prob( x.repeat(px_mean.shape[0], 1, 1)).sum(dim=-1).view( (px_mean.shape[0], -1))) log_pz = (Normal(torch.zeros_like(qz_m), torch.ones_like(qz_m)).log_prob(z).sum(dim=-1)) if not self.linear_encoder: log_qz_given_x = Normal(qz_m, qz_v.sqrt()).log_prob(z).sum(dim=-1) else: log_qz_given_x = MultivariateNormal(qz_m, qz_v).log_prob( z) # No need to sum over latent dim log_pxz = log_px_z + log_pz else: zx = torch.cat([z, x.repeat(px_mean.shape[0], 1, 1)], dim=-1) reshape_dim = x.shape[-1] + z.shape[-1] if not self.linear_encoder: log_qz_given_x = Normal(qz_m, qz_v.sqrt()).log_prob(z).sum(dim=-1) else: log_qz_given_x = MultivariateNormal(qz_m, qz_v).log_prob(z) log_pxz = self.log_normal_full( zx.view((-1, reshape_dim)), torch.zeros_like(zx.view((-1, reshape_dim))), self.log_det_pxz, self.inv_sqrt_pxz, ).view((px_mean.shape[0], -1)) log_ratio = log_pxz - log_qz_given_x if return_full: return log_ratio, log_pxz, log_qz_given_x else: return log_ratio
def act(self, state, memory): action_mean = self._action_mean(state) cov_mat = torch.diag(self.action_var).to(device) dist = MultivariateNormal(action_mean, cov_mat) action = dist.sample() # action = F.softmax(action.reshape(2,-1)).reshape(1,-1) action_logprob = dist.log_prob(action) memory.states.append(state) memory.actions.append(action) memory.logprobs.append(action_logprob) return action.detach()
def test_multivariate_normal() -> None: num_samples = 2000 dim = 2 mu = np.arange(0, dim) / float(dim) L_diag = np.ones((dim, )) L_low = 0.1 * np.ones((dim, dim)) * np.tri(dim, k=-1) L = np.diag(L_diag) + L_low Sigma = L.dot(L.transpose()) distr = MultivariateNormal(loc=torch.Tensor(mu), scale_tril=torch.Tensor(L)) samples = distr.sample((num_samples, )) mu_hat, L_hat = maximum_likelihood_estimate_sgd( MultivariateNormalOutput(dim=dim), samples, init_biases= None, # todo we would need to rework biases a bit to use it in the multivariate case learning_rate=0.01, num_epochs=10, ) distr = MultivariateNormal(loc=torch.tensor(mu_hat), scale_tril=torch.tensor(L_hat)) Sigma_hat = distr.covariance_matrix.numpy() assert np.allclose( mu_hat, mu, atol=0.1, rtol=0.1), f"mu did not match: mu = {mu}, mu_hat = {mu_hat}" assert np.allclose( Sigma_hat, Sigma, atol=0.1, rtol=0.1 ), f"Sigma did not match: sigma = {Sigma}, sigma_hat = {Sigma_hat}"
def act(self, state, memory, stochastic=True): action_mean = self.actor(state) cov_mat = torch.diag(self.action_var).to(device) dist = MultivariateNormal(action_mean, cov_mat) action = dist.sample() action_logprob = dist.log_prob(action) if not stochastic: action = action_mean memory.states.append(state) memory.actions.append(action) memory.logprobs.append(action_logprob) return action.detach()
def _construct_mvn(x: torch.Tensor, w: torch.Tensor, scale=1.0): """ Constructs a multivariate normal distribution of weighted samples. """ mean = (x * w.unsqueeze(-1)).sum(0) centralized = x - mean cov = torch.matmul(w * centralized.t(), centralized) if cov.det() == 0.0: chol = cov.diag().sqrt().diag() else: chol = cov.cholesky() return MultivariateNormal(mean, scale_tril=scale * chol)
def act(self, state, opponent_state): if self.has_continuous_action_space: pre_mean, pre_sigma = self.om(opponent_state) # print('sigma: ', pre_sigma) pre_var = pre_sigma**2 pre_var = pre_var.repeat(1, 2).to(device) pre_mat = torch.diag_embed(pre_var).to(device) pre_dist = MultivariateNormal(pre_mean, pre_mat) pre_action = pre_dist.sample() pre_action = pre_action.clamp(-1, 1) action_mean, action_sigma = self.actor(state, pre_action[0]) action_var = action_sigma**2 action_var = action_var.repeat(1, 2).to(device) cov_mat = torch.diag_embed(action_var).to(device) dist = MultivariateNormal(action_mean, cov_mat) else: action_probs = self.actor(state) dist = Categorical(action_probs) action = dist.sample() action = action.clamp(-1, 1) action_logprob = dist.log_prob(action) return action.detach(), action_logprob.detach(), pre_action.detach()