Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

attention branch - DO NOT ACCEPT PR #160

Open
wants to merge 5 commits into
base: main
Choose a base branch
from
Open

attention branch - DO NOT ACCEPT PR #160

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
9ce004a
attention branch
Aidandos Apr 26, 2023
0695b86
MFOS averaging
Aidandos May 9, 2023
13ce398
hardstop challenge + blindspots
Aidandos May 26, 2023
cbb483a
adding hardstop yamls
Aidandos Jun 1, 2023
d70fea5
eval changes
Aidandos Jun 21, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 24 additions & 0 deletions hardstop_eval_bash.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
#!/bin/bash
###### MFOS AVG ######
python -m pax.experiment -m +experiment/ipd=mfos_att_v_tabular_hardstop_eval ++wandb.log=True ++run_path=ucl-dark/ipd/4ykf9oe8 ++model_path=exp/MFOS-vs-Tabular/run-seed-23-pop-size-1000/2023-05-11_14.58.45.927266/generation_900 ++seed=85768,785678,764578,89678,97869,4567456,856778,3456347,45673,83346 ++stop=1
python -m pax.experiment -m +experiment/ipd=mfos_att_v_tabular_hardstop_eval ++wandb.log=True ++run_path=ucl-dark/ipd/eopf93re ++model_path=exp/MFOS-vs-Tabular/run-seed-65-pop-size-1000/2023-05-11_20.31.48.530245/generation_900 ++seed=85768,785678,764578,89678,97869,4567456,856778,3456347,45673,83346 ++stop=1
python -m pax.experiment -m +experiment/ipd=mfos_att_v_tabular_hardstop_eval ++wandb.log=True ++run_path=ucl-dark/ipd/1sqbd09n ++model_path=exp/MFOS-vs-Tabular/run-seed-47-pop-size-1000/2023-05-11_17.45.03.318240/generation_900 ++seed=85768,785678,764578,89678,97869,4567456,856778,3456347,45673,83346 ++stop=1
python -m pax.experiment -m +experiment/ipd=mfos_att_v_tabular_hardstop_eval ++wandb.log=True ++run_path=ucl-dark/ipd/3n7l8ods ++model_path=exp/MFOS-vs-Tabular/run-seed-8-pop-size-1000/2023-05-11_12.12.19.914211/generation_900 ++seed=85768,785678,764578,89678,97869,4567456,856778,3456347,45673,83346 ++stop=1
python -m pax.experiment -m +experiment/ipd=mfos_att_v_tabular_hardstop_eval ++wandb.log=True ++run_path=ucl-dark/ipd/4mf1ecxq ++model_path=exp/MFOS-vs-Tabular/run-seed-6-pop-size-1000/2023-05-11_09.25.40.656392/generation_900 ++seed=85768,785678,764578,89678,97869,4567456,856778,3456347,45673,83346 ++stop=1

python -m pax.experiment -m +experiment/ipd=mfos_att_v_tabular_hardstop_eval ++wandb.log=True ++run_path=ucl-dark/ipd/4ykf9oe8 ++model_path=exp/MFOS-vs-Tabular/run-seed-23-pop-size-1000/2023-05-11_14.58.45.927266/generation_900 ++seed=85768,785678,764578,89678,97869,4567456,856778,3456347,45673,83346 ++stop=100
python -m pax.experiment -m +experiment/ipd=mfos_att_v_tabular_hardstop_eval ++wandb.log=True ++run_path=ucl-dark/ipd/eopf93re ++model_path=exp/MFOS-vs-Tabular/run-seed-65-pop-size-1000/2023-05-11_20.31.48.530245/generation_900 ++seed=85768,785678,764578,89678,97869,4567456,856778,3456347,45673,83346 ++stop=100
python -m pax.experiment -m +experiment/ipd=mfos_att_v_tabular_hardstop_eval ++wandb.log=True ++run_path=ucl-dark/ipd/1sqbd09n ++model_path=exp/MFOS-vs-Tabular/run-seed-47-pop-size-1000/2023-05-11_17.45.03.318240/generation_900 ++seed=85768,785678,764578,89678,97869,4567456,856778,3456347,45673,83346 ++stop=100
python -m pax.experiment -m +experiment/ipd=mfos_att_v_tabular_hardstop_eval ++wandb.log=True ++run_path=ucl-dark/ipd/3n7l8ods ++model_path=exp/MFOS-vs-Tabular/run-seed-8-pop-size-1000/2023-05-11_12.12.19.914211/generation_900 ++seed=85768,785678,764578,89678,97869,4567456,856778,3456347,45673,83346 ++stop=100
python -m pax.experiment -m +experiment/ipd=mfos_att_v_tabular_hardstop_eval ++wandb.log=True ++run_path=ucl-dark/ipd/4mf1ecxq ++model_path=exp/MFOS-vs-Tabular/run-seed-6-pop-size-1000/2023-05-11_09.25.40.656392/generation_900 ++seed=85768,785678,764578,89678,97869,4567456,856778,3456347,45673,83346 ++stop=100

