Autonomous running train in Pybricks

main.py

from pybricks.hubs import CityHub
from pybricks.pupdevices import DCMotor, ColorDistanceSensor
from pybricks.parameters import Port, Color
from pybricks.tools import wait

hub = CityHub()
motor = DCMotor(Port.A)
color_sensor = ColorDistanceSensor(Port.B)

def wait_for_color(color):
    count = 0
    while count < 10:
        if color_sensor.color() == color:
            count += 1
        else:
            count = 0
        wait(10)

while True:
    motor.dc(20)
    hub.light.on(Color.YELLOW)
    wait_for_color(Color.YELLOW)
    motor.stop()
    wait(1000)

    motor.dc(-20)
    hub.light.on(Color.BLUE)
    wait_for_color(Color.BLUE)
    motor.stop()
    wait(1000)