downfall

The prompt

Build a 3D obstacle course in the Fall Guys style. I am a bean running across raised platforms. The first stretch has spinning hammers and spiked rollers. The next stretch is a floor of tiles that fall after I stand on them. Then more rollers. The last stretch is slopes with bombs dropping onto them. A yellow gate at the end of each stretch sends me to the next one.

I run, strafe, look around, and jump, and I can only jump while I am standing on something. Getting closer to the gate can be worth a little. Falling off, or touching a hammer, roller, or bomb, ends the round and I have lost. Reaching a gate should feel like progress, not like the whole game is over. Keep a long round so a person can cross more than one stretch. Show the course in 3D from behind me, with the score. Different starts can begin on a later stretch. The same start should repeat.

# AI : Bean Run — a 3D Fall Guys style obstacle course.
# A bean runs raised platforms through four stretches (hammers/rollers, falling
# tiles, rollers, bomb slopes), each closed by a yellow gate. Touch a hazard or
# fall off -> lost. Reach a gate -> progress. Reach the last gate -> won.
from __future__ import annotations

import math
from typing import Any, Dict, List, Optional

from helpers import base_observation, key_input, pick_task, rng_for, seed_from_task, task_ids, clamp01

GAME_CLASS = "BeanRun"

# ---- tuning ----------------------------------------------------------------
DT = 1.0 / 60.0
RUN_SPEED = 7.0
JUMP_VY = 4.8
GRAVITY = 9.5
TURN_STEP = 0.22            # radians per look press
DEATH_Y = -3.0
BEAN_R = 0.6
BEAN_H = 1.2
STRETCH_LEN = 26.0
GAP = 4.0
HALF_W = 5.5               # platform half width (bean falls off past this)
TILE_LEN = 3.2
TILE_FALL_TICKS = 14       # ~0.23s of standing before a tile falls
HAMMER_H = 0.75
HAMMER_ARM = 1.5
HAMMER_SPIN = 1.1
ROLLER_R = 1.0
ROLLER_TOP = 0.85
BOMB_DROP = 8.0
BOMB_BOOM_R = 2.2
BOMB_BOOM_TICKS = 9

TASK_IDS = task_ids(8, "s")

MOVE_TOKENS = {
    "up", "down", "left", "right",
    "jump", "turn_left", "turn_right", "stay",
    "forward", "back",
}
ALL_MOVES = [
    "up", "down", "left", "right", "stay",
    "up left", "up right", "down left", "down right",
    "jump", "turn_left", "turn_right",
]


def _inputs() -> List[Dict[str, Any]]:
    return [
        key_input(
            "WASD",
            "run and strafe",
            {
                "w": "up", "ArrowUp": "up",
                "s": "down", "ArrowDown": "down",
                "a": "left", "ArrowLeft": "left",
                "d": "right", "ArrowRight": "right",
            },
            hold=True,
            every_ms=60,
            idle="stay",
        ),
        key_input("Q E", "look around", {"q": "turn_left", "e": "turn_right"}, hold=False),
        key_input("Space", "jump", {" ": "jump"}, hold=False),
    ]


