A sorting job: learn which skill to use when

The quadruped with its OpenMANIPULATOR-X arm gets a real job: put the four cubes on its back, open the door, walk through, close the door behind it, walk to the table and sort the cubes into the boxes of their colour. The robot already can do every part. With Q-learning it learns the order, by trial and error, in about two hundred attempts.

1The job

Open the scene yourself in the quadruped simulator, lane 11: four cubes on the floor, the door with its handle, and behind it the table with four empty sorting boxes.

2Skills and decisions

In tutorials 1–3 a network learned motor control: steering and throttle, ten times a second. A job of two minutes and twenty-odd steps is far too long for that: a robot that has to discover by random twitching that picking up four cubes, opening a door and sorting is rewarded would never get there.

So we split it, as robots (and people) do. The skills are programmed (or learned separately): pick the nearest cube, open the door, walk to the table. The learner only makes the decisions: which skill now? That is task-level or hierarchical reinforcement learning; the skills are also called options or macro-actions.

3What you need

The robotSimulatorAPI repository, Node.js 22+, Chrome and Python 3.10+, as in tutorial 1. No extra Python packages: the Q-table is a plain dictionary. The script is python/sort_rl.py.

node server.mjs                       # terminal 1: the simulators
python python/sort_rl.py baseline   # the right order, by hand
python python/sort_rl.py train      # Q-learning, 400 attempts (~3 min), saves sort_q.json
python python/sort_rl.py eval       # the learned policy, skill by skill

To watch it, start the server with a visible browser and use watch:

node server.mjs --show
python python/sort_rl.py watch --speed 2

A trained Q-table is included: add --model python/models/sort_q.json to eval or watch.

4The skills

Ten skills, built into the quadruped page. Each one first checks whether it can start; if not, it fails at once (and the learner gets a small penalty). You can also call them from your own code: {"action": {"skill": "pick"}}.

These are the same movements as the buttons on the page (Pick, Place on the back, Open and walk through, …), cut into pieces small enough to make decisions between.

5State and Q-learning

The learner doesn't see images or joint angles, only eight small numbers (extra.task in the observation):

state = (lying, zone, door open, door latched, holding a cube,
         cubes on the floor, cubes on the back, cubes sorted)
# zone: start, at_door, doorway, past_door, table
# the start: (0, start, 0, 1, 0, 4, 0, 0)

With so few possible states we don't need a neural network: a Q-table remembers, for every state it has seen, how good each of the ten skills is there, its Q-value: the reward it expects to collect from here on if it chooses that skill. After every skill it updates one number:

Q[s][a] += α · ( r + γ · max(Q[s′]) − Q[s][a] )
# r = reward of this skill, s′ = the state after it, α = 0.3, γ = 0.97

In words: what I got now, plus the best I expect from the new state, compared with what I expected. Choosing is ε-greedy: with probability ε a random skill (exploring), otherwise the skill with the highest Q-value. ε starts at 1 (only random) and drops to 0.05 over the first 60 % of training.

6The reward

A reward only for the finished job would be fair, but the learner would almost never stumble on it (see the trap). So it also gets something for progress:

def progress(task):        # how far the job is
    past = task["zone"] in ("past_door", "table")
    return (4 - floor) * 1 + sorted * 2 + past * 2 + (past and latched) * 2 + at_table * 0.5

r = -0.1                                     # every skill costs time
r -= 0.3 if the skill failed
r += γ · progress(after) − progress(before)  # progress made (0 for after, when the attempt ends)
r += 20 if the job is done                   # all cubes sorted, past the door, door latched

An attempt ends when the job is done, when it has become impossible (past the door with cubes left behind: there is no skill to go back), or after 40 skills.

7Train it

python python/sort_rl.py train
episode   10  ε 0.96  return  -11.2  last 20: mean return  -12.0, job done 0/10  states 16  2 s
…
episode  150  ε 0.38  return    4.9  last 20: mean return    1.2, job done 9/20  states 75  54 s
episode  200  ε 0.17  return   10.5  last 20: mean return    9.9, job done 19/20  states 75  84 s
episode  250  ε 0.05  return   13.6  last 20: mean return   13.2, job done 20/20  states 75  114 s
Saved sort_q.json (79 states) after 204 s of training

Three training runs (--seed 1, 2, 3):

The first 100 attempts look hopeless: mostly random skills that fail. Then the finished jobs come quickly. Q-learning passes the reward back one skill at a time, so the end of the job is learned first and the beginning last.

8The shaping trap

Our first version rewarded progress without the “0 when the attempt ends” rule. It learned this:

  walk_to_door   ok    floor 4 back 0 sorted 0  door latched  (at_door)
  open_door      ok    floor 4 back 0 sorted 0  door open  (at_door)
  walk_through   ok    floor 4 back 0 sorted 0  door open  (past_door)
Learned policy: NOT finished in 3 skills

Straight through the door, cubes left behind. Going through was worth +2 at once, and because the attempt then ends (hopeless), it never had to pay for the cubes it left. That is reward hacking: the robot does exactly what the reward says, not what we meant.

The fix is potential-based shaping (Ng, Harada & Russell, 1999): reward the change of a progress number and count the progress as 0 at the end. Over a whole attempt the shaping then adds up to zero; it only shows the way. Running the trap again yourself is one option away:

python python/sort_rl.py train --shaping naive --model naive.json
python python/sort_rl.py train --shaping none  --model none.json

9What it learned

The Q-table is a JSON file you can read. The best skills (highest Q) in a few states:

Notice that a skill that fails still has a fairly high Q-value: it costs only −0.4 and the robot is still in the same good state afterwards. The difference between the skills is what matters.

And one surprise: with the fourth cube in its gripper, the learned policy stands up first and puts the cube on its back while standing. That saves a little time (107 s instead of 109 s), because the arm moves while the legs stand up. Nobody told it that; it found the shortcut by trying.

10Experiments

11Troubleshooting