Using GNNs for models

[SeaingAnt]

I was trying to use GNN models with topology "from_pos" and I'm struggling to understand where to add position_key. For what I got, I need to change the observation_spec but after that I don't know how to add this extra key in the tensordict. Moreover, it is unclear I could pass this key when using a sequential model where, for example, GNN is in between to MLPs.

What would the right procedure be to do so, in an existing or a custom environment?

[matteobettini]

Thanks a lot for openeing this!

You are right, the current implementation assumes that the position key needs to be in the observation_spec. However, as you point out, this won't work for sequence models where the GNN is not first (as the previous models will process it)

I can fix this making it so that the position key needs not be in the spec but just in the tensordict (similarly to what we do for hidden states of RNNs) so no other model will touch it.

Before this change

Before I make this change, what you can do to use gnn in a multilayer model is

• have gnn as first layer
• make your custom model and create a sequential inside of it

To add the position key to the observation it is sufficient to return it as part of the observation dictinary in your environment (supposing your env supports dictionary observation)

In VMAS you can do

 def observation(self, agent: Agent):
        return {
            "pos": agent.state.pos,
            "other": {"vel": agent.state.vel},
        }

and set position_key="pos"

After this change

After this change I'll make it will be easier.

No need to add it to the observation spec and you can just return it as part of your info

For example, in VMAS:

 def info(self, agent: Agent):
        return {
            "pos": agent.state.pos,
        }

If you tell me what environment you are working with I may be able to help further

[matteobettini]

I have made a draft PR for this at #125

If you want to try it out, you should require no change to the specs, just returning your position key as part of the info

[SeaingAnt]

Thanks a lot.

I imagine that "others" is the agent'name, am I right? Anyway, in pettingzoo adding "pos" in the observation gives an error when resetting pettingzoo env. I couldn't try your PR yet sorry.

I'll let you know next week for that.

[matteobettini]

Thanks a lot.

I imagine that "others" is the agent'name, am I right? Anyway, in pettingzoo adding "pos" in the observation gives an error when resetting pettingzoo env. I couldn't try your PR yet sorry.

I'll let you know next week for that.

In that case “other” was a nested observation but it is just an example so no need to consider it.

ok then i’ll merge the PR, when you have a chance try it by adding “pos” to the info and let me know how it goes.

if you get en error please open another issue with the details for me to reproduce (a simple env) and we’ll get it fixed

[matteobettini]

I had to amend the work in #125 with a new PR #132

The problem is that i discovered that the info dict is almost impossible to access during training.

So the current supported way is to return position and velocity keys as part of the observations (which need to be a dict)

This will however require that GNNs in a sequence model are the first.

You can circumvent this by creating your custom gnn model which does some preprocessing