Add Gymnasium's 11 MuJoCo environments, played from pixels, with a Modal baseline job

examples/mujoco_baseline.py

@@ -0,0 +1,50 @@
+"""A model's baseline on every MuJoCo game (``laya.mujocogames``): a video of its first episode per game, and its
+mean return against random play, doing nothing (always ``NONE``) and the expert, on the same seeded episodes.
+
+    pip install -e . torchvision "gymnasium[mujoco]" "imageio[ffmpeg]" stable-baselines3 sb3-contrib
+    python examples/mujoco_baseline.py --model thaitea/laya-vision \\
+        --revision 8b318c99d7ad3ce19c24369263463882eada9d1e --out mujoco-baseline/
+
+Writes ``<out>/<game>-model.webm`` and ``<out>/baseline.json`` (one row per game, with the model's resolved
+revision), and prints a table. ``modal run modal_app.py::mujoco_eval`` runs every game at once on GPUs. ``--max-steps`` caps every episode (the model runs about 1 s a step on a CPU, and
+most games run to 1000 steps); ``normalized`` = (model - random) / (expert - random), blank where a game has no
+expert. See ``examples/mujoco_video.py`` for one game at a time.
+"""
+import argparse
+import json
+import os
+
+from laya.mujocogames import GAMES, baseline, baseline_table
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--model", default="thaitea/laya-vision", help="checkpoint dir or Hub id")
+    ap.add_argument("--revision", default=None, help="Hub commit to load --model at")
+    ap.add_argument("--device", default=None)
+    ap.add_argument("--games", default=",".join(GAMES))
+    ap.add_argument("--episodes", type=int, default=3, help="model episodes per game (the first is recorded)")
+    ap.add_argument("--baseline-episodes", type=int, default=10, help="random / still / expert episodes per game")
+    ap.add_argument("--seed", type=int, default=200_000)
+    ap.add_argument("--max-steps", type=int, default=300)
+    ap.add_argument("--out", default="mujoco-baseline")
+    args = ap.parse_args()
+
+    from laya.vlm import load_vlm
+
+    agent = load_vlm(args.model, device=args.device, revision=args.revision)
+    os.makedirs(args.out, exist_ok=True)
+    rows = []
+    for game in [g for g in args.games.split(",") if g]:
+        video = os.path.join(args.out, "%s-model.webm" % game)
+        row = baseline(agent, game, video, args.episodes, args.baseline_episodes, args.seed, args.max_steps)
+        row.update(model=args.model, revision=agent.source.get("revision"), video=video)  # the resolved commit
+        rows.append(row)
+        print(json.dumps(row), flush=True)
+        with open(os.path.join(args.out, "baseline.json"), "w") as f:
+            json.dump(rows, f, indent=1)
+    print("\n" + baseline_table(rows))
+
+
+if __name__ == "__main__":
+    main()

examples/mujoco_video.py

