space/control_loop.py

"""30 Hz target-point control and critically damped path smoothing."""
import math

FPS = 30
DEFAULT_PROMPTS = [
    'A person walks forward.',
    'A person runs forward.',
    'A person walks forward in a crouched position.',
    'A person jumps forward repeatedly.',
]
DEFAULT_SPEEDS = [0.8, 2.5, 0.45, 0.8]


def effective_slot(slot, shift):
    slot = int(slot)
    if slot not in range(4):
        raise ValueError('Choose slot 1–4.')
    return 1 if slot == 0 and shift else slot


class SpringPath:
    def __init__(self, frequency=10.):
        self.frequency = frequency
        self.frame = 0
        self.target = [0., 0.]
        self.position = [0., 0.]
        self.velocity = [0., 0.]
        self.heading = 0.

    def step(self, *, x, z, shift, slot, speeds):
        selected = effective_slot(slot, shift)
        dt = 1/FPS
        previous = self.position.copy()
        old_heading = self.heading
        if self.frame:
            norm = max(1., math.hypot(x, z))
            self.target[0] += x/norm * speeds[selected] * dt
            self.target[1] += z/norm * speeds[selected] * dt
            decay = math.exp(-self.frequency*dt)
            for axis in (0, 1):
                offset = self.position[axis] - self.target[axis]
                c = self.velocity[axis] + self.frequency*offset
                self.position[axis] = self.target[axis] + (offset+c*dt)*decay
                self.velocity[axis] = (self.velocity[axis]-self.frequency*c*dt)*decay
            dx, dz = self.position[0]-previous[0], self.position[1]-previous[1]
            if math.hypot(dx,dz) > 1e-5:
                target_heading = math.atan2(dx,dz)
                difference = (target_heading-self.heading+math.pi) % (2*math.pi)-math.pi
                self.heading += max(-math.pi*dt, min(math.pi*dt, difference*(1-math.exp(-10*dt))))
            c, s = math.cos(self.heading), math.sin(self.heading)
            row = [self.heading-old_heading, c*dx-s*dz, s*dx+c*dz]
        else:
            row = [0., 0., 0.]
        point = {'frame':self.frame, 'x':self.position[0], 'z':self.position[1],
                 'heading':self.heading, 'slot':selected, 'base_slot':int(slot),
                 'run':selected==1, 'target_x':self.target[0], 'target_z':self.target[1]}
        self.frame += 1
        return row, point

    def target_point(self, slot):
        return {'x':self.target[0], 'z':self.target[1], 'slot':slot}