Skip to content

Commit fb37935

Browse files
yichao-liangyichao-liang
authored
Add gymnasium-API benchmark wrapper and retire mara_adapter (#21)
* Add gymnasium-API benchmark wrapper and retire mara_adapter shim Wrap predicators' native PyBullet envs in a standard gymnasium.Env so the suite can be consumed as a benchmark independent of the planning framework. Drop the predicators -> mara_robosim adapter layer, which only existed to bridge two near-identical struct types. New: - predicators/envs/gymnasium_wrapper.py: PredicatorsBenchmarkEnv plus register_all_environments(), make(), get_all_env_ids(). Registers 15 envs under mara/*-v0, with a cfg_overrides kwarg for per-make CFG overrides. - notebooks/getting_started.ipynb: interactive walkthrough. - scripts/benchmark_getting_started.py: smoke test that mirrors the notebook and resets every registered env. - tests/envs/test_gymnasium_wrapper.py: 12 unit tests. Removed: - predicators/envs/mara_adapter.py (790 lines). - mara_robosim git dependency from setup.py. - _normalize_env_name_for_gt helper in ground_truth_models, inlined env_name at all 7 callsites. Docs: - predicators/envs/README.md expanded into the benchmark README (install, quick-start, env table, standalone API, walkthroughs). - Top-level README points at the benchmark README. * Rebrand the gym wrapper as MARA RoboSim and audit env status User-facing rename: the Gymnasium wrapper now exports `MARARoboSimEnv` (was `PredicatorsBenchmarkEnv`), the import idiom is `from predicators.envs import gymnasium_wrapper as mara_robosim` (was `as benchmark`), and the smoke script + tests are renamed to `mara_robosim_getting_started.py` and `test_mara_robosim.py`. The underlying file stays `predicators/envs/gymnasium_wrapper.py` so the filename still describes the implementation. Env table: `predicators/envs/README.md` gains three status columns (Tasks / Skills / Demos) populated from a fresh oracle-planning audit: - 4 envs solve oracle test tasks end-to-end (Ants, Blocks, Circuit, Cover). - 8 envs have non-empty option factories but oracle planning currently fails on default config (Balance, Boil, Coffee, Domino, Fan, Float, Grow, Laser). - 3 envs ship empty option sets (Barrier, MagicBin, Switch). Also tweak the top-level README section heading to "MARA RoboSim". * Correct env status columns in MARA RoboSim README Update the Tasks/Skills/Demos columns based on a manual audit: several envs were previously marked ❌ for Demos but actually solve oracle tasks under different configs (Boil, Domino, Fan, Float, Grow); several others have a fixed task generator rather than randomized init/goal sampling (Ants, Barrier, Circuit, Float, Laser, MagicBin, Switch). * Add walkthroughs for getting started with MARA RoboSim environments * Move Walkthroughs section above the env table Promotes the notebook + smoke-script links so readers see them immediately after Quick Start, before scrolling through the 15-row status table.
1 parent 0e7aa3f commit fb37935

10 files changed

Lines changed: 976 additions & 836 deletions

File tree

‎.gitignore‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ logs*
2222
saved_approaches
2323
saved_datasets
2424
scripts/results
25+
scripts/mara_robosim_getting_started_output/
2526
pretrained_model_cache*
2627
tests/datasets/mock_vlm_datasets/cache/
2728
machines.txt

‎README.md‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,6 @@ This codebase implements a framework for *bilevel planning with learned neuro-sy
1111
5. [Embodied Active Learning of Relational State Abstractions for Bilevel Planning](https://arxiv.org/abs/2303.04912). Li, Silver. CoLLAs 2023.
1212
6. [Learning Efficient Abstract Planning Models that Choose What to Predict](https://arxiv.org/abs/2208.07737). Kumar*, McClinton*, Chitnis, Silver, Lozano-Perez, Kaelbling. CoRL 2023.
1313

14-
The codebase is still under active development. **Please contact <tslvr@mit.edu> or <njk@mit.edu> before attempting to use it for your own research.**
15-
1614
### Code Structure
1715

1816
In `predicators/`, the environments are defined in the `envs/` directory, and the approaches (both learning-based and not) are defined in the `approaches/` directory. The core [NSRT learning algorithm](https://arxiv.org/abs/2105.14074) happens in `predicators/nsrt_learning/nsrt_learning_main.py`, which has the following steps:
@@ -26,6 +24,9 @@ Methods for predicate learning are implemented as Approaches (e.g., `predicators
2624

2725
A simple implementation of search-then-sample bilevel planning is provided in `predicators/planning.py`. This implementation uses the "SeSamE" strategy: SEarch-and-SAMple planning, then Execution.
2826

27+
## MARA RoboSim
28+
Predicators ships **MARA RoboSim**, a collection of PyBullet manipulation environments exposed through a standard [Gymnasium](https://gymnasium.farama.org/) API and suitable for world-model learning, causal discovery, and RL research independent of the planning framework. See [`predicators/envs/README.md`](predicators/envs/README.md) for the env list, install instructions, quick-start code, standalone API, and getting-started notebook.
29+
2930
## Installation
3031
* This repository uses Python versions 3.10-3.11. We recommend 3.10.14.
3132
* Run `pip install -e .` to install dependencies.

‎notebooks/getting_started.ipynb‎

Lines changed: 245 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,245 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"metadata": {},
6+
"source": [
7+
"# Getting Started with MARA RoboSim\n",
8+
"\n",
9+
"This notebook walks through MARA RoboSim, the PyBullet manipulation\n",
10+
"suite shipped with [predicators](https://github.com/Learning-and-Intelligent-Systems/predicators):\n",
11+
"discovering available environments, creating one, taking actions,\n",
12+
"and rendering a video — all through the standard Gymnasium API."
13+
]
14+
},
15+
{
16+
"cell_type": "markdown",
17+
"metadata": {},
18+
"source": [
19+
"## Installation\n",
20+
"\n",
21+
"From the predicators repo root:\n",
22+
"\n",
23+
"```bash\n",
24+
"pip install -e .\n",
25+
"```\n",
26+
"\n",
27+
"This installs the agent solvers and MARA RoboSim together. See\n",
28+
"[`predicators/envs/README.md`](../predicators/envs/README.md) for\n",
29+
"more details."
30+
]
31+
},
32+
{
33+
"cell_type": "code",
34+
"execution_count": null,
35+
"metadata": {},
36+
"outputs": [],
37+
"source": [
38+
"import matplotlib\n",
39+
"\n",
40+
"matplotlib.use(\"agg\")\n",
41+
"\n",
42+
"import numpy as np\n",
43+
"import matplotlib.pyplot as plt\n",
44+
"\n",
45+
"from predicators import utils\n",
46+
"from predicators.envs import gymnasium_wrapper as mara_robosim\n",
47+
"\n",
48+
"# MARA RoboSim envs read configuration from predicators.settings.CFG.\n",
49+
"# `reset_config` applies parser defaults so we can use the envs as a\n",
50+
"# library without going through main.py's command-line interface.\n",
51+
"utils.reset_config({\"num_train_tasks\": 1, \"num_test_tasks\": 1})"
52+
]
53+
},
54+
{
55+
"cell_type": "markdown",
56+
"metadata": {},
57+
"source": [
58+
"## Discovering Available Environments\n",
59+
"\n",
60+
"MARA RoboSim provides 15 PyBullet-based robotic manipulation environments.\n",
61+
"Let's register them all and see what's available."
62+
]
63+
},
64+
{
65+
"cell_type": "code",
66+
"execution_count": null,
67+
"metadata": {},
68+
"outputs": [],
69+
"source": [
70+
"mara_robosim.register_all_environments()\n",
71+
"\n",
72+
"env_ids = sorted(mara_robosim.get_all_env_ids())\n",
73+
"print(f\"{len(env_ids)} environments available:\\n\")\n",
74+
"for eid in env_ids:\n",
75+
" print(f\" {eid}\")"
76+
]
77+
},
78+
{
79+
"cell_type": "markdown",
80+
"metadata": {},
81+
"source": [
82+
"## Creating an Environment\n",
83+
"\n",
84+
"We'll use `mara/Blocks-v0`, an environment where a Fetch robot\n",
85+
"must stack and arrange blocks on a tabletop."
86+
]
87+
},
88+
{
89+
"cell_type": "code",
90+
"execution_count": null,
91+
"metadata": {},
92+
"outputs": [],
93+
"source": [
94+
"env = mara_robosim.make(\"mara/Blocks-v0\", render_mode=\"rgb_array\")\n",
95+
"obs, info = env.reset()\n",
96+
"\n",
97+
"frame = env.render()\n",
98+
"if frame is not None:\n",
99+
" plt.imshow(frame)\n",
100+
" plt.axis(\"off\")\n",
101+
" plt.title(\"mara/Blocks-v0\")\n",
102+
" plt.show()\n",
103+
"else:\n",
104+
" print(\"(rendering not available in this configuration)\")"
105+
]
106+
},
107+
{
108+
"cell_type": "markdown",
109+
"metadata": {},
110+
"source": [
111+
"## Exploring the Observation and Action Spaces\n",
112+
"\n",
113+
"MARA RoboSim environments follow the Gymnasium API. Observations and\n",
114+
"actions are continuous-valued numpy arrays."
115+
]
116+
},
117+
{
118+
"cell_type": "code",
119+
"execution_count": null,
120+
"metadata": {},
121+
"outputs": [],
122+
"source": [
123+
"print(\"Observation shape:\", env.observation_space.shape)\n",
124+
"print(\"Action shape: \", env.action_space.shape)\n",
125+
"print()\n",
126+
"\n",
127+
"action = env.action_space.sample()\n",
128+
"obs, reward, terminated, truncated, info = env.step(action)\n",
129+
"print(\"Sample action:\", np.round(action, 3))\n",
130+
"print(\"Reward: \", reward)\n",
131+
"print(\"Terminated: \", terminated)"
132+
]
133+
},
134+
{
135+
"cell_type": "markdown",
136+
"metadata": {},
137+
"source": [
138+
"## Inspecting the Structured State\n",
139+
"\n",
140+
"Beyond the flat observation vector, the `info` dict exposes the full\n",
141+
"object-centric `State` and whether the goal has been reached."
142+
]
143+
},
144+
{
145+
"cell_type": "code",
146+
"execution_count": null,
147+
"metadata": {},
148+
"outputs": [],
149+
"source": [
150+
"state = info[\"state\"]\n",
151+
"print(\"Goal reached:\", info[\"goal_reached\"])\n",
152+
"print()\n",
153+
"print(state.pretty_str())"
154+
]
155+
},
156+
{
157+
"cell_type": "markdown",
158+
"metadata": {},
159+
"source": [
160+
"## Rendering a Video\n",
161+
"\n",
162+
"Let's collect frames from random actions and display them as an animated GIF."
163+
]
164+
},
165+
{
166+
"cell_type": "code",
167+
"execution_count": null,
168+
"metadata": {},
169+
"outputs": [],
170+
"source": [
171+
"from io import BytesIO\n",
172+
"\n",
173+
"from IPython.display import Image\n",
174+
"from PIL import Image as PILImage\n",
175+
"\n",
176+
"obs, info = env.reset()\n",
177+
"frames = []\n",
178+
"frame = env.render()\n",
179+
"if frame is not None:\n",
180+
" frames.append(frame)\n",
181+
"\n",
182+
"for _ in range(50):\n",
183+
" action = env.action_space.sample()\n",
184+
" obs, reward, terminated, truncated, info = env.step(action)\n",
185+
" frame = env.render()\n",
186+
" if frame is not None:\n",
187+
" frames.append(frame)\n",
188+
" if terminated or truncated:\n",
189+
" break\n",
190+
"\n",
191+
"if frames:\n",
192+
" pil_frames = [PILImage.fromarray(f) for f in frames]\n",
193+
" buf = BytesIO()\n",
194+
" pil_frames[0].save(\n",
195+
" buf,\n",
196+
" format=\"GIF\",\n",
197+
" save_all=True,\n",
198+
" append_images=pil_frames[1:],\n",
199+
" duration=100,\n",
200+
" loop=0,\n",
201+
" )\n",
202+
" Image(data=buf.getvalue(), format=\"gif\")\n",
203+
"else:\n",
204+
" print(\"(no frames captured \u2014 rendering may not be available)\")"
205+
]
206+
},
207+
{
208+
"cell_type": "markdown",
209+
"metadata": {},
210+
"source": [
211+
"## Cleanup"
212+
]
213+
},
214+
{
215+
"cell_type": "code",
216+
"execution_count": null,
217+
"metadata": {},
218+
"outputs": [],
219+
"source": [
220+
"env.close()"
221+
]
222+
}
223+
],
224+
"metadata": {
225+
"kernelspec": {
226+
"display_name": "Python 3",
227+
"language": "python",
228+
"name": "python3"
229+
},
230+
"language_info": {
231+
"codemirror_mode": {
232+
"name": "ipython",
233+
"version": 3
234+
},
235+
"file_extension": ".py",
236+
"mimetype": "text/x-python",
237+
"name": "python",
238+
"nbconvert_exporter": "python",
239+
"pygments_lexer": "ipython3",
240+
"version": "3.10.14"
241+
}
242+
},
243+
"nbformat": 4,
244+
"nbformat_minor": 4
245+
}

0 commit comments

Comments
 (0)