Skip to content

Commit 29416a7

Browse files
HaozhiQiHaozhiQi
committed
release
0 parents  commit 29416a7

File tree

181 files changed

+84624
-0
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

181 files changed

+84624
-0
lines changed

‎LICENSE

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
Copyright (c) 2024 Authors of Lessons from Learning to Spin “Pens”
2+
3+
Permission is hereby granted, free of charge, to any person obtaining
4+
a copy of this software and associated documentation files (the
5+
"Software"), to deal in the Software without restriction, including
6+
without limitation the rights to use, copy, modify, merge, publish,
7+
distribute, sublicense, and/or sell copies of the Software, and to
8+
permit persons to whom the Software is furnished to do so, subject to
9+
the following conditions:
10+
11+
The above copyright notice and this permission notice shall be
12+
included in all copies or substantial portions of the Software.
13+
14+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
15+
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
16+
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
17+
NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
18+
LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
19+
OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
20+
WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

‎README.md

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,124 @@
1+
# Lessons from Learning to Spin “Pens”
2+
3+
<p align="center">
4+
<img src="assets/teaser.gif" width="1000"/>
5+
</p>
6+
7+
This repository contains a reference PyTorch implementation of the paper:
8+
9+
<b>Lessons from Learning to Spin “Pens”</b> <br>
10+
[Jun Wang*](https://wang59695487.github.io/),
11+
[Ying Yuan*](https://yingyuan0414.github.io/),
12+
[Haichuan Che*](https://www.linkedin.com/in/haichuan-che-7338721b1/),
13+
[Haozhi Qi*](https://haozhi.io/),
14+
[Yi Ma](http://people.eecs.berkeley.edu/~yima/),
15+
[Jitendra Malik](https://people.eecs.berkeley.edu/~malik/),
16+
[Xiaolong Wang](https://xiaolonw.github.io/) <br>
17+
[[Website](https://penspin.github.io/)]
18+
19+
## Installation
20+
21+
See [installation instructions](docs/install.md).
22+
23+
## Introduction
24+
25+
Our pen spinning method contains the following four steps.
26+
1. Learn a oracle policy with privileged information, point-clouds, and tactile sensor output with RL in simulation.
27+
2. Learn a student policy using the rollout of the oracle policy, also in simulation.
28+
3. Rollout trajectories generated by the oracle policy in a real robot, with initial state distribution matched. The success trajectories are collected while failures are discarded.
29+
4. Finetune the student policy in step 2 with the real-world successful trajectories.
30+
31+
The following session only provides example script of our method. For baselines, checkout [baselines](docs/baseline.md).
32+
33+
## Step 0: Visualize a Pre-trained Oracle Policy
34+
35+
```
36+
cd outputs/AllegroHandHora
37+
gdown 1LCRFE6lvKSUDPpUfEATOmpDUPDbB7n8d
38+
unzip demo.zip -d ./
39+
cd ../../
40+
scripts/vis_teacher.sh demo
41+
```
42+
43+
44+
## Step 1: Oracle Policy training
45+
46+
To train an oracle policy $f$ with RL, run
47+
48+
```
49+
# 0 is GPU is
50+
# 42 is experiment seed
51+
scripts/train_teacher.sh 0 42 output_name
52+
```
53+
54+
After training your oracle policy, you can visualize it as follows:
55+
```
56+
scripts/vis_teacher.sh output_name
57+
```
58+
59+
## Step 2: Student Policy Pretraining
60+
61+
In this section, we train a proprioceptive student policy by distilling from our trained oracle policy $f$.
62+
63+
Note we use the teacher rollout to train student policy, in contrast to DAgger in previous works.
64+
65+
```
66+
scripts/train_student_sim.sh train.ppo.is_demon=True train.demon_path=ORACLE_CHECKPOINT_PATH
67+
```
68+
We have provided a reference teacher checkpoint in [Google Drive](https://drive.google.com/file/d/1LCRFE6lvKSUDPpUfEATOmpDUPDbB7n8d/view?usp=sharing).
69+
70+
## Step 3: Open-Loop Replay in Real Hardware
71+
72+
To generate open-loop replay data for the student policy $\pi$, run
73+
```
74+
python real/robot_controller/teacher_replay.py --data-collect --exp=0 --replay_data_dir=REPLAY_DATA_DIR
75+
```
76+
where `REPLAY_DATA_DIR` is the directory to save the replay data.
77+
78+
Then process the replay data.
79+
80+
## Step 4: Real-world Fine-tuning
81+
82+
To fine-tune the student policy $\pi$ using real data, run
83+
```
84+
scripts/finetune_ppo.sh --real-dataset-folder=REAL_DATA_PATH --checkpoint-path=YOUR_CHECKPOINTPATH
85+
```
86+
87+
## Real Data Download
88+
Please download the real reference data from [Google Drive](https://drive.google.com/drive/folders/1TAMAvqLp3b5vEmdyrdcgW0kBW1GAxoyy?usp=sharing).
89+
```
90+
Real data:
91+
real_data.h5 is in the format of h5 file, which contains the following keys:
92+
-replay_demon_{idx}: the idx-th replay demonstration data
93+
- qpos: the current qpos of the robot
94+
- action: the delta action applied to the robot
95+
- current_target_qpos: the target qpos of the robot
96+
97+
real_data_full.h5 is a full version of real_data.h5, which contains the following keys:
98+
-replay_demon_{idx}: the idx-th replay demonstration data
99+
- qpos: the current qpos of the robot
100+
- action: the delta action applied to the robot
101+
- current_target_qpos: the target qpos of the robot
102+
- rgb_ori: the original rgb image
103+
- rgb_c2d: the rgb image after camera2depth image processing
104+
- depth: the depth image
105+
- pc: the point cloud
106+
- obj_ends: the position of object ends
107+
```
108+
109+
## Acknowledgement
110+
111+
Note: This repository is built based on [Hora](https://github.com/HaozhiQi/hora) and [IsaacGymEnvs](https://github.com/isaac-sim/IsaacGymEnvs).
112+
113+
## Citing
114+
115+
If you find **PenSpin** or this codebase helpful in your research, please consider citing:
116+
117+
```
118+
@article{wang2024penspin,
119+
author={Wang, Jun and Yuan, Ying and Che, Haichuan and Qi, Haozhi and Ma, Yi and Malik, Jitendra and Wang, Xiaolong},
120+
title={Lessons from Learning to Spin “Pens”},
121+
journal={},
122+
year={2024}
123+
}
124+
```

‎assets/ball.urdf

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
<?xml version="1.0"?>
2+
<robot name="ball">
3+
<link name="ball">
4+
<visual>
5+
<origin xyz="0 0 0"/>
6+
<geometry>
7+
<mesh filename="ycb/056_tennis_ball/google_16k/textured.obj" scale="1.1934954497985975 1.1934954497985975 1.1934954497985975"/>
8+
</geometry>
9+
</visual>
10+
<collision>
11+
<origin xyz="0 0 0"/>
12+
<geometry>
13+
<sphere radius="0.04"/>
14+
</geometry>
15+
</collision>
16+
<inertial>
17+
<mass value="0.05" />
18+
<inertia ixx="0.0001" ixy="0.0" ixz="0.0" iyy="0.0001" iyz="0.0" izz="0.0001"/>
19+
</inertial>
20+
</link>
21+
</robot>

‎assets/cube.urdf

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
<?xml version="1.0"?>
2+
<robot name="object">
3+
<link name="object">
4+
<visual>
5+
<origin xyz="0 0 0"/>
6+
<geometry>
7+
<mesh filename="cube_multicolor.obj" scale="0.08 0.08 0.08"/>
8+
</geometry>
9+
</visual>
10+
11+
<collision>
12+
<origin xyz="0 0 0"/>
13+
<geometry>
14+
<box size="0.08 0.08 0.08"/>
15+
</geometry>
16+
</collision>
17+
18+
<inertial>
19+
<mass value="0.05" />
20+
<inertia ixx="0.0001" ixy="0.0" ixz="0.0" iyy="0.0001" iyz="0.0" izz="0.0001"/>
21+
</inertial>
22+
</link>
23+
</robot>

‎assets/cube_multicolor.mtl

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
newmtl red
2+
Ns 10.0
3+
Ka 1.0 1.0 1.0
4+
Kd 1.0 0.0 0.0
5+
Ks 0.125 0.125 0.125
6+
Ke 0.0 0.0 0.0
7+
Ni 1.0
8+
d 1.0
9+
illum 2
10+
11+
newmtl green
12+
Ns 10.0
13+
Ka 1.0 1.0 1.0
14+
Kd 0.0 1.0 0.0
15+
Ks 0.125 0.125 0.125
16+
Ke 0.0 0.0 0.0
17+
Ni 1.0
18+
d 1.0
19+
illum 2
20+
21+
newmtl blue
22+
Ns 10.0
23+
Ka 1.0 1.0 1.0
24+
Kd 0.0 0.0 1.0
25+
Ks 0.125 0.125 0.125
26+
Ke 0.0 0.0 0.0
27+
Ni 1.0
28+
d 1.0
29+
illum 2
30+
31+
newmtl yellow
32+
Ns 10.0
33+
Ka 1.0 1.0 1.0
34+
Kd 1.0 1.0 0.0
35+
Ks 0.125 0.125 0.125
36+
Ke 0.0 0.0 0.0
37+
Ni 1.0
38+
d 1.0
39+
illum 2
40+
41+
newmtl cyan
42+
Ns 10.0
43+
Ka 1.0 1.0 1.0
44+
Kd 0.0 1.0 1.0
45+
Ks 0.125 0.125 0.125
46+
Ke 0.0 0.0 0.0
47+
Ni 1.0
48+
d 1.0
49+
illum 2
50+
51+
newmtl white
52+
Ns 10.0
53+
Ka 1.0 1.0 1.0
54+
Kd 1.0 1.0 1.0
55+
Ks 0.125 0.125 0.125
56+
Ke 0.0 0.0 0.0
57+
Ni 1.0
58+
d 1.0
59+
illum 2

‎assets/cube_multicolor.obj

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
mtllib cube_multicolor.mtl
2+
3+
v -0.5 -0.5 -0.5
4+
v -0.5 -0.5 0.5
5+
v -0.5 0.5 -0.5
6+
v -0.5 0.5 0.5
7+
v 0.5 -0.5 -0.5
8+
v 0.5 -0.5 0.5
9+
v 0.5 0.5 -0.5
10+
v 0.5 0.5 0.5
11+
12+
vn -0.5 -0.5 0.5
13+
vn -0.5 -0.5 -0.5
14+
vn -0.5 0.5 -0.5
15+
vn -0.5 -0.5 -0.5
16+
vn 0.5 -0.5 -0.5
17+
vn -0.5 -0.5 -0.5
18+
19+
usemtl red
20+
f 1//2 7//2 5//2
21+
f 1//2 3//2 7//2
22+
usemtl green
23+
f 1//6 4//6 3//6
24+
f 1//6 2//6 4//6
25+
usemtl blue
26+
f 3//3 8//3 7//3
27+
f 3//3 4//3 8//3
28+
usemtl yellow
29+
f 5//5 7//5 8//5
30+
f 5//5 8//5 6//5
31+
usemtl cyan
32+
f 1//4 5//4 6//4
33+
f 1//4 6//4 2//4
34+
usemtl white
35+
f 2//1 6//1 8//1
36+
f 2//1 8//1 4//1

‎assets/cylinder.urdf

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
<?xml version="1.0"?>
2+
<robot name="object">
3+
<link name="object">
4+
<visual>
5+
<origin xyz="0 0 0"/>
6+
<geometry>
7+
<cylinder radius="0.04" length="0.08"/>
8+
</geometry>
9+
</visual>
10+
11+
<collision>
12+
<origin xyz="0 0 0"/>
13+
<geometry>
14+
<cylinder radius="0.04" length="0.08"/>
15+
</geometry>
16+
</collision>
17+
18+
<inertial>
19+
<mass value="0.05" />
20+
<inertia ixx="0.0001" ixy="0.0" ixz="0.0" iyy="0.0001" iyz="0.0" izz="0.0001"/>
21+
</inertial>
22+
</link>
23+
</robot>

‎assets/cylinder/pencil-5-7/0000.npy

152 Bytes
Binary file not shown.

‎assets/cylinder/pencil-5-7/0000.urdf

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
<robot name="object">
2+
<link name="object">
3+
<visual>
4+
<origin xyz="0 0 0" />
5+
<geometry>
6+
<cylinder radius="0.04" length="0.4" />
7+
</geometry>
8+
</visual>
9+
10+
<collision>
11+
<origin xyz="0 0 0" />
12+
<geometry>
13+
<cylinder radius="0.04" length="0.4" />
14+
</geometry>
15+
</collision>
16+
17+
<inertial>
18+
<mass value="0.05" />
19+
<inertia ixx="0.0001" ixy="0.0" ixz="0.0" iyy="0.0001" iyz="0.0" izz="0.0001" />
20+
</inertial>
21+
</link>
22+
</robot>

‎assets/cylinder/pencil-5-7/0001.npy

152 Bytes
Binary file not shown.

0 commit comments

Comments
 (0)