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

Question

我已经以OpenAi Gym环境的方式实现了自定义环境，在该环境中，我具有可以在边界内的平面上移动的形状（圆形，正方形等）。我想在这个问题上应用诸如REINFORCE之类的强化学习算法，并制作自己的自定义奖励函数。但是我现在在设计神经网络以获取状态并输出动作时遇到了麻烦。

提供一些有关我的状态操作设置的背景信息。对于状态，我使用2D数组描述形状的类型（圆形，正方形等）作为一个热编码数组，并描述位置（边界框边缘的x，y坐标）和形状的旋转（ 0-360），并且当前平面上的每个形状都有一行。对于动作，我有一个1d数组，其中包含要选择的形状的索引以及新的位置和旋转。我首先将2d数组展平为单个数组，然后将其隐藏，然后分支为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来计算对数概率，但是现在我还必须结合分类概率。用单个神经网络预测指数，μ和sigma是否正确？我刚刚进入RL，请先感谢