@@ -0,0 +1,49 @@
+"""Record a MuJoCo episode (``laya.mujocogames``) to a video: the screen, the chosen push and, for a model, its
+probabilities over the pushes.
+
+    pip install -e . torchvision "gymnasium[mujoco]" "imageio[ffmpeg]" stable-baselines3 sb3-contrib
+    python examples/mujoco_video.py --policy expert --out expert.webm
+    python examples/mujoco_video.py --game Hopper --policy random --out random.webm
+    python examples/mujoco_video.py --game HalfCheetah --policy model --model thaitea/laya-vision \\
+        --revision 8b318c99d7ad3ce19c24369263463882eada9d1e --out model.webm
+
+The model sees what a games-suite run sees: the rendered screen with the previous frame ghosted in, and
+``laya.games.control_question``. On a machine with no display, rendering uses EGL (``apt-get install libegl1``) or
+``MUJOCO_GL=osmesa`` (``libosmesa6``). The suffix of ``--out`` picks the container: .webm (VP9) or .mp4 (H.264).
+The video plays in real time (the environment's step length). ``examples/mujoco_baseline.py`` records and scores
+every game.
+"""
+import argparse
+
+from laya.mujocogames import GAMES, MujocoGame, model_chooser, random_policy, record, still_policy
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--game", default="InvertedPendulum", choices=sorted(GAMES))
+    ap.add_argument("--policy", default="expert", choices=("expert", "random", "still", "model"))
+    ap.add_argument("--model", default="thaitea/laya-vision", help="checkpoint dir or Hub id (--policy model)")
+    ap.add_argument("--revision", default=None, help="Hub commit to load --model at")
+    ap.add_argument("--device", default=None)
+    ap.add_argument("--seed", type=int, default=200_000)
+    ap.add_argument("--max-steps", type=int, default=300, help="stop the recording here (0 = the episode's end)")
+    ap.add_argument("--out", default="mujoco.webm")
+    args = ap.parse_args()
+
+    env = MujocoGame(args.game, args.seed)
+    if args.policy == "model":
+        from laya.vlm import load_vlm
+
+        choose = model_chooser(load_vlm(args.model, device=args.device, revision=args.revision), args.game)
+    else:
+        pick = {"random": random_policy(args.seed), "still": still_policy,
+                "expert": lambda e: e.expert()}[args.policy]
+        choose = lambda e: (pick(e), None)  # noqa: E731
+    record(env, choose, args.policy, args.out, args.max_steps)
+    print("%s %s: %d steps, return %.1f, %s -> %s" % (args.game, args.policy, env.steps, env.score,
+                                                    "ended" if env.terminated else "running", args.out))
+    env.close()
+
+
+if __name__ == "__main__":
+    main()

laya/controlgames.py

@@ -52,15 +52,13 @@ def _expert_action(game: str, env, obs) -> int:
 class ControlGame:
     """One seeded episode of a ``GAMES`` environment, stepped by action name."""
 
-    def __init__(self, game: str, seed: int = 0):
-        if game not in GAMES:
-            raise ValueError("unknown control game %r (%s)" % (game, ", ".join(GAMES)))
-        os.environ.setdefault("SDL_VIDEODRIVER", "dummy")  # headless pygame, no audio device
-        os.environ.setdefault("SDL_AUDIODRIVER", "dummy")
-        import gymnasium as gym
+    SPECS = GAMES  # the games this class plays; ``laya.mujocogames.MujocoGame`` swaps in its own
 
-        self.game, self.actions = game, GAMES[game]["actions"]
-        self.env = gym.make(GAMES[game]["env_id"], render_mode="rgb_array")
+    def __init__(self, game: str, seed: int = 0):
+        if game not in self.SPECS:
+            raise ValueError("unknown control game %r (%s)" % (game, ", ".join(self.SPECS)))
+        self.game, self.actions = game, self.SPECS[game]["actions"]
+        self.env = self._make(self.SPECS[game])
         self.obs, _ = self.env.reset(seed=seed)
         self.env.action_space.seed(seed)
         self.score, self.steps, self.terminated, self.truncated = 0.0, 0, False, False
@@ -73,9 +71,20 @@ def done(self) -> bool:
     def expert(self) -> str:
         return self.actions[_expert_action(self.game, self.env, self.obs)]
 
+    def _make(self, spec: Dict):
+        os.environ.setdefault("SDL_VIDEODRIVER", "dummy")  # headless pygame, no audio device
+        os.environ.setdefault("SDL_AUDIODRIVER", "dummy")
+        import gymnasium as gym
+
+        return gym.make(spec["env_id"], render_mode="rgb_array")
+
+    def _env_action(self, action):
+        """What ``env.step`` takes for an action name (a subclass may take other action forms)."""
+        return self.actions.index(action)
+
     def step(self, action: str) -> float:
         self._prev, self._frame = self._frame, None
