Skip to content

rail-berkeley/manipulator_gym

Repository files navigation

Manipulator Gym

⚠️ Still in early development, expect breaking changes ⚠️

This package provides a common gym-like environment for policy to interact with a manipulator robot. The environment is based on the gym interface, and the robots are defined as interfaces. The abstracted robot interfaces is easily swappable, modularized and run distributedly. Common utilities are provided to assist users to run robot policies (e.g. octo) on the robot.

Manipulator Interfaces Description
widowx Interbotix widowx interface in ROS1
widowx_ros2 Interbotix widowx interface in ROS2
widowx_sim Interbotix widowx in a pybullet sim
viperx Interbotix viperx interface in ROS1
service_client client interface to connect a remote interface server

more coming soon.... (etc. mujoco panda)

Other useful features are also provided

  • Eval script to run a robot policy octo
  • VR data collection script to collect data for RLDS training (also support remote data collection)
  • Fine-tuning script to finetune the octo model with the collected RLDS data

Installations:

  • pybullet: used in widowx_sim.
  • octo: for the generalist robot policy
  • agentlace: for distributed robot agent interfaces
  • oxe_envlogger: for logging the data in RLDS format

Quick Start

Code Snippets

To use the gym env, is quite straightforward. Define it as such:

from manipulator_gym.manipulator_env import ManipulatorEnv
from manipulator_gym.interfaces.base_interface import ManipulatorInterface

# 1. define select the robot interface
# ManipulatorInterface is an abstract class, you choose the interfaces from the list above
# e.g. interface = WidowXSimInterface(), ViperXInterface(), etc.
interface = ManipulatorInterface()

# 2. define the gym env, Done!
env = ManipulatorEnv(interface=interface)

Then, you can use the gym env as you would with any other gym env. 🥳 You can also use other interfaces to run different kinds of manipulators, e.g. widowx, viperx, action_client-server, etc.

Alos, we can make the interface to be a server-client setup, to run the policy on a different process, or even machine. Create 2 seperate scripts server.py and client.py:

  1. In server.py. Run the interface as server
from manipulator_gym.interfaces.base_interface import ManipulatorInterface
from manipulator_gym.interfaces.interface_service import ManipulatorInterfaceServer

server = ManipulatorInterfaceServer(manipulator_interface=interface)
server.start()
  1. In slient.py. Then run the env with the client interface
from manipulator_gym.manipulator_env import ManipulatorEnv
from manipulator_gym.interfaces.interface_service import ActionClientInterface

interface = ActionClientInterface(host=FLAGS.ip)
env = ManipulatorEnv(manipulator_interface=interface)

Tada!

Run Examples

First, we show how to run a simple sim env of a widowx robot, while enabling logging of rlds data.

# This runs a sim env of the manipulator
# To logs the data in RLDS format, add: --log_dir log_dir/   (requires oxe_envlogger)
python manipulator_gym/manipulator_env.py --widowx_sim

To visualize logged data

python read_rlds.py --show_img --rlds_dir log_dir/

Also, we can run the octo model to evaluate the robot policy on the robot. This requires the octo model to be installed. Showcase a distributed "robot server" and "policy client" setup

Run the "Robot Server"

python manipulator_server.py --widowx_sim

Run the "Policy Client":

# --text_cond is the task condition, e.g. "put the red marker on the table"
python octo_eval.py --ip IP_ADDRESS --show_img --text_cond "put the banana on the plate"

Communication nodes looks like this:

graph LR
    A[Widowx Sim, server interface]
    A <--agentlace--> B[Gym Env, action client interface <-> Octo Policy]

☝️ This is what you expect to see when running the above commands.


Interfaces

WidowX Sim

This uses the pybullet sim to simulate the widowx robot.

pip install pybullet

ViperX or WidowX

This requires user to download the respective dependencies for the robot, and run the robot interface.

Install the interbotix_ros_arms package for the viper/widowx robot for follow the docs from trossen roboitcs:

Now, it is time to run the policy on the real robot. To enable this, we use the following architecture to abstract out the robot env with the policy. This is enabled by agentlace.

graph LR
    A[Robot Driver] <--ROS Topics--> B[Manipulation server]
    B <--agentlace--> C[Gym Env <-> Octo Policy]
  1. Run the robot driver (ROS2 example)
# ros1: roslaunch interbotix_xsarm_control xsarm_control.launch robot_model:=wx250
# ros2 example:
ros2 launch interbotix_xsarm_control xsarm_control.launch.py robot_model:=wx250s
  1. Run the widowx/viperx server interface
# choose viperx 
python manipulator_server.py --widowx_ros2
  1. Run the octo model on a different machine
python octo_eval.py --ip IP_ADDRESS --show_img --text_cond "PROVIDE YOUR TEXT TASK"

Finetuning pipeline

Data collection and Fine-tuning of Octo model

Generate the log files for the RLDS training, this example uses test_vrmani_env dataset.

# use --test for testing mode, in actual use case with oculus controller
python3 vr_data_collection.py --rlds_output . --test

Note: need to oculus_reader to collect data with occulus vr controller (linear movement with controller, rotation with controller joystick)

(optional) Validate the generated log files by replaying on the robot gym env

python3 read_rlds.py --show_img --rlds_dir PATH_TO_LOGS  --replay 

Now finetune the model using the generated log files

cd octo
python scripts/finetune.py --config=../manipulator_gym/viperx_finetune_config.py --config.pretrained_path=hf://rail-berkeley/octo-small

Others

  • If you wish to directly wrap the gym env for distributed gym inference, you can directly use agentlace's action env wrapper
  • TODO: create util scripts for wandb eval loggings
  • TODO: better interfaces with extend-able sensors and actuators (e.g. camera, bimanual manipulators, etc.)

Notes

This is still in active development. Open issues for wishlist and bugs.

About

No description, website, or topics provided.

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages