需要让action服从action_dist的概率分布得到tenor(1,2)，而不是把action增加维度

如果你需要得到一个形状为[1,2]的张量，并使其服从action_dist的概率分布，可以使用action_dist.probs属性得到概率分布，然后使用torch.multinomial()方法来根据概率分布进行采样。代码如下：

action_dist = torch.distributions.Categorical(probs)
probs = action_dist.probs.unsqueeze(0)  # 将概率分布的形状从[2]变为[1, 2]
action = torch.multinomial(probs, 1)  # 从概率分布中采样1个样本

这里使用了torch.multinomial()方法，它可以根据多项式分布进行采样。在这里，我们将概率分布的形状从[2]变为[1, 2]，然后从中采样1个样本。最终得到的action是一个形状为[1, 1]的张量，其中的元素是0或1，表示两个离散动作的选择情况。如果你需要得到形状为[1, 2]的张量，可以使用torch.squeeze()方法来去除多余的维度。代码如下：

action_dist = torch.distributions.Categorical(probs)
probs = action_dist.probs.unsqueeze(0)  # 将概率分布的形状从[2]变为[1, 2]
action = torch.multinomial(probs, 1)  # 从概率分布中采样1个样本
action = torch.squeeze(action, 1)  # 去除多余维度，得到形状为[1, 2]的张量