Two-Wheel Balancing Robot Kinematics and Control

1        ┌─────┐
2        │ Body│  ← Center of mass (height h)
3        │  •  │
4        └──┬──┘
5           │
6      ─────┴─────  ← Wheel axis
7      ●         ●  ← Wheels (radius r)

  1import numpy as np
  2from scipy.integrate import odeint
  3
  4class BalancingRobot:
  5    def __init__(self, m_body=1.0, m_wheel=0.1, L=0.3, r=0.05,
  6                 I_body=0.05, I_wheel=0.001, wheelbase=0.2):
  7        """
  8        Args:
  9            m_body: Mass of body (kg)
 10            m_wheel: Mass of each wheel (kg)
 11            L: Distance to center of mass (m)
 12            r: Wheel radius (m)
 13            I_body: Body moment of inertia (kg⋅m²)
 14            I_wheel: Wheel moment of inertia (kg⋅m²)
 15            wheelbase: Distance between wheels (m)
 16        """
 17        self.m_b = m_body
 18        self.m_w = m_wheel
 19        self.L = L
 20        self.r = r
 21        self.I_b = I_body
 22        self.I_w = I_wheel
 23        self.wheelbase = wheelbase
 24        self.g = 9.81
 25        
 26        # State: [x, x_dot, theta, theta_dot, psi]
 27        # psi = heading angle
 28        self.state = np.array([0.0, 0.0, 0.0, 0.0, 0.0])
 29    
 30    def dynamics(self, state, t, tau_left, tau_right):
 31        """
 32        Compute state derivatives
 33        
 34        Args:
 35            state: [x, x_dot, theta, theta_dot, psi]
 36            t: Time
 37            tau_left, tau_right: Wheel torques (N⋅m)
 38        
 39        Returns:
 40            State derivatives
 41        """
 42        x, x_dot, theta, theta_dot, psi = state
 43        
 44        # Average torque for forward motion
 45        tau_avg = (tau_left + tau_right) / 2
 46        
 47        # Differential torque for turning
 48        tau_diff = (tau_right - tau_left) / 2
 49        
 50        # Linearized dynamics (valid for small theta)
 51        # Total mass
 52        m_total = self.m_b + 2 * self.m_w
 53        
 54        # Denominator terms
 55        denom1 = self.I_b + self.m_b * self.L**2
 56        denom2 = m_total + 2 * self.I_w / self.r**2
 57        
 58        # Acceleration of body tilt
 59        theta_ddot = (
 60            (m_total * self.g * self.L * np.sin(theta) 
 61             - self.m_b * self.L * np.cos(theta) * 
 62             (tau_avg / self.r - self.m_b * self.L * theta_dot**2 * np.sin(theta)) / denom2)
 63            / (denom1 - self.m_b**2 * self.L**2 * np.cos(theta)**2 / denom2)
 64        )
 65        
 66        # Acceleration of position
 67        x_ddot = (
 68            tau_avg / self.r 
 69            - self.m_b * self.L * (theta_ddot * np.cos(theta) 
 70                                   - theta_dot**2 * np.sin(theta))
 71        ) / denom2
 72        
 73        # Heading rate (differential drive)
 74        psi_dot = tau_diff * self.r / (self.wheelbase * self.I_w)
 75        
 76        return np.array([x_dot, x_ddot, theta_dot, theta_ddot, psi_dot])
 77    
 78    def step(self, tau_left, tau_right, dt):
 79        """
 80        Simulate one time step
 81        
 82        Args:
 83            tau_left, tau_right: Wheel torques
 84            dt: Time step
 85        
 86        Returns:
 87            Updated state
 88        """
 89        t_span = [0, dt]
 90        solution = odeint(self.dynamics, self.state, t_span, 
 91                         args=(tau_left, tau_right))
 92        self.state = solution[-1]
 93        
 94        return self.state.copy()
 95
 96# Example: Free fall (no control)
 97robot = BalancingRobot()
 98robot.state[2] = 0.1  # Initial tilt
 99
