training locomotion policies for a quadruped

Legged locomotion is one of the most researched problems in the physical AI because balancing legged robots and creating generalizable gaits is a much more complex task compared to developing controllers for wheeled robots which makes it harder to use traditional control methods. This is why I wanted to try and train a Reinforcement Policy that enabled a quadruped to walk in challenging terrains.

I was really inspired by a paper I read earlier in the year:

Dynamics and Domain Randomized Gait Modulation with Bezier Curves for Sim-to-Real Legged Locomotion by Rahme Et. Al.

Building off of this I decided to use the URDF from an open source robot, Open Quadruped, designed by Adham Elewaraby.

Overview of the Controller

End-to-End ML systems have been gaining a lot of traction in the field but I decided to go down a different path because I wanted to mix traditional control and learned policies in order to develop the gait. I developed a baseline controller using a trajectory generator into an inverse kinematics module. I then wrapped this in a policy that adjusts the parameters of the trajectory generator to handle more varied terrain in realtime as the controller loop is running.

Baseline Controller

This controller is split into 2 parts: the trajectory generator and the inverse kinematics

Bezier Curve Based Trajectory Generator:
Before I get into the trajectory generator itself, we need to understand walking gaits for quadrupeds. The “gait” is just the movement of the legs that causes the robot to move in whatever direction desired. Generally there is 2 phases to the gait, the swing phase and the support phase. The swing phase is the initial forward movement that occurs in the air and moves the leg forward. The support phase is technically a “backward” movement but is done into the ground so that the leg doesn't move at all and brings the rest of the body back into a neutral position.

As you can see above, after this cycle of Swing Phase into the Support Phase the leg moves forward and returns to its initial configuration. This is the backbone to all gaits in quadruped locomotion. There are many ways to represent these two curves but I decided on doing it with quadratic bezier curves. The gist of bezier curves is essentially the process of stacking linear interpolations on top of each other to form more complex curves. A linear interpolation is the following equation:

$$P(t) = (1 - t)P_0 + tP_1, \quad t \in [0,1]$$

This represents a straight line from P0 to P1. If we want to extend this into a quadratic bezier curve we can essentially nest 2 of these interpolations between 3 points to get a curve between 2 points. That looks like this:

$$\begin{aligned}a(t) &= (1 - t)P_0 + tP_1 \\b(t) &= (1 - t)P_1 + tP_2 \\P(t) &= (1 - t)a(t) + tb(t)\end{aligned}$$

This maps out a curve that looks like this:

That equation for P(t) can be simplified to:

$$P(t) = P_0 + t(-2P_0 + 2P_1) + t^2(P_0 - 2P_1 + P_2)$$

In matrix form:

$$P(t) =\begin{bmatrix}1 & t & t^2\end{bmatrix}\begin{bmatrix}1 & 0 & 0 \\-2 & 2 & 0 \\1 & -2 & 1\end{bmatrix}\begin{bmatrix}P_0 \\P_1 \\P_2\end{bmatrix}$$

This nesting of linear interpolations can be used to create any nth degree bezier curve but the calculation becomes more complex. You can also stack multiple cubic bezier curves together sequentially to form splines but I decided to keep my gait quite simple so I just used quadratic curves.

Now, how do we use these curves for the gaits? We can use these curves for gaits by parameterizing the curves based on some physical parameters:

where $\psi$ is the clearance height, $\delta$ is the ground penetration, and $L_{span}$ is the length of the stride

This means we can just use the matrix form from earlier to represent the swing phase with these 3 (x, z) points:

$$P_0 = (0,0) \quad P_1 = \left(\frac{L_{span}}{2}, \psi\right) \quad P_2 = (L_{span}, 0)$$

and the support phase with these 3 (x, z) points:

$$P_0 = (L_{span}, 0) \quad P_1 = \left(\frac{L_{span}}{2}, -\delta\right) \quad P_2 = (0,0)$$

Now that we have defined the P(t) for both the swing phase and the support phase we can sample points from them based on the “phase” of each leg. We want to parameterize the entire leg trajectory between 0 and 1 so that means 0 to 0.5 would be the swing and then 0.5 to 1 would be the support phase. But each curve is parameterized from 0 to 1 so to normalize the true leg phase to the phase within each curve it looks like this:

$$u =\begin{cases}2\mathrm{leg}_{phase}, & \mathrm{leg}_{phase} < 0.5 \ \mathrm{(swing)} \\2(\mathrm{leg}_{phase} - 0.5), & \mathrm{leg}_{phase} \geq 0.5 \ \mathrm{(support)}\end{cases}$$

or

