Skip to content

Commit 94afccd

Browse files
vateseifvateseif
committed
initial commit
0 parents  commit 94afccd

17 files changed

+1716
-0
lines changed

‎.gitignore

+179
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,179 @@
1+
# Byte-compiled / optimized / DLL files
2+
__pycache__/
3+
*.py[cod]
4+
*$py.class
5+
6+
# C extensions
7+
*.so
8+
9+
# Distribution / packaging
10+
.Python
11+
build/
12+
develop-eggs/
13+
dist/
14+
downloads/
15+
eggs/
16+
.eggs/
17+
lib/
18+
lib64/
19+
parts/
20+
sdist/
21+
var/
22+
wheels/
23+
share/python-wheels/
24+
*.egg-info/
25+
.installed.cfg
26+
*.egg
27+
MANIFEST
28+
29+
# PyInstaller
30+
# Usually these files are written by a python script from a template
31+
# before PyInstaller builds the exe, so as to inject date/other infos into it.
32+
*.manifest
33+
*.spec
34+
35+
# Installer logs
36+
pip-log.txt
37+
pip-delete-this-directory.txt
38+
39+
# Unit test / coverage reports
40+
htmlcov/
41+
.tox/
42+
.nox/
43+
.coverage
44+
.coverage.*
45+
.cache
46+
nosetests.xml
47+
coverage.xml
48+
*.cover
49+
*.py,cover
50+
.hypothesis/
51+
.pytest_cache/
52+
cover/
53+
54+
# Translations
55+
*.mo
56+
*.pot
57+
58+
# Django stuff:
59+
*.log
60+
local_settings.py
61+
db.sqlite3
62+
db.sqlite3-journal
63+
64+
# Flask stuff:
65+
instance/
66+
.webassets-cache
67+
68+
# Scrapy stuff:
69+
.scrapy
70+
71+
# Sphinx documentation
72+
docs/_build/
73+
74+
# PyBuilder
75+
.pybuilder/
76+
target/
77+
78+
# Jupyter Notebook
79+
.ipynb_checkpoints
80+
81+
# IPython
82+
profile_default/
83+
ipython_config.py
84+
85+
# pyenv
86+
# For a library or package, you might want to ignore these files since the code is
87+
# intended to run in multiple environments; otherwise, check them in:
88+
# .python-version
89+
90+
# pipenv
91+
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
92+
# However, in case of collaboration, if having platform-specific dependencies or dependencies
93+
# having no cross-platform support, pipenv may install dependencies that don't work, or not
94+
# install all needed dependencies.
95+
#Pipfile.lock
96+
97+
# poetry
98+
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
99+
# This is especially recommended for binary packages to ensure reproducibility, and is more
100+
# commonly ignored for libraries.
101+
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
102+
#poetry.lock
103+
104+
# pdm
105+
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
106+
#pdm.lock
107+
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
108+
# in version control.
109+
# https://pdm.fming.dev/#use-with-ide
110+
.pdm.toml
111+
112+
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
113+
__pypackages__/
114+
115+
# Celery stuff
116+
celerybeat-schedule
117+
celerybeat.pid
118+
119+
# SageMath parsed files
120+
*.sage.py
121+
122+
# Environments
123+
.env
124+
.venv
125+
env/
126+
venv/
127+
ENV/
128+
env.bak/
129+
venv.bak/
130+
131+
# Spyder project settings
132+
.spyderproject
133+
.spyproject
134+
135+
# Rope project settings
136+
.ropeproject
137+
138+
# mkdocs documentation
139+
/site
140+
141+
# mypy
142+
.mypy_cache/
143+
.dmypy.json
144+
dmypy.json
145+
146+
# Pyre type checker
147+
.pyre/
148+
149+
# pytype static type analyzer
150+
.pytype/
151+
152+
# Cython debug symbols
153+
cython_debug/
154+
155+
# PyCharm
156+
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
157+
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
158+
# and can be added to the global gitignore or merged into this file. For a more nuclear
159+
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
160+
#.idea/
161+
162+
.DS_Store
163+
.zshrc
164+
165+
deletable.*
166+
167+
*.hdf5
168+
169+
# models checkpoint
170+
*.ckpt
171+
172+
# images
173+
*.png
174+
175+
# pickle files
176+
*.pkl
177+
178+
# videos
179+
*.mp4

‎LICENSE

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

‎README.md

