Deploy geoguesser environment with train/eval splits

geoguesser_env/pyproject.toml

@@ -8,7 +8,7 @@ name = "openenv-geoguesser-env"
 # satisfied when an installed distribution matches the version, ignoring the git
 # revision the spec asked for, so a same-version rebuild is invisible to a
 # cached environment: a fixed client sat unused through three job runs.
-version = "0.1.3"
+version = "0.1.4"
 description = "GeoGuessr-style visual geolocation environment for OpenEnv"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -48,3 +48,18 @@ packages = [
     "geoguesser_env.server.render",
 ]
 package-dir = { "geoguesser_env" = ".", "geoguesser_env.server" = "server", "geoguesser_env.server.backends" = "server/backends", "geoguesser_env.server.render" = "server/render" }
+
+# The map renderer and the reverse geocoder read these at runtime, so an install
+# without them is code that cannot pin, guess or name a country -- which is how
+# it failed: `look` worked and `pin` raised FileNotFoundError on the countries
+# outline. 40 MB of coastlines, places, roads, rivers and urban areas.
+#
+# `data/geo/osm_cache` is deliberately not listed. It is 331 MB of fetched
+# Overpass responses that the environment regenerates on demand, and shipping a
+# cache in a wheel would be eight times the size of the data it caches.
+[tool.setuptools.package-data]
+geoguesser_env = [
+    "data/geo/*.geojson",
+    "data/geo/detail/*.json",
+    "openenv.yaml",
+]

geoguesser_env/scripts/render_rollout.py

@@ -64,6 +64,8 @@ DURATIONS = {
     "measure": 2.2,
     "guess": 6.0,
     "forced_guess": 5.0,
+    # Nothing moves and nothing was learnt, so it should register and pass.
+    "unparseable": 1.6,
 }
 DEFAULT_DURATION = 2.4
 
@@ -114,6 +116,42 @@ def interpolate(before: dict, after: dict, steps: int) -> list[dict]:
     return out
 
 
+def walked(turn: dict) -> bool:
+    """Whether a move actually crossed capture points."""
+    before = (turn.get("camera_before") or {}).get("frame_index")
+    after = (turn.get("camera_after") or {}).get("frame_index")
+    return before is not None and after is not None and before != after
+
+
+def walk_states(before: dict, after: dict, steps: int) -> list[dict]:
+    """One state per video frame, stepping through the frames actually walked.
+
+    The camera is held at the pose the move ended on -- a move does not turn the
+    head -- and only the position advances, eased so the walk starts and stops
+    rather than running at a constant rate.
+
+    Args:
+        before (`dict`):
+            Camera state before the move, carrying the starting `frame_index`.
+        after (`dict`):
+            Camera state after it.
+        steps (`int`):
+            How many video frames the segment lasts.
+
+    Returns:
+        `list[dict]`: `steps` states, each a copy of `after` with the
+        `frame_index` of the panorama that should be on screen.
+    """
+    start, end = int(before["frame_index"]), int(after["frame_index"])
+    direction = 1 if end >= start else -1
+    indices = list(range(start, end + direction, direction))
+    out = []
+    for i in range(steps):
+        t = ease(i / max(1, steps - 1))
+        out.append({**after, "frame_index": indices[min(len(indices) - 1, int(t * len(indices)))]})
+    return out
+
+
 def camera_changed(turn: dict) -> bool:
     """Whether the action actually moved the camera, in heading or field of view."""
     before, after = turn.get("camera_before") or {}, turn.get("camera_after") or {}
@@ -161,9 +199,24 @@ def render(args: argparse.Namespace) -> None:
         if line.strip()
     ]
     rows = [r for r in rows if not r.get("error")]
-    if not 0 <= args.episode < len(rows):
-        raise SystemExit(f"--episode must be 0..{len(rows) - 1} ({len(rows)} usable)")
-    record = rows[args.episode]
+    if args.episode_id:
+        # Point at a trace by name rather than by position: an index shifts when
+        # a run is resumed or a file is concatenated, so it is not a stable way
+        # to refer to the episode you actually watched.
+        matches = [r for r in rows if str(r.get("episode_id", "")).startswith(args.episode_id)]
+        if not matches:
+            raise SystemExit(
+                f"no episode id starting {args.episode_id!r} in {args.episodes}. "
+                f"Available: {', '.join(str(r.get('episode_id'))[:12] for r in rows[:8])}..."
+            )
+        if len(matches) > 1:
+            raise SystemExit(f"{args.episode_id!r} matches {len(matches)} episodes; be more specific")
+        record = matches[0]
+    else:
+        if not 0 <= args.episode < len(rows):
+            raise SystemExit(f"--episode must be 0..{len(rows) - 1} ({len(rows)} usable)")
+        record = rows[args.episode]
+    logger.info("episode %s · %s", record.get("episode_id"), record.get("model_name"))
 
     split = record["task"]["split"]
     index_name = INDEX_FOR_SPLIT.get(split)