###### Shaper Nothing #$$$
python -m pax.experiment -m +experiment/ipd=shaper_att_v_tabular_hardstop_eval ++wandb.log=True ++run_path=ucl-dark/ipd/2m3wh5g7 ++model_path=exp/EARL-Shaper-vs-Tabular/run-seed-65-OpenES-pop-size-1000-num-opps-10-att-type-nothing/2023-05-14_22.17.07.592872/generation_900 ++seed=85768,785678,764578,89678,97869,4567456,856778,3456347,45673,83346 ++stop=100
python -m pax.experiment -m +experiment/ipd=shaper_att_v_tabular_hardstop_eval ++wandb.log=True ++run_path=ucl-dark/ipd/1jk5zly5 ++model_path=exp/EARL-Shaper-vs-Tabular/run-seed-47-OpenES-pop-size-1000-num-opps-10-att-type-nothing/2023-05-14_20.46.58.588813/generation_900 ++seed=85768,785678,764578,89678,97869,4567456,856778,3456347,45673,83346 ++stop=100
python -m pax.experiment -m +experiment/ipd=shaper_att_v_tabular_hardstop_eval ++wandb.log=True ++run_path=ucl-dark/ipd/1cvpiolk ++model_path=exp/EARL-Shaper-vs-Tabular/run-seed-23-OpenES-pop-size-1000-num-opps-10-att-type-nothing/2023-05-14_19.16.56.990716/generation_900 ++seed=85768,785678,764578,89678,97869,4567456,856778,3456347,45673,83346 ++stop=100
python -m pax.experiment -m +experiment/ipd=shaper_att_v_tabular_hardstop_eval ++wandb.log=True ++run_path=ucl-dark/ipd/3vml0wjy ++model_path=exp/EARL-Shaper-vs-Tabular/run-seed-6-OpenES-pop-size-1000-num-opps-10-att-type-nothing/2023-05-14_16.16.19.180942/generation_900 ++seed=85768,785678,764578,89678,97869,4567456,856778,3456347,45673,83346 ++stop=100

