API design

[cswinter]

Testing code block link:

incubator/entity_gym/entity_gym/environment.py

Lines 98 to 101 in 2191c9f

	@dataclass
	class SelectEntityAction:
	actions: List[Tuple[EntityID, EntityID]]
	"""Maps each actor to the entity they selected."""

Previous discussion: jeremysalwen/entity-gym#1

[jeremysalwen]

Copying the discussion from jeremysalwen/entity-gym#1 as a series of comments.

incubator/entity_gym/entity_gym/environment.py

Line 17 in 2191c9f

ActionSpace = Union[CategoricalActionSpace, SelectEntityActionSpace]

jeremysalwen:

I assume these two action types are initially created just as a proof of concept (i.e. one action type which is entity independent, and one action type which is entity dependent), but we would later be adding more action types.

A couple questions:

1. Is the intent to have a closed world of action spaces, i.e. the framework defines all possible action-spaces (and or ways to combine them), and the environment implementer is responsible for figuring out how to represent their action space in those terms, or a open world of action spaces, where you could introduce a new action space for your environment, and it will work as long as you define the appropriate encoding/action selection head for it?
2. Do we intend to support complex action spaces, a la Conditional Action Trees (https://arxiv.org/pdf/2104.07294.pdf)?
   2a. If we do, I also had another idea which might be useful here: instead of representing complex actions as a single compound action, we could represent them as a sequence of simple actions. As long as we have a way to decouple the timestep used for exponential discounting from the index of the action, I think this might be a simpler way to handle things.
3. For Stone Ground Hearth Battles, one of the action types we use is a permutation action (for reordering cards). Does this sound like the something you would want to support?

cswinter:

1. Hadn't thought about that yet, sounds like a good idea to add a way for people to define custom action spaces and ways to generate heads.
   YES on 2. and 3. I think the action space is mostly orthogonal to the network architecture so I've just been keeping it extremely minimal for now, but I think we should add support for many more types of actions eventually.

My response:
The way I think about action spaces is in terms of symmetries. In the absence of any symmetries, all action spaces just reduce to categorical spaces. Some of these symmetries are symmetries that describe the action space in isolation (e.g. floating point actions correspond to some sort of translation invariance in the floating point value), but some correspond to symmetries that involve both the action and the observation (e.g. actions which act on an entity correspond to an invariance of permuting the observations and actions of the different entities).

The former type of symmetries are much easier to define and extend. Basically anyone can come along and say "here we need to choose an action which is a point on the surface of a sphere", and define some head to attach to our NN which takes some vector of floating point values, and differentiably converts it into a valid point on the surface of a sphere. They don't need to have any awareness of the structure of the rest of the network, the observation space, etc, in order to do that.

It's the latter type of symmetries that I think are the difficult part. You have already created a demonstration of an API that allows us to express actions which have a single global symmetry (permutation invariance) through SelectEntityActionSpace. I think this is extremely powerful and probably covers 90% of use cases. PermutationActionSpace could also be easily added to this framework. However, in my mind we want to end up with a framework that could be cleanly extended with arbitrary symmetries, (e.g. translational invariance a la convolution).

I think there are still some open questions:

1. How much of the symmetry is actually necessary to capture in order to have a functioning API? i.e. is it possible that the symmetry is only necessary for constructing a good network architecture, not for defining the interface, and so we would essentially be smuggling in the definition of the network architecture through the description of the symmetries?
2. What are even the potential future symmetries that we would be adding? Are we basically done for practical purposes once we get permutation invariance and translation invariance?
3. How practical does this look when you actually turn it into an API? is it not worth it to make things extensible due to how much it will balloon the complexity? (this will possibly be answered when I write some example code)

[cswinter]

I don't have anything particularly intelligent to add at this point, but these are really good questions that I'd like to know the answer to 😄

[jeremysalwen]

incubator/entity_gym/entity_gym/environment.py

Lines 118 to 121 in 2191c9f

	def state_space(cls) -> Dict[str, Entity]:
	"""
	Returns a dictionary mapping the name of observable entities to their type.
	"""

jeremysalwen:

Why do we have a baked in dictionary into our action space? Shouldn't we instead just have our action space be directly an ActionSpace object, where the ActionSpace itself may be a UnionActionSpace[FirstActionSpace, SecondActionSpace,...]? The additional str identifiers everywhere make the API harder to parse. At the least, I think switching to a newtype style type (as you did for EntityID) for all the string identifiers would make the API easier to read.

cswinter:

Good point, also not a fan of stringly typed code. There needs to be some way to tie identify and tie together action spaces/action heads/action masks (e.g. this code). Initially, I had a base class for actions with a name but that turned out to be kind of annoying and I just decided to move it into a dict so it's consistent with the state space as well.
Some alternatives, curious to get your thoughts on those or any ideas you might have:

1. return a List[ActionSpace], instead of a string identifier the primary key is the index
2. return an ActionSpace, and add composite ActionSpaces that allow you to compose base ActionSpaces
   I think we do want some way of associating human-readable names associated with actions for debugging/metrics/manually playable envs. For (1), we could also (optionally) return a corresponding List[str] of labels, and composite ActionSpaces could support something similar.

There's a somewhat related annoyance that for a given environment the actions/action masks/action space have a specific and related structure that isn't captured by the generic interface. So you end up with a bunch of string dict lookups and type asserts, e.g. https://github.com/cswinter/entity-gym/blob/2fa93df66996907b1771291a0d11e8243ab4d857/entity_gym/envs/pick_matching_balls.py#L85-L86.

We could solve this with associate types on Environment/ActionSpace (I think this is not actually expressible in MyPy):

trait ActionSpace {
    type TActionMask: ActionMask;
    type TAction: Action;
}

struct CategoricalActionSpace {
    actions: Vec<String>,
}

impl ActionSpace for CategoricalActionSpace {
    type TActionMask = CategoricalActionMask;
    type TAction = CategoricalAction;
}

trait Environment {
    type TActionSpace: ActionSpace

    fn action_space() -> ActionSpace;

    fn act(&mut self, actions: Self::ActionSpace::TAction) -> (Observation, Self::ActionSpace::ActionMask);
}

Seems like a bad idea unless we want this library to be used only by Haskell fanatics 😛 Composite action spaces would also need something like HList/HDict. Maybe there's a simpler or more Pythonic solution that captures some of the benefits though.

Yeah it does seem like python won't allow you express it directly in the type system. It still might be possible to do something fake, where we use typedefs so it's all strings under the hood (i.e. type checking doesn't check anything), but we use different aliases in different places so to a reader it looks like it's typed?

[cswinter]

Another idea: As part of the observation, the environment returns a mutable object into which the actions can be written. All the type checking that ensures that the actions returned by the policy matches what the environment expects can be done in a generic way by the training loop/environment wrapper. The environment then reads out the actions out of the typed object (which it kept a reference to). Pretty horryfing interface though.

[cswinter]

Related mypy issue, it does sound like this will become possible at some point, but probably not soon: python/mypy#7790
Maybe we can at least future-proof the API so we can switch over to this once it becomes available?
There's also partial workarounds based on TypedDict, not sure if this works in our case but might be worth thinking about a bit more: python/mypy#7856 (comment)

[jeremysalwen]

incubator/entity_gym/entity_gym/environment.py

Lines 93 to 95 in 2191c9f

	class CategoricalAction:
	actions: List[Tuple[EntityID, int]]
	"""Maps each actor to the index of the chosen action."""

jeremysalwen:

The use of EntityID here seems strange, like we are mixing up the entities with the players which should be totally distinct? Is there something else going on here I don't understand? is it because in codecraft you have multiple characters controlled by a single player, and so they are all issued command simultaneously?

cswinter:

There is no multi-agent support currently, observation/action is for a single player. You can sort of do multiple agents by by using a VecEnv where each player is exposed as its own environment, see the multi_snake env. I'm hoping this will be sufficient for now, but I would like to have more comprehensive multi-agent support a la PettingZoo at some point. So the EntityIDs correspond to multiple objects controlled by the same player, e.g. different units in CodeCraft or MicroRTS.

So, if I understand correctly, every action must be associated with an entity? This is what surprised me a bit. How would we handle global actions? Is the idea that we would always have a dummy entity to associate all global actions with?

[cswinter]

From the perspective of the API, global actions can be thought of as syntactic sugar for a dummy Player object. It's possible that some efficiencies can be gained by explicitly marking some actions as always having exactly one actor.

[jeremysalwen]

Thanks for the examples, I think I understand better now.

Feel free to ignore the following musing:
We could describe the action space compositionally:

Categorical[N:int] - Choose one of N actions
Entity[] - Choose one valid entity
Product[A, B] - Choose an action from action space A, and an action from action space B
Exponential[A, N:int] := Product[A, A, ... , A, A], i.e. the product of N copies of A.

We then have

CategoricalAction[N:int]  := Exponential[Categorical[N], NumEntities]
SelectEntityAction := Exponential[Entity[], NumEntities]

i.e. All action spaces are implicitly wrapped in Exponential[X, NumEntities]. I see how this makes sense for a lot of domains where most of the action spaces have this structure. As you point out even if we don't want that structure for our action space, it is easy to use our action mask to essentially hide this structure by having a single "Player" entity.

My concern is that by forcing all actions to have this structure we may be making an interface that becomes annoying to extend in the future, but that is somewhat speculative.

[jeremysalwen]

One thing that does stand out to me is that in code like SelectEntityActionMask(actors=[0,1], mask=[[0,0,0,1,1], [0,0,0,1,1], [0,0,0,1,1]) the actors argument and the mask argument IMO are performing the same function, but this fact is obscured: They are defining a sparse 2d matrix of valid actions. I think it might be better to abstract this out, so we could swap out different representations?

[cswinter]

To be clear, I'm not opposed to having global action. I think it's probably a good idea that will make the API more ergonomic for many environments. Before I was thinking we might want to have something like FullActionSpace { global_actions: ActionSpace, entity_actions: Dict[str, ActionSpace] } but something based on your idea seems like it would be more elegant.

[cswinter]

Re action masks: agreed, not sure what exactly the right abstraction is, but we likely want to support different ways of defining the action masks, both for efficiency reasons and convenience. Note that currently we already have DenseXActionMask types which could have corresponding sparse versions.

[Theomat]

Yeah it does seem like python won't allow you express it directly in the type system. It still might be possible to do something fake, where we use typedefs so it's all strings under the hood (i.e. type checking doesn't check anything), but we use different aliases in different places so to a reader it looks like it's typed?

I am not so keen on the idea of just looking like it is typed, perhaps for speed we can avoid checks at runtime but in a debug case we actively want to check the types are correct and that constraints on tensors are respected (such as dimensional constraints), I believe this would provide a lot more insight when debugging code rather than getting the next error from entity gym because we didn't respect the constraints.

[jeremysalwen]

To be clear, what I was proposing here is to use a typedef e.g. so you write Dict[MyType, Entity] and according to MyPy typechecking it is equivalent to Dict[str, Entity], but the intent is clear to the reader. This would not preclude doing all the same runtime type-checking you would do if you were using Dict[str, Entity] instead.

[Theomat]

But for something like:

def _act(self, actions: Mapping[str, Action]) -> Observation:
        action = actions["Pick Ball"]
        assert isinstance(action, SelectEntityAction)

You would like a check similar to something like this:

def _act(self, actions: Mapping[str, Action]) -> Observation:
        if DEBUG:
             for action_name, action_mask in self._last_step_actions_masks.items():
                  action_mask.check(actions[action_name])
             
        action = actions["Pick Ball"] 

or perhaps this, which ensures that the correct field is present and of the correct type:

class PickBallActions(TypedDict):
        pick_ball: SelectEntityAction

def _act(self, actions: PickBallActions) -> Observation:
        action = actions.pick_ball

[cswinter]

Yeah, it would be nice to have some checking code like this that you can at least enable in debug mode. Ideally, we lift this into the type system, but I haven't been able to come up with a good way to do that yet.