Adding generalized 2D envs #117
Merged
Merged
Conversation
A robocode env that varies the object count per reset, driving kinder's object-centric layer directly so one program spans instances of different sizes. reset() yields an ObjectCentricState at a chosen count; a per-count ConstantObjectKinDEREnv backend cache gives the planner its Box view. Count is inferred by object-name prefix, not type. build_sesame_models gains per-count observation_space/model_kwargs overrides and fails loud on a variable-count env.
Sweep object counts across the eval set so the program and the planner face the same (seed, count) instances, and scale the step budget with the count. Report a by_count breakdown (solve rate per count over all scheduled episodes, crashes and unattempted counted as failures) plus largest_count_all_solved / largest_count_any_solved, and add plot_scaling.py for the program-vs-planner curve.
Rebuild and cache SeSamE models per object count, vectorizing the object-centric state to that count's Box for the planner; thread the pinned count through the per-instance approaches; teach check_action_collision the variable-count env.
For a variable-count env the generated program receives an ObjectCentricState, so the interface note documents that API (get_objects / get) and forbids shape, devectorize, and fixed-count assumptions.
A self-contained ObjectCentricState<->JSON codec (embedding the type hierarchy, preserving is_instance ancestors), a serialize_space branch for the object-centric space, and a stdlib+numpy local mirror the sandbox reconstructs, so a blackboxed program gets the same state object it sees at eval.
Obstruction2D, ClutteredRetrieval2D, ClutteredStorage2D (odd block counts), Motion2D, and StickButton2D, each with feasible design and held-out eval counts.
For a variable-count env the observation is an ObjectCentricState, not a flat vector, so render_state's arbitrary-state mode (a list of floats) and the devectorize/vectorize guidance do not apply. Thread object_centric through build_system_prompt / build_mcp_tool_lines / mcp_tool_descriptions and swap in a seed-mode + render_policy variant.
Replace the side-effect `import kinder_geom2d_env` in VariableObjectCountEnv with an explicit configure_gl_backend() call. That wrapper drives the kinder env classes directly (no kinder.make), so it needs only the GL-backend selection, not register_all_environments() and the gym registry. Factor the GL preamble that the 2D and 3D wrappers each duplicated into the new environments/mujoco_gl.py. This drops the wrong-import-position, unused-import, and broad-except pylint disables from both wrappers and restores normal import order.
merlerm
force-pushed
the
generalized-envs-2d
branch
from
e3e7bba to
e1e2ca9
Compare
Collaborator
Author
|
Some cool very preliminary stuff with this StickButton2D — three-way comparison (identical seed=42 instances, 50 evals each)
|
Import pyplot in normal order and call plt.switch_backend("Agg") after, instead
of matplotlib.use("Agg") wedged between imports. Drops the E402 / wrong-import-position
noqa and pylint disables it forced on every following import, plus a redundant
`import matplotlib`.
Two runs of one approach (e.g. agentic with and without primitives) previously pooled into a single blended line, since collect() keyed only by approach. Track the source dir each run was found under and split a series only when one approach spans several sources, so a genuine three-way comparison shows three lines while seeds of a single run still pool. Make the reusable helpers (collect, series, by_count, plot_environment) public and add tests/experiments/test_plot_scaling.py covering source tracking, label disambiguation, per-count aggregation, and PNG output.
_kinder_check and _variable_count_check differed only in which attribute holds the active kinder backend, so fold them into _kinder_reference_check(env, backend_attr, ...). This also collapses the six per-line type-ignore/disable annotations into one block-scope disable (the type-ignores were unnecessary once the inner env is Any). Add variable-count collision tests (collision, free move, state preservation) so both dispatch paths through the shared routine are covered.
There was a problem hiding this comment.
Pull request overview
This PR introduces support for variable-object-count (generalized) 2D environments by switching observations from fixed-length vectors to object-centric states, and extends the blackbox client/server protocol and experiment tooling to evaluate and plot performance vs. object count.
Changes:
- Add
VariableObjectCountEnvwrapper to exposeObjectCentricStateobservations while still supporting per-count Box views for bilevel planning. - Extend blackbox env client/server to serialize/transport object-centric states (including type hierarchies) and update collision checking + bilevel model building for variable-count instances.
- Add evaluation/reporting utilities for solve-rate-vs-count curves (
summarize_by_count, plotting + result flattening) plus comprehensive tests.
Reviewed changes
Copilot reviewed 31 out of 31 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/utils/test_object_centric_codec.py | New tests covering JSON-safe round-trips and blackbox loop for object-centric state transport. |
| tests/utils/test_episode.py | Adds unit tests for per-count aggregation (summarize_by_count) and threads count into scripted approach calls. |
| tests/utils/test_bilevel.py | Skips env-id inference for variable-count env configs that don’t have env_id. |
| tests/experiments/test_plot_scaling.py | New tests for experiments/plot_scaling.py helpers and PNG output. |
| tests/environments/test_variable_object_count_env.py | New tests validating count pinning, determinism, state restore, and description content. |
| tests/environments/test_collision_checking.py | Adds collision-check coverage for VariableObjectCountEnv dispatch. |
| src/robocode/utils/object_centric_codec.py | New host-side JSON codec and space serializer for ObjectCentricState. |
| src/robocode/utils/episode.py | Adds count pinning to episodes, stores object_count in metrics, and implements summarize_by_count. |
| src/robocode/utils/env_server.py | Registers builtin codecs and adds ObjectCentricStateSpace serialization support. |
| src/robocode/utils/env_server_runtime.py | Allows rendering custom states coming from object-centric payloads. |
| src/robocode/utils/env_client.py | Adds a local object-centric mirror in the sandbox and object-centric-aware state handling. |
| src/robocode/utils/bilevel.py | Updates build_sesame_models to support per-count overrides and reject variable-count primitive misuse. |
| src/robocode/prompts.py | Adds object-centric-specific prompt note and MCP prompt suffix selection. |
| src/robocode/primitives/check_action_collision.py | Extends collision checking dispatch to variable-count envs via backend reference checks. |
| src/robocode/mcp/init.py | Adds object-centric variant of MCP tool descriptions and system suffix. |
| src/robocode/environments/variable_object_count_env.py | New variable-count wrapper env with count inference, per-count model building, and object-centric descriptions. |
| src/robocode/environments/mujoco_gl.py | New helper to configure and lock MuJoCo/PyOpenGL backend once per process. |
| src/robocode/environments/kinder_geom3d_env.py | Refactors GL backend locking to use configure_gl_backend(). |
| src/robocode/environments/kinder_geom2d_env.py | Refactors GL backend locking to use configure_gl_backend(). |
| src/robocode/approaches/bilevel_planning_approach.py | Rebuilds/caches planning models per count and vectorizes object-centric observations for the planner. |
| src/robocode/approaches/base_approach.py | Extends solve_instance contract to accept an optional count pin. |
| src/robocode/approaches/agentic_per_instance_approach.py | Pins eval count per episode and scales horizons for variable-count envs. |
| src/robocode/approaches/agentic_base.py | Detects object-centric observation spaces and switches prompt/tooling accordingly. |
| experiments/run_experiment.py | Sweeps eval_counts for variable-count envs, pins counts per episode, and stores by-count summaries. |
| experiments/plot_scaling.py | New script to plot solve rate and planner degradation vs. object count across runs. |
| experiments/conf/environment/stickbutton2d_generalized.yaml | New variable-count env config for StickButton2D. |
| experiments/conf/environment/obstruction2d_generalized.yaml | New variable-count env config for Obstruction2D. |
| experiments/conf/environment/motion2d_generalized.yaml | New variable-count env config for Motion2D. |
| experiments/conf/environment/clutteredstorage2d_generalized.yaml | New variable-count env config for ClutteredStorage2D. |
| experiments/conf/environment/clutteredretrieval2d_generalized.yaml | New variable-count env config for ClutteredRetrieval2D. |
| experiments/analyze_results.py | Flattens nested by_count solve rates into numeric solve_rate@<count> columns for aggregation. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
…t envs - object_centric_codec: sort objects by name explicitly instead of relying on relational_structs.Object rich comparison, keeping the payload order stable and independent of a dependency's comparison semantics. - summarize_by_count: raise on a scheduled_counts/per_episode length mismatch rather than letting zip silently drop episodes and skew the scheduled denominator. - collision tests: close the variable-count envs, matching the kinder tests.
A variable-object-count env has no single SesameModels bundle (the models bake in the count), so build_sesame_models raised -- leaving the coding agent unable to use the bilevel_models primitive there, a core axis of the generalized head-to-head experiment. Return a VariableCountBilevelModels accessor instead: the program calls models_for_state(state) to get the SesameModels for the current instance's count, built once per count and cached. The primitive description now covers both the fixed-count bundle and the variable-count accessor.
On a variable-count env the observation is an ObjectCentricState, not a flat vector. LLMGenPlanApproach now detects an ObjectCentricStateSpace and its GeneratedApproach interface spec describes the object-centric API (and appends the shared object-centric usage note) instead of calling the state a numpy observation. best_of_k inherits this through the shared prompt path.
A variable-count blackbox observation is a local ObjectCentricState, not a
remote handle, so a program passing it to a remote-module primitive (e.g.
crv_motion_planning) sends it by value as an {__ocs__} leaf. The handle-aware
decode_ref resolved only handle/ndarray/set tags and recursed the codec tag
into a raw dict, so the planner received a dict instead of a state. It now
honors registered codecs the way decode does, covering the deferred non-Box
types as well.
CDL is a prompt strategy (decompose into precondition/subgoal/policy-body behaviors); its Behavior base class takes an opaque state, so nothing structural ties it to flat vectors. Detect an ObjectCentricStateSpace and thread the object_centric flag through the prompt builders, and give the behavior-impl obs guidance object-centric variants (read typed objects via get_objects / get_object_from_name / get, no vector indices) selected for variable-count envs. The fixed-count (vector) prompt is byte-for-byte unchanged.
Adds an end-to-end test that pins object counts through run_episode (including a held-out count) and aggregates by_count over real rollouts, complementing the synthetic-dict summarize_by_count unit tests.
One parametrized test covers Obstruction2D, ClutteredRetrieval2D, ClutteredStorage2D, StickButton2D, and Motion2D: a pinned count is reported and set_state infers it back from the object-name prefix. A one_to_one flag captures that Motion2D makes ~2 obstacle objects per passage (so its prefixed-object count is not the count parameter). A separate test asserts ClutteredStorage2D rejects even block counts.
… wire) Clearbox: crv_motion_planning plans on a multi-object ObjectCentricState pulled from VariableObjectCountEnv, and crv_motion_planning_grasp reaches a geometric outcome over the same state. Blackbox: a local ObjectCentricState passed to the crv remote-module primitive round-trips over the wire and runs the host planner, guarding the decode_ref codec-tag fix. The wire setup is factored into a _blackbox_client helper shared with the existing object-centric loop test.
…unt reporting run_per_instance_eval attached object_count only to solved attempts, so crashed and budget-exhausted entries lacked it. Any consumer that groups per_episode by object_count (e.g. plot_scaling) then dropped those failures, inflating the per-count solve rate once the budget ran out. Every entry now carries its scheduled count, matching the full-scheduled denominator summarize_by_count uses.
A per-instance agentic run was told to solve env.reset(seed=S) but scored on
env.reset(seed=S, options={"object_count": K}) for a variable-count env, so the
agent developed against an unpinned design-count instance rather than the held-out
count it was evaluated on. The prompt now names the exact scored instance, threading
per_instance_count through build_agentic_prompt and the per-instance directive.
The black-box env_client spec always described a numpy observation vector (shape/low/high/dtype, devectorize/vectorize). For a variable-count env the client's observation space is object-centric (types/type_features/get_type) and get_state returns a local ObjectCentricState, so the vector spec pointed the agent at attributes that do not exist. build_agentic_prompt / build_cdl_prompt now select an object-centric interaction spec (read typed objects via get_objects / get / get_object_from_name) when the observation is object-centric; the fixed-count (vector) spec is byte-for-byte unchanged.
The black-box CRV notes told agents to build the planner state with observation_space.devectorize(obs), which does not exist on a variable-count env's object-centric client space (only types/type_features/get_type). Any agent following the note would crash. format_primitives_description is now object-centric-aware and emits notes that pass the ObjectCentricState straight to the host-side CRV planners; agentic, CDL, and genplan thread the flag through.
render_state (reset mode) and render_policy reset with env.reset(seed=seed) and no count, so a held-out per-instance (seed, count) evaluation was visualized as a different design-count instance -- undermining debugging exactly where count pinning matters. Both tools now take an optional object_count, threaded through the local and black-box implementations (env_client -> env server) into the reset.
The object-centric render_state description and system-prompt suffix still listed render_state(seed, label) / "render_state(seed=...), render_policy" without the object_count argument the tools now accept, so an agent in a per-instance variable-count run had no way to learn it could pin the scored held-out (seed, count) instance. Both the render_state description and the object-centric system suffix now advertise object_count for render_state and render_policy.
…rity) The black-box env_client space mirror and the object-centric prompts advertise observation_space.get_type(name), but the real ObjectCentricStateSpace exposes only .types. A program developed black-box (which naturally uses get_type) then crashed at eval with AttributeError: 'ObjectCentricStateSpace' object has no attribute 'get_type'. Attach get_type to the eval-time space (KeyError on an unknown name, matching the client) so test-time and eval-time APIs are identical. Found by an end-to-end object-centric black-box run.
Collaborator
Author
|
I think this is ready for merge. I have checked this as best I can going through all files and testing it end to end. Some points I want to discuss:
Supports all five 2D families and every approach except black-box bilevel_models (deferred by design) and 3D envs (follow-up). All paths are unit-tested. End-to-end (live-agent) validation covers Obstruction2D across the key approaches — generalized agentic (clearbox + black-box), the SeSamE planner head-to-head, and genplan — demonstrating the program-vs-planner scaling claim; the remaining approaches and families are unit-tested but not yet swept end-to-end. |
tomsilver
approved these changes
Jul 9, 2026
tomsilver
left a comment
tomsilver
left a comment
Owner
There was a problem hiding this comment.
not checked carefully but don't want to block!
Collaborator
Author
|
@tomsilver thanks! If you want/have time you can still check with no rush, I am not planning to do much between tonight and tomorrow unfortunately, so it's not super urgent to unblock |
tomsilver
reviewed
Jul 9, 2026
tomsilver
left a comment
tomsilver
left a comment
Owner
There was a problem hiding this comment.
a few thoughts after a slightly close look :)
Lower per-frame resolution keeps long, high-object-count rollouts within memory when capturing gifs; unset, each family keeps its default DPI.
Probing for methods with hasattr/getattr hides renames from mypy and degrades to a silent wrong branch when a probe misses. These six sites all ask the same question -- is this a VariableObjectCountEnv -- so ask it directly; the isinstance check narrows the type, so its methods are verified at type-check time.
to_box() is the single entry point the planner needs, so fold the current_box_space property and the current_box_obs helper (a copy of to_box over the current state) into it. Make _count_for_seed private, since only reset() calls it.
The family description, action-space, and reward prose are built by ConstantObjectKinDEREnv onto its own metadata; the inner object-centric env carries only render metadata. Reading the latter left the Action Space, Reward, and family-description sections of the agent's env card empty. Point the card at the wrapper's metadata (matching KinderGeom2DEnv) and assert the sections are populated so it cannot silently regress.
The reference exemplar drives the observation table; taking it from the smallest design count dropped the count-defining type entirely when that count was 0 (obstruction2d, motion2d), so the agent never saw that type or its features. Use the largest design count, which always contains every type, and drop the unused reference_count knob. Assert the table lists the type so it cannot regress.
Extend the per-family test to check every card section is populated, including the count-defining type row in the observation table. Guards the two card gaps (empty action/reward prose, missing object type) beyond the single obstruction2d case.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
I am working on this, but it's shaping up to be a pretty big PR and this week I will also have rebuttal for ARR, so I wanted to try something different and open a draft PR for now in case anyone else wants to have a look or try it out.
This is currently just claude's first try and I have to test it properly and review everything, starting to do that now