def update_final_weight(self, stage1_input, stage2_input):
current_obs_1st, action_1st, _, discount_1st, next_obs_1st = stage1_input[:
5]
current_obs_2nd, action_2nd, reward_2nd = stage2_input[:3]
next_action_1st = self.policy(next_obs_1st)
discount_1st = tf.expand_dims(discount_1st, axis=1)
instrumental_feature_1st = self.instrumental_feature(
obs=current_obs_1st, action=action_1st, training=False)
instrumental_feature_2nd = self.instrumental_feature(
obs=current_obs_2nd, action=action_2nd, training=False)
target_1st = discount_1st * add_const_col(
self.value_feature(
obs=next_obs_1st, action=next_action_1st, training=True))
target_1st = add_const_col(
self.value_feature(obs=current_obs_1st,
action=action_1st,
training=True)) - self.discount * target_1st
stage1_weight = fit_linear(target_1st, instrumental_feature_1st,
self.stage1_reg)
self.stage1_weight.assign(stage1_weight)
predicted_feature = linear_reg_pred(instrumental_feature_2nd,
stage1_weight)
stage2_weight = fit_linear(tf.expand_dims(reward_2nd, -1),
predicted_feature, self.stage2_reg)
self.value_func.weight.assign(stage2_weight)
return stage1_weight, stage2_weight
def update_value(self, stage1_input, stage2_input):
current_obs_1st, action_1st, _, discount_1st, next_obs_1st = stage1_input[:
5]
current_obs_2nd, action_2nd, reward_2nd = stage2_input[:3]
next_action_1st = self.policy(next_obs_1st)
discount_1st = tf.expand_dims(discount_1st, axis=1)
instrumental_feature_1st = self.instrumental_feature(
obs=current_obs_1st, action=action_1st, training=False)
instrumental_feature_2nd = self.instrumental_feature(
obs=current_obs_2nd, action=action_2nd, training=False)
l2 = snt.regularizers.L2(self.value_reg)
with tf.GradientTape() as tape:
target_1st = discount_1st * add_const_col(
self.value_feature(
obs=next_obs_1st, action=next_action_1st, training=True))
target_1st = add_const_col(
self.value_feature(obs=current_obs_1st,
action=action_1st,
training=True)) - self.discount * target_1st
stage1_weight = fit_linear(target_1st, instrumental_feature_1st,
self.stage1_reg)
predicted_feature = linear_reg_pred(instrumental_feature_2nd,
stage1_weight)
loss = linear_reg_loss(tf.expand_dims(reward_2nd, -1),
predicted_feature, self.stage2_reg)
loss = loss + l2(self.value_feature.trainable_variables)
loss /= action_2nd.shape[0]
gradient = tape.gradient(loss, self.value_feature.trainable_variables)
self._value_func_optimizer.apply(
gradient, self.value_feature.trainable_variables)
return loss
def cal_validation_err(self, valid_input):
"""Return prediction MSE on the validation dataset."""
stage1_weight = self.stage1_weight
stage2_weight = self.value_func.weight
se_sum = 0.
se2_sum = 0.
weight_sum = 0.
for sample in valid_input:
data = sample.data
current_obs, action, reward = data[:3]
d_tm1 = self._get_d_tm1(data)
d_tm1 = tf.expand_dims(d_tm1, axis=1)
instrumental_feature = self.instrumental_feature(obs=current_obs,
action=action,
training=False)
predicted_feature = linear_reg_pred(instrumental_feature,
stage1_weight)
current_feature = add_const_col(
self.value_feature(obs=current_obs,
action=action,
training=True))
predicted_feature = current_feature - d_tm1 * self.discount * predicted_feature
predict = linear_reg_pred(predicted_feature, stage2_weight)
weight = d_tm1 + (1.0 - d_tm1) * tf.convert_to_tensor(
self.d_tm1_weight, dtype=tf.float32)
weight = tf.square(weight)
sq_err = tf.square(tf.expand_dims(reward, -1) - predict)
se_sum += tf.reduce_sum(weight * sq_err)
se2_sum += tf.reduce_sum(weight * tf.square(sq_err))
weight_sum += tf.reduce_sum(weight)
mse = se_sum / weight_sum
mse_err_std = tf.sqrt((se2_sum / weight_sum - mse**2) / weight_sum)
return mse, mse_err_std
def update_value(self, stage1_input, stage2_input):
current_obs_1st, action_1st, _, discount_1st, next_obs_1st = stage1_input[:
5]
d_tm1_1st = self._get_d_tm1(stage1_input)
current_obs_2nd, action_2nd, reward_2nd = stage2_input[:3]
d_tm1_2nd = self._get_d_tm1(stage2_input)
next_action_1st = self.policy(next_obs_1st)
discount_1st = tf.expand_dims(discount_1st, axis=1)
d_tm1_1st = tf.expand_dims(d_tm1_1st, axis=1)
d_tm1_2nd = tf.expand_dims(d_tm1_2nd, axis=1)
instrumental_feature_1st = self.instrumental_feature(
obs=current_obs_1st, action=action_1st, training=False)
instrumental_feature_1st = d_tm1_1st * instrumental_feature_1st
instrumental_feature_2nd = self.instrumental_feature(
obs=current_obs_2nd, action=action_2nd, training=False)
l2 = snt.regularizers.L2(self.value_reg)
with tf.GradientTape() as tape:
# target_1st = discount_1st * self.value_feature(obs=next_obs_1st, action=next_action_1st, training=True)
target_1st = d_tm1_1st * discount_1st * add_const_col(
self.value_feature(
obs=next_obs_1st, action=next_action_1st, training=True))
stage1_weight = fit_linear(target_1st, instrumental_feature_1st,
self.stage1_reg)
predicted_feature = linear_reg_pred(instrumental_feature_2nd,
stage1_weight)
# current_feature = self.value_feature(obs=current_obs_2nd, action=action_2nd, training=True)
current_feature = add_const_col(
self.value_feature(obs=current_obs_2nd,
action=action_2nd,
training=True))
predicted_feature = current_feature - d_tm1_2nd * self.discount * predicted_feature
# loss = linear_reg_loss(tf.expand_dims(reward_2nd, -1), predicted_feature, self.stage2_reg)
weight = d_tm1_2nd + (1.0 - d_tm1_2nd) * tf.convert_to_tensor(
self.d_tm1_weight, dtype=tf.float32)
loss = linear_reg_loss(weight * tf.expand_dims(reward_2nd, -1),
weight * predicted_feature, self.stage2_reg)
loss = loss + l2(self.value_feature.trainable_variables)
loss /= action_2nd.shape[0]
gradient = tape.gradient(loss, self.value_feature.trainable_variables)
self._value_func_optimizer.apply(
gradient, self.value_feature.trainable_variables)
return loss
def update_final_weight(self, stage1_input, stage2_input):
current_obs_1st, action_1st, _, discount_1st, next_obs_1st = stage1_input[:
5]
d_tm1_1st = self._get_d_tm1(stage1_input)
current_obs_2nd, action_2nd, reward_2nd = stage2_input[:3]
d_tm1_2nd = self._get_d_tm1(stage2_input)
next_action_1st = self.policy(next_obs_1st)
discount_1st = tf.expand_dims(discount_1st, axis=1)
d_tm1_1st = tf.expand_dims(d_tm1_1st, axis=1)
d_tm1_2nd = tf.expand_dims(d_tm1_2nd, axis=1)
instrumental_feature_1st = self.instrumental_feature(
obs=current_obs_1st, action=action_1st, training=False)
instrumental_feature_1st = d_tm1_1st * instrumental_feature_1st
instrumental_feature_2nd = self.instrumental_feature(
obs=current_obs_2nd, action=action_2nd, training=False)
# target_1st = discount_1st * self.value_feature(obs=next_obs_1st, action=next_action_1st, training=False)
target_1st = d_tm1_1st * discount_1st * add_const_col(
self.value_feature(
obs=next_obs_1st, action=next_action_1st, training=False))
stage1_weight = fit_linear(target_1st, instrumental_feature_1st,
self.stage1_reg)
self.stage1_weight.assign(stage1_weight)
predicted_feature = linear_reg_pred(instrumental_feature_2nd,
stage1_weight)
# current_feature = self.value_feature(obs=current_obs_2nd, action=action_2nd, training=False)
current_feature = add_const_col(
self.value_feature(obs=current_obs_2nd,
action=action_2nd,
training=False))
# predicted_feature = add_const_col(current_feature) - self.discount * add_const_col(predicted_feature)
predicted_feature = current_feature - d_tm1_2nd * self.discount * predicted_feature
# self.value_func._weight.assign(
# fit_linear(tf.expand_dims(reward_2nd, -1), predicted_feature, self.stage2_reg))
weight = d_tm1_2nd + (1.0 - d_tm1_2nd) * tf.convert_to_tensor(
self.d_tm1_weight, dtype=tf.float32)
stage2_weight = fit_linear(weight * tf.expand_dims(reward_2nd, -1),
weight * predicted_feature, self.stage2_reg)
self.value_func.weight.assign(stage2_weight)
return stage1_weight, stage2_weight
def update_instrumental(self, current_obs, action, reward, discount,
next_obs):
next_action = self.policy(next_obs)
discount = tf.expand_dims(discount, axis=1)
target = discount * add_const_col(
self.value_feature(next_obs, next_action, training=False))
target = add_const_col(self.value_feature(
current_obs, action)) - self.discount * target
l2 = snt.regularizers.L2(self.instrumental_reg)
with tf.GradientTape() as tape:
feature = self.instrumental_feature(obs=current_obs,
action=action,
training=True)
loss = linear_reg_loss(target, feature, self.stage1_reg)
loss = loss + l2(self.instrumental_feature.trainable_variables)
loss /= action.shape[0]
gradient = tape.gradient(loss,
self.instrumental_feature.trainable_variables)
self._instrumental_func_optimizer.apply(
gradient, self.instrumental_feature.trainable_variables)
return loss
def __call__(self, obs, action, training=False):
feature = self._feature(obs, action, training)
return tf.matmul(add_const_col(feature), self._weight)
def __call__(self, obs, action, training=False):
feature = self._net(obs, action)
feature = add_const_col(feature)
return feature