Skip to content

bump studio, load world from local file (GEN-923) #22

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

bump studio, load world from local file (GEN-923) #22

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3bf9730
- bump studio
mhuebert Jan 27, 2025
20f3ff8
ruff
mhuebert Jan 27, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 36 additions & 22 deletions genjax-localization-tutorial/probcomp-localization-tutorial.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,6 @@
import jax
import jax.numpy as jnp
import genjax
from urllib.request import urlopen
from genjax import SelectionBuilder as S
from genjax import ChoiceMapBuilder as C
from genjax.typing import Array, FloatArray, PRNGKey, IntArray
Expand Down Expand Up @@ -197,27 +196,42 @@ def make_world(wall_verts, clutters_vec, start, controls):

def load_world(file_name):
"""
Loads the world configuration from a specified file and constructs the world.
Loads the world configuration from local JSON files and constructs the world.

Args:
- file_name (str): The name of the file containing the world configuration.
- file_name (str): Not used, kept for backwards compatibility

Returns:
- tuple: A tuple containing the world configuration, the initial state, and the total number of control steps.
"""
with urlopen(
"https://raw.githubusercontent.com/probcomp/gen-localization/main/resources/example_20_program.json"
) as url:
data = json.load(url)
# Try to find the JSON files relative to different possible working directories
possible_paths = [
'.', # Current directory
'..', # Parent directory
os.path.dirname(__file__), # Directory containing this file
os.path.dirname(os.path.dirname(__file__)), # Parent of directory containing this file
]

def find_json(filename):
for path in possible_paths:
full_path = os.path.join(path, filename)
if os.path.exists(full_path):
with open(full_path, 'r') as f:
return json.load(f)
raise FileNotFoundError(f"Could not find {filename} in any of: {possible_paths}")

# Load both required JSON files
world_data = find_json('world.json')
program_data = find_json('robot_program.json')

walls_vec = jnp.array(data["wall_verts"])
clutters_vec = jnp.array(data["clutter_vert_groups"])
walls_vec = jnp.array(world_data["wall_verts"])
clutters_vec = jnp.array(world_data["clutter_vert_groups"])
start = Pose(
jnp.array(data["start_pose"]["p"], dtype=float),
jnp.array(data["start_pose"]["hd"], dtype=float),
jnp.array(program_data["start_pose"]["p"], dtype=float),
jnp.array(program_data["start_pose"]["hd"], dtype=float),
)

cs = jnp.array([[c["ds"], c["dhd"]] for c in data["program_controls"]])
cs = jnp.array([[c["ds"], c["dhd"]] for c in program_data["program_controls"]])
controls = Control(cs[:, 0], cs[:, 1])

return make_world(walls_vec, clutters_vec, start, controls)
Expand Down Expand Up @@ -426,7 +440,7 @@ def pose_plot(p, fill: str | Any = "black", **opts):
strokeWidth=2,
stroke="#ccc",
),
{"margin": 0, "inset": 50, "width": 500, "axis": None, "aspectRatio": 1},
{"margin": 0, "inset": 50, "maxWidth": 500, "axis": None, "aspectRatio": 1},
Plot.domain([0, 20]),
)
# Plot the world with walls only
Expand Down Expand Up @@ -807,7 +821,7 @@ def animate_path_with_confidence(path: Pose, motion_settings: dict):
+ Plot.color_map({"some pose": "green", "with heading modified": "red"})
+ Plot.title("Modifying a heading")
)
| html("span.tc", f"score ratio: {rotated_trace_weight_diff}")
| html(["span.tc", f"score ratio: {rotated_trace_weight_diff}"])
)

# %% [markdown]
Expand All @@ -834,7 +848,7 @@ def animate_path_with_confidence(path: Pose, motion_settings: dict):
for pose in path_from_trace(trace)
]
+ Plot.color_map({"some path": "green", "with heading modified": "red"})
) | html("span.tc", f"score ratio: {rotated_first_step_weight_diff}")
) | html(["span.tc", f"score ratio: {rotated_first_step_weight_diff}"])

# %% [markdown]
# ### Ideal sensors
Expand Down Expand Up @@ -1127,7 +1141,7 @@ def plt(readings):
return Plot.new(
plot_base or Plot.domain([0, 20]),
plot_sensors(pose, readings),
{"width": 400, "height": 400},
{"maxWidth": 400, "aspectRatio": 1},
)

return plt(readings1) & plt(readings2)
Expand Down Expand Up @@ -1169,13 +1183,13 @@ def plt(readings):
sample, log_weight = model_importance(
sub_key, constraints_low_deviation, (motion_settings_low_deviation,)
)
animate_full_trace(sample) | html("span.tc", f"log_weight: {log_weight}")
animate_full_trace(sample) | html(["span.tc", f"log_weight: {log_weight}"])
# %%
key, sub_key = jax.random.split(key)
sample, log_weight = model_importance(
sub_key, constraints_high_deviation, (motion_settings_high_deviation,)
)
animate_full_trace(sample) | html("span.tc", f"log_weight: {log_weight}")
animate_full_trace(sample) | html(["span.tc", f"log_weight: {log_weight}"])
# %% [markdown]
# A trace resulting from a call to `importance` is structurally indistinguishable from one drawn from `simulate`. But there is a key situational difference: while `get_score` always returns the frequency with which `simulate` stochastically produces the trace, this value is **no longer equal to** the frequency with which the trace is stochastically produced by `importance`. This is both true in an obvious and less relevant sense, as well as true in a more subtle and extremely germane sense.
#
Expand Down Expand Up @@ -1256,9 +1270,9 @@ def constraint_from_path(path):
Plot.Row(
*[
(
html("div.f3.b.tc", title)
html(["div.f3.b.tc", title])
| animate_full_trace(trace, frame_key="frame")
| html("span.tc", f"score: {score:,.2f}")
| html(["span.tc", f"score: {score:,.2f}"])
)
for (title, trace, motion_settings, score) in [
[
Expand Down Expand Up @@ -1521,8 +1535,8 @@ class SequentialImportanceSampling(Generic[StateT, ControlT]):
Given:
- a functional wrapper for the importance method of a generative function
- an initial state of type StateT, which should be a PyTree $z_0$
- a vector of control inputs, also a PyTree $u_i, of shape $(T, \ldots)$
- an array of observations $y_i$, also of shape $(T, \ldots)$
- a vector of control inputs, also a PyTree $u_i, of shape $(T, \\ldots)$
- an array of observations $y_i$, also of shape $(T, \\ldots)$
perform the inference technique known as Sequential Importance Sampling.

The signature of the GFI importance method is
Expand Down
23 changes: 13 additions & 10 deletions poetry.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ package-mode = false
python = ">=3.11,<3.13"
jupytext = "^1.16.1"
genjax = {version = "0.7.0.post4.dev0+eacb241e" , source = "gcp" }
genstudio = {version = "2024.10.1", source = "gcp"}
genstudio = "2025.1.12"
ipykernel = "^6.29.3"
matplotlib = "^3.8.3"
anywidget = "^0.9.7"
Expand Down
Loading