This package implements conv-autoencoder based visuo-motor imitation learning using pytorch. This package focuses on extensibility. One can define custom data types (vector, image and composite-image) and also map/inverse from these custom types to feature vector that fed into the LSTM. Alongside the imitation learning framework, this package provides two demo using pybullet and rlbench to show example usage. Please try kinematic pybullet demo and dynamic rlbench demo. The results of demo is available on google drive.
The ros wrapper of this package can be found in mohou_ros. Although mohou_ros currently supports only PR2 robot, many useful utilities and scripts for working with a real robot are included.
one of result of applying this framework to rlbench's task
For stable version:
pip3 install mohouFor beta from source:
git clone https://github.com/HiroIshida/mohou.git
cd mohou
pip install -e . After the pip install you are ready to start pybullet demo! We also provide rlbench_demo. As for rlbench demo, additional installation step of pyrep and rlbench is required. See https://github.com/stepjam/RLBench for the detail.
First, the important concept of the mohou package is "project". Each "project" has each directory and the directory contains everything, e.g. dataset, trained models, visualization results. Thanks to this concept, hard-coding the file path of TrainCache and EpisodeBundle, and many other stuff can be avoided. The use of the concept of project enables easy loading many objects. For example, EpisodeBundle which is a bundle of episodic sequential data is
can be dumped and loaded by
EpisodeBundle.dump(project_path)
EpisodeBundle.load(project_path)
Except the visualization stuff, the pipeline consists of 1) generation of dataset, 2) training autoencoder, 3) trainling lstm, 4) execution using the trained policy. For example, in the pybullet_demo.sh, kuka_reaching.py, python3 -m mohou.script.train_autoencoder, python3 -m mohou.script.train_lstm, and python3 $example_path/kuka_reaching.py --feedback corresponds to the four core steps. The result of all trained model is saved in {project_path}/models directory.
Note that step 1 and step 4 must vary according to the problem and kuka_reaching.py is just an example. That is, if you use the real robot, you must write own dataset collection program and execution program.
Other than the above steps, the software provide visualization method for autoencoder and lstm training reuslts, which are saved in {project_path}/autoencoder_result and {project_path}/lstm_result respectively.
The visualization of autoencoder result is done by
python3 -m mohou.script.visualize_autoencoder_result # plus optionswhich plots comparison of original and reconstruction images side-by-side.
The visualization of lstm result is done by
python3 -m mohou.script.visualize_lstm_result # plus optionswhich plots the result of LSTM prediction of images as gif files and joint angles as png files. In the prediction, we first feeds some sequential state to LSTM and then propagate without feeding any extra images.
These visualization is extremely important to know the training quality. Based on the visualization result, you can decide increase the number of episode data or increase the training epoch.
The teaching data must be saved as a EpisodeBundle. EpisodeBundle consists of multiple EpisodeData. EpisodeData consists of multiple ElementSequence. And, ElementSequence consists of sequence of each elements like AngleVector and RGBImage. The pseudo-code of data collection looks like the below.
import numpy as np
from mohou.types import AngleVector, ElementSequence, EpisodeBundle, EpisodeData, RGBImage
def obtain_rgb_from_camera() -> np.ndarray: # type: ignore
# implement by your self
pass
def obtain_joint_configuration() -> np.ndarray: # type: ignore
# implement by your self
pass
def create_episode_data() -> EpisodeData:
n_step = 100
rgb_list = []
av_list = []
for _ in range(n_step):
rgb_numpy = obtain_rgb_from_camera()
av_numpy = obtain_joint_configuration()
rgb = RGBImage(rgb_numpy, dtype=np.uint8) # type: ignore
av = AngleVector(av_numpy)
rgb_list.append(rgb)
av_list.append(av)
# convering list of element to a sequence type
rgb_sequence = ElementSequence(rgb_list)
av_sequence = ElementSequence(av_list)
return EpisodeData.from_seq_list([rgb_sequence, av_sequence])
n_episode = 20
episode_list = []
for _ in range(n_episode):
episode_list.append(create_episode_data())
chunk = EpisodeBundle.from_data_list(episode_list)
chunk.dump(project_path)Of course you can make a EpisodeBundle consists of DepthImage and GripperState or other your custom type.
A pseudo-code for execution can be written as below:
from mohou.default import create_default_propagator
from mohou.propagator import Propagator
from mohou.types import AngleVector, ElementDict, RGBImage
# create_default_propagator functions automatically resolve the autoencoder and lstm model path
# given the project_path, and then create the propagator.
propagator: Propagator = create_default_propagator(your_project_path)
while True:
# Observation using real/simulated robot
rgb: RGBImage = obtain_rgb_from_camera() # define the function by yourself
av: AngleVector = obtain_joint_configuration() # define the function by yourself
elem_dict = ElementDict((rgb, av))
propagator.feed(elem_dict)
# If your fed elem_dict contains RGBImage and AngleVector, then propagated
# elem_dict_pred also has RGBImage and AngleVector
elem_dict_pred: ElementDict = propagator.predict(n_prop=1)[0]
# Get specific element by providing the elemen type as a key
rgb_pred = elem_dict_pred[AngleVector]
av_pred = elem_dict_pred[RGBImage]
# send command
send_next_angle_vector(av_pred) # define by yourselfThe following figure show the type hierarchy. In the mohou framework, only the leaf types (filled by grey) can be instantiated. In most case, users would create custom type by inheriting from either CompositeImageBase, or VectorBase or PrimitiveImageBase. For the detail, please refere to mohou/types.py for how the built-in concrete types such as RGBDImage, RGBImage and AngleVector are defined.
When you make a new PR, you need to check that the tests passed and formatting is correct.
The test for mohou software consists of 3 steps: 1) static type check by mypy, 2) unit test by pytest, 3) integration test and 4) regression test.
To running these tests, you need install mypy and pytest by
pip3 install pytest mypy # or use --user option
Then, do the following
python3 -m mypy .
python3 -m pytest -v -s tests/unittest/
./tests/integration_test.sh
./tests/regression_test.pymohou code follows black standard style. Additionally, we use isort and flake8 to check if the code is following pep standard. Basically, what you have to do for formatting is running
./format.shTo running the format command, you need to install
pip3 install black isort flake8 autoflake