$$\begin{aligned}u_{swing} &= 2leg_{phase}, \ \mathrm{for}\ leg_{phase} \in [0,0.5) \\u_{support} &= 2(leg_{phase} - 0.5), \ \mathrm{for}\ leg_{phase} \in [0.5,1)\end{aligned}$$

where $u$ is the input into the curves.

Now so far these curves have only been in the x-z plane and only really support forward movement. We have represented P(t) with (x, z) points which really were (x, 0, z) points but we can extend P(t) to include lateral movement with a quick change:

$$\begin{aligned}&\text{Let the original point from the curve be } P_b(t) = (x_b, z_b). \\&\text{If we want to move with the translation angle } \rho: \\&P(t) = (x_b\cos(\rho), x_b\sin(\rho), z_b).\end{aligned}$$

Yaw control can also be added by defining a new bezier curve for this that rotates the robot in the desired direction. This is done by taking in an angular velocity parameter and the desired gait period and then defining this new P_yaw(t) and adding it to the original P(t). This is first done by defining a dx and dy for the leg:

$$\begin{aligned}\theta &= T_{gait} \cdot \omega \\d_x &= x\cos\theta - y\sin\theta - x \\d_y &= x\sin\theta + y\cos\theta - y\end{aligned}$$

x and y in this equation are based on the neutral positions for each leg which are defined as this, where L and W are the length and width of the robot:

$$\begin{aligned}\mathrm{front\text{-}left}: &\quad \left(\frac{L}{2}, \frac{W}{2}\right)&\qquad\mathrm{front\text{-}right}: &\quad \left(\frac{L}{2}, -\frac{W}{2}\right) \\\mathrm{back\text{-}left}: &\quad \left(-\frac{L}{2}, \frac{W}{2}\right)&\qquad\mathrm{back\text{-}right}: &\quad \left(-\frac{L}{2}, -\frac{W}{2}\right)\end{aligned}$$

The new yaw swing phase is defined by these 3 points:

$$Y_0 = (0,0,0), \quad Y_1 = \left(\frac{dx}{2}, \frac{dy}{2}, \psi\right), \quad Y_2 = (dx,dy,0)$$

And the support phase by these 3 points:

$$Y_0 = (dx,dy,0), \quad Y_1 = \left(\frac{dx}{2}, \frac{dy}{2}, -\delta\right), \quad Y_2 = (0,0,0)$$

Now the new trajectory curve is defined by:

$$P(t) = P(t) + Y(t)$$

This is the basis of the bezier curve trajectory generator as this allows us to sample from the trajectory and get points along the path from the initial position of the gait to the final position of the gait with reference to the initial position.

Inverse Kinematics Module:
Because I control the quadruped via each leg individually, I derived the inverse kinematics for each leg individually with reference from its hip. Here is the full derivation:

Putting it Together
Now that we defined both the trajectory generator and the inverse kinematics we just need to put them together in a loop. Because we are creating a walking gait it mean that not all 4 legs can be off the ground at the time, this mean we need to offset their gaits in diagonal pairs and we do this by just adding a 0.5 offset to one of the pairings. This means while 1 pair is in the swing phase (in the air) the other pair is in the support phase (digging into the ground) and this creates stability during the gait. This means the the leg phase for each leg becomes:

$$leg_{phase} = (global_{phase} + phase_{offset}) \bmod 1$$

Now with this we can generate a set of the positions for each leg of the robot, but these positions are only relative to the starting position of the trajectory which is not how the inverse kinematics works so we cant just plug these values into the IK module we need to transform them by a fixed transform which is the x, y, z offset of the foot to the hip in the neutral stance for each leg. When we add this to the points generated by the trajectory generator it allows the IK module to correctly determine the joint angles needed to accomplish the desired trajectory. Finally the last thing needed to make this controller work effectively was to add PD control that ensures the heading of the robot stays fixed to a desired angle because any mass imbalances can cause undesired torques to affect the robot and this helps to correct for it.

Designing the Environment

Since I was utilizing reinforcement learning the most involved portion of the project was designing the environment in MujoCo.

MujoCo environments are based on XML and thankfully the general physical design of the world in simulation wasn’t too complex since it was really just a terrain that the robot would walk in. All I had to do was convert the robot urdf to a MJCF xml file using the standard converter. I had to make small tweaks to the collision meshes on the feet but other than that it was quite standard.

Domain and Dynamics Randomization:
The most interesting thing I gained from the paper was their usage of dynamics and domain randomization within the environment. Domain in this sense referring to the actual environment the robot is in, and Dynamics is referring to the physics governing the environment, things like friction, body masses, and gravity. The best way I could explain the benefit of this is to think of it like this: There is a fundamental different between the simulated world and the real world so training a policy in a singular simulated world would allow for those differences to really show up when deploying the policy in the real world. By randomizing the domain and dynamics you are essentially training the robot to be adaptable across thousands of potential worlds and this makes it so that bringing the robot to the real world is just 1 more of those worlds its already been able to adapt to.

