CartPole: Learn Swing-Up and Balance
Start with CartPole before Canyon Run. The cart applies one horizontal force to swing a hanging pole up and then balance it. The physics, learner, random-number state and episode resets all run in Modelica. The viewer displays their results; it does not select weights.
Train in the Editor
Press Run below. The initial pole angle is π − 0.01 radians: 0.57° off straight down, with the cart and pole at rest. This small asymmetry is explicit; exactly symmetric down-start performance is not claimed. All seven policy gains begin at zero.
The example starts in Realtime. Choose Fast to learn sooner and switch back to watch the learned policy. It freezes its gains after five simulated seconds of balance. The cart centre must stay within ±2 m. Attempting to leave a rail ends and penalizes that episode, then resets the plant.
// rumoca-live-scenario: ../repo-examples/reinforcement_learning/cartpole/rumoca-scenario.toml
Drag the 3D view to rotate; right-drag to pan; scroll or pinch to zoom. Reset view restores the camera. R restarts training. After learning, Bump left/right requests an actual ±20 N force for 0.3 simulated seconds. A disturbance can still cause failure: the plant resets, retaining the frozen gains. The model bounds the combined policy and bump force to ±30 N.
The diagrams show actual learned gains and observed episode rewards. Fast mode can skip intermediate episodes in the displayed plot; training still evaluates each episode. Realtime pacing and simulation time are separate from rendering.
A Small Policy with an Explicit Physics Prior
A single tanh neuron combines seven clipped features. Three pump the pole’s energy while it hangs down; four emphasize balance near the upright position. The features depend on cart position/velocity, sin/cos of the pole angle and angular velocity. This physical feature map is fixed and documented in the source. All seven gains are learned; no stabilizing gains are preset. This is a small policy-search example, not an end-to-end deep actor.
The learner retains the two best policies across episodes and refits a diagonal Gaussian every 32 evaluations, with a minimum standard deviation of 1. Each candidate gets at most 20 simulated seconds. Its reward favors pole height and sustained balance, penalizes energy error, cart displacement and force, and subtracts 50 on a rail failure. Success requires cos(angle) > 0.98, |angular velocity| < 1 rad/s and |cart position| < 1.2 m for 250 control ticks. The control period is 0.02 s. A one-sample observation/weight delay and the one-candidate distribution-refit delay are explicit in the source.
The shared NeuralNetwork package supplies the dense layer,
tanh, clipping and Gaussian samples. PolicySearch.mo owns candidate sampling,
elite retention and refitting; CartPoleSwingUp.mo owns the task and force.
One checked native seed-1 run evaluated 84 policies before freezing. A separate 30-second run from the same near-down start, with learning disabled, held the balance criterion from 1.92–30.00 s: 28.08 s. Maximum absolute cart position was 1.498 m, with no rail contacts or resets and unchanged weights. Six separate boundary controls checked both signs, inward noncontacts and outward starts at the limits. These are nominal development results, not held-out, multiple-seed, disturbance or Gymnasium benchmark acceptance.
Native equivalent, from the repository root:
rumoca sim examples/reinforcement_learning/cartpole/CartPoleSwingUp.mo \
--source-root examples/reinforcement_learning/NeuralNetwork.mo \
--source-root examples/reinforcement_learning/cartpole/PolicySearch.mo \
--source-root examples/reinforcement_learning/cartpole/CartPole.mo \
--model CartPoleSwingUp --solver rk-like
Read the Modelica policy-search implementation
within;
package PolicySearch "Diagonal cross-entropy search over neural policy weights"
function candidate "Draw one policy; the caller owns the random state"
input Real mean[:];
input Real deviation[size(mean,1)];
input Integer seeds[size(mean,1)];
output Real weights[size(mean,1)];
output Integer nextSeeds[size(mean,1)];
protected
Real noise;
algorithm
for i in 1:size(mean,1) loop
(noise,nextSeeds[i]) := NeuralNetwork.normal(seeds[i]);
weights[i] := mean[i]+deviation[i]*noise;
end for;
end candidate;
function retain "Replace only the worst elite when this episode is better"
input Real elite[:, :];
input Real scores[size(elite,1)];
input Real weights[size(elite,2)];
input Real score;
output Real updated[size(elite,1),size(elite,2)];
output Real updatedScores[size(elite,1)];
protected
Integer worst;
algorithm
updated := elite;
updatedScores := scores;
worst := 1;
for i in 2:size(scores,1) loop
if scores[i]<scores[worst] then
worst := i;
end if;
end for;
if score>scores[worst] then
for j in 1:size(weights,1) loop
updated[worst,j] := weights[j];
end for;
updatedScores[worst] := score;
end if;
end retain;
function fit "Refit each weight's Gaussian from the retained elite policies"
input Real elite[:, :];
output Real mean[size(elite,2)];
output Real deviation[size(elite,2)];
protected
Real variance;
algorithm
for j in 1:size(elite,2) loop
mean[j] := 0;
for i in 1:size(elite,1) loop
mean[j] := mean[j]+elite[i,j]/size(elite,1);
end for;
variance := 0;
for i in 1:size(elite,1) loop
variance := variance+(elite[i,j]-mean[j])^2/size(elite,1);
end for;
deviation[j] := max(1.0,sqrt(variance));
end for;
end fit;
model Learner "Episode-level search, separate from the fast plant/controller clock"
parameter Integer randomSeeds[7];
input Boolean ended;
input Boolean reached;
input Real score;
output Boolean trained(start=false,fixed=true);
output Integer evaluations(start=0,fixed=true);
output Integer generation(start=1,fixed=true);
output Real weights[7](each start=0,each fixed=true);
protected
Integer rng[7](start=randomSeeds,each fixed=true);
Real mean[7](each start=0,each fixed=true);
Real deviation[7](each start=1,each fixed=true);
Real elite[2,7](each start=0,each fixed=true);
Real eliteScores[2](each start=-1e9,each fixed=true);
Real retained[2,7](each start=0,each fixed=true);
Real scores[2](each start=-1e9,each fixed=true);
algorithm
when ended then
trained := pre(trained) or reached;
if not pre(trained) then
evaluations := pre(evaluations)+1;
(retained,scores) := retain(pre(elite),pre(eliteScores),pre(weights),score);
if mod(evaluations,32)==0 then
(mean,deviation) := fit(retained);
generation := pre(generation)+1;
else
mean := pre(mean);
deviation := pre(deviation);
generation := pre(generation);
end if;
elite := retained;
eliteScores := scores;
if trained then
weights := pre(weights);
rng := pre(rng);
else
// The first candidate after a refit uses the prior distribution;
// this explicit event delay keeps search updates separate from sampling.
(weights,rng) := candidate(pre(mean),pre(deviation),pre(rng));
end if;
else
evaluations := pre(evaluations);
generation := pre(generation);
retained := pre(retained);
scores := pre(scores);
mean := pre(mean);
deviation := pre(deviation);
elite := pre(elite);
eliteScores := pre(eliteScores);
weights := pre(weights);
rng := pre(rng);
end if;
end when;
end Learner;
end PolicySearch;
Compare the Python Upright-Balance Route
This older comparison teaches upright balancing from a small initial tilt; it does not perform the swing-up task above. It uses a four-gain policy and a separate cross-entropy search. The checked replay below belongs to that separate run; it is not presented as new training.
Loading recorded native result…
This animation replays the checked native evaluation with learning disabled. The chart shows the actual policy-search history. Use the command below to train a fresh policy.
Run the Python Learner
From the repository root, with a native rumoca executable on your path:
python3 examples/reinforcement_learning/cartpole/train.py \
--rumoca rumoca --output cartpole-run
If you built the CLI locally, use --rumoca target/release/rumoca instead.
The output directory must be new. The default run uses seed 1, eight
generations of 16 candidate policies, and a 180-second wall-time limit.
The program first measures a zero-force baseline, then prints the best episode
length and population mean after each generation. It saves the learned gains,
training history and a final evaluation in cartpole-run/result.json.
cartpole-run/flight.csv records that final flight’s states and force, with
learning disabled.
In a checked native run, the zero-force policy lasted 34 steps. The best policy reached the 500-step limit in generation 2 and again in its final frozen evaluation. The population mean rose from 44.5 steps in generation 1 to 479.9 in generation 8. The run took about 9 seconds on one core; elapsed time depends on the machine and compiler build.
What Is Learned?
The policy is a linear combination of the four normalized states, passed
through tanh to bound the force to ±10 N:
scales = [2.4, 2.0, math.radians(12), 2.0]
force = 10 * math.tanh(sum(w*x/s for w, x, s in zip(weights, state, scales)))
The four weights start unknown. The learner uses cross-entropy policy
search:
- Sample 16 sets of gains from a Gaussian distribution.
- Run one episode for each set and score it by how many steps it lasts.
- Keep the best quarter and fit the next Gaussian to their gains.
- Repeat, then evaluate the best policy without changing its gains.
This is reinforcement learning by direct policy search: episode reward selects the gains. It does not use a hand-tuned balancing controller, a value network, physics derivatives or backpropagation. It is smaller than the actor-critic methods used for more complex continuous-control tasks.
Each control step advances the Modelica plant by 0.02 s. An episode ends if the cart leaves ±2.4 m or the pole tilts beyond ±12 degrees; otherwise it is truncated after 500 steps, or 10 s. Every step, including a terminating step, contributes one point. The old episode is scored before resetting the simulator. The policy search retains its gains and random-number state across resets.
Physics and Scope
The model is an upright, frictionless cart with a hinged rigid pole. Its parameters are a 1 kg cart, a 0.1 kg pole, a 0.5 m pivot-to-centre-of-mass distance and gravity of 9.8 m/s². The cart starts at rest, with a 0.03 rad pole tilt. Rumoca integrates the equations; the Python program supplies force and reads the four states.
The equations and failure limits follow the familiar CartPole problem. This example uses a continuous bounded force, a fixed initial state and Rumoca’s solver, so its scores are not Gymnasium CartPole benchmark results. Balancing from this one start also does not prove robustness to disturbances.
Try changing initialAngle in the model, increasing the training population,
or evaluating the saved gains from a different starting state. Keep those
evaluation results separate from the episodes used to select the gains.
Read the complete Python learner
"""Learn four CartPole feedback gains from episode reward, using only stdlib."""
import argparse
import csv
import json
import math
import os
from pathlib import Path
import random
import selectors
import shutil
import subprocess
import tempfile
import time
class Session:
def __init__(self, binary, scratch, deadline):
self.deadline = deadline
self.time = 0.0
self.buffer = b''
self.log = (scratch / 'stderr.log').open('wb')
self.child = subprocess.Popen([binary, 'sim', str(Path(__file__).with_name('CartPole.mo')),
'-m', 'CartPole', '--solver', 'rk-like', '--serve-stdio',
'--cache-dir', str(scratch / 'cache')], stdin=subprocess.PIPE,
stdout=subprocess.PIPE, stderr=self.log)
self.selector = selectors.DefaultSelector()
self.selector.register(self.child.stdout, selectors.EVENT_READ)
try:
hello = self.read()
if hello.get('event') != 'hello' or hello.get('protocol_version') != 1:
raise ValueError('unsupported session handshake')
except Exception:
self.close()
raise
def read(self):
while b'\n' not in self.buffer:
remaining = self.deadline - time.monotonic()
if remaining <= 0 or not self.selector.select(remaining):
raise TimeoutError('CartPole session deadline')
data = os.read(self.child.stdout.fileno(), 65536)
if not data:
raise RuntimeError('session ended; inspect stderr.log')
self.buffer += data
if len(self.buffer) > 1024*1024:
raise ValueError('oversized session response')
line, self.buffer = self.buffer.split(b'\n', 1)
return json.loads(line)
def call(self, command, expected='ok'):
if time.monotonic() >= self.deadline:
raise TimeoutError('CartPole session deadline')
self.child.stdin.write((json.dumps(command) + '\n').encode())
self.child.stdin.flush()
event = self.read()
observed = event.get('time')
if event.get('event') != expected:
raise ValueError('unexpected session response: ' + str(event))
if type(observed) not in (float, int) or not math.isfinite(observed):
raise ValueError('invalid response time')
return event
def state(self):
event = self.call({'command': 'state'}, 'state')
if abs(event['time'] - self.time) > 1e-9:
raise ValueError('stale observation')
values = [event['values'][name] for name in ['x', 'v', 'theta', 'omega']]
if not all(type(v) in (int, float) and math.isfinite(v) for v in values):
raise ValueError('invalid physical state')
return values
def set_force(self, force):
if not math.isfinite(force) or abs(force) > 10:
raise ValueError('force outside policy bounds')
event = self.call({'command': 'set_inputs', 'inputs': [['force', force]]})
if abs(event['time'] - self.time) > 1e-9:
raise ValueError('input update advanced time')
def reset(self):
self.set_force(0.0)
event = self.call({'command': 'reset'})
if event['time'] != 0:
raise ValueError('reset did not return time zero')
self.time = 0.0
return self.state()
def step(self, force):
self.set_force(force)
target = self.time + .02
event = self.call({'command': 'advance_to', 'time': target})
if abs(event['time'] - target) > 1e-9:
raise ValueError('step time mismatch')
self.time = event['time']
return self.state()
def close(self):
try:
if self.child.poll() is None:
try:
self.child.stdin.write(b'{"command":"close"}\n')
self.child.stdin.flush()
self.child.wait(timeout=3)
except (BrokenPipeError, subprocess.TimeoutExpired):
pass
finally:
if self.child.poll() is None:
self.child.kill()
self.child.wait(timeout=3)
self.selector.close()
self.log.close()
return self.child.returncode
def force_for(weights, state):
scales = [2.4, 2.0, math.radians(12), 2.0]
return 10 * math.tanh(sum(w*x/s for w, x, s in zip(weights, state, scales)))
def episode(session, weights, trace=None):
state = session.reset()
for step in range(1, 501):
force = force_for(weights, state)
state = session.step(force)
failed = abs(state[0]) > 2.4 or abs(state[2]) > math.radians(12)
if trace is not None:
trace.writerow([session.time, *state, force, failed])
if failed:
return step
return 500
def search(session, rng, generations, population):
mean, deviation = [0.0]*4, [1.0]*4
best, best_score, history = mean[:], 0, []
for generation in range(generations):
candidates = [[rng.gauss(m, s) for m, s in zip(mean, deviation)]
for _ in range(population)]
scored = [(episode(session, weights), weights) for weights in candidates]
scored.sort(key=lambda item: item[0], reverse=True)
if scored[0][0] > best_score:
best_score, best = scored[0]
elites = [weights for _, weights in scored[:max(2, population//4)]]
mean = [sum(w[j] for w in elites)/len(elites) for j in range(4)]
deviation = [max(.1, math.sqrt(sum((w[j]-mean[j])**2 for w in elites)/len(elites)))
for j in range(4)]
row = {'generation': generation+1, 'best_steps': best_score,
'population_mean_steps': sum(score for score, _ in scored)/population}
history.append(row)
print(json.dumps(row), flush=True)
return best, history
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--rumoca', default='rumoca')
parser.add_argument('--output', type=Path, default=Path('cartpole-run'))
parser.add_argument('--generations', type=int, default=8)
parser.add_argument('--population', type=int, default=16)
parser.add_argument('--seed', type=int, default=1)
parser.add_argument('--wall-seconds', type=int, default=180)
args = parser.parse_args()
if args.population < 4 or min(args.generations, args.wall_seconds) < 1:
parser.error('positive generations/wall budget and population >=4 required')
args.output.mkdir(exist_ok=False)
result = {'completed': False, 'seed': args.seed, 'method': 'cross-entropy policy search',
'scope': 'fixed-start instructional CartPole, not Gymnasium benchmark acceptance'}
started = time.monotonic()
with tempfile.TemporaryDirectory(prefix='cartpole-') as scratch_name:
scratch, session = Path(scratch_name), None
try:
session = Session(args.rumoca, scratch, started + args.wall_seconds)
result['zero_policy_steps'] = episode(session, [0.0]*4)
weights, history = search(session, random.Random(args.seed), args.generations, args.population)
with (args.output/'flight.csv').open('w', newline='') as handle:
writer = csv.writer(handle)
writer.writerow(['time', 'x', 'v', 'theta', 'omega', 'force', 'failed'])
result['frozen_steps'] = episode(session, weights, writer)
result.update(completed=True, weights=weights, training=history)
except Exception as error:
result['error'] = repr(error)
raise
finally:
if session is not None:
result['child_exit_code'] = session.close()
if (scratch/'stderr.log').exists():
shutil.copy2(scratch/'stderr.log', args.output/'stderr.log')
result['elapsed_seconds'] = time.monotonic() - started
(args.output/'result.json').write_text(json.dumps(result, indent=2)+'\n')
print(json.dumps(result), flush=True)
if __name__ == '__main__':
main()