def _make_env_dataloader(
self,
env_factory: Callable[[], gym.Env],
batch_size: Optional[int],
num_workers: Optional[int] = None,
seed: Optional[int] = None,
max_steps: Optional[int] = None,
max_episodes: Optional[int] = None,
) -> GymDataLoader:
""" Helper function for creating a (possibly vectorized) environment.
"""
logger.debug(
f"batch_size: {batch_size}, num_workers: {num_workers}, seed: {seed}"
)
env: Union[gym.Env, gym.vector.VectorEnv]
if batch_size is None:
env = env_factory()
else:
env = make_batched_env(
env_factory,
batch_size=batch_size,
num_workers=num_workers,
# TODO: Still debugging shared memory + custom spaces (e.g. Sparse).
shared_memory=False,
)
## Apply the "post-batch" wrappers:
# from sequoia.common.gym_wrappers import ConvertToFromTensors
# TODO: Only the BaselineMethod requires this, we should enable it only
# from the BaselineMethod, and leave it 'off' by default.
if self.add_done_to_observations:
env = AddDoneToObservation(env)
# # Convert the samples to tensors and move them to the right device.
# env = ConvertToFromTensors(env)
# env = ConvertToFromTensors(env, device=self.config.device)
# Add a wrapper that converts numpy arrays / etc to Observations/Rewards
# and from Actions objects to numpy arrays.
env = TypedObjectsWrapper(
env,
observations_type=self.Observations,
rewards_type=self.Rewards,
actions_type=self.Actions,
)
# Create an IterableDataset from the env using the EnvDataset wrapper.
dataset = EnvDataset(env, max_steps=max_steps, max_episodes=max_episodes,)
# Create a GymDataLoader for the EnvDataset.
env_dataloader = GymDataLoader(dataset)
if batch_size and seed:
# Seed each environment with its own seed (based on the base seed).
env.seed([seed + i for i in range(env_dataloader.num_envs)])
else:
env.seed(seed)
return env_dataloader
def test_loss_is_nonzero_at_episode_end_iterate(batch_size: int):
""" Test that when *iterating* through the env (active-dataloader style),
when the episode ends, a non-zero loss is returned by the output head.
"""
with gym.make("CartPole-v0") as temp_env:
temp_env = AddDoneToObservation(temp_env)
obs_space = temp_env.observation_space
action_space = temp_env.action_space
reward_space = getattr(
temp_env, "reward_space", spaces.Box(*temp_env.reward_range, shape=())
)
env = gym.vector.make("CartPole-v0", num_envs=batch_size, asynchronous=False)
env = AddDoneToObservation(env)
env = ConvertToFromTensors(env)
env = EnvDataset(env)
head = PolicyHead(
# observation_space=obs_space,
input_space=obs_space["x"],
action_space=action_space,
reward_space=reward_space,
hparams=PolicyHead.HParams(accumulate_losses_before_backward=False),
)
env.seed(123)
non_zero_losses = 0
for i, obs in zip(range(100), env):
print(i, obs)
x = obs["x"]
done = obs["done"]
representations = x
assert isinstance(x, Tensor)
assert isinstance(done, Tensor)
observations = ContinualRLSetting.Observations(
x=x,
done=done,
# info=info,
)
head_output = head.forward(observations, representations=representations)
actions = head_output.actions.numpy().tolist()
# actions = np.zeros(batch_size, dtype=int).tolist()
rewards = env.send(actions)
# print(f"Step {i}, obs: {obs}, done: {done}")
assert isinstance(representations, Tensor)
forward_pass = ForwardPass(
observations=observations,
representations=representations,
actions=head_output,
)
rewards = ContinualRLSetting.Rewards(rewards)
loss = head.get_loss(forward_pass, actions=head_output, rewards=rewards)
print("loss:", loss)
for env_index, env_is_done in enumerate(observations.done):
if env_is_done:
print(f"Episode ended for env {env_index} at step {i}")
assert loss.total_loss != 0.0
non_zero_losses += 1
break
else:
print(f"No episode ended on step {i}, expecting no loss.")
assert loss.total_loss == 0.0
assert non_zero_losses > 0
def test_loss_is_nonzero_at_episode_end(batch_size: int):
""" Test that when stepping through the env, when the episode ends, a
non-zero loss is returned by the output head.
"""
with gym.make("CartPole-v0") as temp_env:
temp_env = AddDoneToObservation(temp_env)
obs_space = temp_env.observation_space
action_space = temp_env.action_space
reward_space = getattr(
temp_env, "reward_space", spaces.Box(*temp_env.reward_range, shape=())
)
env = gym.vector.make("CartPole-v0", num_envs=batch_size, asynchronous=False)
env = AddDoneToObservation(env)
env = ConvertToFromTensors(env)
env = EnvDataset(env)
head = PolicyHead(
input_space=obs_space.x,
action_space=action_space,
reward_space=reward_space,
hparams=PolicyHead.HParams(accumulate_losses_before_backward=False),
)
# TODO: Simulating as if the output head were attached to a BaselineModel.
PolicyHead.base_model_optimizer = torch.optim.Adam(head.parameters(), lr=1e-3)
head.train()
env.seed(123)
obs = env.reset()
# obs = torch.as_tensor(obs, dtype=torch.float32)
done = torch.zeros(batch_size, dtype=bool)
info = np.array([{} for _ in range(batch_size)])
loss = None
non_zero_losses = 0
encoder = nn.Linear(4, 4)
encoder.train()
for i in range(100):
representations = encoder(obs["x"])
observations = ContinualRLSetting.Observations(
x=obs["x"],
done=done,
# info=info,
)
head_output = head.forward(observations, representations=representations)
actions = head_output.actions.numpy().tolist()
# actions = np.zeros(batch_size, dtype=int).tolist()
obs, rewards, done, info = env.step(actions)
done = torch.as_tensor(done, dtype=bool)
rewards = ContinualRLSetting.Rewards(rewards)
assert len(info) == batch_size
print(f"Step {i}, obs: {obs}, done: {done}, info: {info}")
forward_pass = ForwardPass(
observations=observations,
representations=representations,
actions=head_output,
)
loss = head.get_loss(forward_pass, actions=head_output, rewards=rewards)
print("loss:", loss)
assert observations.done is not None
for env_index, env_is_done in enumerate(observations.done):
if env_is_done:
print(f"Episode ended for env {env_index} at step {i}")
assert loss.loss != 0.0
non_zero_losses += 1
break
else:
print(f"No episode ended on step {i}, expecting no loss.")
assert loss is None or loss.loss == 0.0
assert non_zero_losses > 0