-        self.obs, reward, self.terminated, self.truncated, _ = self.env.step(self.actions.index(action))
+        self.obs, reward, self.terminated, self.truncated, _ = self.env.step(self._env_action(action))
         self.score += float(reward)
         self.steps += 1
         return float(reward)
@@ -100,24 +109,25 @@ def close(self) -> None:
         self.env.close()
 
 
-def play_episodes(game: str, policy, episodes: int, seed: int = 0, max_steps: int = 0) -> Dict:
+def play_episodes(game: str, policy, episodes: int, seed: int = 0, max_steps: int = 0, make=ControlGame) -> Dict:
     """Play ``episodes`` seeded episodes (seed ``seed + i``); ``policy(env) -> action name``. ``max_steps`` caps
-    an episode below the environment's own limit (0 = no extra cap)."""
+    an episode below the environment's own limit (0 = no extra cap). ``make`` is the game class (``ControlGame``
+    or a subclass such as ``laya.mujocogames.MujocoGame``)."""
     from collections import Counter
 
     counts, eps = Counter(), []
     for i in range(episodes):
-        env = ControlGame(game, seed + i)
+        env = make(game, seed + i)
         while not env.done and not (max_steps and env.steps >= max_steps):
             a = policy(env)
-            counts[a] += 1
+            counts.update(a.values() if isinstance(a, dict) else [a])  # per-joint actions: count each joint's level
             env.step(a)
         eps.append({"score": round(env.score, 3), "steps": env.steps, "terminated": env.terminated})
         env.close()
-    scores = [e["score"] for e in eps]
+    scores, solved = [e["score"] for e in eps], make.SPECS[game]["solved"]  # None: no threshold
     return {"game": game, "episodes": episodes, "seed": seed, "actions": dict(counts), "results": eps,
             "mean_score": float(np.mean(scores)), "std_score": float(np.std(scores)),
-            "solved_rate": float(np.mean([s >= GAMES[game]["solved"] for s in scores])),
+            "solved_rate": None if solved is None else float(np.mean([s >= solved for s in scores])),
             "mean_steps": float(np.mean([e["steps"] for e in eps]))}
 
 

laya/games.py

@@ -85,19 +85,77 @@ def snake_question() -> Dict:
     "MountainCar": "The car is too weak to drive straight up; rock back and forth to build speed and reach the "
                    "flag on the right hill.",
     "LunarLander": "Fire the engines to land the lander gently and upright between the two flags.",
+    "InvertedPendulum": "A pole is hinged on a cart that slides along a rail; push the cart left or right to keep "
+                        "the pole upright.",
+    "InvertedDoublePendulum": "Two poles are hinged end to end on a cart that slides along a rail; push the cart "
+                              "left or right to keep both poles upright.",
+    "Reacher": "A two-joint arm lies flat on a table; move its fingertip onto the red target and keep it there "
+               "with little effort.",
+    "Pusher": "A seven-joint arm stands by a table; push the white cylinder onto the red goal.",
+    "Swimmer": "A three-segment swimmer lies in a thick fluid; wriggle its two joints to swim forward.",
+    "Hopper": "A one-legged robot stands upright; hop forward as fast as you can without falling.",
+    "Walker2d": "A two-legged robot stands upright; walk forward as fast as you can without falling.",
+    "HalfCheetah": "A two-legged running robot; run forward as fast as you can.",
+    "Ant": "A four-legged robot; walk forward as fast as you can without flipping over.",
+    "Humanoid": "A humanoid robot stands upright; walk forward as fast as you can without falling.",
+    "HumanoidStandup": "A humanoid robot lies on the ground; get up and raise your head as high as you can.",
 }
+
+# MuJoCo joint-torque games (``laya.mujocogames``): the actuated joints in the environment's actuator order. Each
+# joint gets its own ``choice`` every step (``mujoco_questions``), over ``TORQUE_LEVELS``.
+TORQUE_JOINTS = {
+    "Reacher": ("shoulder", "elbow"),
+    "Pusher": ("shoulder pan", "shoulder lift", "upper arm roll", "elbow flex", "forearm roll", "wrist flex",
+               "wrist roll"),
+    "Swimmer": ("front joint", "back joint"),
+    "Hopper": ("thigh", "knee", "foot"),
+    "Walker2d": ("right thigh", "right knee", "right foot", "left thigh", "left knee", "left foot"),
+    "HalfCheetah": ("back thigh", "back shin", "back foot", "front thigh", "front shin", "front foot"),
+    "Ant": ("leg 4 hip", "leg 4 ankle", "leg 1 hip", "leg 1 ankle", "leg 2 hip", "leg 2 ankle", "leg 3 hip",
+            "leg 3 ankle"),
+    "Humanoid": ("abdomen y", "abdomen z", "abdomen x", "right hip x", "right hip z", "right hip y", "right knee",
+                 "left hip x", "left hip z", "left hip y", "left knee", "right shoulder 1", "right shoulder 2",
+                 "right elbow", "left shoulder 1", "left shoulder 2", "left elbow"),
+}
+TORQUE_JOINTS["HumanoidStandup"] = TORQUE_JOINTS["Humanoid"]
+TORQUE_LEVELS = {"STRONG_NEG": "full torque the negative way", "NEG": "half torque the negative way",
+                 "NONE": "no torque", "POS": "half torque the positive way",
+                 "STRONG_POS": "full torque the positive way"}
+
+
+def joint_key(joint: str) -> str:
+    """The answer key for a joint: ``"right knee"`` -> ``"right_knee"``."""
+    return joint.replace(" ", "_")
+
+
+def mujoco_questions(game: str) -> Dict[str, Dict]:
+    """One ``choice`` per joint of a ``TORQUE_JOINTS`` game, all asked about the same frame. The game and its goal come
+    first and the joint last, so the questions share a prefix (``laya.vlm`` computes it once)."""
+    return {joint_key(j): {
+        "type": "choice",
+        "instructions": "You are playing the control task %s. %s A faint copy shows where things were one step "
+                        "earlier. How should you torque the %s now?" % (game, CONTROL_GOALS[game], j),
+        "criteria": dict(TORQUE_LEVELS),
+    } for j in TORQUE_JOINTS[game]}
+
+
 CONTROL_ACTIONS = {
     "CartPole": {"LEFT": "push the cart left", "RIGHT": "push the cart right"},
     "Acrobot": {"CLOCKWISE": "twist the lower link clockwise", "NONE": "do nothing",
                 "COUNTERCLOCKWISE": "twist the lower link counter-clockwise"},
     "MountainCar": {"LEFT": "accelerate left", "NONE": "do not accelerate", "RIGHT": "accelerate right"},
     "LunarLander": {"NOOP": "do nothing", "LEFT_ENGINE": "fire the left orientation engine",
                     "MAIN_ENGINE": "fire the main engine", "RIGHT_ENGINE": "fire the right orientation engine"},
+    "InvertedPendulum": {"LEFT": "push the cart left", "NONE": "do not push", "RIGHT": "push the cart right"},
+    "InvertedDoublePendulum": {"HARD_LEFT": "push the cart left hard", "LEFT": "push the cart left gently",
+                               "NONE": "do not push", "RIGHT": "push the cart right gently",
+                               "HARD_RIGHT": "push the cart right hard"},
 }
 
 
 def control_question(game: str) -> Dict:
-    """The question for a ``laya.controlgames`` game; the screen ghosts the previous frame to show motion."""
+    """The question for a ``laya.controlgames`` game or a MuJoCo pendulum; the screen ghosts the previous frame to show
+    motion. The MuJoCo torque games ask one question per joint instead (``mujoco_questions``)."""
     return {"action": {
         "type": "choice",
         "instructions": "You are playing the control task %s. %s A faint copy shows where things were one step "