Сам разобрался.
В итоге всё сводится к тому, чтобы в своей какой-то функции обработать ответы сети, оценить их и самостоятельно высчитать ошибку. Потом взять любой ответ сети, на его основе создать свой правильный ответ, в сравнении с которым будет ошибка нужного нам размера, и потом это всё запихнуть в функцию потерь. Вся фишка в том, что к ответам сети привязаны графы их получения, т.е. в тензоре ответа сети есть вся последовательность как он получен. И на основании правильного ответа, ответа сети и этого графа, привязанного к ответу сети, выполняется обратное распространение ошибки. Нет графа - нет обучения. ))
Ещё один момент - если в сети используется какой-то не типовой не дифференцируемый слой, то штатная оптимизация тоже не будет работать, но при этом никаких ошибок не покажет.
В моём случае с выходной активацией softmax это оказалось не очень удачным вариантом, потому что допустим при ошибке 0,2 я не смог придумать как правильно создать целевой ответ.
Условный код как это работает:
optimizer.zero_grad()
answers = agent.forward(train_data)
#вычисляем ошибку на основе ответов сети в какой-то своей функции f().
#Получаем например 0,2 - типа ошибка 20%
nn_error = f(answers)
#берём первый из ответов сети, который содержит граф расчётов
nn_ans = answers[0]
#тут нам нужно создать правильный ответ - возможны вариант и нужно придумать правильно
#чтобы получить правильный можно как прибавить 20%, так и вычесть. Я прибавил.
target_ans = nn_ans * (1 + nn_error)
#функцию потерь можно сделать как свою, так и использовать штатную
f_loss = nn.L1Loss()
loss = f_loss(nn_ans, target_ans)
loss.backward()
optimizer.step()
#print(list(agent.parameters()))
