Hide demo rows and fix world-model agents

Author: biru-codeastromer

agents/worldmodel_agents/imagination_mpc.py

@@ -1,14 +1,25 @@
 from __future__ import annotations
 
-import copy
-
 import numpy as np
+import torch
 from worldmodel_models.registry import create_world_model
 from worldmodel_planners.mpc_cem import MPCCEMPlanner
 
 from worldmodel_agents.base import AgentConfig, BaseAgent
 
 
+def _clone_state(value):
+    if isinstance(value, torch.Tensor):
+        return value.detach().clone()
+    if isinstance(value, dict):
+        return {key: _clone_state(item) for key, item in value.items()}
+    if isinstance(value, list):
+        return [_clone_state(item) for item in value]
+    if isinstance(value, tuple):
+        return tuple(_clone_state(item) for item in value)
+    return value
+
+
 class ImaginationMPCAgent(BaseAgent):
     def __init__(self, config: AgentConfig | None = None):
         super().__init__(config=config)
@@ -40,7 +51,7 @@ def rollout_fn(state, action_seq):
         result = self.planner.plan(
             root_state=self.latent,
             rollout_fn=rollout_fn,
-            clone_state_fn=copy.deepcopy,
+            clone_state_fn=_clone_state,
         )
         self.last_imagined_transitions = result.imagined_transitions
         self.last_planner_trace = result.trace

agents/worldmodel_agents/search_mcts.py

@@ -1,14 +1,25 @@
 from __future__ import annotations
 
-import copy
-
 import numpy as np
+import torch
 from worldmodel_models.registry import create_world_model
 from worldmodel_planners.mcts import MCTSPlanner
 
 from worldmodel_agents.base import AgentConfig, BaseAgent
 
 
+def _clone_state(value):
+    if isinstance(value, torch.Tensor):
+        return value.detach().clone()
+    if isinstance(value, dict):
+        return {key: _clone_state(item) for key, item in value.items()}
+    if isinstance(value, list):
+        return [_clone_state(item) for item in value]
+    if isinstance(value, tuple):
+        return tuple(_clone_state(item) for item in value)
+    return value
+
+
 class SearchMCTSAgent(BaseAgent):
     """Minimal MuZero-style skeleton: learned model + MCTS planning."""
 
@@ -42,7 +53,7 @@ def transition_fn(state, action):
         result = self.planner.plan(
             root_state=self.latent,
             transition_fn=transition_fn,
-            clone_state_fn=copy.deepcopy,
+            clone_state_fn=_clone_state,
         )
 
         self.last_imagined_transitions = result.imagined_transitions

server/worldmodel_server/main.py

@@ -263,13 +263,16 @@ def leaderboard(
     track: str = Query(default="test"),
     env: str | None = Query(default=None),
     agent: str | None = Query(default=None),
+    include_demo: bool = Query(default=False),
     session: Session = Depends(get_session),
 ):
     q = select(RunEntry).where(RunEntry.status == "uploaded", RunEntry.track == track)
     if env:
         q = q.where(RunEntry.env == env)
     if agent:
         q = q.where(RunEntry.agent == agent)
+    if not include_demo:
+        q = q.where(RunEntry.created_by != "demo-seed")
 
     rows = session.scalars(q.order_by(desc(RunEntry.created_at))).all()
     out: list[LeaderboardRow] = []

tests/test_agent_registry.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from worldmodel_agents.registry import create_agent
+from worldmodel_gym.envs.registry import make_env
+
+
+def test_search_mcts_can_act_on_memory_maze_observation():
+    env = make_env("memory_maze", obs_mode="both", max_steps=8)
+    obs, info = env.reset(seed=123)
+    agent = create_agent("search_mcts")
+    agent.reset(seed=123)
+
+    action = agent.act(obs, info)
+
+    assert isinstance(action, int)
+
+
+def test_imagination_mpc_can_act_on_switch_quest_observation():
+    env = make_env("switch_quest", obs_mode="both", max_steps=8)
+    obs, info = env.reset(seed=123)
+    agent = create_agent("imagination_mpc")
+    agent.reset(seed=123)
+
+    action = agent.act(obs, info)
+
+    assert isinstance(action, int)

tests/test_server_app.py

@@ -44,14 +44,26 @@ def test_seed_demo_data_populates_leaderboard(tmp_path, monkeypatch):
     app = modules["worldmodel_server.main"].app
 
     with TestClient(app) as client:
-        response = client.get("/api/leaderboard?track=test")
+        response = client.get("/api/leaderboard?track=test&include_demo=true")
 
     assert response.status_code == 200
     rows = response.json()
     assert len(rows) >= 2
     assert any(row["agent"] == "demo-mpc" for row in rows)
 
 
+def test_public_leaderboard_hides_seeded_demo_rows_by_default(tmp_path, monkeypatch):
+    modules = load_test_modules(monkeypatch, tmp_path, seed_demo=True)
+    app = modules["worldmodel_server.main"].app
+
+    with TestClient(app) as client:
+        response = client.get("/api/leaderboard?track=test")
+
+    assert response.status_code == 200
+    rows = response.json()
+    assert rows == []
+
+
 def test_api_key_can_create_and_upload_run(tmp_path, monkeypatch):
     modules = load_test_modules(monkeypatch, tmp_path)
     app = modules["worldmodel_server.main"].app

worldmodels/worldmodel_models/common.py

@@ -33,8 +33,11 @@ def __init__(self, config: ModelConfig):
         super().__init__()
         self.config = config
         self.device = torch.device(config.device)
+        self.optimizer: torch.optim.Optimizer | None = None
+
+    def initialize_optimizer(self) -> None:
         self.to(self.device)
-        self.optimizer = torch.optim.Adam(self.parameters(), lr=config.lr)
+        self.optimizer = torch.optim.Adam(self.parameters(), lr=self.config.lr)
 
     def _obs_tensor(self, obs) -> torch.Tensor:
         arr = to_numpy_obs(obs)

worldmodels/worldmodel_models/deterministic.py

@@ -24,6 +24,7 @@ def __init__(self, config: ModelConfig | None = None):
         self.obs_head = torch.nn.Linear(c.latent_dim, c.obs_dim)
         self.reward_head = torch.nn.Linear(c.latent_dim, 1)
         self.done_head = torch.nn.Linear(c.latent_dim, 1)
+        self.initialize_optimizer()
 
     def init_state(self, batch_size: int = 1) -> dict[str, torch.Tensor]:
         latent = torch.zeros(

worldmodels/worldmodel_models/stochastic.py

@@ -22,6 +22,7 @@ def __init__(self, config: ModelConfig | None = None):
         self.reward_head = torch.nn.Linear(c.latent_dim, 1)
         self.done_head = torch.nn.Linear(c.latent_dim, 1)
         self.obs_head = torch.nn.Linear(c.latent_dim, c.obs_dim)
+        self.initialize_optimizer()
 
     def init_state(self, batch_size: int = 1) -> dict[str, torch.Tensor]:
         h = torch.zeros((batch_size, self.config.latent_dim), device=self.device)