100trajectory = [robot.state.copy()]
101dt = 0.01
102
103for _ in range(100):
104    robot.step(tau_left=0, tau_right=0, dt=dt)
105    trajectory.append(robot.state.copy())
106
107trajectory = np.array(trajectory)
108
109import matplotlib.pyplot as plt
110time = np.arange(len(trajectory)) * dt
111
112plt.figure(figsize=(12, 4))
113plt.subplot(131)
114plt.plot(time, trajectory[:, 0])
115plt.xlabel('Time (s)')
116plt.ylabel('Position (m)')
117plt.grid(True)
118
119plt.subplot(132)
120plt.plot(time, np.degrees(trajectory[:, 2]))
121plt.xlabel('Time (s)')
122plt.ylabel('Tilt Angle (deg)')
123plt.grid(True)
124
125plt.subplot(133)
126plt.plot(time, trajectory[:, 1])
127plt.xlabel('Time (s)')
128plt.ylabel('Velocity (m/s)')
129plt.grid(True)
130
131plt.tight_layout()
132plt.show()

  1class PIDController:
  2    def __init__(self, Kp, Ki, Kd, setpoint=0):
  3        self.Kp = Kp
  4        self.Ki = Ki
  5        self.Kd = Kd
  6        self.setpoint = setpoint
  7        
  8        self.integral = 0
  9        self.prev_error = 0
 10    
 11    def update(self, measurement, dt):
 12        """Compute control output"""
 13        error = self.setpoint - measurement
 14        
 15        # Proportional
 16        P = self.Kp * error
 17        
 18        # Integral (with anti-windup)
 19        self.integral += error * dt
 20        self.integral = np.clip(self.integral, -10, 10)
 21        I = self.Ki * self.integral
 22        
 23        # Derivative
 24        D = self.Kd * (error - self.prev_error) / dt
 25        self.prev_error = error
 26        
 27        return P + I + D
 28    
 29    def reset(self):
 30        """Reset controller state"""
 31        self.integral = 0
 32        self.prev_error = 0
 33
 34class BalancingController:
 35    def __init__(self, robot):
 36        self.robot = robot
 37        
 38        # Balance controller (tilt angle)
 39        self.balance_pid = PIDController(
 40            Kp=50.0,   # Proportional gain
 41            Ki=5.0,    # Integral gain
 42            Kd=5.0,    # Derivative gain
 43            setpoint=0.0
 44        )
 45        
 46        # Position controller
 47        self.position_pid = PIDController(
 48            Kp=2.0,
 49            Ki=0.1,
 50            Kd=3.0,
 51            setpoint=0.0
 52        )
 53    
 54    def compute_control(self, target_position=0, target_heading=0):
 55        """
 56        Compute wheel torques for balance and position control
 57        
 58        Returns:
 59            (tau_left, tau_right)
 60        """
 61        x, x_dot, theta, theta_dot, psi = self.robot.state
 62        
 63        # Position error
 64        position_error = target_position - x
 65        
 66        # Desired tilt angle for position control
 67        theta_desired = self.position_pid.update(x, dt=0.01)
 68        theta_desired = np.clip(theta_desired, -0.3, 0.3)  # Limit tilt
 69        
 70        # Balance control
 71        self.balance_pid.setpoint = theta_desired
 72        tau_balance = self.balance_pid.update(theta, dt=0.01)
 73        
 74        # Heading control (simple proportional)
 75        heading_error = target_heading - psi
 76        heading_error = np.arctan2(np.sin(heading_error), 
 77                                   np.cos(heading_error))
 78        tau_heading = 0.5 * heading_error
 79        
 80        # Combine
 81        tau_left = tau_balance - tau_heading
 82        tau_right = tau_balance + tau_heading
 83        
 84        # Limit torques
 85        tau_max = 2.0  # N⋅m
 86        tau_left = np.clip(tau_left, -tau_max, tau_max)
 87        tau_right = np.clip(tau_right, -tau_max, tau_max)
 88        
 89        return tau_left, tau_right
 90
 91# Example: Balance with position control
 92robot = BalancingRobot()
 93robot.state[2] = 0.1  # Initial tilt
 94
 95controller = BalancingController(robot)
 96
 97trajectory = [robot.state.copy()]
 98dt = 0.01
 99
