# How to not flatten action mask with Dict observation

**URL:** <https://discuss.ray.io/t/how-to-not-flatten-action-mask-with-dict-observation/5727>\
**Category:** RLlib\
**Created:** [April 8, 2022, 1:54pm UTC](https://discuss.ray.io/t/how-to-not-flatten-action-mask-with-dict-observation/5727 "2022-04-08T13:54:40Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![fedetask](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/fedetask/32/2233_2.png) [@fedetask](https://discuss.ray.io/u/fedetask)\
**Post date:** [April 8, 2022, 1:54pm UTC](https://discuss.ray.io/t/how-to-not-flatten-action-mask-with-dict-observation/5727/1 "2022-04-08T13:54:40Z")

</div>

**How severe does this issue affect your experience of using Ray?**

- Medium: It contributes to significant difficulty to complete my task, but I can work around it.

I want to perform action masking in my multi-agent environment using DQN with a custom model.  
My multi-agent environment produces observations as follows:

```python
observation = {
    agent_1_key: {
        'observation': {...} # Dictionary with observations for agent 1
        'action_mask': [1, 1, 0, ...] # Action mask for agent 1
    },
    agent_2_key: {
        'observation': {...} # Dictionary with observations for agent 2
        'action_mask': [1, 1, 0, ...] # Action mask for agent 2
    }
    ... # For each agent
}

```

My understanding is that, in a custom model, the `input_dict` of the `forward()` method will contain `'obs'` (the actual observation) and `'obs_flat'` (the flattened observation tensor).

However, if I use `'obs_flat'` as input to my model, it will contain also the `"action_mask"` field, and I don’t want that. On the other hand, I cannot use `"obs"` as input since it is a dictionary.

Do I need to modify the `"obs_flat"` tensor, flatten the `"obs"` dictionary with the `"action_mask"` fields removed, or is there another way to do the job ?

Thank you in advance