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