The way I went about implementing this was in 2 ways:

• Domain: I added a height map to the xml file which essentially took a set of x, y, z values that described the amplitudes of noise to add to the floor height and I randomly sampled values for this according to a gaussian distribution with a mean of 0.
• Dynamics: In a similar fashion to above I randomly sampled values to adjust the friction between the feet and the ground, as well as the body masses of the various links within the robot. I centred the distribution around whatever values were in the original xml and then changed them before the rollout occurred

Observation Space:
The robot is fitted with an IMU so the observation space was defined as a 10 dimensional space with the following quantities which were the same used by the paper:

1. Pitch
2. Roll
3. Yaw
4–6. Linear Accelerations: x, y, z
7–10. Phase of trajectory for each leg: FL, FR, BL, BR

Action Space:
The policy essentially modulate the baseline controller loop so the action space is 14 dimensions for both positional offsets and parameter modulations into the trajectory generator:

1–12. dx, dy, dz for each leg end-effector
13. clearance height parameter
14. ground penetration parameter

Reward Function
This was probably the part I spent the most time trying to get right because it starts off quite arbitrarily and requires tuning to get stable training. At first I tried using the reward structure from the original paper:

$$r_t = \Delta_x - 10(|r| + |p|) - 0.03\sum |\omega|$$

where $\Delta_{x}$ is the forward movement of the robot, $r$ is the roll angle, $p$ is the pitch angle, and $\omega$ is the vector of all angular velocities. Overall this reward formulation incentivized forward movement and then punished any other angular movement.

When I tried this however, it did not work very well and I am assuming this has to do with differences in the physics environment that I used compared to the paper. I had used MujoCo and they had used PyBullet. After a couple more iterations, some researching, and more testing I ended up using this reward function:

$$r_t = r_{forward} + 0.2r_{yaw} + 0.2r_{lateral} + 0.5r_{posture} + 0.1r_{shake} + r_{backward}$$

where

$$\begin{aligned}r_{forward} &= \exp\left(-\frac{(v_x - v_x^*)^2}{0.1}\right) \\r_{yaw} &= -|\omega_z - \omega_z^*| \\r_{lateral} &= -|v_y| \\r_{posture} &= -(|\phi| + |\theta|) \\r_{shake} &= -|a_z| \\r_{backward} &=\begin{cases}-2|v_x|, & \mathrm{if}\ v_x < 0 \\0, & \mathrm{otherwise}\end{cases}\end{aligned}$$

I decided to instead of just plainly rewarding forward movement I wanted to instead incentivize the robot to follow the target velocity and this ended up performing much better. I did the same thing with the yaw rate. The lateral penalty just penalizes any movement to the right because training happened with only forward movement and the prior reward never penalized this. The posture penalty penalizes any pitch or roll in the robot similar to the original reward function. The shake penalizes any vertical shake which incentivizes stability when walking across varied terrain. Finally I added the backwards penalty because there were weird glitches in the physics of the environment where the robot would jump backwards so I just wanted to heavily penalize this behaviour.

Model Architecture

Because this model would be ran on physical robots I wanted it to be as lightweight as possible while still being powerful which is why I chose a simple fully connected architecture with 10 neuron input layer, 64 hidden units in a single layer, and then the output layer of 14 neurons.

Training

The training for this was quite straightforward because I utilized stable_baselines3 for the training algorithm since I was more interested in the building up the environment than implementing the algorithms from scratch. I used PPO because it such a reliable algorithm and is sample efficient. The last detail about training that would also help to explain the reward function above is that the policy was trained on a fixed gait. By that I mean the parameters of the gait like how fast it was, the fact it was only forward movement with no yaw, and the stride length were all kept fixed throughout the training because the policy generalizes well to different gaits when deployed in the real world according to the original paper.

Results

Here are 2 videos showing the robot in the varied terrain conditions before and after wrapping with a policy

Future Work

First and foremost I would love to deploy this policy on the physical quadruped which is something I am currently working on but the electrical for it is proving to be much harder than expected.

Beyond deploying on a physical robot however, there are a couple features I believe I could add in the future to improve the robustness of the policy or make the controller more effective:

• More gait types: gallops (4 legs in the air) , bounds (back together and front together), pace (side pairs together)
• More complex and accurate terrain types: stairs, rocks and real obstacles, ramps
• More complex trajectories other than simple quadratic curves like combining cubic curves to form complex splines