lib/cartpole.viv

#| The cart-pole of Barto, Sutton and Anderson (1983): a cart of mass CART_M on a frictionless
#| track, a pole of mass POLE_M and half-length POLE_L hinged on it. The state is
#| [x, xdot, theta, thetadot], with theta = 0 the pole upright and theta > 0 leaning to the right.
#| Use it with `import cartpole`: a controller returns a force, cartpole_step advances the
#| state, cartpole_draw puts the cart, the pole and its bob on the stage.

CART_M = 1.0
POLE_M = 0.1
POLE_L = 0.5
GRAVITY = 9.81

#| The accelerations [xdd, thetadd] for a state and a horizontal force on the cart.
def cartpole_accel(s, force, gravity = GRAVITY) {
    th = s[2]
    thd = s[3]
    total = CART_M + POLE_M
    sn = sin(th)
    cs = cos(th)
    tmp = (force + POLE_M * POLE_L * thd * thd * sn) / total
    thdd = (gravity * sn - cs * tmp) / (POLE_L * (4.0 / 3.0 - POLE_M * cs * cs / total))
    xdd = tmp - POLE_M * POLE_L * thdd * cs / total
    return [xdd, thdd]
}

#| The time derivative of the state.
def cartpole_rate(s, force) {
    a = cartpole_accel(s, force)
    return [s[1], a[0], s[3], a[1]]
}

#| One step of dt by fourth-order Runge-Kutta, the force held for the step.
def cartpole_step(s, force, dt) {
    k1 = cartpole_rate(s, force)
    k2 = cartpole_rate(s + 0.5 * dt * k1, force)
    k3 = cartpole_rate(s + 0.5 * dt * k2, force)
    k4 = cartpole_rate(s + dt * k3, force)
    return s + (dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4)
}

#| The mechanical energy, with the pole's rotational inertia (1/3) m L^2 about the hinge.
def cartpole_energy(s) {
    xd = s[1]
    th = s[2]
    thd = s[3]
    bob_vx = xd + POLE_L * thd * cos(th)
    bob_vy = -POLE_L * thd * sin(th)
    kinetic = 0.5 * CART_M * xd * xd + 0.5 * POLE_M * (bob_vx * bob_vx + bob_vy * bob_vy) + 0.5 * (POLE_M * POLE_L * POLE_L / 3.0) * thd * thd
    return kinetic + POLE_M * GRAVITY * POLE_L * cos(th)
}

#| Draw the cart, the pole and its bob for a state, on a stage whose window is in meters.
def cartpole_draw(s) {
    x = s[0]
    th = s[2]
    tip_x = x + 2.0 * POLE_L * sin(th)
    tip_y = 0.12 + 2.0 * POLE_L * cos(th)
    rect("cart", x, 0.06, 0.36, 0.14, 0.0)
    segment("pole", x, 0.12, tip_x, tip_y)
    circle("bob", tip_x, tip_y, 0.05)
}

test "with no force the cart-pole conserves energy: a free swing from 60 degrees drifts by less than a part in a million" {
    s = [0.0, 0.0, 1.0471975511965976, 0.0]
    e0 = cartpole_energy(s)
    steps = 400.0
    for k to steps {
        s = cartpole_step(s, 0.0, 0.005)
    }
    assert_near(cartpole_energy(s), e0, 1.0e-6 * abs(e0))
}
test "hanging straight down it stays there; balanced upright it falls" {
    down = [0.0, 0.0, 3.141592653589793, 0.0]
    steps = 200.0
    for k to steps {
        down = cartpole_step(down, 0.0, 0.01)
    }
    assert_near(down[2], 3.141592653589793, 1.0e-9)
    up = [0.0, 0.0, 0.01, 0.0]
    for k to steps {
        up = cartpole_step(up, 0.0, 0.01)
    }
    assert_gt(abs(up[2]), 0.5)
}

# SPDX-License-Identifier: MIT OR Apache-2.0
# SPDX-FileCopyrightText: 2026 The Viv authors