Action Mask API Design

[Bam4d]

In the Environment class we are asked to return an Observation object which has the following structure:

@dataclass
class Observation:
    entities: Dict[str, np.ndarray]
    """Maps each entity type to an array with the features for each observed entity of that type."""
    ids: Sequence[EntityID]
    """
    Maps each entity index to an opaque identifier used by the environment to
    identify that entity.
    """
    action_masks: Mapping[str, ActionMask]
    """Maps each action to an action mask."""
    reward: float
    done: bool
    end_of_episode_info: Optional[EpisodeStats] = None

In particular there is the action mask action_masks: Mapping[str, ActionMask] which provides a map of "action name" -> "action mask"

The ActionMask itself defines a 2d array defining which actor entity (row) can perform which action (column).

I think there is a case where this is not compatible with certain environments.

If an environment has a number of actor entities which define actions with the name "move", but those entities have different definitions under the hood. For example: one entity can move up,down,left,right.... but another can only rotate left, rotate right and move forwards.

Will we need to prefix all actions wtih some sort of entity_type, so that actions with the same names do not conflict if their inputs are different?

Something like this:
entity1_move -> up, down, left, right
entity2_move -> rotate left, rotate right, move forward

This feels like a hacky solution, so want to open this question to more ideas.

[Bam4d]

Another solution would be to flatten all actions for all entities and mask accordingly, but this leaves the arrays very sparse and feels wasteful.

"move" -> up, down, left, right, rotate left, rotate right, move forward

[cswinter]

Something like this:
entity1_move -> up, down, left, right
entity2_move -> rotate left, rotate right, move forward

This feels like a hacky solution, so want to open this question to more ideas.

This strikes me as the correct solution. The entity1_move and entity2_move are really two very distinct kinds of actions that should have separate action heads. The str key of action_masks is really more of a unique identifier rather than necessarily the canonical human readable "name" of the action. We could still have other places where both of these actions are surfaced with the same "Move" name.