Mirrors_Framework/deepmind-research
1
0
Fork 0
mirror of https://github.com/google-deepmind/deepmind-research.git synced 2026-10-06 06:59:41 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add GPE/GPI experiments.

Browse Source 
PiperOrigin-RevId: 323750949
This commit is contained in:
Shaobo Hou
2020-07-29 14:36:59 +01:00
committed by Diego de Las Casas
parent 59c0cf5044
commit a24bda5ed0
20 changed files with 1732 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files
option_keyboard/configs.py
+25 -1
View File
@@ -21,7 +21,7 @@ def get_task_config():
return dict(
arena_size=11,
num_channels=2,
max_num_steps=50, # 5o for the actual task.
max_num_steps=50, # 50 for the actual task.
num_init_objects=10,
object_priors=[0.5, 0.5],
egocentric=True,
@@ -39,3 +39,27 @@ def get_pretrain_config():
egocentric=True,
default_w=(1, 1),
)
def get_fig4_task_config():
return dict(
arena_size=11,
num_channels=2,
max_num_steps=50, # 50 for the actual task.
num_init_objects=10,
object_priors=[0.5, 0.5],
egocentric=True,
default_w=(1, -1),
)
def get_fig5_task_config(default_w):
return dict(
arena_size=11,
num_channels=2,
max_num_steps=50, # 50 for the actual task.
num_init_objects=10,
object_priors=[0.5, 0.5],
egocentric=True,
default_w=default_w,
)
option_keyboard/environment_wrappers.py
+171 -2
View File
@@ -24,6 +24,10 @@ import dm_env
import numpy as np
import tensorflow.compat.v1 as tf
import tensorflow_hub as hub
import tree
from option_keyboard import smart_module
class EnvironmentWithLogging(dm_env.Environment):
@@ -93,8 +97,9 @@ class EnvironmentWithKeyboard(dm_env.Environment):
self._keyboard(tf.expand_dims(obs_ph, axis=0))[0], [obs_ph])
session.run(tf.global_variables_initializer())
saver = tf.train.Saver(var_list=keyboard.variables)
saver.restore(session, keyboard_ckpt_path)
if keyboard_ckpt_path:
saver = tf.train.Saver(var_list=keyboard.variables)
saver.restore(session, keyboard_ckpt_path)
def _compute_reward(self, option, obs):
return np.sum(self._options_np[option] * obs["cumulants"])
@@ -152,6 +157,102 @@ class EnvironmentWithKeyboard(dm_env.Environment):
return getattr(self._env, name)
class EnvironmentWithKeyboardDirect(dm_env.Environment):
"""Wraps an environment with a keyboard.
This is different from EnvironmentWithKeyboard as the actions space is not
discretized.
TODO(shaobohou) Merge the two implementations.
"""
def __init__(self,
env,
keyboard,
keyboard_ckpt_path,
additional_discount,
call_and_return=False):
self._env = env
self._keyboard = keyboard
self._discount = additional_discount
self._call_and_return = call_and_return
obs_spec = self._extract_observation(env.observation_spec())
obs_ph = tf.placeholder(shape=obs_spec.shape, dtype=obs_spec.dtype)
option_ph = tf.placeholder(
shape=(keyboard.num_cumulants,), dtype=tf.float32)
gpi_action = self._keyboard.gpi(obs_ph, option_ph)
session = tf.Session()
self._gpi_action = session.make_callable(gpi_action, [obs_ph, option_ph])
self._keyboard_action = session.make_callable(
self._keyboard(tf.expand_dims(obs_ph, axis=0))[0], [obs_ph])
session.run(tf.global_variables_initializer())
if keyboard_ckpt_path:
saver = tf.train.Saver(var_list=keyboard.variables)
saver.restore(session, keyboard_ckpt_path)
def _compute_reward(self, option, obs):
assert option.shape == obs["cumulants"].shape
return np.sum(option * obs["cumulants"])
def reset(self):
return self._env.reset()
def step(self, option):
"""Take a step in the keyboard, then the environment."""
step_count = 0
option_step = None
while True:
obs = self._extract_observation(self._env.observation())
action = self._gpi_action(obs, option)
action_step = self._env.step(action)
step_count += 1
if option_step is None:
option_step = action_step
else:
new_discount = (
option_step.discount * self._discount * action_step.discount)
new_reward = (
option_step.reward + new_discount * action_step.reward)
option_step = option_step._replace(
observation=action_step.observation,
reward=new_reward,
discount=new_discount,
step_type=action_step.step_type)
if action_step.last():
break
# Terminate option.
if self._compute_reward(option, action_step.observation) > 0:
break
if not self._call_and_return:
break
return option_step
def action_spec(self):
return dm_env.specs.BoundedArray(shape=(self._keyboard.num_cumulants,),
dtype=np.float32,
minimum=-1.0,
maximum=1.0,
name="action")
def _extract_observation(self, obs):
return obs["arena"]
def observation_spec(self):
return self._env.observation_spec()
def __getattr__(self, name):
return getattr(self._env, name)
def _discretize_actions(num_actions_per_dim,
action_space_dim,
min_val=-1.0,
@@ -188,3 +289,71 @@ def _discretize_actions(num_actions_per_dim,
logging.info("Discretized actions: %s", discretized_actions)
return discretized_actions
class EnvironmentWithLearnedPhi(dm_env.Environment):
"""Wraps an environment with learned phi model."""
def __init__(self, env, model_path):
self._env = env
create_ph = lambda x: tf.placeholder(shape=x.shape, dtype=x.dtype)
add_batch = lambda x: tf.expand_dims(x, axis=0)
# Make session and callables.
with tf.Graph().as_default():
model = smart_module.SmartModuleImport(hub.Module(model_path))
obs_spec = env.observation_spec()
obs_ph = tree.map_structure(create_ph, obs_spec)
action_ph = tf.placeholder(shape=(), dtype=tf.int32)
phis = model(tree.map_structure(add_batch, obs_ph), add_batch(action_ph))
self.num_phis = phis.shape.as_list()[-1]
self._last_phis = np.zeros((self.num_phis,), dtype=np.float32)
session = tf.Session()
self._session = session
self._phis_fn = session.make_callable(
phis[0], tree.flatten([obs_ph, action_ph]))
self._session.run(tf.global_variables_initializer())
def reset(self):
self._last_phis = np.zeros((self.num_phis,), dtype=np.float32)
return self._env.reset()
def step(self, action):
"""Take action in the environment and do some logging."""
phis = self._phis_fn(*tree.flatten([self._env.observation(), action]))
step = self._env.step(action)
if step.first():
phis = self._phis_fn(*tree.flatten([self._env.observation(), action]))
step = self._env.step(action)
step.observation["cumulants"] = phis
self._last_phis = phis
return step
def action_spec(self):
return self._env.action_spec()
def observation(self):
obs = self._env.observation()
obs["cumulants"] = self._last_phis
return obs
def observation_spec(self):
obs_spec = self._env.observation_spec()
obs_spec["cumulants"] = dm_env.specs.BoundedArray(
shape=(self.num_phis,),
dtype=np.float32,
minimum=-1e9,
maximum=1e9,
name="collected_resources")
return obs_spec
def __getattr__(self, name):
return getattr(self._env, name)
option_keyboard/experiment.py
+19 -12
View File
@@ -18,37 +18,44 @@
from absl import logging
import numpy as np
def _ema(base, val, decay=0.995):
return base * decay + (1 - decay) * val
def run(environment, agent, num_episodes, report_every=200):
def run(env, agent, num_episodes, report_every=200, num_eval_reps=1):
"""Runs an agent on an environment.
Args:
environment: The environment.
env: The environment.
agent: The agent.
num_episodes: Number of episodes to train for.
report_every: Frequency at which training progress are reported (episodes).
num_eval_reps: Number of eval episodes to run per training episode.
"""
train_returns = []
train_return_ema = 0.
eval_returns = []
eval_return_ema = 0.
for episode_id in range(num_episodes):
# Run a training episode.
train_episode_return = run_episode(environment, agent, is_training=True)
train_episode_return = run_episode(env, agent, is_training=True)
train_returns.append(train_episode_return)
train_return_ema = _ema(train_return_ema, train_episode_return)
# Run an evaluation episode.
eval_episode_return = run_episode(environment, agent, is_training=False)
eval_returns.append(eval_episode_return)
for _ in range(num_eval_reps):
eval_episode_return = run_episode(env, agent, is_training=False)
eval_returns.append(eval_episode_return)
eval_return_ema = _ema(eval_return_ema, eval_episode_return)
if ((episode_id + 1) % report_every) == 0:
logging.info(
"Episode %s, avg train return %.3f, avg eval return %.3f",
episode_id + 1,
np.mean(train_returns[-report_every:]),
np.mean(eval_returns[-report_every:]),
)
logging.info("Episode %s, avg train return %.3f, avg eval return %.3f",
episode_id + 1, train_return_ema, eval_return_ema)
if hasattr(agent, "get_logs"):
logging.info("Episode %s, agent logs: %s", episode_id + 1,
agent.get_logs())
def run_episode(environment, agent, is_training=False):
option_keyboard/gpe_gpi_experiments/eval_keyboard_fig5.py
+144
View File
@@ -0,0 +1,144 @@
# Lint as: python3
# pylint: disable=g-bad-file-header
# pylint: disable=line-too-long
# Copyright 2020 DeepMind Technologies Limited. All Rights Reserved.
#
# 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
#
# http://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.
# ============================================================================
r"""Run an experiment.
This script generates the raw data for the polar plots used to visualise how
well a trained keyboard covers the space of w.
For example, train 3 separate keyboards with different base policies:
python3 train_keyboard.py --logtostderr --policy_weights_name=12
python3 train_keyboard.py --logtostderr --policy_weights_name=34
python3 train_keyboard.py --logtostderr --policy_weights_name=5
Then generate the polar plot data as follows:
python3 eval_keyboard_fig5a.py --logtostderr \
--keyboard_paths=/tmp/option_keyboard/keyboard_12/tfhub,/tmp/option_keyboard/keyboard_34/tfhub,/tmp/option_keyboard/keyboard_5/tfhub \
--num_episodes=1000
Example outout:
[[ 0.11 0.261 -0.933 ]
[ 1.302 3.955 0.54 ]
[ 2.398 4.434 1.2105359 ]
[ 3.459 4.606 2.087 ]
[ 4.09026795 4.60911325 3.06106882]
[ 4.55499485 4.71947818 3.8123229 ]
[ 4.715 4.835 4.395 ]
[ 4.75743564 4.64095528 4.46330207]
[ 4.82518207 4.71232378 4.56190708]
[ 4.831 4.7155 4.5735 ]
[ 4.78074425 4.6754641 4.58312762]
[ 4.70154374 4.5416429 4.47850417]
[ 4.694 4.631 4.427 ]
[ 4.25085125 4.56606664 3.68157677]
[ 3.61726795 4.4838453 2.68154403]
[ 2.714 4.43 1.554 ]
[ 1.69 4.505 0.9635359 ]
[ 0.894 4.043 0.424 ]
[ 0.099 0.349 0.055 ]]
"""
from absl import app
from absl import flags
import numpy as np
import tensorflow.compat.v1 as tf
import tensorflow_hub as hub
from option_keyboard import configs
from option_keyboard import environment_wrappers
from option_keyboard import experiment
from option_keyboard import scavenger
from option_keyboard import smart_module
from option_keyboard.gpe_gpi_experiments import regressed_agent
FLAGS = flags.FLAGS
flags.DEFINE_integer("num_episodes", 1000, "Number of training episodes.")
flags.DEFINE_list("keyboard_paths", [], "Path to keyboard model.")
def evaluate_keyboard(keyboard_path):
"""Evaluate a keyboard."""
angles_to_sweep = np.deg2rad(np.linspace(-90, 180, num=19, endpoint=True))
weights_to_sweep = np.stack(
[np.cos(angles_to_sweep),
np.sin(angles_to_sweep)], axis=-1)
weights_to_sweep /= np.sum(
np.maximum(weights_to_sweep, 0.0), axis=-1, keepdims=True)
weights_to_sweep = np.clip(weights_to_sweep, -1000, 1000)
tf.logging.info(weights_to_sweep)
# Load the keyboard.
keyboard = smart_module.SmartModuleImport(hub.Module(keyboard_path))
# Create the task environment.
all_returns = []
for w_to_sweep in weights_to_sweep.tolist():
base_env_config = configs.get_fig5_task_config(w_to_sweep)
base_env = scavenger.Scavenger(**base_env_config)
base_env = environment_wrappers.EnvironmentWithLogging(base_env)
# Wrap the task environment with the keyboard.
with tf.variable_scope(None, default_name="inner_loop"):
additional_discount = 0.9
env = environment_wrappers.EnvironmentWithKeyboardDirect(
env=base_env,
keyboard=keyboard,
keyboard_ckpt_path=None,
additional_discount=additional_discount,
call_and_return=False)
# Create the player agent.
agent = regressed_agent.Agent(
batch_size=10,
optimizer_name="AdamOptimizer",
# Disable training.
optimizer_kwargs=dict(learning_rate=0.0,),
init_w=w_to_sweep)
returns = []
for _ in range(FLAGS.num_episodes):
returns.append(experiment.run_episode(env, agent))
tf.logging.info(f"Task: {w_to_sweep}, mean returns over "
f"{FLAGS.num_episodes} episodes is {np.mean(returns)}")
all_returns.append(returns)
return all_returns, weights_to_sweep
def main(argv):
del argv
all_returns = []
for keyboard_path in FLAGS.keyboard_paths:
returns, _ = evaluate_keyboard(keyboard_path)
all_returns.append(returns)
print("Results:")
print(np.mean(all_returns, axis=-1).T)
if __name__ == "__main__":
tf.disable_v2_behavior()
app.run(main)
option_keyboard/gpe_gpi_experiments/regressed_agent.py
+95
View File
@@ -0,0 +1,95 @@
# Lint as: python3
# pylint: disable=g-bad-file-header
# Copyright 2020 DeepMind Technologies Limited. All Rights Reserved.
#
# 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
#
# http://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.
# ============================================================================
"""Regressed agent."""
import numpy as np
import tensorflow.compat.v1 as tf
class Agent():
"""A DQN Agent."""
def __init__(
self,
batch_size,
optimizer_name,
optimizer_kwargs,
init_w,
):
"""A simple DQN agent.
Args:
batch_size: Size of update batch.
optimizer_name: Name of an optimizer from tf.train
optimizer_kwargs: Keyword arguments for the optimizer.
init_w: The initial cumulant weight.
"""
self._batch_size = batch_size
self._init_w = np.array(init_w)
self._replay = []
# Regress w by gradient descent, could also use closed-form solution.
self._n_cumulants = len(init_w)
self._regressed_w = tf.get_variable(
"regressed_w",
dtype=tf.float32,
initializer=lambda: tf.to_float(init_w))
cumulants_ph = tf.placeholder(
shape=(None, self._n_cumulants), dtype=tf.float32)
rewards_ph = tf.placeholder(shape=(None,), dtype=tf.float32)
predicted_rewards = tf.reduce_sum(
tf.multiply(self._regressed_w, cumulants_ph), axis=-1)
loss = tf.reduce_sum(tf.square(predicted_rewards - rewards_ph))
with tf.variable_scope("optimizer"):
self._optimizer = getattr(tf.train, optimizer_name)(**optimizer_kwargs)
train_op = self._optimizer.minimize(loss)
# Make session and callables.
session = tf.Session()
self._update_fn = session.make_callable(train_op,
[cumulants_ph, rewards_ph])
self._action = session.make_callable(self._regressed_w.read_value(), [])
session.run(tf.global_variables_initializer())
def step(self, timestep, is_training=False):
"""Select actions according to epsilon-greedy policy."""
del timestep
if is_training:
# Can also just use random actions at environment level.
return np.random.uniform(low=-1.0, high=1.0, size=(self._n_cumulants,))
return self._action()
def update(self, step_tm1, action, step_t):
"""Takes in a transition from the environment."""
del step_tm1, action
transition = [
step_t.observation["cumulants"],
step_t.reward,
]
self._replay.append(transition)
if len(self._replay) == self._batch_size:
batch = list(zip(*self._replay))
self._update_fn(*batch)
self._replay = [] # Just a queue.
def get_logs(self):
return dict(regressed=self._action())
option_keyboard/gpe_gpi_experiments/run_dqn_fig4b.py
+64
View File
@@ -0,0 +1,64 @@
# Lint as: python3
# pylint: disable=g-bad-file-header
# Copyright 2020 DeepMind Technologies Limited. All Rights Reserved.
#
# 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
#
# http://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.
# ============================================================================
"""Run an experiment.
Run a q-learning agent on task (1, -1).
"""
from absl import app
from absl import flags
import tensorflow.compat.v1 as tf
from option_keyboard import configs
from option_keyboard import dqn_agent
from option_keyboard import environment_wrappers
from option_keyboard import experiment
from option_keyboard import scavenger
FLAGS = flags.FLAGS
flags.DEFINE_integer("num_episodes", 10000, "Number of training episodes.")
def main(argv):
del argv
# Create the task environment.
env_config = configs.get_fig4_task_config()
env = scavenger.Scavenger(**env_config)
env = environment_wrappers.EnvironmentWithLogging(env)
# Create the flat agent.
agent = dqn_agent.Agent(
obs_spec=env.observation_spec(),
action_spec=env.action_spec(),
network_kwargs=dict(
output_sizes=(64, 128),
activate_final=True,
),
epsilon=0.1,
additional_discount=0.9,
batch_size=10,
optimizer_name="AdamOptimizer",
optimizer_kwargs=dict(learning_rate=3e-4,))
experiment.run(env, agent, num_episodes=FLAGS.num_episodes)
if __name__ == "__main__":
tf.disable_v2_behavior()
app.run(main)
option_keyboard/gpe_gpi_experiments/run_dqn_fig5.py
+66
View File
@@ -0,0 +1,66 @@
# Lint as: python3
# pylint: disable=g-bad-file-header
# Copyright 2020 DeepMind Technologies Limited. All Rights Reserved.
#
# 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
#
# http://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.
# ============================================================================
"""Run an experiment.
Run a q-learning agent on a task.
"""
from absl import app
from absl import flags
import tensorflow.compat.v1 as tf
from option_keyboard import configs
from option_keyboard import dqn_agent
from option_keyboard import environment_wrappers
from option_keyboard import experiment
from option_keyboard import scavenger
FLAGS = flags.FLAGS
flags.DEFINE_integer("num_episodes", 10000, "Number of training episodes.")
flags.DEFINE_list("test_w", [], "The w to test.")
def main(argv):
del argv
# Create the task environment.
test_w = [float(x) for x in FLAGS.test_w]
env_config = configs.get_fig5_task_config(test_w)
env = scavenger.Scavenger(**env_config)
env = environment_wrappers.EnvironmentWithLogging(env)
# Create the flat agent.
agent = dqn_agent.Agent(
obs_spec=env.observation_spec(),
action_spec=env.action_spec(),
network_kwargs=dict(
output_sizes=(64, 128),
activate_final=True,
),
epsilon=0.1,
additional_discount=0.9,
batch_size=10,
optimizer_name="AdamOptimizer",
optimizer_kwargs=dict(learning_rate=3e-4,))
experiment.run(env, agent, num_episodes=FLAGS.num_episodes)
if __name__ == "__main__":
tf.disable_v2_behavior()
app.run(main)
option_keyboard/gpe_gpi_experiments/run_regressed_w_fig4.py
+92
View File
@@ -0,0 +1,92 @@
# Lint as: python3
# pylint: disable=g-bad-file-header
# Copyright 2020 DeepMind Technologies Limited. All Rights Reserved.
#
# 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
#
# http://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.
# ============================================================================
r"""Run an experiment.
Run GPE/GPI on task (1, -1) with w obtained by regression.
For example, first train a keyboard:
python3 train_keyboard.py -- --logtostderr --policy_weights_name=12 \
--export_path=/tmp/option_keyboard/keyboard
Then, evaluate the keyboard with w by regression.
python3 run_regressed_w_fig4.py -- --logtostderr \
--keyboard_path=/tmp/option_keyboard/keyboard_12/tfhub
"""
from absl import app
from absl import flags
import numpy as np
import tensorflow.compat.v1 as tf
import tensorflow_hub as hub
from option_keyboard import configs
from option_keyboard import environment_wrappers
from option_keyboard import experiment
from option_keyboard import scavenger
from option_keyboard import smart_module
from option_keyboard.gpe_gpi_experiments import regressed_agent
FLAGS = flags.FLAGS
flags.DEFINE_integer("num_episodes", 1000, "Number of training episodes.")
flags.DEFINE_string("keyboard_path", None, "Path to keyboard model.")
def main(argv):
del argv
# Load the keyboard.
keyboard = smart_module.SmartModuleImport(hub.Module(FLAGS.keyboard_path))
# Create the task environment.
base_env_config = configs.get_fig4_task_config()
base_env = scavenger.Scavenger(**base_env_config)
base_env = environment_wrappers.EnvironmentWithLogging(base_env)
# Wrap the task environment with the keyboard.
additional_discount = 0.9
env = environment_wrappers.EnvironmentWithKeyboardDirect(
env=base_env,
keyboard=keyboard,
keyboard_ckpt_path=None,
additional_discount=additional_discount,
call_and_return=False)
# Create the player agent.
agent = regressed_agent.Agent(
batch_size=10,
optimizer_name="AdamOptimizer",
optimizer_kwargs=dict(learning_rate=1e-1,),
init_w=np.random.normal(size=keyboard.num_cumulants) * 0.1,
)
experiment.run(
env,
agent,
num_episodes=FLAGS.num_episodes,
report_every=2,
num_eval_reps=100)
if __name__ == "__main__":
tf.disable_v2_behavior()
app.run(main)
option_keyboard/gpe_gpi_experiments/run_regressed_w_with_phi_fig4c.py
+105
View File
@@ -0,0 +1,105 @@
# Lint as: python3
# pylint: disable=g-bad-file-header
# Copyright 2020 DeepMind Technologies Limited. All Rights Reserved.
#
# 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
#
# http://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.
# ============================================================================
r"""Run an experiment.
Run GPE/GPI on task (1, -1) with a learned phi model and w by regression.
For example, first train a phi model with 3 dimenional phi:
python3 train_phi_model.py -- --logtostderr --use_random_tasks \
--export_path=/tmp/option_keyboard/phi_model_3d --num_phis=3
Then train a keyboard:
python3 train_keyboard_with_phi.py -- --logtostderr \
--export_path=/tmp/option_keyboard/keyboard_3d \
--phi_model_path=/tmp/option_keyboard/phi_model_3d \
--num_phis=2
Finally, evaluate the keyboard with w by regression.
python3 run_regressed_w_with_phi_fig4b.py -- --logtostderr \
--phi_model_path=/tmp/option_keyboard/phi_model_3d \
--keyboard_path=/tmp/option_keyboard/keyboard_3d/tfhub
"""
from absl import app
from absl import flags
import numpy as np
import tensorflow.compat.v1 as tf
import tensorflow_hub as hub
from option_keyboard import configs
from option_keyboard import environment_wrappers
from option_keyboard import experiment
from option_keyboard import scavenger
from option_keyboard import smart_module
from option_keyboard.gpe_gpi_experiments import regressed_agent
FLAGS = flags.FLAGS
flags.DEFINE_integer("num_episodes", 1000, "Number of training episodes.")
flags.DEFINE_string("phi_model_path", None, "Path to phi model.")
flags.DEFINE_string("keyboard_path", None, "Path to keyboard model.")
def main(argv):
del argv
# Load the keyboard.
keyboard = smart_module.SmartModuleImport(hub.Module(FLAGS.keyboard_path))
# Create the task environment.
base_env_config = configs.get_fig4_task_config()
base_env = scavenger.Scavenger(**base_env_config)
base_env = environment_wrappers.EnvironmentWithLogging(base_env)
base_env = environment_wrappers.EnvironmentWithLearnedPhi(
base_env, FLAGS.phi_model_path)
# Wrap the task environment with the keyboard.
additional_discount = 0.9
env = environment_wrappers.EnvironmentWithKeyboardDirect(
env=base_env,
keyboard=keyboard,
keyboard_ckpt_path=None,
additional_discount=additional_discount,
call_and_return=False)
# Create the player agent.
agent = regressed_agent.Agent(
batch_size=10,
optimizer_name="AdamOptimizer",
optimizer_kwargs=dict(learning_rate=1e-1,),
init_w=np.random.normal(size=keyboard.num_cumulants) * 0.1,
)
experiment.run(
env,
agent,
num_episodes=FLAGS.num_episodes,
report_every=2,
num_eval_reps=100)
if __name__ == "__main__":
tf.disable_v2_behavior()
app.run(main)
option_keyboard/gpe_gpi_experiments/run_true_w_fig4.py
+92
View File
@@ -0,0 +1,92 @@
# Lint as: python3
# pylint: disable=g-bad-file-header
# Copyright 2020 DeepMind Technologies Limited. All Rights Reserved.
#
# 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
#
# http://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.
# ============================================================================
r"""Run an experiment.
Run GPE/GPI on task (1, -1) with the groundtruth w.
For example, first train a keyboard:
python3 train_keyboard.py -- --logtostderr --policy_weights_name=12
Then, evaluate the keyboard with groundtruth w.
python3 run_true_w_fig4.py -- --logtostderr \
--keyboard_path=/tmp/option_keyboard/keyboard_12/tfhub
"""
from absl import app
from absl import flags
import numpy as np
import tensorflow.compat.v1 as tf
import tensorflow_hub as hub
from option_keyboard import configs
from option_keyboard import environment_wrappers
from option_keyboard import experiment
from option_keyboard import scavenger
from option_keyboard import smart_module
from option_keyboard.gpe_gpi_experiments import regressed_agent
FLAGS = flags.FLAGS
flags.DEFINE_integer("num_episodes", 1000, "Number of training episodes.")
flags.DEFINE_string("keyboard_path", None, "Path to keyboard model.")
def main(argv):
del argv
# Load the keyboard.
keyboard = smart_module.SmartModuleImport(hub.Module(FLAGS.keyboard_path))
# Create the task environment.
base_env_config = configs.get_fig4_task_config()
base_env = scavenger.Scavenger(**base_env_config)
base_env = environment_wrappers.EnvironmentWithLogging(base_env)
# Wrap the task environment with the keyboard.
additional_discount = 0.9
env = environment_wrappers.EnvironmentWithKeyboardDirect(
env=base_env,
keyboard=keyboard,
keyboard_ckpt_path=None,
additional_discount=additional_discount,
call_and_return=False)
# Create the player agent.
agent = regressed_agent.Agent(
batch_size=10,
optimizer_name="AdamOptimizer",
# Disable training.
optimizer_kwargs=dict(learning_rate=0.0,),
init_w=[1., -1.])
returns = []
for _ in range(FLAGS.num_episodes):
returns.append(experiment.run_episode(env, agent))
tf.logging.info("#" * 80)
tf.logging.info(
f"Avg. return over {FLAGS.num_episodes} episodes is {np.mean(returns)}")
tf.logging.info("#" * 80)
if __name__ == "__main__":
tf.disable_v2_behavior()
app.run(main)
option_keyboard/gpe_gpi_experiments/run_true_w_fig6.py
+93
View File
@@ -0,0 +1,93 @@
# Lint as: python3
# pylint: disable=g-bad-file-header
# Copyright 2020 DeepMind Technologies Limited. All Rights Reserved.
#
# 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
#
# http://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.
# ============================================================================
r"""Run an experiment.
Run GPE/GPI on the "balancing" task with a fixed w
For example, first train a keyboard:
python3 train_keyboard.py -- --logtostderr --policy_weights_name=12
Then, evaluate the keyboard with a fixed w.
python3 run_true_w_fig4.py -- --logtostderr \
--keyboard_path=/tmp/option_keyboard/keyboard_12/tfhub
"""
from absl import app
from absl import flags
import numpy as np
import tensorflow.compat.v1 as tf
import tensorflow_hub as hub
from option_keyboard import configs
from option_keyboard import environment_wrappers
from option_keyboard import experiment
from option_keyboard import scavenger
from option_keyboard import smart_module
from option_keyboard.gpe_gpi_experiments import regressed_agent
FLAGS = flags.FLAGS
flags.DEFINE_integer("num_episodes", 1000, "Number of training episodes.")
flags.DEFINE_string("keyboard_path", None, "Path to keyboard model.")
flags.DEFINE_list("test_w", [], "The w to test.")
def main(argv):
del argv
# Load the keyboard.
keyboard = smart_module.SmartModuleImport(hub.Module(FLAGS.keyboard_path))
# Create the task environment.
base_env_config = configs.get_task_config()
base_env = scavenger.Scavenger(**base_env_config)
base_env = environment_wrappers.EnvironmentWithLogging(base_env)
# Wrap the task environment with the keyboard.
additional_discount = 0.9
env = environment_wrappers.EnvironmentWithKeyboardDirect(
env=base_env,
keyboard=keyboard,
keyboard_ckpt_path=None,
additional_discount=additional_discount,
call_and_return=False)
# Create the player agent.
agent = regressed_agent.Agent(
batch_size=10,
optimizer_name="AdamOptimizer",
# Disable training.
optimizer_kwargs=dict(learning_rate=0.0,),
init_w=[float(x) for x in FLAGS.test_w])
returns = []
for _ in range(FLAGS.num_episodes):
returns.append(experiment.run_episode(env, agent))
tf.logging.info("#" * 80)
tf.logging.info(
f"Avg. return over {FLAGS.num_episodes} episodes is {np.mean(returns)}")
tf.logging.info("#" * 80)
if __name__ == "__main__":
tf.disable_v2_behavior()
app.run(main)
option_keyboard/gpe_gpi_experiments/train_keyboard.py
+65
View File
@@ -0,0 +1,65 @@
# Lint as: python3
# pylint: disable=g-bad-file-header
# Copyright 2020 DeepMind Technologies Limited. All Rights Reserved.
#
# 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
#
# http://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.
# ============================================================================
"""Train a keyboard."""
from absl import app
from absl import flags
import numpy as np
import tensorflow.compat.v1 as tf
from option_keyboard import keyboard_utils
FLAGS = flags.FLAGS
flags.DEFINE_integer("num_pretrain_episodes", 20000,
"Number of pretraining episodes.")
flags.DEFINE_string("export_path", None,
"Where to save the keyboard checkpoints.")
flags.DEFINE_string("policy_weights_name", None,
"A string repsenting the policy weights.")
def main(argv):
del argv
all_policy_weights = {
"1": [1., 0.],
"2": [0., 1.],
"3": [1., -1.],
"4": [-1., 1.],
"5": [1., 1.],
}
if FLAGS.policy_weights_name:
policy_weights = np.array(
[all_policy_weights[v] for v in FLAGS.policy_weights_name])
num_episodes = ((FLAGS.num_pretrain_episodes // 2) *
max(2, len(policy_weights)))
export_path = FLAGS.export_path + "_" + FLAGS.policy_weights_name
else:
policy_weights = None
num_episodes = FLAGS.num_pretrain_episodes
export_path = FLAGS.export_path
keyboard_utils.create_and_train_keyboard(
num_episodes=num_episodes,
policy_weights=policy_weights,
export_path=export_path)
if __name__ == "__main__":
tf.disable_v2_behavior()
app.run(main)
option_keyboard/gpe_gpi_experiments/train_keyboard_with_phi.py
+49
View File
@@ -0,0 +1,49 @@
# Lint as: python3
# pylint: disable=g-bad-file-header
# Copyright 2020 DeepMind Technologies Limited. All Rights Reserved.
#
# 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
#
# http://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.
# ============================================================================
"""Train a keyboard."""
from absl import app
from absl import flags
import numpy as np
import tensorflow.compat.v1 as tf
from option_keyboard import keyboard_utils
FLAGS = flags.FLAGS
flags.DEFINE_integer("num_pretrain_episodes", 20000,
"Number of pretraining episodes.")
flags.DEFINE_integer("num_phis", None, "Size of phi")
flags.DEFINE_string("phi_model_path", None,
"Where to load the phi model checkpoints.")
flags.DEFINE_string("export_path", None,
"Where to save the keyboard checkpoints.")
def main(argv):
del argv
keyboard_utils.create_and_train_keyboard_with_phi(
num_episodes=FLAGS.num_pretrain_episodes,
phi_model_path=FLAGS.phi_model_path,
policy_weights=np.eye(FLAGS.num_phis, dtype=np.float32),
export_path=FLAGS.export_path)
if __name__ == "__main__":
tf.disable_v2_behavior()
app.run(main)
option_keyboard/gpe_gpi_experiments/train_phi_model.py
+270
View File
File diff suppressed because it is too large Load Diff
option_keyboard/keyboard_agent.py
+23 -7
View File
@@ -16,10 +16,14 @@
# ============================================================================
"""Keyboard agent."""
import os
import numpy as np
import sonnet as snt
import tensorflow.compat.v1 as tf
from option_keyboard import smart_module
class Agent():
"""An Option Keyboard Agent."""
@@ -51,6 +55,7 @@ class Agent():
optimizer_kwargs: Keyword arguments for the optimizer.
"""
tf.logging.info(policy_weights)
self._policy_weights = tf.convert_to_tensor(
policy_weights, dtype=tf.float32)
self._current_policy = None
@@ -61,13 +66,16 @@ class Agent():
self._n_actions = action_spec.num_values
self._n_policies, self._n_cumulants = policy_weights.shape
self._network = OptionValueNet(
self._n_policies,
self._n_cumulants,
self._n_actions,
network_kwargs=network_kwargs,
)
def create_network():
return OptionValueNet(
self._n_policies,
self._n_cumulants,
self._n_actions,
network_kwargs=network_kwargs,
)
self._network = smart_module.SmartModuleExport(create_network)
self._replay = []
obs_spec = self._extract_observation(obs_spec)
@@ -103,6 +111,12 @@ class Agent():
td_error = tf.stop_gradient(c_t + g * qa_t) - qa_tm1
loss = tf.reduce_sum(tf.square(td_error) / 2)
# Dummy calls to keyboard for SmartModule
_ = self._network.gpi(o_tm1[0], c_t[0])
_ = self._network.num_cumulants
_ = self._network.num_policies
_ = self._network.num_actions
with tf.variable_scope("optimizer"):
self._optimizer = getattr(tf.train, optimizer_name)(**optimizer_kwargs)
train_op = self._optimizer.minimize(loss)
@@ -155,7 +169,9 @@ class Agent():
def export(self, path):
tf.logging.info("Exporting keyboard to %s", path)
self._saver.save(self._session, path)
self._network.export(
os.path.join(path, "tfhub"), self._session, overwrite=True)
self._saver.save(self._session, os.path.join(path, "checkpoints"))
class OptionValueNet(snt.AbstractModule):
option_keyboard/keyboard_utils.py
+88
View File
@@ -0,0 +1,88 @@
# Lint as: python3
# pylint: disable=g-bad-file-header
# Copyright 2020 DeepMind Technologies Limited. All Rights Reserved.
#
# 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
#
# http://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.
# ============================================================================
"""Keyboard utils."""
import numpy as np
from option_keyboard import configs
from option_keyboard import environment_wrappers
from option_keyboard import experiment
from option_keyboard import keyboard_agent
from option_keyboard import scavenger
def create_and_train_keyboard(num_episodes,
policy_weights=None,
export_path=None):
"""Train an option keyboard."""
if policy_weights is None:
policy_weights = np.eye(2, dtype=np.float32)
env_config = configs.get_pretrain_config()
env = scavenger.Scavenger(**env_config)
env = environment_wrappers.EnvironmentWithLogging(env)
agent = keyboard_agent.Agent(
obs_spec=env.observation_spec(),
action_spec=env.action_spec(),
policy_weights=policy_weights,
network_kwargs=dict(
output_sizes=(64, 128),
activate_final=True,
),
epsilon=0.1,
additional_discount=0.9,
batch_size=10,
optimizer_name="AdamOptimizer",
optimizer_kwargs=dict(learning_rate=3e-4,))
if num_episodes:
experiment.run(env, agent, num_episodes=num_episodes)
agent.export(export_path)
return agent
def create_and_train_keyboard_with_phi(num_episodes,
phi_model_path,
policy_weights,
export_path=None):
"""Train an option keyboard."""
env_config = configs.get_pretrain_config()
env = scavenger.Scavenger(**env_config)
env = environment_wrappers.EnvironmentWithLogging(env)
env = environment_wrappers.EnvironmentWithLearnedPhi(env, phi_model_path)
agent = keyboard_agent.Agent(
obs_spec=env.observation_spec(),
action_spec=env.action_spec(),
policy_weights=policy_weights,
network_kwargs=dict(
output_sizes=(64, 128),
activate_final=True,
),
epsilon=0.1,
additional_discount=0.9,
batch_size=10,
optimizer_name="AdamOptimizer",
optimizer_kwargs=dict(learning_rate=3e-4,))
if num_episodes:
experiment.run(env, agent, num_episodes=num_episodes)
agent.export(export_path)
return agent
option_keyboard/requirements.txt
+3
View File
@@ -1,6 +1,9 @@
absl-py
dm-env==1.2
dm-sonnet==1.34
dm-tree
numpy==1.16.4
tensorflow==1.13.2
tensorflow_hub==0.7.0
tensorflow_probability==0.6.0
wrapt
option_keyboard/run_ok.py
+20 -36
View File
@@ -16,60 +16,44 @@
# ============================================================================
"""Run an experiment."""
import os
from absl import app
from absl import flags
import numpy as np
import tensorflow.compat.v1 as tf
import tensorflow_hub as hub
from option_keyboard import configs
from option_keyboard import dqn_agent
from option_keyboard import environment_wrappers
from option_keyboard import experiment
from option_keyboard import keyboard_agent
from option_keyboard import keyboard_utils
from option_keyboard import scavenger
from option_keyboard import smart_module
FLAGS = flags.FLAGS
flags.DEFINE_integer("num_episodes", 10000, "Number of training episodes.")
flags.DEFINE_integer("num_pretrain_episodes", 20000,
"Number of pretraining episodes.")
def _train_keyboard(num_episodes):
"""Train an option keyboard."""
env_config = configs.get_pretrain_config()
env = scavenger.Scavenger(**env_config)
env = environment_wrappers.EnvironmentWithLogging(env)
agent = keyboard_agent.Agent(
obs_spec=env.observation_spec(),
action_spec=env.action_spec(),
policy_weights=np.array([
[1.0, 0.0],
[0.0, 1.0],
]),
network_kwargs=dict(
output_sizes=(64, 128),
activate_final=True,
),
epsilon=0.1,
additional_discount=0.9,
batch_size=10,
optimizer_name="AdamOptimizer",
optimizer_kwargs=dict(learning_rate=3e-4,))
experiment.run(env, agent, num_episodes=num_episodes)
return agent
flags.DEFINE_string("keyboard_path", None, "Path to pretrained keyboard model.")
def main(argv):
del argv
# Pretrain the keyboard and save a checkpoint.
pretrain_agent = _train_keyboard(num_episodes=FLAGS.num_pretrain_episodes)
keyboard_ckpt_path = "/tmp/option_keyboard/keyboard.ckpt"
pretrain_agent.export(keyboard_ckpt_path)
if FLAGS.keyboard_path:
keyboard_path = FLAGS.keyboard_path
else:
with tf.Graph().as_default():
export_path = "/tmp/option_keyboard/keyboard"
_ = keyboard_utils.create_and_train_keyboard(
num_episodes=FLAGS.num_pretrain_episodes, export_path=export_path)
keyboard_path = os.path.join(export_path, "tfhub")
# Load the keyboard.
keyboard = smart_module.SmartModuleImport(hub.Module(keyboard_path))
# Create the task environment.
base_env_config = configs.get_task_config()
@@ -80,11 +64,11 @@ def main(argv):
additional_discount = 0.9
env = environment_wrappers.EnvironmentWithKeyboard(
env=base_env,
keyboard=pretrain_agent.keyboard,
keyboard_ckpt_path=keyboard_ckpt_path,
keyboard=keyboard,
keyboard_ckpt_path=None,
n_actions_per_dim=3,
additional_discount=additional_discount,
call_and_return=True)
call_and_return=False)
# Create the player agent.
agent = dqn_agent.Agent(
option_keyboard/scavenger.py
+20 -4
View File
@@ -56,7 +56,8 @@ class Scavenger(auto_reset_environment.Base):
num_init_objects=15,
object_priors=None,
egocentric=True,
rewarder=None):
rewarder=None,
aux_tasks_w=None):
self._arena_size = arena_size
self._num_channels = num_channels
self._max_num_steps = max_num_steps
@@ -64,12 +65,13 @@ class Scavenger(auto_reset_environment.Base):
self._egocentric = egocentric
self._rewarder = (
getattr(this_module, rewarder)() if rewarder is not None else None)
self._aux_tasks_w = aux_tasks_w
if object_priors is None:
self._object_priors = np.ones(num_channels) / num_channels
else:
assert len(object_priors) == num_channels
self._object_priors = np.array(object_priors)
self._object_priors = np.array(object_priors) / np.sum(object_priors)
if default_w is None:
self._default_w = np.ones(shape=(num_channels,))
@@ -203,10 +205,15 @@ class Scavenger(auto_reset_environment.Base):
collected_resources = np.copy(self._prev_collected).astype(np.float32)
return dict(
obs = dict(
arena=arena,
cumulants=collected_resources,
)
if self._aux_tasks_w is not None:
obs["aux_tasks_reward"] = np.dot(
np.array(self._aux_tasks_w), self._prev_collected).astype(np.float32)
return obs
def observation_spec(self):
arena = dm_env.specs.BoundedArray(
@@ -222,10 +229,19 @@ class Scavenger(auto_reset_environment.Base):
maximum=1e9,
name="collected_resources")
return dict(
obs_spec = dict(
arena=arena,
cumulants=collected_resources,
)
if self._aux_tasks_w is not None:
obs_spec["aux_tasks_reward"] = dm_env.specs.BoundedArray(
shape=(len(self._aux_tasks_w),),
dtype=np.float32,
minimum=-1e9,
maximum=1e9,
name="aux_tasks_reward")
return obs_spec
def action_spec(self):
return dm_env.specs.DiscreteArray(num_values=len(Action), name="action")
option_keyboard/smart_module.py
+228
View File
File diff suppressed because it is too large Load Diff