def test_cartpole_policy_model():
flat_action_space = flat_structured_space(_mock_action_spaces_dict())
distribution_mapper = DistributionMapper(action_space=flat_action_space,
distribution_mapper_config={})
action_logits_shapes = {step_key: {action_head: distribution_mapper.required_logits_shape(action_head)
for action_head in _mock_action_spaces_dict()[step_key].spaces.keys()}
for step_key in _mock_action_spaces_dict().keys()}
obs_shapes = observation_spaces_to_in_shapes(_mock_observation_spaces_dict())
policy = CustomComplexPolicyNet(obs_shapes[0], action_logits_shapes[0], non_lin='torch.nn.ReLU',
hidden_units=[128])
critic = CustomComplexCriticNet(obs_shapes[0], non_lin='torch.nn.ReLU',
hidden_units=[128])
obs_np = _mock_observation_spaces_dict()[0].sample()
obs = {k: torch.from_numpy(v) for k, v in obs_np.items()}
actions = policy(obs)
values = critic(obs)
assert 'action_move' in actions
assert 'action_use' in actions
assert 'value' in values
def test_distribution_mapper():
""" distribution test """
# action space
act_space = spaces.Dict(
spaces={
"selection":
spaces.Discrete(10),
"order":
spaces.MultiBinary(15),
"scale_input":
spaces.Box(shape=(5, ), low=0, high=100, dtype=np.float64),
"order_by_weight":
spaces.Box(shape=(5, ), low=0, high=100, dtype=np.float64)
})
# default config
config = [{
"action_space":
spaces.Box,
"distribution":
"maze.distributions.squashed_gaussian.SquashedGaussianProbabilityDistribution"
}, {
"action_head":
"order_by_weight",
"distribution":
"maze.distributions.beta.BetaProbabilityDistribution"
}]
# initialize distribution mapper
distribution_mapper = DistributionMapper(action_space=act_space,
distribution_mapper_config=config)
repr(distribution_mapper)
# assign action heads to registered distributions
logits_dict = dict()
for action_head in act_space.spaces.keys():
logits_shape = distribution_mapper.required_logits_shape(action_head)
logits_tensor = torch.from_numpy(np.random.randn(*logits_shape))
torch_dist = distribution_mapper.action_head_distribution(
action_head=action_head, logits=logits_tensor, temperature=1.0)
logits_dict[action_head] = logits_tensor
# check if distributions are correctly assigned
if action_head == "selection":
assert isinstance(torch_dist, CategoricalProbabilityDistribution)
elif action_head == "order":
assert isinstance(torch_dist, BernoulliProbabilityDistribution)
elif action_head == "scale_input":
assert isinstance(torch_dist,
SquashedGaussianProbabilityDistribution)
elif action_head == "order_by_weight":
assert isinstance(torch_dist, BetaProbabilityDistribution)
# test dictionary distribution mapping
dict_dist = distribution_mapper.logits_dict_to_distribution(
logits_dict=logits_dict, temperature=1.0)
assert isinstance(dict_dist, DictProbabilityDistribution)
def main(n_epochs) -> None:
"""Trains the cart pole environment with the ES implementation.
"""
env = GymMazeEnv(env="CartPole-v0")
distribution_mapper = DistributionMapper(action_space=env.action_space,
distribution_mapper_config={})
obs_shapes = observation_spaces_to_in_shapes(env.observation_spaces_dict)
action_shapes = {
step_key: {
action_head: distribution_mapper.required_logits_shape(action_head)
for action_head in env.action_spaces_dict[step_key].spaces.keys()
}
for step_key in env.action_spaces_dict.keys()
}
# initialize policies
policies = [
PolicyNet(obs_shapes=obs_shapes[0],
action_logits_shapes=action_shapes[0],
non_lin=nn.SELU)
]
# initialize optimizer
policy = TorchPolicy(networks=list_to_dict(policies),
distribution_mapper=distribution_mapper,
device="cpu")
shared_noise = SharedNoiseTable(count=1_000_000)
algorithm_config = ESAlgorithmConfig(n_rollouts_per_update=100,
n_timesteps_per_update=0,
max_steps=0,
optimizer=Adam(step_size=0.01),
l2_penalty=0.005,
noise_stddev=0.02,
n_epochs=n_epochs,
policy_wrapper=None)
trainer = ESTrainer(algorithm_config=algorithm_config,
torch_policy=policy,
shared_noise=shared_noise,
normalization_stats=None)
setup_logging(job_config=None)
maze_rng = np.random.RandomState(None)
# run with pseudo-distribution, without worker processes
trainer.train(ESDummyDistributedRollouts(
env=env,
n_eval_rollouts=10,
shared_noise=shared_noise,
agent_instance_seed=MazeSeeding.generate_seed_from_random_state(
maze_rng)),
model_selection=None)
def test_dummy_model_with_dummy_network():
"""
Unit test for the DummyStructuredEnvironment
"""
maze_env = build_dummy_maze_env()
# init the distribution_mapper with the flat action space
distribution_mapper_config = [{
"action_space":
spaces.Box,
"distribution":
"maze.distributions.squashed_gaussian.SquashedGaussianProbabilityDistribution"
}]
distribution_mapper = DistributionMapper(
action_space=maze_env.action_space,
distribution_mapper_config=distribution_mapper_config)
obs_shapes = observation_spaces_to_in_shapes(
maze_env.observation_spaces_dict)
dummy_actor = DummyPolicyNet(
obs_shapes=obs_shapes[0],
action_logits_shapes={
key: distribution_mapper.required_logits_shape(key)
for key in maze_env.action_space.spaces.keys()
},
non_lin=nn.Tanh)
dummy_critic = DummyValueNet(obs_shapes=obs_shapes[0], non_lin=nn.Tanh)
obs_np = maze_env.reset()
obs = {k: torch.from_numpy(v) for k, v in obs_np.items()}
for i in range(100):
logits_dict = dummy_actor(obs)
prob_dist = distribution_mapper.logits_dict_to_distribution(
logits_dict=logits_dict, temperature=1.0)
sampled_actions = prob_dist.sample()
obs_np, _, _, _ = maze_env.step(sampled_actions)
obs = {k: torch.from_numpy(v) for k, v in obs_np.items()}
_ = dummy_critic(obs)
maze_env.close()
def required_model_output_shape(action_space: gym.spaces.Dict,
model_config: Dict) -> int:
"""Returns the required logits shape (network output shape) for a given action head.
:param action_space: The action space of the env.
:param model_config: The rllib model config.
:return: The number of the flattened output.
"""
# Retrieve the distribution_mapper_config from the model config
method_distribution_mapper_config = \
model_config['custom_model_config']['maze_model_composer_config']['distribution_mapper_config']
# Build the distribution mapper
method_distribution_mapper = DistributionMapper(
action_space,
distribution_mapper_config=method_distribution_mapper_config)
# Compute the flattened number of logits
num_outputs = sum([
np.prod(
method_distribution_mapper.required_logits_shape(action_head))
for action_head in method_distribution_mapper.action_space.spaces
])
return num_outputs
def test_cartpole_policy_model():
env = GymMazeEnv(env='CartPole-v0')
observation_spaces_dict = env.observation_spaces_dict
action_spaces_dict = env.action_spaces_dict
flat_action_space = flat_structured_space(action_spaces_dict)
distribution_mapper = DistributionMapper(action_space=flat_action_space,
distribution_mapper_config={})
action_logits_shapes = {
step_key: {
action_head: distribution_mapper.required_logits_shape(action_head)
for action_head in action_spaces_dict[step_key].spaces.keys()
}
for step_key in action_spaces_dict.keys()
}
obs_shapes = observation_spaces_to_in_shapes(observation_spaces_dict)
policy = CustomPlainCartpolePolicyNet(obs_shapes[0],
action_logits_shapes[0],
hidden_layer_0=16,
hidden_layer_1=32,
use_bias=True)
critic = CustomPlainCartpoleCriticNet(obs_shapes[0],
hidden_layer_0=16,
hidden_layer_1=32,
use_bias=True)
obs_np = env.reset()
obs = {k: torch.from_numpy(v) for k, v in obs_np.items()}
actions = policy(obs)
values = critic(obs)
assert 'action' in actions
assert 'value' in values
def train_function(n_epochs: int, epoch_length: int, deterministic_eval: bool,
eval_repeats: int, distributed_env_cls,
split_rollouts_into_transitions: bool) -> SAC:
"""Implements the lunar lander continuous env and performs tests on it w.r.t. the sac trainer.
"""
# initialize distributed env
env_factory = lambda: GymMazeEnv(env="LunarLanderContinuous-v2")
# initialize the env and enable statistics collection
eval_env = distributed_env_cls([env_factory for _ in range(2)],
logging_prefix='eval')
env = env_factory()
# init distribution mapper
distribution_mapper = DistributionMapper(
action_space=env.action_space,
distribution_mapper_config=[{
'action_space':
'gym.spaces.Box',
'distribution':
'maze.distributions.squashed_gaussian.SquashedGaussianProbabilityDistribution'
}])
action_shapes = {
step_key: {
action_head:
tuple(distribution_mapper.required_logits_shape(action_head))
for action_head in env.action_spaces_dict[step_key].spaces.keys()
}
for step_key in env.action_spaces_dict.keys()
}
obs_shapes = observation_spaces_to_in_shapes(env.observation_spaces_dict)
# initialize policies
policies = {
ii: PolicyNet(obs_shapes=obs_shapes[ii],
action_logits_shapes=action_shapes[ii],
non_lin=nn.Tanh)
for ii in obs_shapes.keys()
}
for key, value in env.action_spaces_dict.items():
for act_key, act_space in value.spaces.items():
obs_shapes[key][act_key] = act_space.sample().shape
# initialize critic
critics = {
ii: QCriticNetContinuous(obs_shapes[ii],
non_lin=nn.Tanh,
action_spaces_dict=env.action_spaces_dict)
for ii in obs_shapes.keys()
}
# initialize optimizer
algorithm_config = SACAlgorithmConfig(
n_rollout_steps=5,
lr=0.001,
entropy_coef=0.2,
gamma=0.99,
max_grad_norm=0.5,
batch_size=100,
num_actors=2,
tau=0.005,
target_update_interval=1,
entropy_tuning=False,
device='cpu',
replay_buffer_size=10000,
initial_buffer_size=100,
initial_sampling_policy={
'_target_': 'maze.core.agent.random_policy.RandomPolicy'
},
rollouts_per_iteration=1,
split_rollouts_into_transitions=split_rollouts_into_transitions,
entropy_coef_lr=0.0007,
num_batches_per_iter=1,
n_epochs=n_epochs,
epoch_length=epoch_length,
rollout_evaluator=RolloutEvaluator(eval_env=eval_env,
n_episodes=eval_repeats,
model_selection=None,
deterministic=deterministic_eval),
patience=50,
target_entropy_multiplier=1.0)
actor_policy = TorchPolicy(networks=policies,
distribution_mapper=distribution_mapper,
device='cpu')
replay_buffer = UniformReplayBuffer(
buffer_size=algorithm_config.replay_buffer_size, seed=1234)
SACRunner.init_replay_buffer(
replay_buffer=replay_buffer,
initial_sampling_policy=algorithm_config.initial_sampling_policy,
initial_buffer_size=algorithm_config.initial_buffer_size,
replay_buffer_seed=1234,
split_rollouts_into_transitions=split_rollouts_into_transitions,
n_rollout_steps=algorithm_config.n_rollout_steps,
env_factory=env_factory)
distributed_actors = DummyDistributedWorkersWithBuffer(
env_factory=env_factory,
worker_policy=actor_policy,
n_rollout_steps=algorithm_config.n_rollout_steps,
n_workers=algorithm_config.num_actors,
batch_size=algorithm_config.batch_size,
rollouts_per_iteration=algorithm_config.rollouts_per_iteration,
split_rollouts_into_transitions=split_rollouts_into_transitions,
env_instance_seeds=list(range(algorithm_config.num_actors)),
replay_buffer=replay_buffer)
critics_policy = TorchStepStateActionCritic(
networks=critics,
num_policies=1,
device='cpu',
only_discrete_spaces={0: False},
action_spaces_dict=env.action_spaces_dict)
learner_model = TorchActorCritic(policy=actor_policy,
critic=critics_policy,
device='cpu')
# initialize trainer
sac = SAC(learner_model=learner_model,
distributed_actors=distributed_actors,
algorithm_config=algorithm_config,
evaluator=algorithm_config.rollout_evaluator,
model_selection=None)
# train agent
sac.train(n_epochs=algorithm_config.n_epochs)
return sac
def perform_test_maze_rllib_action_distribution(batch_dim: int):
""" distribution test """
random.seed(42)
np.random.seed(42)
torch.manual_seed(42)
# action space
act_space = spaces.Dict(spaces=dict(
sorted({
"selection":
spaces.Discrete(10),
"scale_input":
spaces.Box(shape=(5, ), low=0, high=100, dtype=np.float64),
"order_by_weight":
spaces.Box(shape=(5, ), low=0, high=100, dtype=np.float64)
}.items())))
# default config
config = [{
"action_space":
spaces.Box,
"distribution":
"maze.distributions.squashed_gaussian.SquashedGaussianProbabilityDistribution"
}, {
"action_head":
"order_by_weight",
"distribution":
"maze.distributions.beta.BetaProbabilityDistribution"
}]
# initialize distribution mapper
distribution_mapper = DistributionMapper(action_space=act_space,
distribution_mapper_config=config)
num_outputs = sum([
np.prod(distribution_mapper.required_logits_shape(action_head))
for action_head in distribution_mapper.action_space.spaces
])
model_config = {
'custom_model_config': {
'maze_model_composer_config': {
'distribution_mapper_config': config
}
}
}
assert num_outputs == MazeRLlibActionDistribution.required_model_output_shape(
act_space, model_config)
# assign action heads to registered distributions
logits_dict = dict()
for action_head in act_space.spaces.keys():
logits_shape = distribution_mapper.required_logits_shape(action_head)
if batch_dim > 0:
logits_shape = (batch_dim, *logits_shape)
logits_tensor = torch.from_numpy(np.random.randn(*logits_shape))
logits_dict[action_head] = logits_tensor
flat_input = torch.cat([tt for tt in logits_dict.values()], dim=-1)
if batch_dim == 0:
flat_input = flat_input.unsqueeze(0)
fake_model = FakeRLLibModel(distribution_mapper)
rllib_dist = MazeRLlibActionDistribution(flat_input,
fake_model,
temperature=0.5)
# test dictionary distribution mapping
maze_dist = distribution_mapper.logits_dict_to_distribution(
logits_dict=logits_dict, temperature=0.5)
for action_head in act_space.spaces.keys():
maze_distribution = maze_dist.distribution_dict[action_head]
maze_rllib_distribution = rllib_dist.maze_dist.distribution_dict[
action_head]
if hasattr(maze_distribution, 'logits'):
assert torch.allclose(maze_distribution.logits,
maze_rllib_distribution.logits)
if hasattr(maze_distribution, 'low'):
assert torch.allclose(maze_distribution.low,
maze_rllib_distribution.low)
assert torch.allclose(maze_distribution.high,
maze_rllib_distribution.high)
test_action_maze = maze_dist.sample()
test_action_rllib = rllib_dist.sample()
for action_head in act_space.spaces.keys():
assert test_action_maze[action_head].shape == test_action_rllib[
action_head].shape[int(batch_dim == 0):]
maze_action = maze_dist.deterministic_sample()
rllib_action = rllib_dist.deterministic_sample()
for action_head in act_space.spaces.keys():
assert torch.all(maze_action[action_head] == rllib_action[action_head])
maze_action = convert_to_torch(maze_action,
device=None,
cast=torch.float64,
in_place=True)
rllib_action = convert_to_torch(rllib_action,
device=None,
cast=torch.float64,
in_place=True)
# This un-sqeeze is preformed by rllib before passing an action to log p
for action_head in act_space.spaces.keys():
if len(rllib_action[action_head].shape) == 0:
rllib_action[action_head] = rllib_action[action_head].unsqueeze(0)
logp_maze_dict = maze_dist.log_prob(maze_action)
action_concat = torch.cat(
[v.unsqueeze(-1) for v in logp_maze_dict.values()], dim=-1)
logp_maze = torch.sum(action_concat, dim=-1)
logp_rllib = rllib_dist.logp(rllib_action)
if batch_dim == 0:
logp_rllib = logp_rllib[0]
assert torch.equal(logp_maze, logp_rllib)
logp_rllib_2 = rllib_dist.sampled_action_logp()
if batch_dim == 0:
logp_rllib_2 = logp_rllib_2[0]
assert torch.equal(logp_maze, logp_rllib_2)
maze_entropy = maze_dist.entropy()
rllib_entropy = rllib_dist.entropy()
if batch_dim == 0:
rllib_entropy = rllib_entropy[0]
assert torch.equal(maze_entropy, rllib_entropy)
logits_dict2 = dict()
for action_head in act_space.spaces.keys():
logits_shape = distribution_mapper.required_logits_shape(action_head)
if batch_dim > 0:
logits_shape = (batch_dim, *logits_shape)
logits_tensor = torch.from_numpy(np.random.randn(*logits_shape))
logits_dict2[action_head] = logits_tensor
flat_input = torch.cat([tt for tt in logits_dict2.values()], dim=-1)
if batch_dim == 0:
flat_input = flat_input.unsqueeze(0)
fake_model = FakeRLLibModel(distribution_mapper)
rllib_dist_2 = MazeRLlibActionDistribution(flat_input,
fake_model,
temperature=0.5)
# test dictionary distribution mapping
maze_dist_2 = distribution_mapper.logits_dict_to_distribution(
logits_dict=logits_dict2, temperature=0.5)
maze_kl = maze_dist.kl(maze_dist_2)
rllib_kl = rllib_dist.kl(rllib_dist_2)
if batch_dim == 0:
rllib_kl = rllib_kl[0]
assert torch.equal(maze_kl, rllib_kl)
nn.Linear(in_features=obs_shapes[OBSERVATION_NAME][0],
out_features=16), nn.Tanh(),
nn.Linear(in_features=16,
out_features=action_logits_shapes[ACTION_NAME][0]))
def forward(self, in_dict: Dict[str,
torch.Tensor]) -> Dict[str, torch.Tensor]:
""" forward pass. """
return {ACTION_NAME: self.net(in_dict[OBSERVATION_NAME])}
# init default distribution mapper
distribution_mapper = DistributionMapper(
action_space=spaces.Dict(spaces={ACTION_NAME: spaces.Discrete(2)}),
distribution_mapper_config={})
# request required action logits shape and init a policy net
logits_shape = distribution_mapper.required_logits_shape(ACTION_NAME)
policy_net = PolicyNet(obs_shapes={OBSERVATION_NAME: (4, )},
action_logits_shapes={ACTION_NAME: logits_shape})
# compute action logits (here from random input)
logits_dict = policy_net({OBSERVATION_NAME: torch.randn(4)})
# init action sampling distribution from model output
dist = distribution_mapper.logits_dict_to_distribution(logits_dict,
temperature=1.0)
# sample action (e.g., {my_action: 1})
action = dist.sample()