def offline_gym(
env_name: str,
pkl_path: str,
num_train_transitions: int,
max_steps: Optional[int],
seed: Optional[int] = None,
):
"""
Generate samples from a DiscreteRandomPolicy on the Gym environment and
saves results in a pandas df parquet.
"""
initialize_seed(seed)
env = EnvFactory.make(env_name)
replay_buffer = ReplayBuffer.create_from_env(
env=env, replay_memory_size=num_train_transitions, batch_size=1)
fill_replay_buffer(env, replay_buffer, num_train_transitions)
if isinstance(env.action_space, gym.spaces.Discrete):
is_discrete_action = True
else:
assert isinstance(env.action_space, gym.spaces.Box)
is_discrete_action = False
df = replay_buffer_to_pre_timeline_df(is_discrete_action, replay_buffer)
logger.info(f"Saving dataset with {len(df)} samples to {pkl_path}")
df.to_pickle(pkl_path)
def run_test_offline(
env_name: str,
max_steps: Optional[int],
model: ModelManager__Union,
replay_memory_size: int,
num_batches_per_epoch: int,
num_train_epochs: int,
passing_score_bar: float,
num_eval_episodes: int,
use_gpu: bool,
):
env = EnvFactory.make(env_name)
env.seed(SEED)
env.action_space.seed(SEED)
normalization = build_normalizer(env)
logger.info(f"Normalization is: \n{pprint.pformat(normalization)}")
manager = model.value
trainer = manager.initialize_trainer(
use_gpu=use_gpu,
reward_options=RewardOptions(),
normalization_data_map=normalization,
)
# first fill the replay buffer to burn_in
replay_buffer = ReplayBuffer.create_from_env(
env=env,
replay_memory_size=replay_memory_size,
batch_size=trainer.minibatch_size,
)
# always fill full RB
fill_replay_buffer(
env=env, replay_buffer=replay_buffer, desired_size=replay_memory_size
)
device = torch.device("cuda") if use_gpu else None
# pyre-fixme[6]: Expected `device` for 2nd param but got `Optional[torch.device]`.
trainer_preprocessor = make_replay_buffer_trainer_preprocessor(trainer, device, env)
writer = SummaryWriter()
with summary_writer_context(writer):
for epoch in range(num_train_epochs):
logger.info(f"Evaluating before epoch {epoch}: ")
eval_rewards = evaluate_cem(env, manager, max_steps, 1)
for _ in tqdm(range(num_batches_per_epoch)):
train_batch = replay_buffer.sample_transition_batch_tensor()
preprocessed_batch = trainer_preprocessor(train_batch)
trainer.train(preprocessed_batch)
logger.info(f"Evaluating after training for {num_train_epochs} epochs: ")
eval_rewards = evaluate_cem(env, manager, max_steps, num_eval_episodes)
mean_rewards = np.mean(eval_rewards)
assert (
mean_rewards >= passing_score_bar
), f"{mean_rewards} doesn't pass the bar {passing_score_bar}."
def run_test(
env: str,
model: ModelManager__Union,
replay_memory_size: int,
train_every_ts: int,
train_after_ts: int,
num_episodes: int,
max_steps: Optional[int],
last_score_bar: float,
):
env = EnvFactory.make(env)
env.seed(SEED)
env.action_space.seed(SEED)
normalization = build_normalizer(env)
logger.info(f"Normalization is {normalization}")
manager = model.value
trainer = manager.initialize_trainer(
use_gpu=False,
reward_options=RewardOptions(),
normalization_data_map=normalization,
)
policy = manager.create_policy()
replay_buffer = ReplayBuffer.create_from_env(
env=env,
replay_memory_size=replay_memory_size,
batch_size=trainer.minibatch_size,
)
post_step = train_with_replay_buffer_post_step(
replay_buffer=replay_buffer,
trainer=trainer,
training_freq=train_every_ts,
batch_size=trainer.minibatch_size,
replay_burnin=train_after_ts,
)
agent = Agent.create_for_env(env,
policy=policy,
post_transition_callback=post_step)
reward_history = []
for i in range(num_episodes):
logger.info(f"running episode {i}")
ep_reward = run_episode(env=env, agent=agent, max_steps=max_steps)
reward_history.append(ep_reward)
assert reward_history[-1] >= last_score_bar, (
f"reward after {len(reward_history)} episodes is {reward_history[-1]},"
f"less than < {last_score_bar}...\n"
f"Full reward history: {reward_history}")
return reward_history
def train_seq2reward(
env: gym.Env,
trainer: Seq2RewardTrainer,
trainer_preprocessor,
num_train_transitions: int,
seq_len: int,
batch_size: int,
num_train_epochs: int,
# for optional validation
test_replay_buffer=None,
):
train_replay_buffer = ReplayBuffer.create_from_env(
env=env,
replay_memory_size=num_train_transitions,
batch_size=batch_size,
stack_size=seq_len,
return_everything_as_stack=True,
)
fill_replay_buffer(env, train_replay_buffer, num_train_transitions)
num_batch_per_epoch = train_replay_buffer.size // batch_size
logger.info("Made RBs, starting to train now!")
state_dim = env.observation_space.shape[0]
for epoch in range(num_train_epochs):
for i in range(num_batch_per_epoch):
batch = train_replay_buffer.sample_transition_batch(batch_size=batch_size)
preprocessed_batch = trainer_preprocessor(batch)
adhoc_action_padding(preprocessed_batch, state_dim=state_dim)
losses = trainer.train(preprocessed_batch)
print_seq2reward_losses(epoch, i, losses)
# validation
if test_replay_buffer is not None:
with torch.no_grad():
trainer.seq2reward_network.eval()
test_batch = test_replay_buffer.sample_transition_batch(
batch_size=batch_size
)
preprocessed_test_batch = trainer_preprocessor(test_batch)
adhoc_action_padding(preprocessed_test_batch, state_dim=state_dim)
valid_losses = trainer.get_loss(preprocessed_test_batch)
print_seq2reward_losses(epoch, "validation", valid_losses)
trainer.seq2reward_network.train()
return trainer
def create_embed_rl_dataset(
env: gym.Env,
memory_network: MemoryNetwork,
num_state_embed_transitions: int,
batch_size: int,
seq_len: int,
hidden_dim: int,
use_gpu: bool,
):
assert isinstance(env.action_space, gym.spaces.Discrete)
assert isinstance(env.observation_space, gym.spaces.Box)
assert len(env.observation_space.shape) == 1
logger.info("Starting to create embedded RL Dataset!")
# seqlen+1 because MDNRNN embeds the first seq_len steps and then
# the embedded state will be concatenated with the last step
# Ie.. (o1,o2,...,on) -> RNN -> h1,h2,...,hn
# and we set s_{n+1} = [o_{n+1}, h_n]
embed_env = StateEmbedEnvironment(
gym_env=env, mdnrnn=memory_network, max_embed_seq_len=seq_len + 1
)
# now create a filled replay buffer of embeddings
# new obs shape dim = state_dim + hidden_dim
embed_rb = ReplayBuffer.create_from_env(
env=embed_env,
replay_memory_size=num_state_embed_transitions,
batch_size=batch_size,
stack_size=1,
)
fill_replay_buffer(
env=embed_env, replay_buffer=embed_rb, desired_size=num_state_embed_transitions
)
batch = embed_rb.sample_transition_batch_tensor(
batch_size=num_state_embed_transitions
)
state_min = min(batch.state.min(), batch.next_state.min()).item()
state_max = max(batch.state.max(), batch.next_state.max()).item()
logger.info(
f"Finished making embed dataset with size {embed_rb.size}, "
f"min {state_min}, max {state_max}"
)
return embed_rb, state_min, state_max
def train_mdnrnn(
env: gym.Env,
trainer: MDNRNNTrainer,
trainer_preprocessor,
num_train_transitions: int,
seq_len: int,
batch_size: int,
num_train_epochs: int,
# for optional validation
test_replay_buffer=None,
):
train_replay_buffer = ReplayBuffer.create_from_env(
env=env,
replay_memory_size=num_train_transitions,
batch_size=batch_size,
stack_size=seq_len,
return_everything_as_stack=True,
)
fill_replay_buffer(env, train_replay_buffer, num_train_transitions)
num_batch_per_epoch = train_replay_buffer.size // batch_size
logger.info("Made RBs, starting to train now!")
for epoch in range(num_train_epochs):
for i in range(num_batch_per_epoch):
batch = train_replay_buffer.sample_transition_batch_tensor(
batch_size=batch_size
)
preprocessed_batch = trainer_preprocessor(batch)
losses = trainer.train(preprocessed_batch)
print_mdnrnn_losses(epoch, i, losses)
# validation
if test_replay_buffer is not None:
with torch.no_grad():
trainer.memory_network.mdnrnn.eval()
test_batch = test_replay_buffer.sample_transition_batch_tensor(
batch_size=batch_size
)
preprocessed_test_batch = trainer_preprocessor(test_batch)
valid_losses = trainer.get_loss(preprocessed_test_batch)
print_mdnrnn_losses(epoch, "validation", valid_losses)
trainer.memory_network.mdnrnn.train()
return trainer
def test_create_from_recsim_interest_evolution(self):
env_config = {
"num_candidates": 20,
"slate_size": 3,
"resample_documents": False,
"seed": 1,
}
env = interest_evolution.create_environment(env_config)
replay_buffer = ReplayBuffer.create_from_env(env,
replay_memory_size=100,
batch_size=10,
store_log_prob=True)
obs = env.reset()
observation = obs["user"]
action = env.action_space.sample()
log_prob = -1.0
doc_features = np.stack(list(obs["doc"].values()), axis=0)
next_obs, reward, terminal, _env = env.step(action)
response = next_obs["response"]
click = np.array([r["click"] for r in response])
response_quality = np.stack([r["quality"] for r in response], axis=0)
repsonse_cluster_id = np.array([r["cluster_id"] for r in response])
response_watch_time = np.stack([r["watch_time"] for r in response],
axis=0)
response_liked = np.array([r["liked"] for r in response])
replay_buffer.add(
observation,
action,
reward,
terminal,
mdp_id=0,
sequence_number=0,
doc=doc_features,
response_click=click,
response_cluster_id=repsonse_cluster_id,
response_quality=response_quality,
response_liked=response_liked,
response_watch_time=response_watch_time,
log_prob=log_prob,
)
def run_test(
env_name: str,
model: ModelManager__Union,
replay_memory_size: int,
train_every_ts: int,
train_after_ts: int,
num_train_episodes: int,
max_steps: Optional[int],
passing_score_bar: float,
num_eval_episodes: int,
use_gpu: bool,
):
env = EnvFactory.make(env_name)
env.seed(SEED)
env.action_space.seed(SEED)
normalization = build_normalizer(env)
logger.info(f"Normalization is: \n{pprint.pformat(normalization)}")
manager = model.value
try:
# pyre-fixme[16]: `Env` has no attribute `state_feature_config_provider`.
manager.state_feature_config_provider = env.state_feature_config_provider
logger.info(
f"Using environment's state_feature_config_provider.\n"
f"{manager.state_feature_config_provider}"
)
except AttributeError:
logger.info("state_feature_config_provider override not applicable")
trainer = manager.initialize_trainer(
use_gpu=use_gpu,
reward_options=RewardOptions(),
normalization_data_map=normalization,
)
training_policy = manager.create_policy(serving=False)
replay_buffer = ReplayBuffer.create_from_env(
env=env,
replay_memory_size=replay_memory_size,
batch_size=trainer.minibatch_size,
)
device = torch.device("cuda") if use_gpu else torch.device("cpu")
# first fill the replay buffer to burn_in
train_after_ts = max(train_after_ts, trainer.minibatch_size)
fill_replay_buffer(
env=env, replay_buffer=replay_buffer, desired_size=train_after_ts
)
post_step = train_with_replay_buffer_post_step(
replay_buffer=replay_buffer,
env=env,
trainer=trainer,
training_freq=train_every_ts,
batch_size=trainer.minibatch_size,
device=device,
)
agent = Agent.create_for_env(
env, policy=training_policy, post_transition_callback=post_step, device=device
)
writer = SummaryWriter()
with summary_writer_context(writer):
train_rewards = []
for i in range(num_train_episodes):
trajectory = run_episode(
env=env, agent=agent, mdp_id=i, max_steps=max_steps
)
ep_reward = trajectory.calculate_cumulative_reward()
train_rewards.append(ep_reward)
logger.info(
f"Finished training episode {i} (len {len(trajectory)})"
f" with reward {ep_reward}."
)
logger.info("============Train rewards=============")
logger.info(train_rewards)
logger.info(f"average: {np.mean(train_rewards)};\tmax: {np.max(train_rewards)}")
# Check whether the max score passed the score bar; we explore during training
# the return could be bad (leading to flakiness in C51 and QRDQN).
assert np.max(train_rewards) >= passing_score_bar, (
f"max reward ({np.max(train_rewards)})after training for "
f"{len(train_rewards)} episodes is less than < {passing_score_bar}.\n"
)
serving_policy = manager.create_policy(serving=True)
agent = Agent.create_for_env_with_serving_policy(env, serving_policy)
eval_rewards = evaluate_for_n_episodes(
n=num_eval_episodes, env=env, agent=agent, max_steps=max_steps
).squeeze(1)
logger.info("============Eval rewards==============")
logger.info(eval_rewards)
logger.info(f"average: {np.mean(eval_rewards)};\tmax: {np.max(eval_rewards)}")
assert np.mean(eval_rewards) >= passing_score_bar, (
f"Predictor reward is {np.mean(eval_rewards)},"
f"less than < {passing_score_bar}.\n"
)
def run_test(
env_name: str,
model: ModelManager__Union,
replay_memory_size: int,
train_every_ts: int,
train_after_ts: int,
num_train_episodes: int,
max_steps: Optional[int],
passing_score_bar: float,
num_eval_episodes: int,
use_gpu: bool,
):
env = EnvFactory.make(env_name)
env.seed(SEED)
env.action_space.seed(SEED)
normalization = build_normalizer(env)
logger.info(f"Normalization is: \n{pprint.pformat(normalization)}")
manager = model.value
trainer = manager.initialize_trainer(
use_gpu=use_gpu,
reward_options=RewardOptions(),
normalization_data_map=normalization,
)
training_policy = manager.create_policy(serving=False)
replay_buffer = ReplayBuffer.create_from_env(
env=env,
replay_memory_size=replay_memory_size,
batch_size=trainer.minibatch_size,
)
device = torch.device("cuda") if use_gpu else None
# first fill the replay buffer to burn_in
train_after_ts = max(train_after_ts, trainer.minibatch_size)
fill_replay_buffer(env=env,
replay_buffer=replay_buffer,
desired_size=train_after_ts)
post_step = train_with_replay_buffer_post_step(
replay_buffer=replay_buffer,
env=env,
trainer=trainer,
training_freq=train_every_ts,
batch_size=trainer.minibatch_size,
device=device,
)
agent = Agent.create_for_env(
env,
policy=training_policy,
post_transition_callback=post_step,
# pyre-fixme[6]: Expected `Union[str, torch.device]` for 4th param but got
# `Optional[torch.device]`.
device=device,
)
writer = SummaryWriter()
with summary_writer_context(writer):
train_rewards = []
for i in range(num_train_episodes):
trajectory = run_episode(env=env,
agent=agent,
mdp_id=i,
max_steps=max_steps)
ep_reward = trajectory.calculate_cumulative_reward()
train_rewards.append(ep_reward)
logger.info(
f"Finished training episode {i} with reward {ep_reward}.")
assert train_rewards[-1] >= passing_score_bar, (
f"reward after {len(train_rewards)} episodes is {train_rewards[-1]},"
f"less than < {passing_score_bar}...\n"
f"Full reward history: {train_rewards}")
logger.info("============Train rewards=============")
logger.info(train_rewards)
serving_policy = manager.create_policy(serving=True)
agent = Agent.create_for_env_with_serving_policy(env, serving_policy)
eval_rewards = evaluate_for_n_episodes(n=num_eval_episodes,
env=env,
agent=agent,
max_steps=max_steps).squeeze(1)
assert np.mean(eval_rewards) >= passing_score_bar, (
f"Predictor reward is {np.mean(eval_rewards)},"
f"less than < {passing_score_bar}...\n"
f"Full eval rewards: {eval_rewards}.")
logger.info("============Eval rewards==============")
logger.info(eval_rewards)
def train_mdnrnn_and_compute_feature_stats(
env_name: str,
model: ModelManager__Union,
num_train_transitions: int,
num_test_transitions: int,
seq_len: int,
batch_size: int,
num_train_epochs: int,
use_gpu: bool,
saved_mdnrnn_path: Optional[str] = None,
):
""" Train MDNRNN Memory Network and compute feature importance/sensitivity. """
env: gym.Env = EnvFactory.make(env_name)
env.seed(SEED)
manager = model.value
trainer = manager.initialize_trainer(
use_gpu=use_gpu,
reward_options=RewardOptions(),
normalization_data_map=build_normalizer(env),
)
device = "cuda" if use_gpu else "cpu"
trainer_preprocessor = make_replay_buffer_trainer_preprocessor(trainer, device, env)
test_replay_buffer = ReplayBuffer.create_from_env(
env=env,
replay_memory_size=num_test_transitions,
batch_size=batch_size,
stack_size=seq_len,
return_everything_as_stack=True,
)
fill_replay_buffer(env, test_replay_buffer, num_test_transitions)
if saved_mdnrnn_path is None:
# train from scratch
trainer = train_mdnrnn(
env=env,
trainer=trainer,
trainer_preprocessor=trainer_preprocessor,
num_train_transitions=num_train_transitions,
seq_len=seq_len,
batch_size=batch_size,
num_train_epochs=num_train_epochs,
test_replay_buffer=test_replay_buffer,
)
else:
# load a pretrained model, and just evaluate it
trainer.memory_network.mdnrnn.load_state_dict(torch.load(saved_mdnrnn_path))
with torch.no_grad():
trainer.memory_network.mdnrnn.eval()
test_batch = test_replay_buffer.sample_transition_batch_tensor(
batch_size=test_replay_buffer.size
)
preprocessed_test_batch = trainer_preprocessor(test_batch)
feature_importance = calculate_feature_importance(
env=env,
trainer=trainer,
use_gpu=use_gpu,
test_batch=preprocessed_test_batch,
)
feature_sensitivity = calculate_feature_sensitivity(
env=env,
trainer=trainer,
use_gpu=use_gpu,
test_batch=preprocessed_test_batch,
)
trainer.memory_network.mdnrnn.train()
return feature_importance, feature_sensitivity
def run_test(
env: str,
model: ModelManager__Union,
replay_memory_size: int,
train_every_ts: int,
train_after_ts: int,
num_train_episodes: int,
max_steps: Optional[int],
passing_score_bar: float,
num_eval_episodes: int,
use_gpu: bool,
):
env = EnvFactory.make(env)
env.seed(SEED)
env.action_space.seed(SEED)
normalization = build_normalizer(env)
logger.info(f"Normalization is {normalization}")
manager = model.value
trainer = manager.initialize_trainer(
use_gpu=use_gpu,
reward_options=RewardOptions(),
normalization_data_map=normalization,
)
replay_buffer = ReplayBuffer.create_from_env(
env=env,
replay_memory_size=replay_memory_size,
batch_size=trainer.minibatch_size,
)
device = torch.device("cuda") if use_gpu else None
post_step = train_with_replay_buffer_post_step(
replay_buffer=replay_buffer,
trainer=trainer,
training_freq=train_every_ts,
batch_size=trainer.minibatch_size,
replay_burnin=train_after_ts,
device=device,
)
training_policy = manager.create_policy(serving=False)
agent = Agent.create_for_env(env,
policy=training_policy,
post_transition_callback=post_step,
device=device)
train_rewards = []
for i in range(num_train_episodes):
ep_reward = run_episode(env=env, agent=agent, max_steps=max_steps)
train_rewards.append(ep_reward)
logger.info(f"Finished training episode {i} with reward {ep_reward}.")
assert train_rewards[-1] >= passing_score_bar, (
f"reward after {len(train_rewards)} episodes is {train_rewards[-1]},"
f"less than < {passing_score_bar}...\n"
f"Full reward history: {train_rewards}")
logger.info("============Train rewards=============")
logger.info(train_rewards)
def gym_to_reagent_serving(
obs: np.array) -> Tuple[torch.Tensor, torch.Tensor]:
obs_tensor = torch.tensor(obs).float().unsqueeze(0)
presence_tensor = torch.ones_like(obs_tensor)
return (obs_tensor, presence_tensor)
serving_policy = manager.create_policy(serving=True)
agent = Agent.create_for_env(env,
policy=serving_policy,
obs_preprocessor=gym_to_reagent_serving)
eval_rewards = []
for i in range(num_eval_episodes):
ep_reward = run_episode(env=env, agent=agent, max_steps=max_steps)
eval_rewards.append(ep_reward)
logger.info(f"Finished eval episode {i} with reward {ep_reward}.")
assert np.mean(eval_rewards) >= passing_score_bar, (
f"Predictor reward is {np.mean(eval_rewards)},"
f"less than < {passing_score_bar}...\n"
f"Full eval rewards: {eval_rewards}.")
logger.info("============Eval rewards==============")
logger.info(eval_rewards)
def train_seq2reward_and_compute_reward_mse(
env_name: str,
model: ModelManager__Union,
num_train_transitions: int,
num_test_transitions: int,
seq_len: int,
batch_size: int,
num_train_epochs: int,
use_gpu: bool,
saved_seq2reward_path: Optional[str] = None,
):
""" Train Seq2Reward Network and compute reward mse. """
env: gym.Env = EnvFactory.make(env_name)
env.seed(SEED)
manager = model.value
trainer = manager.initialize_trainer(
use_gpu=use_gpu,
reward_options=RewardOptions(),
normalization_data_map=build_normalizer(env),
)
device = "cuda" if use_gpu else "cpu"
# pyre-fixme[6]: Expected `device` for 2nd param but got `str`.
trainer_preprocessor = make_replay_buffer_trainer_preprocessor(trainer, device, env)
test_replay_buffer = ReplayBuffer.create_from_env(
env=env,
replay_memory_size=num_test_transitions,
batch_size=batch_size,
stack_size=seq_len,
return_everything_as_stack=True,
)
fill_replay_buffer(env, test_replay_buffer, num_test_transitions)
if saved_seq2reward_path is None:
# train from scratch
trainer = train_seq2reward(
env=env,
trainer=trainer,
trainer_preprocessor=trainer_preprocessor,
num_train_transitions=num_train_transitions,
seq_len=seq_len,
batch_size=batch_size,
num_train_epochs=num_train_epochs,
test_replay_buffer=test_replay_buffer,
)
else:
# load a pretrained model, and just evaluate it
trainer.seq2reward_network.load_state_dict(torch.load(saved_seq2reward_path))
state_dim = env.observation_space.shape[0]
with torch.no_grad():
trainer.seq2reward_network.eval()
test_batch = test_replay_buffer.sample_transition_batch(
batch_size=test_replay_buffer.size
)
preprocessed_test_batch = trainer_preprocessor(test_batch)
adhoc_action_padding(preprocessed_test_batch, state_dim=state_dim)
losses = trainer.get_loss(preprocessed_test_batch)
detached_losses = losses.cpu().detach().item()
trainer.seq2reward_network.train()
return detached_losses