Flatten observation space (dictionary) in parametric actions

RLlib
1

Hi,

I am new to RLLib and not sure what is the best way to do build my parametric action model.
I followed the paramtertic_actions_models example the difference is that my ‘real_obs’ space is not a Box or Multibinary but a spaceDict, thus the action_embed_model.forward fails when looking for these obs_flat and they are a dictionary instead.

Here the code snippet from the current parametric action model example

class TorchParametricActionsModel(DQNTorchModel):
    """PyTorch version of above ParametricActionsModel."""

    def __init__(self,
                 obs_space,
                 action_space,
                 num_outputs,
                 model_config,
                 name,
                 true_obs_shape=(4, ),
                 action_embed_size=2,
                 **kw):
        DQNTorchModel.__init__(self, obs_space, action_space, num_outputs,
                               model_config, name, **kw)

        self.action_embed_model = TorchFC(
            Box(-1, 1, shape=true_obs_shape), action_space, action_embed_size,
            model_config, name + "_action_embed")

    def forward(self, input_dict, state, seq_lens):
        # Extract the available actions tensor from the observation.
        avail_actions = input_dict["obs"]["avail_actions"]
        action_mask = input_dict["obs"]["action_mask"]

        # Compute the predicted action embedding
        action_embed, _ = self.action_embed_model({
            "obs": input_dict["obs"]["cart"]
        })

        # Expand the model output to [BATCH, 1, EMBED_SIZE]. Note that the
        # avail actions tensor is of shape [BATCH, MAX_ACTIONS, EMBED_SIZE].
        intent_vector = torch.unsqueeze(action_embed, 1)

        # Batch dot product => shape of logits is [BATCH, MAX_ACTIONS].
        action_logits = torch.sum(avail_actions * intent_vector, dim=2)

        # Mask out invalid actions (use -inf to tag invalid).
        # These are then recognized by the EpsilonGreedy exploration component
        # as invalid actions that are not to be chosen.
        inf_mask = torch.clamp(torch.log(action_mask), FLOAT_MIN, FLOAT_MAX)

        return action_logits + inf_mask, state

    def value_function(self):
        return self.action_embed_model.value_function()

The difference here is that in my env the returned input_dict["obs"]["cart"] gym space is dictionary instead of a Box

action_embed, _ = self.action_embed_model({
            "obs": input_dict["obs"]["cart"]
        })

Any advice, hints will be welcome.

Thanks,

H

1 Like
2

Hi @hjarraya,

Have a look here and see if that works for you.

1 Like
3

Hi thanks for the answer,
I found it easier to flatten observations this way

def flatten_obs(obs_dict: Dict):
    tensors = []
    for key in obs_dict.keys():
        tensors.append(obs_dict[key])
    return torch.cat(tensors, dim=1)