class CustomTorchRPGModel(TorchModelV2, nn.Module):
"""Example of interpreting repeated observations."""
def __init__(self, obs_space, action_space, num_outputs, model_config,
name):
super().__init__(obs_space, action_space, num_outputs, model_config,
name)
nn.Module.__init__(self)
self.model = TorchFCNet(obs_space, action_space, num_outputs,
model_config, name)
def forward(self, input_dict, state, seq_lens):
# The unpacked input tensors, where M=MAX_PLAYERS, N=MAX_ITEMS:
# {
# 'items', <torch.Tensor shape=(?, M, N, 5)>,
# 'location', <torch.Tensor shape=(?, M, 2)>,
# 'status', <torch.Tensor shape=(?, M, 10)>,
# }
print("The unpacked input tensors:", input_dict["obs"])
print()
print("Unbatched repeat dim", input_dict["obs"].unbatch_repeat_dim())
print()
print("Fully unbatched", input_dict["obs"].unbatch_all())
print()
return self.model.forward(input_dict, state, seq_lens)
def value_function(self):
return self.model.value_function()
class ShallowModel(TorchModelV2):
def __init__(self, obs_space, action_space, num_outputs, model_config,
name):
super().__init__(obs_space, action_space, num_outputs, model_config,
name)
self.model = FullyConnectedNetwork(obs_space, action_space,
num_outputs, model_config, name)
self.register_variables(self.model.variables())
def forward(self, input_dict, state, seq_lens):
return self.model.forward(input_dict, state, seq_lens)
def value_function(self):
return self.model.value_function()
class CustomFCModel(TorchModelV2, nn.Module):
"""Example of interpreting repeated observations."""
def __init__(self, obs_space: gym.spaces.Space,
action_space: gym.spaces.Space, num_outputs: int,
model_config, name: str, **customized_model_kwargs):
super(CustomFCModel, self).__init__(
obs_space=obs_space,
action_space=action_space,
num_outputs=num_outputs,
model_config=model_config,
name=name,
)
nn.Module.__init__(self)
if "social_vehicle_config" in model_config["custom_model_config"]:
social_vehicle_config = model_config["custom_model_config"][
"social_vehicle_config"]
else:
social_vehicle_config = customized_model_kwargs[
"social_vehicle_config"]
social_vehicle_encoder_config = social_vehicle_config["encoder"]
social_feature_encoder_class = social_vehicle_encoder_config[
"social_feature_encoder_class"]
social_feature_encoder_params = social_vehicle_encoder_config[
"social_feature_encoder_params"]
self.social_feature_encoder = (social_feature_encoder_class(
**social_feature_encoder_params) if social_feature_encoder_class
else None)
self.model = TorchFCNet(obs_space, action_space, num_outputs,
model_config, name)
def forward(self, input_dict, state, seq_lens):
low_dim_state = input_dict["obs"]["low_dim_states"]
social_vehicles_state = input_dict["obs"]["social_vehicles"]
social_feature = []
if self.social_feature_encoder is not None:
social_feature, _ = self.social_feature_encoder(
social_vehicles_state)
else:
social_feature = [e.reshape(1, -1) for e in social_vehicles_state]
input_dict["obs"]["social_vehicles"] = (torch.cat(
social_feature, 0) if len(social_feature) > 0 else [])
return self.model.forward(input_dict, state, seq_lens)
def value_function(self):
return self.model.value_function()