def softmax_like(env, *, trajectory_model, agent_model, log=False): """softmax_like :param env: OpenAI Gym environment :param trajectory_model: trajectory probabilistic program :param agent_model: agent's probabilistic program :param log: boolean; if True, print log info """ Qs = torch.as_tensor([ infer_Q( env, action, trajectory_model=trajectory_model, agent_model=agent_model, log=log, ) for action in range(env.action_space.n) ]) action_logits = args.alpha * Qs action_dist = Categorical(logits=action_logits) if log: print('policy:') print( tabulate( [action_dist.probs.tolist()], headers=env.actions, tablefmt='fancy_grid', )) return action_dist.sample()
def sample_TRG_sentence(self, src_lengths): num_words = self.model.trg_embed.weight.size(0) prob = 1. / num_words distr = Categorical( probs=torch.tensor([prob for _ in range(num_words)])) trgs = [[self.sos_index] for _ in src_lengths] for i, s in enumerate(src_lengths): for _ in range(s): trgs[i] = trgs[i] + [distr.sample().item()] return src_lengths.new_tensor(trgs).long()
def _tmc_mixture_sample(msg): dist, num_samples = msg["fn"], msg["infer"].get("num_samples") # find batch dims that aren't plate dims batch_shape = [1] * len(dist.batch_shape) for f in msg["cond_indep_stack"]: if f.vectorized: batch_shape[f.dim] = f.size if f.size > 0 else dist.batch_shape[ f.dim] batch_shape = tuple(batch_shape) # sample a batch sample_shape = (num_samples, ) fat_sample = dist(sample_shape=sample_shape) # TODO thin before sampling assert fat_sample.shape == sample_shape + dist.batch_shape + dist.event_shape assert any(d > 1 for d in fat_sample.shape) target_shape = (num_samples, ) + batch_shape + dist.event_shape # if this site has any possible ancestors, sample ancestor indices uniformly thin_sample = fat_sample if thin_sample.shape != target_shape: index = [Ellipsis] + [slice(None)] * (len(thin_sample.shape) - 1) squashed_dims = [] for squashed_dim, squashed_size in zip( range(1, len(thin_sample.shape)), thin_sample.shape[1:]): if squashed_size > 1 and (target_shape[squashed_dim] == 1 or squashed_dim == 0): # uniformly sample one ancestor per upstream particle population ancestor_dist = Categorical(logits=torch.zeros( (squashed_size, ), device=thin_sample.device)) ancestor_index = ancestor_dist.sample( sample_shape=(num_samples, )) index[squashed_dim] = ancestor_index squashed_dims.append(squashed_dim) thin_sample = Vindex(thin_sample)[tuple(index)] for squashed_dim in squashed_dims: thin_sample = thin_sample.unsqueeze(squashed_dim) assert thin_sample.shape == target_shape return thin_sample
def addWord(lex): newLex = lex.duplicate() x = Categorical(probs=torch.tensor(world.mis)) pair = world.word_object_pairs[int(x.sample())] newLex.addPair(pair[0], pair[1]) return newLex