Add SERL launcher with a Franka Arm from state example

[Leo428]

Hi all,

This PR includes:

• the edgeml launcher codes from youliang (with unused parts removed)
• replay buffer implementations for fast online RL training. This part is taken from jaxrl-5, which we used in the ICRA paper. The memory-efficient buffer supports faster CPU-to-GPU transfer by packing obs and next obs for pixel images and performs pre-fetching to speed up training further. However, in this PR, the test example is with state-only training.
• an example of Franka Arm lifting a cube in simulation, with separate actor and learner nodes. This example should start picking up a cube around 8 minutes (20k steps) and should converge around 15 minutes on a 4090 machine.
• I also removed the isort hook in pre-commit due to a format conflict with black formatter hook. I checked this with youliang already.

Thank you so much for reviewing!

Notes:

• I followed the installation instructions for jaxrl-minimal. But I found that I have to upgrade the chex and tensorflow versions to run training correctly. Maybe jaxrl-minimal needs a pip requirement update at some point?
• I tested the code with https://github.com/rail-berkeley/jaxrl_minimal/tree/serl_dev. This branch is based on and identical to the main branch. I will add DrQ and encoder-related changes later this week once this PR is through.

[youliangtan]

Thanks for this PR. Some high-level thoughs.

Install edgeml from the original repo. This can easily be done by including the git path and commit hash in setup.py. This will automatically pull and install the code during pip3 install -e . in serl_launcher. The main idea here is that edgeml will get maintained overtime, and we try to prevent duplication of code when possible.

After these changes, i can try run the simulation code.

[youliangtan]

serl_launcher package

[Leo428]

Cool, I can change that.

[youliangtan]

Since Iam maintaining edgeml, so it is better not to duplicate code

    install_requires=[
        ....... TODO DEPENDENCIES
        "'edgeml @ git+https://github.com/youliangtan/edgeml.git@COMMIT_HASH',
    ],

[Leo428]

Would this still work if I added additional code for an efficient online RL replay buffer in the /data folder?

Also, for SERL's use cases, some of the features in edgeml are not used. Shall we meet some times tomorrow? I'd love to hear your suggestions in person.

[youliangtan]

yes, you will just add it to serl_launcher/data and the custom code will import edgeml as library.

It is okay that features are not used. Think edgeml as a numpy library, you only use classes that you required. If you make custom changes to edgeml, either open a PR, or just have the code in serl_launcher.

[youliangtan]

I see no diff in serl_dev and main in jaxrl_minimal, is that what it suppose to?

[Leo428]

This branch is created for future use, we are going to add in encoder stuff, DrQ, and VICE.

[Leo428]

eventually, we should merge changes back to the main branch though. This is for our convenience when testing and running code for now.

[youliangtan]

This seems to be similar to the one in jaxrl_m. If there's diff, place it in serl_dev, or keep it here for now (both options are good)

[Leo428]

Right, it can be in either places. I'm not placing it in jaxrl_minimal for now because I'm not sure if anyone else's code depends on the specifics of the dataset and replay_buffer implementation in jaxrl_minimal. Although I highly recommend this version of the memory efficient replay buffer as it is fast and compact when doing image-based online RL training.

[youliangtan]

Since this package is now call serl_launcher, the path should be:

serl_launcher/serl_launcher/data/memory_efficient_replay_buffer.py

[Leo428]

Hi @youliangtan,

I made the changes per your requests.

Some side notes: When running jaxrl-minimal, you might run into some bugs related to chex and tensorflow versions. You can run pip install --upgrade chex and pip install --upgrade tensorflow to fix those errors.

[youliangtan]

I managed to test-run the environment and run the actor and learner bash scripts.

Did minor code cleanups, and also updated the Readme doc. srry for the nitpicking, but is a good practice to use code block rather than code line in markdown (easier for users to copy)

0170394

Note: need to update chex to 0.1.85, will need to update jaxrl_m soon

[youliangtan]

Overall this works well. I set it up and ran it with a clean env. Havnt try to run long enough since i only tried it on a CPU setup. Can do it for the next iteration

[Leo428]

Nice! with CPU it might take a bit longer to converge.

Thank you for the code improvements!