Discrete RL Updates take longer than Continuous updates

[chrisw-shadow]

Hi there,

I've noticed that the RL updates take significantly longer when using a2c_discrete versus using a2c_continuous, with all else being as similar as possible. This is somewhat unexpected. Do you guys have any inclination as to what causes it? Is this a known issue?

Thank you 🙏

[chrisw-shadow]

Some more details: for a continuous action space, the RL update takes about 2ms but for discrete action space we're looking at around 12-14ms.

This seems to be tied to the number of DoFs that we are producing actions for. We are providing the action space as a tuple of discrete spaces. Is this correct? It appears that the actual MultiDiscrete space from gym is not supported

[Denys88]

Hi, Ill take a look. Overall sampling from categorical distribution is slower than Normal but 2 ms 12-14ms is a huge. Ill take a look.
What environment did you use to compare?

[chrisw-shadow]

Hey, we are using a custom environment in a custom sim so I won't be able to share that unfortunately. But, if my hypothesis is correct you should be able to observe this in any environment with a Multi-Discrete action space (or rather a tuple of Discrete spaces as it seems to be set up in rl games)

I think the slow down comes from this kinda thing:

rl_games/rl_games/algos_torch/network_builder.py

Line 467 in e087014

logits = [logit(out) for logit in self.logits]

Where each DoF is being processed sequentially in the forward pass and probably in the RL update too. I confirmed this by looking at the pytorch profile and we are making 12 calls to functions like multinomial and logsumexp in sequence for our robot with 12 DoFs. The continuous case, in contrast, makes 1 call to normal (which I assume is the continuous equivalent to multinomial).

sorry for the horrible screenshot.

[Denys88]

Ah I see, it might be expected in this case ( at least with my approach) theoretically if you can guarantee that all you distributions have the same amount of classes we can make it in one call. Ill take a look what could be done here.

[Denys88]

Working on demo for you :) should be ready today/tomorrow I hope you will be able to profile with your env.

[Denys88]

@chrisw-shadow could you try this branch #331 ?

[chrisw-shadow]

@chrisw-shadow could you try this branch #331 ?

hey ! thanks for this, I'll aim to give it a try some time this week and I'll get back to you with some results