+46
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
# Imiation Learning for 250$ robot arm
2+
This repository contains a re-adapatation of [Action Chunking Transformer](https://github.com/tonyzhaozh/act/tree/main) that works for this [low-cost robot](https://github.com/AlexanderKoch-Koch/low_cost_robot) design (250$).
3+
4+
We are sharing the repo so anyone (non-experts included) can train a robot policy after a few teleoperated demonstraions.
5+
6+
## AI training
7+
### Setup
8+
Create conda environment
9+
~~~
10+
conda create --name act python=3.9
11+
conda activate act
12+
~~~
13+
14+
Install torch (for reference we add the versions we are using)
15+
~~~
16+
conda install pytorch==1.13.1 torchvision==0.14.1
17+
~~~
18+
19+
You can now install the requirements:
20+
~~~
21+
pip install -r requirements.txt
22+
~~~
23+
24+
Go to `config/config.py` and change the paths of the ports that connect leader and follower robots to your computer.
25+
26+
### Data collection
27+
In order to collect data simply run:
28+
~~~
29+
python record_episodes.py
30+
~~~
31+
You can select how many episodes you intend collecting in `config/config.py`. Change the value of `'num_episodes'` (set to 1 by default). You should be turning on the volume of your pc-- data for each episode will be recorded after you hear "Go" and it will stop when you hear "Stop".
32+
33+
You will also need to connect a camera to your computer and point it towards the robot while collecting the data via teleoperation. It's important the camera doesn't move otherwise evaluation of the policy is likely to fail.
34+
35+
### Train policy
36+
We slightly re-adapt [Action Chunking Tranfosrmer](https://github.com/tonyzhaozh/act/tree/main) to account for our setup. To start training simply run:
37+
~~~
38+
python train.py
39+
~~~
40+
The policy will be saved in `checkpoints/`.
41+
42+
### Evaluate policy
43+
Make sure to keep the same setup while you were collecting the data. To evaluate the policy simply run:
44+
~~~
45+
python evaluate.py
46+
~~~

‎cam.ipynb

+96
Large diffs are not rendered by default.

‎checkpoints/.gitkeep

Whitespace-only changes.

‎config/config.py

+66
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
import os
2+
# fallback to cpu if mps is not available for specific operations
3+
os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = "1"
4+
import torch
5+
6+
# data directory
7+
DATA_DIR = 'data/'
8+
9+
# checkpoint directory
10+
CHECKPOINT_DIR = 'checkpoints/'
11+
12+
# device
13+
device = 'cpu'
14+
if torch.cuda.is_available(): device = 'cuda'
15+
if torch.backends.mps.is_available(): device = 'mps'
16+
os.environ['DEVICE'] = device
17+
18+
# robot port names
19+
ROBOT_PORTS = {
20+
'leader': '/dev/tty.usbmodem57380045221',
21+
'follower': '/dev/tty.usbmodem57380046991'
22+
}
23+
24+
25+
# task config (you can add new tasks)
26+
TASK_CONFIG = {
27+
'sort':{
28+
'dataset_dir': os.path.join(DATA_DIR, 'sort'),
29+
'num_episodes': 20,
30+
'episode_len': 300,
31+
'state_dim': 5,
32+
'action_dim': 5,
33+
'cam_width': 640,
34+
'cam_height': 480,
35+
'camera_names': ['front']
36+
}
37+
}
38+
39+
40+
# policy config
41+
POLICY_CONFIG = {
42+
'lr': 1e-5,
43+
'device': device,
44+
'num_queries': 100,
45+
'kl_weight': 10,
46+
'hidden_dim': 512,
47+
'dim_feedforward': 3200,
48+
'lr_backbone': 1e-5,
49+
'backbone': 'resnet18',
50+
'enc_layers': 4,
51+
'dec_layers': 7,
52+
'nheads': 8,
53+
'camera_names': ['front'],
54+
'policy_class': 'ACT',
55+
'temporal_agg': False
56+
}
57+
58+
# training config
59+
TRAIN_CONFIG = {
60+
'seed': 42,
61+
'num_epochs': 2000,
62+
'batch_size_val': 8,
63+
'batch_size_train': 8,
64+
'eval_ckpt_name': 'policy_last.ckpt',
65+
'checkpoint_dir': os.path.join(CHECKPOINT_DIR, 'sort')
66+
}

‎data/.gitkeep

Whitespace-only changes.

0 commit comments

Comments
 (0)