100for i in range(500):
101    # Target position moves
102    target_pos = 0.5 * np.sin(i * dt * 0.5)
103    
104    tau_l, tau_r = controller.compute_control(target_position=target_pos)
105    robot.step(tau_l, tau_r, dt)
106    trajectory.append(robot.state.copy())
107
108trajectory = np.array(trajectory)
109
110# Plot results
111fig, axes = plt.subplots(2, 2, figsize=(12, 10))
112time = np.arange(len(trajectory)) * dt
113
114axes[0, 0].plot(time, trajectory[:, 0], label='Actual')
115axes[0, 0].plot(time, 0.5 * np.sin(time * 0.5), '--', label='Target')
116axes[0, 0].set_ylabel('Position (m)')
117axes[0, 0].legend()
118axes[0, 0].grid(True)
119
120axes[0, 1].plot(time, np.degrees(trajectory[:, 2]))
121axes[0, 1].set_ylabel('Tilt Angle (deg)')
122axes[0, 1].grid(True)
123
124axes[1, 0].plot(time, trajectory[:, 1])
125axes[1, 0].set_xlabel('Time (s)')
126axes[1, 0].set_ylabel('Velocity (m/s)')
127axes[1, 0].grid(True)
128
129axes[1, 1].plot(time, trajectory[:, 3])
130axes[1, 1].set_xlabel('Time (s)')
131axes[1, 1].set_ylabel('Angular Velocity (rad/s)')
132axes[1, 1].grid(True)
133
134plt.tight_layout()
135plt.show()

 1from scipy.linalg import solve_continuous_are
 2
 3def compute_lqr_gain(robot):
 4    """
 5    Compute LQR gain matrix
 6    
 7    Returns:
 8        K: Feedback gain matrix
 9    """