class BeanRun:
    def __init__(self) -> None:
        self.reset("")

    # ---- world construction -------------------------------------------------
    def _build(self, task_id: str) -> None:
        seed = seed_from_task(task_id)
        rng = rng_for(task_id)
        canonical = pick_task(task_id, TASK_IDS, "")
        start = seed % 4
        self.task_id = canonical or ""
        self.start_stretch = start
        self.total_gates = 4

        # stretch kinds, in order
        kinds = ["hammers", "tiles", "rollers", "bombs"]
        self.stretches: List[Dict[str, Any]] = []
        z = 0.0
        for i, kind in enumerate(kinds):
            self.stretches.append({
                "index": i,
                "kind": kind,
                "z0": z,
                "z1": z + STRETCH_LEN,
                "gate": z + STRETCH_LEN,
                "height": 0.0,
            })
            z += STRETCH_LEN + GAP

        # floor segments (tiles are per-tile alive, others solid)
        self.segments: List[Dict[str, Any]] = []
        self.hammers: List[Dict[str, Any]] = []
        self.rollers: List[Dict[str, Any]] = []
        self.tile_segments: List[Dict[str, Any]] = []
        self.bombs: List[Dict[str, Any]] = []
        self._bomb_schedule: List[float] = []   # upcoming spawn times (absolute t)

        for st in self.stretches:
            z0, z1 = st["z0"], st["z1"]
            kind = st["kind"]
            if kind == "tiles":
                ntiles = int((STRETCH_LEN - 1.0) // TILE_LEN)
                for t in range(ntiles):
                    tz0 = z0 + t * TILE_LEN
                    tz1 = min(tz0 + TILE_LEN, z1)
                    seg = {
                        "x0": -HALF_W, "x1": HALF_W, "z0": tz0, "z1": tz1,
                        "height": 0.0, "kind": "tile", "state": "alive",
                        "ticks": 0,
                    }
                    self.segments.append(seg)
                    self.tile_segments.append(seg)
            else:
                self.segments.append({
                    "x0": -HALF_W, "x1": HALF_W, "z0": z0, "z1": z1,
                    "height": 0.0, "kind": kind, "state": "solid", "ticks": -1,
                })

            if kind == "hammers":
                zs = [z0 + 6.0, z0 + 13.0, z0 + 20.0]
                for j, cz in enumerate(zs):
                    self.hammers.append({
                        "x": 0.0, "z": cz, "angle": rng.random() * 6.28,
                        "spin": HAMMER_SPIN * (1 if j % 2 == 0 else -1),
                        "arm": HAMMER_ARM, "height": HAMMER_H,
                    })
            elif kind == "rollers":
                centers = [z0 + 8.0, z0 + 17.0]
                for j, cz in enumerate(centers):
                    amp = rng.uniform(3.0, 4.5)
                    self.rollers.append({
                        "x": 0.0, "z": cz, "base_z": cz,
                        "amp": amp, "phase": rng.random() * 6.28,
                        "r": ROLLER_R, "top": ROLLER_TOP,
                    })
            elif kind == "bombs":
                # slopes are drawn tilted; physics floor stays level.
                n = 6
                for b in range(n):
                    when = self._start_time_for_self() + 3.0 + b * 3.4
                    self._bomb_schedule.append(when)

        # bomb spawn positions derived from schedule at spawn time
        self._bomb_seed = rng
        self._bomb_times: List[float] = list(self._bomb_schedule)

    def _start_time_for_self(self) -> float:
        # helper only during build; base time is 0
        return 0.0

    # ---- reset / step -------------------------------------------------------
    def reset(self, task_id: str = "") -> Dict[str, Any]:
        self._build(task_id)
        seed = seed_from_task(task_id)
        self.t = 0.0
        self.alive = True
        self.done = False
        self.won = False
        start = self.start_stretch
        self.stretch = start
        st = self.stretches[start]
        # bean spawns a little behind the start of its stretch, center, on the floor
        spawn_z = max(0.0, st["z0"] + 1.5)
        self.x = 0.0
        self.y = 0.0
        self.z = spawn_z
        self.vy = 0.0
        self.yaw = 0.0
        self.grounded = True
        self.gates = 0
        self.score = 0.0
        self.message = ""
        self.prev_z = self.z
        self._age_ticks = 0
        self._bomb_schedule = list(self._bomb_times)
        self.bombs = []
        self._last_move = None
        return self._observe(0.0)

    def _observe(self, reward: float) -> Dict[str, Any]:
        return base_observation(
            reward=reward,
            done=self.done,
            legal_moves=ALL_MOVES,
            last_move=self._last_move,
            bean={
                "x": round(self.x, 3),
                "y": round(self.y, 3),
                "z": round(self.z, 3),
                "yaw": round(self.yaw, 4),
                "grounded": self.grounded,
            },
            alive=self.alive,
            won=self.won,
            score=round(self.score, 3),
            gates=self.gates,
            total_gates=self.total_gates,
            t=round(self.t, 3),
            stretch=self.stretch,
            stretch_label=["Hammers", "Falling tiles", "Rollers", "Bomb slopes"][min(self.stretch, 3)],
            message=self.message,
            inputs=_inputs(),
            segments=[
                {
                    "x0": s["x0"], "x1": s["x1"], "z0": s["z0"], "z1": s["z1"],
                    "height": s["height"], "kind": s["kind"], "state": s["state"],
                }
                for s in self.segments
            ],
            hammers=[
                {"x": h["x"], "z": h["z"], "angle": h["angle"], "arm": h["arm"], "height": h["height"]}
                for h in self.hammers
            ],
            rollers=[
                {"x": r["x"], "z": r["z"], "r": r["r"], "top": r["top"]}
                for r in self.rollers
            ],
            bombs=[
                {"x": b["x"], "z": b["z"], "y": b["y"], "state": b["state"]}
                for b in self.bombs
            ],
            gate_list=[
                {"z": st["gate"], "passed": i < self.gates or (i < self.stretch)}
                for i, st in enumerate(self.stretches)
            ],
        )

    def _ground_at(self, x: float, z: float) -> Optional[float]:
        for s in self.segments:
            if s["state"] == "solid" or s["state"] == "alive" or s["state"] == "falling":
                if s["z0"] <= z <= s["z1"] and s["x0"] <= x <= s["x1"]:
                    return s["height"]
        return None

    def step(self, action):
        if self.done:
            return self._observe(0.0), True, None
        move = getattr(action, "move", "") or ""
        return self._tick(move), True, None

    def _tick(self, move: str) -> Dict[str, Any]:
        self._last_move = move or None
        reward = 0.0
        if self.done:
            return self._observe(0.0)

        self.t += DT
        self._age_ticks += 1
        tokens = [t for t in (move or "").split() if t]
        jump = False
        fwd = back = left = right = False

        for tok in tokens:
            if tok == "jump":
                jump = True
            elif tok in ("turn_left",):
                self.yaw -= TURN_STEP
            elif tok in ("turn_right",):
                self.yaw += TURN_STEP
            elif tok in ("up", "forward"):
                fwd = True
            elif tok == "down":
                back = True
            elif tok == "left":
                left = True
            elif tok == "right":
                right = True
            elif tok == "back":
                back = True

        # facing / right vectors
        fx, fz = math.sin(self.yaw), math.cos(self.yaw)
        rx, rz = math.cos(self.yaw), -math.sin(self.yaw)
        dvx = fx * (fwd - back) + rx * (right - left)
        dvz = fz * (fwd - back) + rz * (right - left)
        if abs(dvx) > 1e-9 or abs(dvz) > 1e-9:
            n = math.hypot(dvx, dvz)
            dvx /= n
            dvz /= n

        if jump and self.grounded:
            self.vy = JUMP_VY
            self.grounded = False

        # integrate
        self.x += dvx * RUN_SPEED * DT
        self.z += dvz * RUN_SPEED * DT
        self.vy -= GRAVITY * DT
        self.y += self.vy * DT

        ground = self._ground_at(self.x, self.z)
        if ground is not None and self.vy <= 0.0 and self.y <= ground:
            self.y = ground
            self.vy = 0.0
            self.grounded = True
        elif ground is None:
            self.grounded = False

        # progress reward toward the gate is worth a little
        reward += clamp01((self.z - self.prev_z) * (RUN_SPEED * DT) / STRETCH_LEN) * 0.04
        self.prev_z = self.z

        self._update_hazards()
        self._update_bombs()
        self._update_tiles()
        self._collide()

        if self.done:
            return self._observe(reward)

        # gate crossing
        st = self.stretches[self.stretch]
        if self.z >= st["gate"]:
            self.gates += 1
            reward += 0.30
            if self.stretch + 1 >= self.total_gates:
                self.won = True
                self.done = True
                self.score += 1.0
                self.message = "You reached the finish gate. You won!"
                return self._observe(reward + 1.0)
            else:
                self.stretch += 1
                self.score += 0.3
                nxt = self.stretches[self.stretch]
                self.z = nxt["z0"] + 1.0
                self.y = nxt["height"]
                self.vy = 0.0
                self.grounded = True
                self.prev_z = self.z
                self.message = "Gate cleared! On to the next stretch."
        else:
            self.message = ""

        self.score += reward
        return self._observe(reward)

    # ---- dynamics -----------------------------------------------------------
    def _update_hazards(self) -> None:
        for h in self.hammers:
            h["angle"] += h["spin"] * DT
        for r in self.rollers:
            r["x"] = r.get("amp", 3.0) * math.sin(r.get("phase", 0.0) + self.t * 1.1)

    def _update_bombs(self) -> None:
        # spawn due bombs
        due = [i for i, when in enumerate(self._bomb_schedule) if self.t >= when]
        for i in due:
            bx = self._bomb_seed.uniform(-HALF_W + 0.8, HALF_W - 0.8)
            bz = self._bomb_seed.choice(
                [self.stretches[-1]["z0"] + 4.0, self.stretches[-1]["z0"] + 12.0,
                 self.stretches[-1]["z0"] + 20.0]
            )
            self.bombs.append({"x": bx, "z": bz, "y": BOMB_DROP, "vy": 0.0,
                               "state": "fall", "ticks": 0})
        for i in sorted(due, reverse=True):
            self._bomb_schedule.pop(i)

        for b in self.bombs:
            if b["state"] == "fall":
                b["vy"] -= GRAVITY * 0.55 * DT
                b["y"] += b["vy"] * DT
                if b["y"] <= 0.4:
                    b["y"] = 0.4
                    b["state"] = "boom"
                    b["ticks"] = BOMB_BOOM_TICKS
            elif b["state"] == "boom":
                b["ticks"] -= 1
                if b["ticks"] <= 0:
                    b["state"] = "gone"
        self.bombs = [b for b in self.bombs if b["state"] != "gone"]

    def _update_tiles(self) -> None:
        for seg in self.tile_segments:
            if seg["state"] == "falling":
                seg["ticks"] -= 1
                if seg["ticks"] <= 0:
                    seg["state"] = "gone"
            if seg["state"] == "alive" and self.grounded and not self.done:
                if seg["z0"] <= self.z <= seg["z1"] and seg["x0"] <= self.x <= seg["x1"]:
                    seg["state"] = "falling"
                    seg["ticks"] = TILE_FALL_TICKS

    # ---- collision ----------------------------------------------------------
    def _collide(self) -> None:
        if self.y < DEATH_Y:
            self._lose("You fell off the course.")
            return
        for h in self.hammers:
            ca = math.cos(h["angle"])
            sa = math.sin(h["angle"])
            dx = self.x - h["x"]
            dz = self.z - h["z"]
            radial = abs(dx * ca + dz * sa)
            dperp = abs(dx * sa - dz * ca)
            if radial < h["arm"] and dperp < 0.22 + BEAN_R and self.y < h["height"] + 0.2:
                self._lose("A spinning hammer knocked you off.")
                return
        for r in self.rollers:
            d = math.hypot(self.x - r["x"], self.z - r["z"])
            if d < r["r"] + BEAN_R * 0.6 and self.y < r["top"] + 0.15:
                self._lose("A spiked roller flattened you.")
                return
        for b in self.bombs:
            if b["state"] == "fall":
                d = math.hypot(self.x - b["x"], self.z - b["z"])
                if d < 0.9 and abs(b["y"] - (self.y + BEAN_H * 0.5)) < 1.0:
                    self._lose("A bomb landed right on you.")
                    return
            elif b["state"] == "boom":
                d = math.hypot(self.x - b["x"], self.z - b["z"])
                if d < BOMB_BOOM_R:
                    self._lose("A bomb exploded too close.")
                    return

    def _lose(self, msg: str) -> None:
        self.alive = False
        self.done = True
        self.message = msg
        self.score -= 1.0