diff --git a/config/value_dice.json5 b/config/value_dice.json5 index 439eef0..a2289e6 100644 --- a/config/value_dice.json5 +++ b/config/value_dice.json5 @@ -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, diff --git a/scratch/johannes/valuedice.py b/scratch/johannes/valuedice.py new file mode 100644 index 0000000..82311b4 --- /dev/null +++ b/scratch/johannes/valuedice.py @@ -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) \ No newline at end of file diff --git a/src/value_dice/value_dice.py b/src/value_dice/value_dice.py index d9e29d6..574c588 100644 --- a/src/value_dice/value_dice.py +++ b/src/value_dice/value_dice.py @@ -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)