huangfu/InteractionImitation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

MInor change in value dice loss, activate print statements, only do EITHER value OR policy update for each batch

Browse Source
This commit is contained in:
Johannes Fischer
2021-08-06 18:52:17 +02:00
parent 025c71767f
commit 88b4466e57
1 changed files with 27 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
src/value_dice/value_dice.py
View File

@@ -1,3 +1,5 @@
import numpy as np
import torch import torch
import torch.nn as nn import torch.nn as nn
from torch.utils.data import DataLoader from torch.utils.data import DataLoader
@@ -213,7 +215,7 @@ def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs):
# nonlinear loss # nonlinear loss
value_diff = value - discount * value_next value_diff = value - discount * value_next
nonlinear_loss = torch.logsumexp(value_diff, dim=0) nonlinear_loss = torch.logsumexp(value_diff, dim=0) - np.log(len(value_diff))
loss = nonlinear_loss - linear_loss loss = nonlinear_loss - linear_loss
return loss return loss
@@ -243,22 +245,33 @@ def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs):
policy_loss = -loss policy_loss = -loss
value_loss = loss value_loss = loss
# compute loss and step optimizer # # compute loss and step optimizer
# policy_optimizer.zero_grad()
# value_optimizer.zero_grad()
# policy_loss.backward(retain_graph=True)
# value_loss.backward()
# clip_grad_norm_(policy.policy.parameters(), clip_grad_norm)
# clip_grad_norm_(policy.value.parameters(), clip_grad_norm)
# policy_optimizer.step()
# value_optimizer.step()
if batch_idx % 2 == 0:
policy_optimizer.zero_grad() policy_optimizer.zero_grad()
value_optimizer.zero_grad() policy_loss.backward()
policy_loss.backward(retain_graph=True)
value_loss.backward()
clip_grad_norm_(policy.policy.parameters(), clip_grad_norm) clip_grad_norm_(policy.policy.parameters(), clip_grad_norm)
clip_grad_norm_(policy.value.parameters(), clip_grad_norm)
policy_optimizer.step() policy_optimizer.step()
else:
value_optimizer.zero_grad()
value_loss.backward()
clip_grad_norm_(policy.value.parameters(), clip_grad_norm)
value_optimizer.step() value_optimizer.step()
epoch_loss += loss.item() / len(train_dataset) epoch_loss += loss.item() / len(train_dataset)
# if i % print_epoch_every == 0: if i % print_epoch_every == 0:
# print('Epoch: {}, Training Loss: {}'.format(i, epoch_loss)) print('Epoch: {}, Training Loss: {}'.format(i, epoch_loss))
# measure cv loss # measure cv loss
if i % cv_every == 0: if i % cv_every == 0:
@@ -280,8 +293,8 @@ def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs):
if i % cv_every == 0: if i % cv_every == 0:
writer.add_scalar('cv loss', cv_loss, i) writer.add_scalar('cv loss', cv_loss, i)
# if i % print_cv_every == 0: if i % print_cv_every == 0:
# print('Epoch: {}, CV Loss: {}'.format(i, cv_loss)) print('Epoch: {}, CV Loss: {}'.format(i, cv_loss))
policy.save_model(filestr) policy.save_model(filestr)