Я делаю проект обучения с подкреплением и пытаюсь получить тензор, который представляет ожидаемое вознаграждение за все заданные действия. У меня есть длинный тензор выбранных действий размером batch со значениями либо ноль, либо единица (два возможных действия). У меня есть тензор ожидаемых вознаграждений за каждое действие размером batch * action_size, и я хочу, чтобы тензор размера batch.
Например, если размер партии был 4, тогда у меня есть
action = tensor([1,0,0,1])
expectedReward = tensor([[3,7],[5,9],[-1,12],[0,1]])
и я хочу
rewardForActions = tensor([7,5,-1,1])
Я думал, этот ответит на мой вопрос, но это совсем не то же самое, потому что, если бы я пошел с этим решением, он бы в итоге получим тензор 4 * 4, выбирая из каждой строки 4 раза, а не один раз.
Есть идеи?