10    # Linearized system matrices
11    m_total = robot.m_b + 2 * robot.m_w
12    denom = robot.I_b + robot.m_b * robot.L**2
13    
14    # State matrix A
15    A = np.array([
16        [0, 1, 0, 0],
17        [0, 0, -robot.m_b * robot.g * robot.L / m_total, 0],
18        [0, 0, 0, 1],
19        [0, 0, m_total * robot.g * robot.L / denom, 0]
20    ])
21    
22    # Input matrix B
23    B = np.array([
24        [0],
25        [1 / (m_total * robot.r)],
26        [0],
27        [-1 / (denom * robot.r)]
28    ])
29    
30    # Cost matrices
31    Q = np.diag([10.0, 1.0, 100.0, 10.0])  # State cost
32    R = np.array([[0.1]])  # Control cost
33    
34    # Solve Riccati equation
35    P = solve_continuous_are(A, B, Q, R)
36    
37    # Compute gain
38    K = np.linalg.inv(R) @ B.T @ P
39    
40    return K
41
42class LQRBalancingController:
43    def __init__(self, robot):
44        self.robot = robot
45        self.K = compute_lqr_gain(robot)
46        print(f"LQR Gains: {self.K}")
47    
48    def compute_control(self, target_state=None):
49        """
50        Compute control using LQR
51        
52        Args:
53            target_state: Desired [x, x_dot, theta, theta_dot]
54        
55        Returns:
56            (tau_left, tau_right)
57        """
58        if target_state is None:
59            target_state = np.array([0, 0, 0, 0])
60        
61        # State error
62        current_state = self.robot.state[:4]
63        error = current_state - target_state
64        
65        # LQR control law: u = -K * error
66        tau = -self.K @ error
67        tau = tau[0]  # Extract scalar
68        
69        # Apply to both wheels (no turning)
70        tau_left = tau
71        tau_right = tau
72        
73        # Limit
74        tau_max = 2.0
75        tau_left = np.clip(tau_left, -tau_max, tau_max)
76        tau_right = np.clip(tau_right, -tau_max, tau_max)
77        
78        return tau_left, tau_right
79
80# Example with LQR
81robot = BalancingRobot()
82robot.state[2] = 0.15  # Larger initial tilt
83
84controller = LQRBalancingController(robot)
85
86trajectory = [robot.state.copy()]
87dt = 0.01
88
89for _ in range(500):
90    tau_l, tau_r = controller.compute_control()
91    robot.step(tau_l, tau_r, dt)
92    trajectory.append(robot.state.copy())
93
94trajectory = np.array(trajectory)
95print(f"Final tilt: {np.degrees(trajectory[-1, 2]):.2f} degrees")

 1def forward_kinematics(self, omega_left, omega_right):
 2    """
 3    Compute robot velocity from wheel velocities
 4    
 5    Args:
 6        omega_left, omega_right: Wheel angular velocities (rad/s)
 7    
 8    Returns:
 9        (v, omega): Linear and angular velocity
10    """
11    # Linear velocity
12    v = self.r * (omega_left + omega_right) / 2
13    
14    # Angular velocity (heading rate)
15    omega = self.r * (omega_right - omega_left) / self.wheelbase
16    
17    return v, omega
18
19# Add to BalancingRobot class
20BalancingRobot.forward_kinematics = forward_kinematics

 1def inverse_kinematics(self, v_desired, omega_desired):
 2    """
 3    Compute wheel velocities for desired motion
 4    
 5    Args:
 6        v_desired: Desired linear velocity (m/s)
 7        omega_desired: Desired angular velocity (rad/s)
 8    
 9    Returns:
10        (omega_left, omega_right): Wheel angular velocities
11    """
12    omega_left = (v_desired - omega_desired * self.wheelbase / 2) / self.r
13    omega_right = (v_desired + omega_desired * self.wheelbase / 2) / self.r
14    
15    return omega_left, omega_right
16
17# Add to BalancingRobot class
18BalancingRobot.inverse_kinematics = inverse_kinematics

 1class IMUSensor:
 2    def __init__(self, noise_gyro=0.01, noise_accel=0.1):
 3        self.noise_gyro = noise_gyro
 4        self.noise_accel = noise_accel
 5    
 6    def read(self, true_theta, true_theta_dot):
 7        """Simulate IMU readings with noise"""
 8        # Gyroscope (angular velocity)
 9        gyro = true_theta_dot + np.random.normal(0, self.noise_gyro)
10        
11        # Accelerometer (tilt angle estimate)
12        accel_theta = true_theta + np.random.normal(0, self.noise_accel)
13        
14        return gyro, accel_theta
15
16class ComplementaryFilter:
17    def __init__(self, alpha=0.98):
18        """
19        Args:
20            alpha: Filter coefficient (0.95-0.99 typical)
21        """
22        self.alpha = alpha
23        self.theta_fused = 0
24    
25    def update(self, gyro, accel_theta, dt):
26        """
27        Fuse gyro and accelerometer data
28        
29        Args:
30            gyro: Angular velocity from gyroscope
31            accel_theta: Angle from accelerometer
32            dt: Time step
33        """
34        # Integrate gyro
35        theta_gyro = self.theta_fused + gyro * dt
36        
37        # Complementary filter
38        self.theta_fused = self.alpha * theta_gyro + (1 - self.alpha) * accel_theta
39        
40        return self.theta_fused
41
42# Example
43imu = IMUSensor()
44filter = ComplementaryFilter(alpha=0.98)
45
46true_theta = 0.1
47true_theta_dot = 0.5
48
49for _ in range(100):
50    gyro, accel = imu.read(true_theta, true_theta_dot)
51    theta_estimate = filter.update(gyro, accel, dt=0.01)
52    
53    # Update true values (simple dynamics)
54    true_theta += true_theta_dot * 0.01
55
56print(f"True: {true_theta:.3f}, Estimated: {theta_estimate:.3f}")

 1def apply_motor_limits(self, tau_left, tau_right, tau_max=2.0, 
 2                      omega_max=50.0):
 3    """
 4    Apply motor torque and velocity limits
 5    
 6    Returns:
 7        Limited torques
 8    """
 9    # Torque limits
