<jupyter_start><jupyter_text>***Copyright 2020 Google LLC.***Licensed under the Apache License, Version 2.0 (the "License");<jupyter_code>#@title Default title text
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import numpy as np
import tensorflow.compat.v2 as tf
import matplotlib.pyplot as plt
tf.enable_v2_behavior()
# Mode as a function of observation
def f(s):
  return np.sin(s*2*np.pi/100.)/2.
N = 100
s = np.random.uniform(-100, 100, size=N)  # observations between -100 and 100
a = np.random.uniform(-1, 1, size=N)  # Actions between -1 and 1
P = 0.2
y = -100*np.abs(a - f(s))**P
y /= np.max(np.abs(y))
print(np.max(y))
print(np.min(y))
plt.scatter(s, a, c=y)
plt.plot(np.sort(s), f(np.sort(s)))
plt.plot()<jupyter_output><empty_output><jupyter_text>ExplanationObservations $s$ are scalar values between -100 and 100. Actions $a$ are scalar values between -1 and 1.The plot above shows $(s, a)$ pairs, with their corresponding targets $y$ as color (dark is low, light is hight). The high target regions follow the curve $f$, which gives the mode (argmax action) as a function of $s$.The goal is to recover $f$ from the given data points.<jupyter_code>data = (s[:, np.newaxis], a[:, np.newaxis], y[:, np.newaxis])
s_features = tf.constant(np.linspace(-100, 100, 50)[np.newaxis, :], dtype=tf.float32)

hidden_widths = [1000, 500]
model = tf.keras.Sequential(
    [tf.keras.layers.Lambda(lambda x: tf.exp(-(x - s_features)**2/2000))]
    + [tf.keras.layers.Dense(w, activation='relu') for w in hidden_widths]
    + [tf.keras.layers.Dense(1, activation=None)]
)<jupyter_output><empty_output><jupyter_text>What loss functions best recover the curve $f$ from our dataset?<jupyter_code># loss A
#    ||h(s) - a|^p - R|^q
# This is danabo's mode regression loss

p = 0.1
q = 1/P
# p = q = 2.0
def loss(model, s, a, y):
  reg = tf.linalg.global_norm(model.trainable_variables)
  return tf.reduce_mean(tf.abs(-tf.abs(model(s)-a)**p - y)**q) + 0.003*reg
# loss B
#    |h(s) - a|^p * exp(R/tau)
# This is one of Dale's surrogate loss, specifically dot-product loss.

p = 1.0
tau = 1/10.
def loss(model, s, a, y):
  reg = tf.linalg.global_norm(model.trainable_variables)
  target = tf.cast(tf.exp(y/tau), tf.float32)
  return tf.reduce_mean(tf.abs(model(s)-a)**p * target) + 0.0005*reg
np.var(s)
# Initialize model

device_string = '/device:GPU:0'
# device_string = '/device:TPU:0'
# device_string = ''  # CPU

with tf.device(device_string):
  model(data[0])
  print(loss(model, *data).numpy())  # Initialize model

optimizer = tf.keras.optimizers.Adam(learning_rate=0.0001)
def sample_batch(batch_size, *args):
  assert args
  idx = np.random.choice(args[0].shape[0], batch_size)
  return tuple([arg[idx] for arg in args])
for i in range(10000):
  # batch = sample_batch(100, *data)
  batch = data
  optimizer.minimize(lambda: loss(model, *batch), model.trainable_variables)
  if i % 100 == 0:
    print(i, '\t', loss(model, *data).numpy())<jupyter_output><empty_output><jupyter_text>Test recovery of $f$.<jupyter_code>X = np.linspace(-100, 100, 200)[:, np.newaxis]
Y = model(X).numpy()
plt.plot(X, Y)
plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15162355531095212, 'fraction_numerical': 0.0341221794166208, 'mean_word_length': 4.6269349845201235, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 17, 'https://': 1, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '441384', 'license': 'Apache-2.0', 'n_tokens_mistral': 1414, 'n_tokens_neox': 1316, 'n_words': 382, 'path': 'action_gap_rl/notebooks/mode_regression_nn.ipynb', 'repo_name': 'ojInc/google-research', 'size': '3634'}