python -m pax.experiment -m +experiment/ipd=shaper_att_v_tabular_hardstop_eval ++wandb.log=True ++run_path=ucl-dark/ipd/2m3wh5g7 ++model_path=exp/EARL-Shaper-vs-Tabular/run-seed-65-OpenES-pop-size-1000-num-opps-10-att-type-nothing/2023-05-14_22.17.07.592872/generation_900 ++seed=85768,785678,764578,89678,97869,4567456,856778,3456347,45673,83346 ++stop=1
python -m pax.experiment -m +experiment/ipd=shaper_att_v_tabular_hardstop_eval ++wandb.log=True ++run_path=ucl-dark/ipd/1jk5zly5 ++model_path=exp/EARL-Shaper-vs-Tabular/run-seed-47-OpenES-pop-size-1000-num-opps-10-att-type-nothing/2023-05-14_20.46.58.588813/generation_900 ++seed=85768,785678,764578,89678,97869,4567456,856778,3456347,45673,83346 ++stop=1
python -m pax.experiment -m +experiment/ipd=shaper_att_v_tabular_hardstop_eval ++wandb.log=True ++run_path=ucl-dark/ipd/1cvpiolk ++model_path=exp/EARL-Shaper-vs-Tabular/run-seed-23-OpenES-pop-size-1000-num-opps-10-att-type-nothing/2023-05-14_19.16.56.990716/generation_900 ++seed=85768,785678,764578,89678,97869,4567456,856778,3456347,45673,83346 ++stop=1
python -m pax.experiment -m +experiment/ipd=shaper_att_v_tabular_hardstop_eval ++wandb.log=True ++run_path=ucl-dark/ipd/3vml0wjy ++model_path=exp/EARL-Shaper-vs-Tabular/run-seed-6-OpenES-pop-size-1000-num-opps-10-att-type-nothing/2023-05-14_16.16.19.180942/generation_900 ++seed=85768,785678,764578,89678,97869,4567456,856778,3456347,45673,83346 ++stop=1
22 changes: 22 additions & 0 deletions pax/agents/mfos_ppo/networks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -146,6 +146,28 @@ def forward_fn(
network = hk.without_apply_rng(hk.transform(forward_fn))
return network, hidden_state

def make_mfos_avg_network(num_actions: int, hidden_size: int):
hidden_state = jnp.zeros((1, 3 * hidden_size))

def forward_fn(
inputs: jnp.ndarray,
state: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray],
) -> Tuple[Tuple[jnp.ndarray, jnp.ndarray], jnp.ndarray]:
mfos = ActorCriticMFOS(num_actions, hidden_size)
hidden_t, hidden_a, hidden_v = jnp.split(state, 3, axis=-1)
avg_hidden_t = jnp.mean(hidden_t, axis=0, keepdims=True).repeat(state.shape[0], axis=0)
avg_hidden_a = jnp.mean(hidden_a, axis=0, keepdims=True).repeat(state.shape[0], axis=0)
avg_hidden_v = jnp.mean(hidden_v, axis=0, keepdims=True).repeat(state.shape[0], axis=0)
hidden_t = 0.5*hidden_t + 0.5*avg_hidden_t
hidden_a = 0.5*hidden_a + 0.5*avg_hidden_a
hidden_v = 0.5*hidden_v + 0.5*avg_hidden_v
state = jnp.concatenate([hidden_t, hidden_a, hidden_v], axis=-1)
logits, values, state = mfos(inputs, state)
return (logits, values), state

network = hk.without_apply_rng(hk.transform(forward_fn))
return network, hidden_state


def make_mfos_ipditm_network(
num_actions: int, hidden_size: int, output_channels, kernel_shape
Expand Down
21 changes: 15 additions & 6 deletions pax/agents/mfos_ppo/ppo_gru.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
from pax.agents.mfos_ppo.networks import (
make_mfos_ipditm_network,
make_mfos_network,
make_mfos_avg_network,
)
from pax.utils import TrainingState, get_advantages

Expand Down Expand Up @@ -63,7 +64,6 @@ def __init__(
obs_spec: Tuple,
batch_size: int = 2000,
num_envs: int = 4,
num_steps: int = 500,
num_minibatches: int = 16,
num_epochs: int = 4,
clip_value: bool = True,
Expand Down Expand Up @@ -479,8 +479,8 @@ def prepare_batch(

# Other useful hyperparameters
self._num_envs = num_envs # number of environments
self._num_steps = num_steps # number of steps per environment
self._batch_size = int(num_envs * num_steps) # number in one batch
# self._num_steps = num_steps # number of steps per environment
# self._batch_size = int(num_envs * num_steps) # number in one batch
self._num_minibatches = num_minibatches # number of minibatches
self._num_epochs = num_epochs # number of epochs to use sample
self._gru_dim = gru_dim
Expand Down Expand Up @@ -566,15 +566,25 @@ def make_mfos_agent(
agent_args.output_channels,
agent_args.kernel_shape,
)
elif args.env_id == "iterated_matrix_game":
if args.att_type=='att':
raise ValueError("Attention not supported")
elif args.att_type=='avg':
network, initial_hidden_state = make_mfos_avg_network(
action_spec, agent_args.hidden_size
)
elif args.att_type=='nothing':
network, initial_hidden_state = make_mfos_network(
action_spec, agent_args.hidden_size
)
else:
raise ValueError("Unsupported environment")

gru_dim = initial_hidden_state.shape[1]

# Optimizer
transition_steps = (
num_iterations,
*agent_args.num_epochs * agent_args.num_minibatches,
num_iterations * agent_args.num_epochs * agent_args.num_minibatches
)

if agent_args.lr_scheduling:
Expand Down Expand Up @@ -609,7 +619,6 @@ def make_mfos_agent(
obs_spec=obs_spec,
batch_size=None,
num_envs=args.num_envs,
num_steps=args.num_steps,
num_minibatches=agent_args.num_minibatches,
num_epochs=agent_args.num_epochs,
clip_value=agent_args.clip_value,
Expand Down
Loading