如何处理多输出神经网络的损失函数和对数概率？

Question

我已经实现了一个OpenAi健身房环境风格的自定义环境，在这个环境中，我可以在边界内的平面上移动形状（圆、正方形等）。我想在这个问题上应用强化学习算法，比如REINFORCE，并定制了自己的奖励函数。但是我现在在设计神经网络以获取一个状态并输出一个动作时遇到了麻烦。你知道吗

提供一些关于我的状态操作设置的背景信息。对于这个状态，我使用一个2d数组来描述形状的类型（圆、正方形等）作为一个热编码数组和位置（边界框边缘的x、y坐标）以及形状的旋转（0-360），并且为平面上当前的每个形状都有一行。对于这个动作，我有一个一维数组，其中包含要选择的形状索引以及新的位置和旋转。我首先将二维数组展平为单个数组，然后有一个隐藏层，然后分支为3个输出：连续输出的概率分布的Mu和Sigma，以及选择形状索引的概率。你知道吗

class Model(tf.keras.Model):
    def __init__(self, action_pair, input_shape):
        super().__init__('mlp_policy')

        self.input = kl.Input(shape=input_shape)
        self.flat = kl.Flatten()(self.input)
        self.hidden = kl.Dense(128, activation='relu')(self.flat)
        # action_pair is a tuple with the size of the categorical output and the continuous output
        self.index = kl.Dense(action_pair[0], activation='softmax', name='index')(self.hidden)
        self.mu_branch = kl.Dense(action_pair[1], activation='tanh', name='mu')(self.hidden)
        self.sigma_branch = kl.Dense(action_pair[1], activation='softplus', name='sigma')(self.hidden)


    def call(self, inputs):

        x = self.input(inputs)
        x = self.flat(x)
        hid = self.hidden(x)

        index = self.index(hid)
        mu = self.mu_branch(hid)
        sigma = self.sigma_branch(hid)

        return index_prob, mu, sigma

我不知道我应该如何处理损失函数的对数概率，因为通常我会使用mu和sigma来计算对数概率，但是现在我还必须合并分类概率。用一个神经网络预测指数mu和sigma是否正确？我刚进入RL，提前谢谢