Image Based Reinforcement Learning
By: Karanbir Chahal
Tags: Reinforcement Learning, Sim2real

This blog post describes the project which I worked on for my Fall 2020 Deep Reinforcement Learning class. The topic of my project was to explore learning RL policies from pixels, these policies were trained to solve the OpenAI Fetch Robotic environments. My work wanted to explore the possible answers to the following questions:

What are effective RL algorithms to train these OpenAI algorithms ? What are the differences in the learning dynamics of state based and image based learning ?

Before we dive in the tldr of this blog post would be:

My experiences training the Asymmetric Actor Critic (using DDPG and HER) algorithm. Exploring distillation as an alternate route to Image based HER.

The codebase which implements and runs all of these experiments is given here. It implements Asymmetric Actor Critic, Hindsight Experience Replay, Domain Randomization and Student/Teacher distillation.

The Tasks

Fetch Reach

Fetch Pick and Place

Fetch Push

The first task, Fetch Reach involves training a robotic arm to reach to a particular point in 3D space. The second task, Fetch Pick and Place involves training the arm to pick a rectangular object and place it at a particular point in 3D space. The last task, Fetch Push involves pushing the object to a particular point in space.

At each time step, these OpenAI environments offer image based, state based observations, desired goals and achieved goals as part of their step function. This is a sparse reward setting, hence a 1 is rewarded at the end of a successful episode otherwise the agent receives a 0.

The state of these environments range from a 10 dimensional vector to a 15 dimensional vector which represent the various positions of the robot joints and the position of the object and the goal. The goal states are a 3 dimensional vector which indicate the location of the goal in 3D space.

The Algorithm

The algorithm mentioned in the Asymmetric Actor Critic Paper was used to solve this task. The benefits of using this algorithm are given as follows:

While we have access to the full state of the system in a simulation, training a state based policy is not amenable to transfer in the real world. A policy that uses image observations is much easier to transfer to the real world. However, training an image based RL algorithm is very tough, often times leading to the policy failing to converge. The authors of the Asym Actor Critic paper discovered that training an off policy actor critic RL algorithm with a critic trained on states(one that has been trained on states) together with an image based actor is enough for the policy to converge (solve the task). Hence, the name Asymmetric Actor Critic. This policy only requires image observations as input and can hence be transferred to the real world with domain randomization.

The algorithm takes an off the shelf off-policy actor critic algorithm (I used DDPG) and trains the critic who’s input is state observations and where the actor takes image observations as input. Since, it is a sparse reward environment, effective exploration is sorely needed. Hence, the Hindsight Experience Replay (HER) algorithm is also used. HER relabels a rollout performed in the environment by switching the desired goal with the observed goal at the end of the episode. With HER, the policy gets a reward at the end of an episode and retrieves some useful signal from a high percentage of rollouts. However, 20% of rollouts performed in the environment are not relabeled by HER and are kept as is. Another key part of this algorithm is to use domain randomization on the simulator to enable transfer to the real robot. We implement all facets of this algorithm.

Experiments

First, training a DDPG+HER policy on state based version of the environment was performed to understand the performance of the policy using well crafted state vectors describing the state of the environment.

We observe that the FetchReach environment is by far the easiest environment to solve taking the least amount of time. The other 2 environments take significantly more time to train with Fetch Push taking 50 times the samples and Fetch Pick and Place taking around 80 times the samples needed for Fetch Reach. However, state based training is fairly quick taking over 3 hours while using 4 processes running in parallel performing Synchronous SGD to optimize the policy.

Next, we wanted to understand the performance of training an image only based DDPG+HER policy. After training it on Fetch Reach (the easiest of the environments), we were unable to make the policy to converge further reinforcing that symmetric Image based RL is hard.

Comparison between perf of state based, asym actor/critic and sym image based algorithms

Finally, we compare the performance of state based and asymmetric image based algorithms for the Fetch Push and Fetch Pick and Place environments. We observe that performance is comparable between the two.

