[BUG] Actor SPS includes learning elements

EdanToledo / Stoix

🏛️A research-friendly codebase for fast experimentation of single-agent reinforcement learning in JAX • End-to-End JAX RL

Apache License 2.0

239 stars 24 forks source link

[BUG] Actor SPS includes learning elements #130

Open thomashirtz opened 1 week ago

thomashirtz commented 1 week ago

Describe the bug

When logging the acting, the SPS is around the learn function, therefore the acting sps is actually logging the acting time.

See https://github.com/EdanToledo/Stoix/pull/129#discussion_r1838576663

EdanToledo commented 1 week ago

I think the only thing to do is figure out if its even possible to have timings inside a jitted function. If its not possible then we are unfortunately bound by jax. I'll look into this soon.

thomashirtz commented 1 week ago

Yeah that is what I was thinking I guess that if we have the evaluation time and the training time it is quite good already. Although if we can extract the actor time, that would be nice to also be able to extract the env step SPS