@@ -200,7 +253,20 @@ def render(args: argparse.Namespace) -> None:
     )
 
     size = (args.width_pano, args.width_pano)
+    # Matched to the composition's left panel, not to 16:9. The panel is about
+    # 1.42:1 once the rail and the bars are taken out, so a 16:9 frame had its
+    # sides cropped away by object-fit -- throwing away exactly the extra width
+    # the wide render exists to provide.
+    wide_size = (args.wide_width, int(round(args.wide_width / args.wide_aspect)))
+    wide_dir = out_dir / "wide"
+    mini_dir = out_dir / "minimap"
+    wide_dir.mkdir(parents=True, exist_ok=True)
+    mini_dir.mkdir(parents=True, exist_ok=True)
     counter = 0
+    wide_counter = 0
+    mini_counter = 0
+    # Minimaps repeat whenever the pin set is unchanged, which is most turns.
+    minimap_cache: dict[tuple, str] = {}
     segments = []
     # The last panorama the agent saw, kept so a map segment can sit over the
     # street view it was placed from -- the way the human UI's map inset does --
@@ -216,15 +282,75 @@ def render(args: argparse.Namespace) -> None:
         image.convert("RGB").save(frames_dir / name, quality=args.quality)
         return f"frames/{name}"
 
+    def write_wide(image) -> str:
+        """Save one wide-context frame. Kept in its own sequence so the two
+        panes can be composed side by side without the numbering interleaving."""
+        nonlocal wide_counter
+        wide_counter += 1
+        name = f"{wide_counter:05d}.jpg"
+        image.convert("RGB").save(wide_dir / name, quality=args.quality)
+        return f"wide/{name}"
+
+    def write_minimap(image) -> str:
+        """Save one minimap still."""
+        nonlocal mini_counter
+        mini_counter += 1
+        name = f"{mini_counter:05d}.png"
+        image.convert("RGB").save(mini_dir / name)
+        return f"minimap/{name}"
+
+    def minimap_for(pins: list[tuple[float, float]]) -> str | None:
+        """A wide map of what the agent has pinned so far, and nothing more.
+
+        Deliberately never passed `truth`: the minimap is on screen from the
+        first turn, so drawing the answer there would spoil every episode before
+        the agent has said anything. A run with no pins yet gets no minimap.
+        """
+        if not pins:
+            return None
+        key = tuple(pins)
+        if key not in minimap_cache:
+            focus = pins[-1]
+            minimap_cache[key] = write_minimap(
+                render_map(
+                    list(pins),
+                    focus=focus,
+                    span_deg=args.minimap_span,
+                    dpi=max(70, args.map_dpi // 2),
+                )
+            )
+        return minimap_cache[key]
+
     for turn in record["turns"]:
         action = turn.get("action") or {}
-        kind = str(action.get("action") or "reset")
+        # An unparseable reply has no action at all, and defaulting that to
+        # "reset" captioned the segment "dropped somewhere on Earth" -- telling
+        # a viewer the episode restarted when the model had simply written
+        # something the parser could not read. The distinction matters: one is
+        # the environment acting, the other is the model failing.
+        raw_kind = action.get("action") or action.get("op")
+        if raw_kind is None:
+            kind = "reset" if turn["turn"] == 0 else "unparseable"
+        else:
+            # Mirror the environment's own aliases so a video labels the action
+            # the environment performed, not the spelling the model used.
+            kind = ALIASES.get(str(raw_kind).lower(), str(raw_kind).lower())
+        if kind == "unparseable" and args.skip_unparsed:
+            # A reply the harness could not read may say more about the parser
+            # than about the model, and the environment did nothing with the
+            # turn either way. Left in, the segment spends two seconds telling
+            # the viewer about our own tooling. The turn is still in the trace;
+            # pass --keep-unparsed to show it.
+            logger.info("skipping turn %s: reply was not parsed", turn["turn"])
+            continue
         seconds = DURATIONS.get(kind, DEFAULT_DURATION)
         if kind in {"look", "pan", "zoom"} and not camera_changed(turn):
             seconds = min(seconds, STATIC_DURATION)
         total = max(1, int(round(seconds * args.fps)))
         paths: list[str] = []
+        wide_paths: list[str] = []
         cameras: list[list[float]] = []
+        pins_now = [tuple(pt) for pt in (turn.get("pins") or [])]
 
         if turn.get("image_kind") == "map":
             # Maps are single stills: re-render at video resolution from the pins
@@ -258,19 +384,62 @@ def render(args: argparse.Namespace) -> None:
                     focus = pins[-1] if pins else None
             else:
                 focus = pins[-1] if pins else None
-            image = render_map(
-                pins,
-                focus=focus,
-                span_deg=span,
-                dpi=args.map_dpi,
-                truth=truth,
-            )
-            map_aspect = image.width / image.height
-            paths = [write(image)]
+
+            if truth is not None and pins and args.reveal_frames > 1:
+                # The reveal is animated rather than cut to. A single still
+                # showed the guess, the answer and the line between them all at
+                # once, which lands as a fact rather than as a result: the
+                # viewer has no moment of "where is it, then". So the map opens
+                # tight on the guess and widens until the true location comes
+                # into frame, and the truth marker is withheld until the
+                # framing is wide enough to contain it.
+                start_span = max(0.35, span * 0.18)
+                reveal_at = max(1, int(args.reveal_frames * 0.42))
+                for i in range(args.reveal_frames):
+                    t = ease(i / (args.reveal_frames - 1))
+                    step_focus = (
+                        pins[0][0] + (focus[0] - pins[0][0]) * t,
+                        pins[0][1] + (focus[1] - pins[0][1]) * t,
+                    )
+                    image = render_map(
+                        pins,
+                        focus=step_focus,
+                        span_deg=start_span + (span - start_span) * t,
+                        dpi=args.reveal_dpi,
+                        truth=truth if i >= reveal_at else None,
+                    )
+                    map_aspect = image.width / image.height
+                    paths.append(write(image))
+                # The last frame at full resolution, because it is held for
+                # several seconds and is the shot people screenshot.
+                final = render_map(
+                    pins, focus=focus, span_deg=span, dpi=args.map_dpi, truth=truth
+                )
+                map_aspect = final.width / final.height
+                paths.append(write(final))
+            else:
+                image = render_map(
+                    pins,
+                    focus=focus,
+                    span_deg=span,
+                    dpi=args.map_dpi,
+                    truth=truth,
+                )
+                map_aspect = image.width / image.height
+                paths = [write(image)]
         else:
             before, after = turn["camera_before"], turn["camera_after"]
             moving = kind in {"look", "pan", "zoom"} and camera_changed(turn)
-            states = interpolate(before, after, total) if moving else [after]
+            if kind == "move" and walked(turn):
+                # A real walk, not a cut. The capture points are about 3.3 m
+                # apart, so a 34 m move crosses roughly ten panoramas -- every
+                # one of which exists and can be rendered at the same heading.
+                # Holding each for a few video frames gives the forward motion
+                # Street View shows when you step down a road, and every frame
+                # is something the environment could actually have returned.
+                states = walk_states(before, after, total)
+            else:
+                states = interpolate(before, after, total) if moving else [after]
             for state in states:
                 if state.get("frame_index") is None:
                     continue
@@ -284,6 +453,20 @@ def render(args: argparse.Namespace) -> None:
                 if image.size != size:
                     image = image.resize(size)
                 paths.append(write(image))
+                # The same moment at a wider field of view. Rendered from the
+                # same pose, so the two panes cannot disagree about where the
+                # agent is looking -- the left one simply shows more of the
+                # street than the model was given.
+                wide = backend.render_view(
+                    task,
+                    int(state["frame_index"]),
+                    float(state["heading_deg"]),
+                    float(state["pitch_deg"]),
+                    min(120.0, float(state["fov_deg"]) * args.wide_factor),
+                )
+                if wide.size != wide_size:
+                    wide = wide.resize(wide_size)
+                wide_paths.append(write_wide(wide))
                 # The camera the frame was rendered with, so the HUD can read
                 # live through a pan instead of jumping to the end state the
                 # moment the shot begins.
@@ -319,6 +502,13 @@ def render(args: argparse.Namespace) -> None:
                 # Only set on a map segment: the street view it was placed from.
                 "backdrop": backdrop if turn.get("image_kind") == "map" else None,
                 "frames": paths,
+                # The same pose at a wider field of view, for the left canvas.
+                # Empty on a map segment, where there is nothing to widen.
+                "wide": wide_paths,
+                # What the agent has pinned by the end of this turn, so the
+                # minimap and the bubbles agree about the state of the guess.
+                "pins": [[round(a, 4), round(b, 4)] for a, b in pins_now],
+                "minimap": minimap_for(pins_now),
                 "cameras": cameras,
                 # A single still is held for the whole segment; a pan has one
                 # frame per video frame.
@@ -353,6 +543,8 @@ def render(args: argparse.Namespace) -> None:
         # letterboxing it: React sizes the card to the image instead of fitting
         # the image into a card the wrong shape.
         "mapAspect": map_aspect,
+        "theme": args.theme,
+        "wideAspect": args.wide_aspect,
         "episode": {
             "episode_id": record["episode_id"],
             "model_name": record["model_name"],
@@ -387,8 +579,20 @@ def render(args: argparse.Namespace) -> None:
 
 #: Uppercase verb for the on-screen action badge, in the human UI's own words
 #: (the control pad says LOOK / WALK / ZOOM, not `look` / `move` / `zoom`).
+# The spellings a model may emit for an action the environment accepts. Kept in
+# step with `ACTION_ALIASES` in examples/geoguesser_llm_rollout.py.
+ALIASES = {
+    "place_pin": "pin",
+    "submit_guess": "guess",
+    "final_guess": "guess",
+    "reverse_geocode": "view_map",
+    "turn": "pan",
+    "rotate": "pan",
+}
+
 VERBS = {
     "reset": "START",
+    "unparseable": "UNREAD",
     "look": "LOOK",
     "pan": "LOOK",
     "zoom": "ZOOM",
@@ -450,6 +654,8 @@ def action_caption(kind: str, action: dict, turn: dict) -> str:
     separately, and a viewer three seconds into a muted video needs to know
     "walking forward" before they need to know the argument names.
     """
+    if kind == "unparseable":
+        return "reply could not be parsed as an action"
     before = turn.get("camera_before") or {}
     after = turn.get("camera_after") or {}
     if kind == "reset":
@@ -576,6 +782,42 @@ def main() -> None:
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument("episodes", type=pathlib.Path)
     parser.add_argument("--episode", type=int, default=0)
+    parser.add_argument(
+        "--episode-id",
+        default=None,
+        help="Select by episode_id (a unique prefix is enough) instead of by "
+        "position, which shifts when a run is resumed.",
+    )
+    parser.add_argument(
+        "--theme",
+        choices=("dark", "light"),
+        default="dark",
+        help="Composition palette. Light reads better on a timeline; dark reads "
+        "better as an instrument.",
+    )
+    parser.add_argument(
+        "--wide-factor",
+        type=float,
+        default=2.0,
+        help="How much wider than the agent's own field of view the left canvas "
+        "renders. A pan then scrolls through a scene instead of cutting between "
+        "two crops, which is what the environment actually models.",
+    )
+    parser.add_argument("--wide-width", type=int, default=1408)
+    parser.add_argument(
+        "--wide-aspect",
+        type=float,
+        default=1.42,
+        help="Aspect of the left canvas in the composition. Anything wider is "
+        "cropped away by object-fit rather than shown.",
+    )
+    parser.add_argument(
+        "--minimap-span",
+        type=float,
+        default=60.0,
+        help="Half-width in degrees of the minimap. Wide enough to place a "
+        "continent without revealing anything the agent has not pinned.",
+    )
     parser.add_argument("--fps", type=int, default=30)
     # 4:5 rather than 16:9. The clip is made for a phone timeline, where a
     # portrait frame occupies about twice the screen a landscape one does, and
@@ -590,6 +832,21 @@ def main() -> None:
     )
     parser.add_argument("--quality", type=int, default=88)
     parser.add_argument("--map-dpi", type=int, default=160)
+    parser.add_argument(
+        "--reveal-frames",
+        type=int,
+        default=30,
+        help="Frames in the guess reveal, widening from the guess until the "
+        "true location comes into frame. 1 disables the animation.",
+    )
+    parser.add_argument(
+        "--reveal-dpi",
+        type=int,
+        default=100,
+        help="Resolution of the reveal's intermediate frames. Lower than "
+        "--map-dpi because they are on screen for a thirtieth of a second "
+        "each, while the final frame is rendered at full resolution.",
+    )
     parser.add_argument(
         "--reply-chars",
         type=int,
@@ -605,6 +862,15 @@ def main() -> None:
         help="View size the server used, for the fidelity check.",
     )
     parser.add_argument("--skip-verify", action="store_true")
+    parser.add_argument(
+        "--keep-unparsed",
+        dest="skip_unparsed",
+        action="store_false",
+        default=True,
+        help="Show turns whose reply the harness could not parse. Skipped by "
+        "default: such a turn may be a limitation of the parser rather than of "
+        "the model, and the environment did nothing with it either way.",
+    )
     parser.add_argument("--cache", type=pathlib.Path, default=ROOT / "data" / "panos")
     parser.add_argument("--tasks-dir", type=pathlib.Path, default=ROOT / "tasks")
     parser.add_argument(