araffin / sbx

SBX: Stable Baselines Jax (SB3 + Jax)
MIT License
328 stars 32 forks source link

[Question] framestack and train_freq for sbx #53

Open Jackflyingzzz opened 2 weeks ago

Jackflyingzzz commented 2 weeks ago

Hi there!

Thank you for developing SBX! I'm currently working with SB3 for real-time robot control and was wondering if SBX supports the framestack using DummyVecEnv wrapper? Additionally, can SBX algorithms like crossQ perform update/gradient steps per episode instead of per RL step by adjusting the train_freq parameter (e.g., from (5, "step") to (2, "episode"))? This feature is crucial for my real-time RL loop.

Thanks again for your help!

Best regards, Jacky

araffin commented 2 weeks ago

Hello,

real-time robot control and was wondering if SBX supports the framestack using DummyVecEnv wrapper? A

VecFrameStack, yes

Additionally, can SBX algorithms like crossQ perform update/gradient steps per episode instead of per RL step by adjusting the train_freq parameter (e.g., from (5, "step") to (2,

It should work but it might want to re-compile the training loop if the number of step changes. Also, this jax version should be fast enough to use with robot (I did use it at 30Hz with 10 gradient steps per step with the DroQ configuration of SAC).

Jackflyingzzz commented 1 week ago

Hello,

real-time robot control and was wondering if SBX supports the framestack using DummyVecEnv wrapper? A

VecFrameStack, yes

Additionally, can SBX algorithms like crossQ perform update/gradient steps per episode instead of per RL step by adjusting the train_freq parameter (e.g., from (5, "step") to (2,

It should work but it might want to re-compile the training loop if the number of step changes. Also, this jax version should be fast enough to use with robot (I did use it at 30Hz with 10 gradient steps per step with the DroQ configuration of SAC).

wow glad to hear that, thanks for sharing!