Incorporate SerialEnv and introduct multistep policy logic

[alexander-soare]

This PR has two unrelated features but it was more sensible to make sure they work together.

1. Incorporate SerialEnv. For rollout, this runs the policy in batch mode, but the environment still runs in a sequential batch. For PushT this still provides a 10x speedup on a machine with an RTX 3090 and 16 CPU threads. The goal is to swap this out with ParallelEnv.

2. Creates a base class for policies thereby enabling a base forward method that handles multi-step policies. This is for correct handling of the step outputs (reward, terminated, truncated etc). For example, see: https://github.com/users/Cadene/projects/1?pane=issue&itemId=56225488

[Cadene]

Thanks! Very useful. Dont hesitate to ping me on discord if you need insights to pass the unit tests ;)
Best

[Cadene]

In the description could you add a minimal code to run the tests so that we can breakpoint easily to better understand the code
Also, the main issue that we are trying to address here is probably the miscalculation of the reward, no? Maybe we should add this in the description !
Thanks for your contribution ;)

[Cadene]

Let's iterate one more them and let's merge this very cool feature!
Could you add the time that you saved by adding SerialEnv?
Not sure we have time for this, but ideally we should make sure by training a diffusion policy on Pusht that we can reproduce results obtained before this PR and post them in the PR description.
Thanks!

[alexander-soare]

In the description could you add a minimal code to run the tests so that we can breakpoint easily to better understand the code Also, the main issue that we are trying to address here is probably the miscalculation of the reward, no? Maybe we should add this in the description ! Thanks for your contribution ;)

You mean the PR description? I'm confused, as the way to run the tests is the same as for all other tests.

I'll link the relevant issue.

[Cadene]

Left some comments ;) Almost ready to merge

[alexander-soare]

@Cadene bty. Can you please check the previously unresolved comments as well?

[alexander-soare]

@Cadene 🏓. I removed abstractmethods as you suggested. Also, please take a look at what I did re. online training. There's definitely a way to do it properly but I'm short-circuiting for now to get this PR through.