[Feature] multiagent data standardization: PPO advantages

[matteobettini]

Hello!

Context

So, in this library we frequently apply standardizations over data (also called normalization) where we subtract the mean from the data and divide it by a standard deviation.

A few places where we do this that come to mind are in the transforms (where it can be used over observations and rewards) and in the advantages in ppo. Please let me know about other places if you can think of more.

Issue

The issue is that normally the mean and std statistics are computed over all the dimensions. But in multiagent settings, where for example a reward tensor could have an "agent" dimension, this is not the proper way to do it as different agents might have different reward distributions and applying standardization naively across them will create issues.

Example

For example consider the case where we have 2 agents each with 2 rewards (agent_1_r = [100,200], agent_2_r = [0,-1]), structured in a 2x2 tensor r = torch.tensor([[100,200],[0,-1]])

If we standardize using mean(r) and std(r), we get r = torch.tensor([[0.3,1.5],[-0.9,-.91]]), which is bad

If we standardize using mean(r, dim=1) and std(r,dim=1), hence excluding the agent dim from the statistics, we get r = torch.tensor([[-1,1],[1,-1]]), which is better as it does not cross pollute over agents

Proposal

We should adapt all places in the library where standardization is used to take a list of dimensions to exlude from the statistics. This way, multiagent users can provide the agent dimension as input and use the standardization feature

What I did here

I did this for the advantages in PPO

To do it, I made a nice standardization util function that takes as input a list of dimensions to exlude from the standardization (a more flexible implementation of a functional layer norm)

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/2677

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.