Project 01 · robot learning
Training an SO-101 arm to pick and place
SO-101 Arms Project
The first project I completed this summer was based on training four Action Chunking with Transformers (ACT) policies to accomplish a pick-and-place task. ACT is an imitation-learning approach that uses camera images and robot joint positions to predict short sequences of actions. In order to perform this task, I used the SO-101 leader and follower arms, Codex, RunPod, Hugging Face, and multiple cameras. The process of building a policy for this particular robot went as follows. I would first collect a dataset using the leader arm to guide the follower arm, hence mirroring the task I wanted the follower arm to complete.
The dataset would be a collection of 30 episodes, including two different camera angles: one from the wrist of the follower arm and one from an overhead view. Along with these videos, the robot would also track each motor's joint position while it was moving. It was able to do this by using built-in joint encoders inside the servo motors. Hugging Face would match the joint positions and video feeds together in order to get the most accurate dataset. After each attempt, the ball would be reset into its original position along with the arm, while the video would be stored locally in a file on Codex. At the end of the dataset, the file would be uploaded to Hugging Face. Once the dataset was completed, I would rent a GPU on RunPod and then upload all of the joint data and video data in order to train the policy inside of the GPU. My Codex would connect to the GPU via SSH and then train the policy. Each training run I did would take around 2–3 hours to complete. For the SO-101 Arm project, I trained a total of 4 policies this way. The first policy I trained ran a simple pick-and-place task with the ball and arm starting in the same starting position for each episode. This policy was trained on a total of 30 episodes. Once the policy was trained, I began its evaluation. For each policy's evaluation, I ran 20 trials and would record if the trial was a success or a failure and, if it was a failure, what happened during the trial. Being able to write down and record what happened during the trials if there was a failure allowed me to go back and look at my results and adapt my training/data collection for future policies to make them more successful. The first policy had a success rate of 5%.
The second policy I trained ran the same pick-and-place task, but this time the ball was rotated through 10 different starting positions throughout the episodes, and the arm started in a variety of different positions. Due to the complexity of this policy, I trained it on a hundred-episode dataset instead of the regular thirty. The point of this policy was to see if the SO-101 arm would be able to learn/adapt to the new situations it was put into and still complete the task it was given, unlike the first policy, which used simply one spot and one arm position. Due to the complexity of the task and all of the different starting positions that were used, the evaluation process had to be changed. This time, we used 40 trials, and in addition to that change, I also recorded the position that the ball started in and the position the arm started in. The second policy had a success rate of 7.5%.
The third policy I trained took the first two datasets and combined them in order to see if the results from the first evaluation could be improved with the added 100 episodes to its training. Like the first evaluation, I used 20 trials for this evaluation. Unfortunately, this policy was not a success, as my success rate was 0%.
After the three evaluations of the two different types of datasets, I had noticed a few specific problems during my trials. The arm, when approaching the ball, would consistently approach the ball at a low angle. Because of this low trajectory, when it would go to open its claw, it would knock the ball off its spot and would be unable to then adapt and pick the ball up. Additionally, there were multiple occasions where the arm would stutter throughout its movements, especially at the beginning of each episode. It would stutter back and forth, like it was being indecisive on what to do. So, for the fourth policy that I trained, when collecting episodes, I made sure to try and make really smooth and clean movements in order to eliminate the jittering. On top of this, I emphasized an overhead approach in order to obtain a high trajectory with the arm. Once the dataset was complete, I trained the policy on RunPod and then began the evaluation. I made one minor adjustment to this evaluation. I extended each trial run to 30 seconds instead of the original 15. I did this because I felt that, in the previous evaluations, if the robot made one small mistake, with only 15 seconds, it did not have time to try and adapt and correct its mistake. The goal for this policy was to run 20 trials during this evaluation and try and get a success rate of above 70%. I exceeded this goal with a success rate of 85%.
But during this evaluation, when I had my three failures occur, they occurred in the span of four trials, and during all four of those trials, a ton of stuttering occurred. So, mid-evaluation, I had Codex evaluate the policy's performance and find out why the stuttering was occurring. Eventually, it came to a conclusion that, in order to get the best performance of the policy, I should connect the policy to a GPU on RunPod and then run it through there. Once this was completed, the policy was back up and running with no errors at all and completed the rest of the 20 trials without any problems or stuttering, just like the trials before this four-trial spell. After the 20 trials were done, I wanted to see if I could get three additional trials as successes to make up for the failed trials. I was able to complete that with a hundred percent success rate. So, even though I wasn't able to get a true 100% success rate, knowing that I could have if I had a clean policy running for all 20 trials felt like a huge success for me.
Results
Success rates for policies using the same task definition.
With such a high success rate for this policy, I wanted to run further tests on it to see just how good it was. I wanted to see if I turned the orientation of the ball, flipped it over, etc., if the arm would still be able to complete the task. I ran a small five-trial evaluation, in which each time the ball would be in a new orientation. The evaluation had a 100% success rate. With all the continued success, I wanted to test the arm once more. I wanted to see what would happen if I used a different object from the brown ball and if I used no object at all. For the different object, I used a circular black object with a flat bottom that was smaller than the brown ball, and it had a much harder texture than the brown ball. I tried two trials on the different-object test and one trial on the no-object test. Although all of them ended in failures, these results were able to tell me a lot. For the black-object test, the arm performed the test for the most part successfully. It was when it was moving to the bowl with the black object that it sensed that it didn't like something, and it dropped the object. It reset back to its original starting position after and then just stayed there for the duration of the test. It performed the same series of events both times. For the final test with no object, the arm made multiple attempts to grab at the spot where it thought an object would be, but after it was unable to grab anything, it got very confused. It tried to drop off nothing in the bowl and then made some more attempts back at the spot, but after multiple attempts of nothing happening, it finally returned to its home position. Both of these tests showed our policy was not able to initially recognize that something was off in the situation it was given. Additionally, we see that for both of these tests, the policy repeats the motion as if the original object was there. So, instead of being able to recognize at first that something was off in the situation and the object had changed, the policy has simply memorized the path of motion. But we do see, as the arm is moving towards the bowl, it may have detected something was wrong, but ultimately gets confused.
The next part of this project takes a little bit of a different turn. For this specific policy I trained, I first used a pre-trained model called Pi0.5. Pi0.5 is a vision-language-action model that enables robots to understand instructions and perform physical tasks in new environments. Using the Pi0.5 model, I post-trained, or fine-tuned, the policy using my thirty episodes from my V4 policy. The goal of this new fine-tuned policy was to try and see if I could improve my previous evaluation tests. The first test was a regular pick-and-place task from one spot, the second test involved using four different spots, 2 inches in each direction (left, right, up, and down) from the original spot, and the third test was from the original spot but used a different object. Unfortunately, due to time constraints on this project, I was only able to run the first test and half of the second test. The post-training of the Pi0.5 policy with the 30 episodes from V4 took roughly 10 hours to complete. The results of the first evaluation were very inconclusive, unlike the results of the V4 policy. My success rate was 30%—much lower than my 85% success rate on just the V4 policy.
For the second evaluation test, my success rate was 0%. Although I only completed half the test due to time constraints, I believe, due to what was occurring during the failed attempts, that the rest of the tests would have also failed, as on all the failed attempts, the arm kept trying to grab at the original spot. That being said, it would be incorrect for me to make a full scientific conclusion here, as the evaluation was not fully finished. My original hypothesis was that each evaluation's success rate would greatly improve; this was not the case whatsoever. I believe this happened because, for such a large model as Pi0.5, only thirty episodes of post-training just aren't enough to create a new policy that will consistently perform. Pi0.5 succeeded in 6 of 20 trials, but the successes were inconsistent. In addition, for the second evaluation, we only post-trained on one specific task, so even with all of Pi0.5's initial training, it isn't enough to create results, as during testing the policy seemed very reliant on the original spot.
Task definitions
Success rates for the same task under different definitions.
Download evaluation workbookRegular Pick and Place from One Spot
Multiple Starting Arm Spots and Starting Ball Positions
V4: 20 trials · Pi0.5 regular task: 20 trials · V2: 40 trials · Pi0.5 multi-position: 10 trials
Overall, the experience I gained using AI to train models, build datasets, and go through the process of tweaking datasets in order to build stronger policies based on the results of my prior policies was an important learning experience and, I believe, will really help me in future jobs and lab work.