10    tau_left = np.clip(tau_left, -tau_max, tau_max)
11    tau_right = np.clip(tau_right, -tau_max, tau_max)
12    
13    # Velocity limits (if needed)
14    # This would require tracking wheel velocities
15    
16    return tau_left, tau_right

 1def simulate_balancing_robot(duration=10.0, dt=0.01, 
 2                            controller_type='lqr'):
 3    """Complete simulation with visualization"""
 4    robot = BalancingRobot()
 5    robot.state[2] = 0.1  # Initial tilt
 6    
 7    if controller_type == 'lqr':
 8        controller = LQRBalancingController(robot)
 9    else:
10        controller = BalancingController(robot)
11    
12    trajectory = []
13    controls = []
14    
15    t = 0
16    while t < duration:
17        # Compute control
18        if controller_type == 'lqr':
19            tau_l, tau_r = controller.compute_control()
20        else:
21            target_pos = 0.5 * np.sin(t * 0.5)
22            tau_l, tau_r = controller.compute_control(target_position=target_pos)
23        
24        # Step simulation
25        robot.step(tau_l, tau_r, dt)
26        
27        trajectory.append(robot.state.copy())
28        controls.append([tau_l, tau_r])
29        
30        t += dt
31    
32    trajectory = np.array(trajectory)
33    controls = np.array(controls)
34    time = np.arange(len(trajectory)) * dt
35    
36    # Plot
37    fig, axes = plt.subplots(3, 2, figsize=(14, 12))
38    
39    axes[0, 0].plot(time, trajectory[:, 0])
40    axes[0, 0].set_ylabel('Position (m)')
41    axes[0, 0].set_title('Position')
42    axes[0, 0].grid(True)
43    
44    axes[0, 1].plot(time, np.degrees(trajectory[:, 2]))
45    axes[0, 1].set_ylabel('Tilt (deg)')
46    axes[0, 1].set_title('Tilt Angle')
47    axes[0, 1].grid(True)
48    
49    axes[1, 0].plot(time, trajectory[:, 1])
50    axes[1, 0].set_ylabel('Velocity (m/s)')
51    axes[1, 0].set_title('Linear Velocity')
52    axes[1, 0].grid(True)
53    
54    axes[1, 1].plot(time, trajectory[:, 3])
55    axes[1, 1].set_ylabel('Angular Vel (rad/s)')
56    axes[1, 1].set_title('Tilt Rate')
57    axes[1, 1].grid(True)
58    
59    axes[2, 0].plot(time, controls[:, 0], label='Left')
60    axes[2, 0].plot(time, controls[:, 1], label='Right')
61    axes[2, 0].set_xlabel('Time (s)')
62    axes[2, 0].set_ylabel('Torque (N⋅m)')
63    axes[2, 0].set_title('Control Torques')
64    axes[2, 0].legend()
65    axes[2, 0].grid(True)
66    
67    axes[2, 1].plot(trajectory[:, 0], np.degrees(trajectory[:, 2]))
68    axes[2, 1].set_xlabel('Position (m)')
69    axes[2, 1].set_ylabel('Tilt (deg)')
70    axes[2, 1].set_title('Phase Plot')
71    axes[2, 1].grid(True)
72    
73    plt.tight_layout()
74    plt.show()
75    
76    return trajectory, controls
77
78# Run simulation
79trajectory, controls = simulate_balancing_robot(duration=10.0, 
80                                                controller_type='lqr')