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

Add debug output and original implementation

Browse Source
This commit is contained in:
Johannes Fischer
2021-09-08 13:52:21 +02:00
parent 1fd0a71646
commit 26cde18be3
3 changed files with 149 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
config/value_dice.json5
View File

@@ -9,12 +9,12 @@
deepsets: {
input_dim: 6, // number of relative state vars for others
phi: {
hidden_n: 2,
hidden_n: 1,
hidden_dim: 20,
},
latent_dim: 20,
rho: {
hidden_n: 2,
hidden_n: 1,
hidden_dim: 10,
},
output_dim: 10
@@ -43,12 +43,12 @@
deepsets: {
input_dim: 6, // number of relative state vars for others
phi: {
hidden_n: 2,
hidden_n: 1,
hidden_dim: 20,
},
latent_dim: 20,
rho: {
hidden_n: 2,
hidden_n: 1,
hidden_dim: 10,
},
output_dim: 10
@@ -70,13 +70,13 @@
},
policy_optim: {
optimizer: 'adam',
lr: 1e-3,
weight_decay: 0.1,
lr: 1e-0,
weight_decay: 0.01,
},
value_optim: {
optimizer: 'adam',
lr: 1e-3,
weight_decay: 0.1,
lr: 1e-6,
weight_decay: 0.01,
},
train_epochs: 200,
train_batch_size: 32,

121
scratch/johannes/valuedice.py Normal file
View File

@@ -0,0 +1,121 @@
def weighted_softmax(x, weights, axis=0):
x = x - tf.reduce_max(x, axis=axis)
return weights * tf.exp(x) / tf.reduce_sum(
weights * tf.exp(x), axis=axis, keepdims=True)
@tf.function
def update(self,
expert_dataset_iter,
policy_dataset_iter,
discount,
replay_regularization=0.05,
nu_reg=10.0):
"""A function that updates nu network.
When replay regularization is non-zero, it learns
(d_pi * (1 - replay_regularization) + d_rb * replay_regulazation) /
(d_expert * (1 - replay_regularization) + d_rb * replay_regulazation)
instead.
Args:
expert_dataset_iter: An tensorflow graph iteratable over expert data.
policy_dataset_iter: An tensorflow graph iteratable over training policy
data, used for regularization.
discount: An MDP discount.
replay_regularization: A fraction of samples to add from a replay buffer.
nu_reg: A grad penalty regularization coefficient.
"""
(expert_states, expert_actions,
expert_next_states) = expert_dataset_iter.get_next()
expert_initial_states = expert_states
rb_states, rb_actions, rb_next_states, _, _ = policy_dataset_iter.get_next(
)[0]
with tf.GradientTape(
watch_accessed_variables=False, persistent=True) as tape:
tape.watch(self.actor.variables)
tape.watch(self.nu_net.variables)
_, policy_next_actions, _ = self.actor(expert_next_states)
# _, rb_next_actions, rb_log_prob = self.actor(rb_next_states)
_, policy_initial_actions, _ = self.actor(expert_initial_states)
Inputs for the linear part of DualDICE loss.
expert_init_inputs = tf.concat(
[expert_initial_states, policy_initial_actions], 1)
expert_inputs = tf.concat([expert_states, expert_actions], 1)
expert_next_inputs = tf.concat([expert_next_states, policy_next_actions],
1)
rb_inputs = tf.concat([rb_states, rb_actions], 1)
rb_next_inputs = tf.concat([rb_next_states, rb_next_actions], 1)
expert_nu_0 = self.nu_net(expert_init_inputs)
expert_nu = self.nu_net(expert_inputs)
expert_nu_next = self.nu_net(expert_next_inputs)
rb_nu = self.nu_net(rb_inputs)
rb_nu_next = self.nu_net(rb_next_inputs)
expert_diff = expert_nu - discount * expert_nu_next
rb_diff = rb_nu - discount * rb_nu_next
linear_loss_expert = tf.reduce_mean(expert_nu_0 * (1 - discount))
linear_loss_rb = tf.reduce_mean(rb_diff)
rb_expert_diff = tf.concat([expert_diff, rb_diff], 0)
rb_expert_weights = tf.concat([
tf.ones(expert_diff.shape) * (1 - replay_regularization),
tf.ones(rb_diff.shape) * replay_regularization
], 0)
rb_expert_weights /= tf.reduce_sum(rb_expert_weights)
non_linear_loss = tf.reduce_sum(
tf.stop_gradient(
weighted_softmax(rb_expert_diff, rb_expert_weights, axis=0)) *
rb_expert_diff)
linear_loss = (
linear_loss_expert * (1 - replay_regularization) +
linear_loss_rb * replay_regularization)
loss = (non_linear_loss - linear_loss)
alpha = tf.random.uniform(shape=(expert_inputs.shape[0], 1))
nu_inter = alpha * expert_inputs + (1 - alpha) * rb_inputs
nu_next_inter = alpha * expert_next_inputs + (1 - alpha) * rb_next_inputs
nu_inter = tf.concat([nu_inter, nu_next_inter], 0)
with tf.GradientTape(watch_accessed_variables=False) as tape2:
tape2.watch(nu_inter)
nu_output = self.nu_net(nu_inter)
nu_grad = tape2.gradient(nu_output, [nu_inter])[0] + EPS
nu_grad_penalty = tf.reduce_mean(
tf.square(tf.norm(nu_grad, axis=-1, keepdims=True) - 1))
nu_loss = loss + nu_grad_penalty * nu_reg
pi_loss = -loss + keras_utils.orthogonal_regularization(self.actor.trunk)
nu_grads = tape.gradient(nu_loss, self.nu_net.variables)
pi_grads = tape.gradient(pi_loss, self.actor.variables)
self.nu_optimizer.apply_gradients(zip(nu_grads, self.nu_net.variables))
self.actor_optimizer.apply_gradients(zip(pi_grads, self.actor.variables))
del tape
self.avg_nu_expert(expert_nu)
self.avg_nu_rb(rb_nu)
self.nu_reg_metric(nu_grad_penalty)
self.avg_loss(loss)
self.avg_actor_loss(pi_loss)
self.avg_actor_entropy(-rb_log_prob)

24
src/value_dice/value_dice.py
View File

@@ -215,9 +215,16 @@ def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs):
# nonlinear loss
value_diff = value - discount * value_next
nonlinear_loss = torch.logsumexp(value_diff, dim=0) - np.log(len(value_diff))
# print(value_diff)
# nonlinear_loss = torch.logsumexp(value_diff, dim=0) #- np.log(len(value_diff))
nonlinear_loss = torch.log(torch.mean(torch.exp(value_diff), dim=0))
loss = nonlinear_loss - linear_loss
print("Loss report:")
print("Linear: {}".format(linear_loss.item()))
print("Nonlinear: {}".format(nonlinear_loss.item()))
print("Total: {}".format(loss.item()))
return loss
@@ -237,7 +244,6 @@ def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs):
# train
epoch_loss = 0
for (batch_idx, batch) in enumerate(training_loader):
loss = f_value_dice_loss(batch)
# In original implementation policy is regularized with orthogonal regularization,
@@ -260,12 +266,22 @@ def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs):
if batch_idx % 2 == 0:
policy_optimizer.zero_grad()
policy_loss.backward()
clip_grad_norm_(policy.policy.parameters(), clip_grad_norm)
# clip_grad_norm_(policy.policy.parameters(), clip_grad_norm)
policy_optimizer.step()
grad_list = torch.cat([torch.flatten(p.grad) for p in policy.policy.parameters()])
torch.mean(grad_list)
print("gradient stats:")
print(torch.mean(grad_list))
print(torch.std(grad_list))
print(torch.min(grad_list))
print(torch.max(grad_list))
# print(policy.policy.head.layers[0].weight.grad)
# print(policy.policy.head.layers[0].bias.grad)
else:
value_optimizer.zero_grad()
value_loss.backward()
clip_grad_norm_(policy.value.parameters(), clip_grad_norm)
# clip_grad_norm_(policy.value.parameters(), clip_grad_norm)
value_optimizer.step()
epoch_loss += loss.item() / len(train_dataset)