Fetch Push Asym Actor Critic

I also trained a Fetch Reach policy with Domain Randomization and was able to make it converge. This policy is amenable to be transferred to a real robot.

Fetch Reach Asym Actor Critic + Domain Randomization

An implementation detail about implementing Asymmetric Actor Critic is that it is quite hard to create a purely image based input of a relabeled trajectory for the actor whilst using HER. This is because when we relabel the trajectory, we have to give a reward to an agent at the end of the rollout even when it has not done anything useful. Hence, this would entail moving the goal to the position of the end effector at the end of the robot for the Fetch Reach task. This is quite simple to do. However, for the Fetch Push environments this would entail moving the position of the goal to the position of the box. Hence, we would ideally need to use a rollout where the robot has touched the objects and the objects should not have fallen off the table. Otherwise, naive relabeling doesn't necessarily provide good learning signal as it is entirely possible that the robot arm might get a positive reward even when it doesn’t interact with the box. This is fine as it provides signal that the goal should be at the box’s position. However, the experience in the replay buffer with the robot interacting with the object is infinitesimally small if no human engineering is put into the HER relabeling operation. Additionally, this would require a different set of relabeling requirements for novel environments. In training these algorithms, I observed that pure image based actor training without relabeling tricks doesn’t converge for Fetch Push. To solve for this, we input the 3D achieved goal state vector along with the image based observation and are able to make the policy converge. This raises an opportunity for potential improvement by rethinking the benefits of HER and what could be an easier alternative to get this exploration.

Asym Actor Critic Fetch Push Loss curves

Asym Actor Critic Fetch Pick & Place Loss Curves

Exploring Distillation

One idea is to use a state based policy trained in the simulator to act as a teacher to the student teacher paradigm of model distillation.

In our setting, we use the state based actor as the teacher and train a student policy that minimizes the distance between the outputs of the teacher with the outputs of the student using the mean squared error loss.

This works out well if we sample expert (teacher) samples and the student samples with equal probability in the rollout. The reason for that is that the expert would only go over a small subset of state space hence leading to the student policy to fail when it sees states that it has never seen in training. I was able to achieve a policy that learns the most important parts of the tasks but might lack in the preciseness of a RL trained policy. It is useful to note that we have essentially boiled the image based policy learning to a supervised learning problem to a large degree in a sense (it would be fully supervised (i.i.d) if we only sampled expert actions). I think one reason as to why the distilled policy learns faster as it is trying to predict a fixed target instead of a moving/non stationary target which is the case with most RL algorithms. Here is a video of a distilled agent performing the Fetch Push task:

Distilled Fetch Push Student

Here is a plot comparing the performance of FetchPush among the 3 approaches used

Fetch Push Comparison

Fetch Pick And Place

As we can see, it is really astounding how fast a behavior cloned/distilled policy trains using image input. It proves to us that building better image features is not the bottleneck for these environments but good high quality exploration is a much more effective way to increase sample efficiency. The limitations of the distillation method from what I have observed are as follows:

The distilled policy albeit learning quickly doesn’t approach the asymptotic performance of a state based RL policy. Methods that cleverly combine the benefits of both approaches could help alleviate this problem. The distilled policy assumes that we have access to states in a simulation environment and the image based policy is bounded by the performance of a state based policy.

Observations

I noticed that state based RL policies achieve a higher reward than image based based policies- the difference is minute only around a 0.1 difference in mean reward and I think it is because the images are blurry and the robot arrives at the vague position of the goal but doesn’t go to the exact center at times which decreases it’s score. The state based policies also train faster than image based policies but the difference is not huge. The image based policy might take around 5–10 epochs more (40,000 steps — 80,000 steps) to train for environments like FetchPush and FetchPickAndPlace. This difference is even smaller in Fetch Reach(8,000–16,000) steps. I believe that this is the amount of samples needed to learn a representation comparable to the state based one.

Future Directions and Limitations

In my time training these models, I have a few observations that could lead to potential gains in performance/sample efficiency: