Implementation of Voxel-based vision for SERL

[nisutte]

Hi all

This is a pull request as mentioned here. The goal was to implement voxel-based vision for a box picking task with a UR5 robot arm.

New implementations:

• serl_robot_infra/robot_controllers/ur5_controller.py UR5 Impedance controller
• serl_robot_infra/ur_env A new environment for picking boxes
• examples/box_picking_drq New experiment files for the same purpose
• serl_launcher/vision/resnet_v1_18.py A pre-trained ResNet18 model by Pytorch trained on ImagNet.
• serl_launcher/vision/voxel_grid_encoders.py Implementation of a 3D-conv based Voxel encoder.
• serl_launcher/wrappers Some new wrappers for: (1) Scaling the observation (2) Logging the observation statistics to wandb.

Modifications to SERL:

These modifications improved the performance for my case, you can choose which you want to include / change.

• serl_launcher/agents/continuous/sac.py Used original jax_rl implementation with backup_entropy enabled. Also, the target entropy is set to -action_space instead of -action_space/2.
• serl_launcher/agents/continuous/drq.py Added classmethod create_voxel_drq with voxel encoders. Additionally, implemented the 3D augmentation technique for voxels analogous to batched_random_crop for images.
• serl_launcher/common/encoding.py Added a Masked EncodingWrapper to easily ignore certain proprioceptive states (in our experiments, forces/torques were too noisy and decreased the performance).
• serl_launcher/vision/resnet_v1.py Dropout bugfix and added an optional num_kp parameter, which controls the number of keypoints.

Additional library requirements

• ur-rtde Library implements API for Universal Robots RTDE realtime interface.
• open3d Used to visualize the voxel perception of the end effector.
• clu Common Loop Utils for a quick parameter overview of the model. (Is not needed, but useful)

Videos showing the voxel-based policies can be seen here: Videos

[jianlanluo]

Hi, thanks for the work, since this is quite different from the main branch, I think it would make more sense if you push to a different branch on its own and write a documentation how to use it?