Extension of What's the Move? Hybrid Imitation Learning via Salient Points (SPHINX) to mobile manipulation setting.
git clone https://github.com/priyasundaresan/mink.gitFirst create a conda env:
conda env create -f mac_env.yml Then, source set_env.sh to activate the tidybot conda env and set the PYTHONPATH appropriately.
# NOTE: run this once per shell before running any script from this repo
source set_env.shThis is a good sanity check to test if the repo works locally for you, in the conda environment you created.
Remember to run source set_env.sh once per shell before running any script from this repo.
To test whole-body teleoperation using simple mouse click-drag interactions, you can use the following script:
source set_env.sh
mjpython interactive_scripts/teleop_mouse.pyYou will see a little red interaction cube appear at the end effector.
You can Double Click to select it, then use Ctrl + Right Click and Drag (two fingers down on trackpad) to move it positionally, and Ctrl + Left Click and Drag to control orientation (one finger down on trackpad).
To test out whole-body teleoperation using the iPhone as the teleoperation device, you can use:
source set_env.sh
mjpython interactive_scripts/teleop_phone.pyThis will print out something like Starting server at 10.30.163.179:5001. Next, make sure your iPhone is connected to the same Wi-Fi network as your laptop, and that the app XRBrowser is installed. Open XRBrowser, and go to the IP address printed out by the script. You can then use your iPhone to teleoperate the robot. A few quick notes:
- Make sure your iPhone has Portrait Lock Orientaton ON, and start with top of phone facing up (towards your face) not forward (towards your toes).
- Click
Start episodeto start data collection, similarlyEnd episodeto finish an episode. - Your actions will only be mirrored when you are pressing the screen.
- Swipe up/down to open/close the gripper. NOTE: the simulated gripper is currently a bit finicky.
- You'll definitely want to practice before collecting any useful datasets! :)
Create a new data folder, download the data from here, and move it into data (NOTE, it is also available on sc cluster node: /iliad/u/priyasun/mink/data). See instructions at the bottom if you want to collect your own dataset.
You can run python dataset_utils/waypoint_dataset.py to load the cube dataset, and save some visualizations.
To train Mobile-SPHINX, run the following (training logs and eval success rates will be logged to Weights & Biases).
To train the cube and/or cabinet task:
# cube
python scripts/train_waypoint.py --config_path cfgs/waypoint/cube.yaml
# drawer
python scripts/train_waypoint.py --config_path cfgs/waypoint/cabinet.yamlUse --save_dir PATH to specify where to store the logs and models.
Use --use_wb 0 to disable logging to W&B (this is useful when debugging, to avoid saving unnecessary logs).
Assuming the resulting checkpoints are saved to exps/waypoint/cube, to eval the waypoint policy, you can run the following.
If you have access to a workstation (with GPU and display), run:
python scripts/eval_waypoint.py --model exps/waypoint/cube/ema.pt --env_cfg envs/cfgs/cube.yaml Otherwise, if you are evaluating on the cluster or some machine without a display (or over SSH):
MUJOCO_GL=egl python scripts/eval_waypoint.py --model exps/waypoint/cube/ema.pt --env_cfg envs/cfgs/cube.yaml --headless This will by default run 20 rollouts and save videos to the folder rollouts. For easier viewing, you can then use python common_utils/display_rollouts.py to creat a grid of all the rollout videos in a .html file that can easily be viewed in any browser.
Note:
--record 0 will run the rollouts without saving videos (faster if you don't care about visualizing)
Remember to run source set_env.sh once per shell before running any script from this repo.
This part walks through how to collect data for a task from scratch. You can use this general workflow to collect data for & train other custom tasks.
Run the following:
source set_env.sh
mjpython interactive_scripts/record_sim.py --env_cfg envs/cfgs/cube.yaml- Open XRBrowser, and go to the IP address printed out by the script, and hit
Start episode. - Wait for the simulator window to load, then begin teleoperation.
- Once done, you can click
End episode.- If
is_successis implemented, you should also see some feedback in Terminal when you have successfully teleoperated the task (after which you should hitEnd episode).
- If
- After you see
Done savingin Terminal, you can clickResetto begin the next episode. - In general, wait for the simulator to load before teleoperating, and if the robot is not responsive to your iPhone actions, just try refreshing the page.
Each teleoperated episode will be saved as an npz to dev1 as follows:
dev1/
└── demo00000.npz
└── demo00001.npz
...
NOTE: If you do mess up a demo after starting an episode and click End episode, you will need to manually delete the last recorded npz file. Every time you run the script record_sim.py, it will start saving from the last recorded demo index if there is one (i.e. if you just recorded demo00004.npz and quit, then re-run, it will save from demo00005.npz).
To sanity check how your recorded demos look, we provide a script that loads demos from dev1 and replays them.
python interactive_scripts/replay_sim.py --env_cfg envs/cfgs/cube.yamlOnce happy with the demos recorded in dev1, we need to post-process them into a SPHINX-compatible format (e.g., with mode labels and salient point annotations).
The first step is to break up each demo temporally into waypoint and dense modes.
Run the following script, which will load each demo in dev1, visualize it, and allow you to temporally annotate modes.
python dataset_utils/annotate_modes.py- Go to
http://127.0.0.1:5000in your browser. - Use the blue circular cursor to scroll through the frames of the first demo, and
Shift+Clickto specify a waypoint at that frame.Deletewill remove the last waypoint (if you mess up)- Try to use a consistent strategy and number of waypoints across demos. For
cube, I typically use 3 waypoints: one at the frame where the gripper 'approaches' the cube, one when it 'grasps', and one frame towards the very end of the demo (to 'lift').
- When you're happy labeling that demo, then go to Terminal and press
Enter. - Refresh the page to load the next demo.
- If the script for some reason crashes/hangs, interrupt it, re-run, and go back to the URL & refresh. It should load the most recent un-annotated demo.
After this step, you will have a new folder dev1_relabeled which contains all the demos, now annotated with modes.
The purpose of the above was to temporally relabel demos into dense/waypoint modes. The last step is to label salient points for the extracted waypoint observations above. Run:
python dataset_utils/annotate_salient_points.py- In the window that appears, you will see the point cloud of the first extracted waypoint timestep from the first demo in
dev1_relabeled- You can drag the point cloud around with just
Clickinteractions and zoom in using the trackpad to get a better view of where you want to put a salient point
- You can drag the point cloud around with just
Shift+Clickto label a salient point (a colored sphere will appear)- You can re-click if you mess up, just note that only the last click will be recorded.
- Press
qorEscto go to the next obs.
After this step, dev1_relabeled contains all the demos, now annotated with both modes and salient points!
Rename this folder to whatever you want and put it in data. See above for how to train/eval the policy on the task for which you just collected data.
To add a custom task, you need to do the following:
- Create XML's:
- In
interactive_scripts/stanford_tidybot, create two files:<task_name>.xmlandtidybot_<task_name>.xml(you can basically just copy overcube.xmlandtidybot_cube.xml, replacing with your object assets in<task_name>.xml).
- In
- Create a new
env_cfg:envs/cfgs/<task_name>.yamland give it a name in thetaskfield - Register the task in
envs/mj_env.py- In the
__init__method, updatexml_filebased ontask - Update
reset_task - Update
is_success
- In the
- Finally, you can try
mjpython interactive_scripts/record_sim.py --env_cfg envs/cfgs/<task_name>.xml- NOTE: If you get the following error:
ValueError: Error: keyframe 0: invalid qpos size, expected length X, it means thehomekeyframe oftidybot_<task_name>.xmlis not the right dimension. This keyframe represents the home configuration of the whole robot, plus thefree jointsof whatever assets are in the scene (i.e., there are 7 free joints for thecubetask, so thehomekeyframe is padded with0.6 0 0 0 0 0 0representing that the cube should be initially positioned at[0.6, 0, 0]and a "zero" quaternion -- you can set this to something else if you want different initial object poses). - NOTE: You can make
reset_taskdo nothing at first, andis_successtrivially beFalseduring debugging. I typically do this just so I can first focus on loading the scene properly, and once I am able to successfully teleoperate the task, I work backwards to figure out what initial scene randomizations are reasonable (forreset_task) and what the task success condition should be (foris_success, used to evaluate if rollouts are successful).
- NOTE: If you get the following error:
- If all of this works as expected, you should be able to just pass the
env_cfgyou created whenever you are collecting demos withrecord_sim.pyor evaluating policies with the eval scripts inscripts/eval_<waypoint/hybrid/dense>.py.